mongochunks.rb

require "moped"
require "pp"

class Mongochunks
  attr_accessor :interval, :config_session, :database_session, :sort_by, :stat
  
  def initialize(config_host: "localhost:27017", host: "localhost:27017", database: "test", sort_by: :total, interval: 1)
    self.interval = interval
    self.config_session = Moped::Session.new([config_host])
    self.config_session.use(:config)
    self.database_session = Moped::Session.new([host])
    self.database_session.use(database)
    self.sort_by = sort_by
    self.chunks
  end
  
  def chunks
    return @chunks unless @chunks.nil?
    
    @chunks = {}
    self.config_session[:chunks].find.each do |chunk|
      @chunks[chunk["ns"]] ||= {}
      @chunks[chunk["ns"]][:shard_key] ||= chunk["min"].keys.first
      @chunks[chunk["ns"]][:chunks] ||= []
      
      @chunks[chunk["ns"]][:chunks] << {
        id: chunk["_id"],
        shard: chunk["shard"],
        min: chunk["min"][@chunks[chunk["ns"]][:shard_key]],
        max: chunk["max"][@chunks[chunk["ns"]][:shard_key]]
      }
    end
    
    @chunks
  end
  
  def find_chunks(record)
    collection = self.chunks[record["ns"]]
    # TODO: Inserts?
    return false if !collection || !record["query"]
    
    ids = case record["query"][collection[:shard_key]].class.to_s
    when "Moped::BSON::ObjectId"
      [record["query"][collection[:shard_key]]]
    when "Hash"
      query = record["query"][collection[:shard_key]]
      if query["$in"]
        query["$in"]
      end
    end
    return false unless ids
    
    chunks = ids.map do |id|
      collection[:chunks].find do |chunk|
        true if
          (
            chunk[:min] != Moped::BSON::MinKey &&
            chunk[:max] != Moped::BSON::MaxKey &&
            id >= chunk[:min] && id < chunk[:max]
          ) ||
          (chunk[:min] == Moped::BSON::MinKey && id < chunk[:max]) ||
          (chunk[:max] == Moped::BSON::MaxKey && id >= chunk[:min])
      end
    end
    
    chunks
  end
  
  def print_stat
    stat = self.stat
    stat.each do |collection, chunks|
      stat[collection] = chunks.sort_by { |_, stats| stats[self.sort_by] }
      # chunks.each do |id, stats|
      #   pp self.chunks[collection][:chunks].find { |x| x[:id] == id }
      # end
    end
    
    puts ""
    puts Time.new.utc
    pp stat
    puts ""
  end
  
  def fire
    while true do
      self.stat = {}
      
      offset = Time.new - self.interval
      records = database_session[:"system.profile"].find(ts: { "$gte" => offset }).to_a
      records.each do |record|
        # # FIX
        # record["ns"].gsub!("reader_test", "reader")
        chunks = self.find_chunks(record)
        next unless chunks
        
        chunks.each do |chunk|
          self.stat[record["ns"]] ||= {}
          self.stat[record["ns"]][chunk[:id]] ||= { query: 0, update: 0, insert: 0, delete: 0, total: 0 }
          self.stat[record["ns"]][chunk[:id]][record["op"].to_sym] += 1
          self.stat[record["ns"]][chunk[:id]][:total] += 1
        end
      end
      
      self.print_stat
      sleep interval
    end
  end
end

chunks = Mongochunks.new(config_host: "mongo-config1.theoldreader.intra:27019", host: "mongo.theoldreader.intra:10001", database: "reader")
# chunks = Mongochunks.new(database: "reader_test")
chunks.fire