btrfs auto rebalancing

#!/usr/bin/ruby

# Maximum free space wasted as percentage
MAX_WASTED_RATIO = 0.3
# Maximum unallocated space above which no check is needed
UNALLOCATED_THRESHOLD_RATIO = 0.3
# Target max waste when rebalancing (not much less than max_wasted)
TARGET_WASTED_RATIO = 0.25
MAX_FAILURES = 20
MAX_REBALANCES = 100

require 'open3'

def filesystems
  mounts = []
  IO.popen("mount") do |io|
    io.each_line do |line|
      if match = line.match(/^(.*) on (.*) type btrfs/)
        mounts << [ match[1], match[2] ] unless mounts.map(&:first).include?(match[1])
      end
    end
  end
  return mounts.map { |m| m[1] }
end

class Btrfs
  def initialize(mountpoint)
    @mountpoint = mountpoint
    refresh_usage
  end

  def refresh_usage
    IO.popen("btrfs fi usage --raw '#{@mountpoint}'") do |io|
      io.each_line do |line|
        case line
        when /Device size:\s*(\d*)$/
          @device_size = Regexp.last_match[1].to_i
        when /Device allocated:\s*(\d*)$/
          @allocated = Regexp.last_match[1].to_i
        when /Device unallocated:\s*(\d*)$/
          @unallocated = Regexp.last_match[1].to_i
        when /Free \(estimated\):\s*(\d*)\s*\(min: \d*\)$/
          @free = Regexp.last_match[1].to_i
        when /Data ratio:\s*(\d+\.\d+)$/
          @ratio = Regexp.last_match[1].to_f
        end
      end
    end
  end

  def rebalance_if_needed
    if rebalance_needed?
      log("waste %.2f%%" % (free_wasted * 100))
      log "unallocated: #{@unallocated}"
      log "ratio:       #{@ratio}"
      log "free:        #{@free}"
      rebalance
    end
  rescue => ex
    log "can't rebalance, error: #{ex}"
  end

  def rebalance
    failures = 0
    count = 0
    log "First quick pass in case removing empty chunks would be enough"
    balance_usage(0)
    # Try to guess best usage_target
    usage_target = ((1 - TARGET_WASTED_RATIO) * 50).to_i
    while (free_wasted > target_wasted) && (count < MAX_REBALANCES)
      log "rebalance with usage: #{usage_target}"
      count += 1
      fail "not enough space, remove snapshots or files" if usage_target == 0
      exit_status = balance_usage(usage_target)
      if exit_status.exitstatus != 0
        failures += 1
        if usage_target == 0
          fail "can't retry rebalancing, usage_target reached 0 already"
        end
        if failures >= MAX_FAILURES
          fail "too many rebalance failures: MAX_FAILURES"
        end
        usage_target -= 1
      else
        failures = 0
        usage_target += 1
      end
      refresh_usage
      log("waste %.2f%%" % (free_wasted * 100))
    end
  end

  def balance_usage(target)
    output, status = Open3.capture2e("btrfs fi balance start -dusage=#{target} #{@mountpoint}")
    log output
    return status
  end

  def log(msg)
    puts "#{@mountpoint}: #{msg}"
  end

  def rebalance_needed?
    (unallocated_ratio < UNALLOCATED_THRESHOLD_RATIO) && (free_wasted > max_wasted)
  end

  def free_wasted
    1 - (@unallocated.to_f / (@free * @ratio))
  end
  def unallocated_ratio
    @unallocated.to_f / @device_size
  end
  def max_wasted
    MAX_WASTED_RATIO
  end

  def target_wasted
    TARGET_WASTED_RATIO
  end
end

def rebalance_fs
  filesystems.each do |fs|
    Btrfs.new(fs).rebalance_if_needed
  end
end

rebalance_fs