ruby_mongo_sharding.rb

ruby_mongo_sharding.rb

# -*- coding: utf-8 -*-
require 'pp'
require 'fileutils'
require 'mongo'
require 'yaml'

module Sample
  class Mongo
    BASE_DATA_PATH = '/tmp/mongodb'
    LOG_FILE_PATH = "#{BASE_DATA_PATH}/log"

    MONGO_D_PORT = 10000
    MONGO_C_PORT = 10010
    MONGO_S_PORT = 27018
    CHUNK_SIZE   = 20 # in MB
    SHARDS_NUM = 2
    DB_NAME = 'sample'

    def self.configure(&blocks)
      self.new(&blocks)
    end

    def initialize(&blocks)
      @connection = nil
      @workers = []
      start(&blocks)
    end

    def connection
      @connection ||= ::Mongo::Connection.new("localhost", MONGO_S_PORT)
    end

    private
    def start(&blocks)
      begin
        start_mongo
        configure_sharding
        blocks.call(self) if block_given?
        loop {}
      rescue => e
        puts e
        shutdown
      end
    end

    # starting mongo
    def start_mongo
      trap do
        bootstrap
        starting_shard
        starting_mongo_c
        starting_mongo_s
      end
    end

    # configure sharding
    def configure_sharding
      db_admin = connection.db("admin")
      (1..SHARDS_NUM).each do |i|
        path = "localhost:#{MONGO_D_PORT + i}"
        output "Add Shard: #{path}" do
          db_admin.command addshard: path
        end
      end

      output 'List Shards' do
        db_admin.command listshards: 1
      end

      output "Enable Sharding: #{DB_NAME}" do
        db_admin.command enablesharding: DB_NAME
      end

      YAML.load_file(File.expand_path(File.dirname(__FILE__) + '/shard_key.yml')).each do |col, skey|
        hash = skey.each_with_object({}) {|attr, _hash| _hash[attr] = 1}
        db_admin.command shardcollection: "#{DB_NAME}.#{col}", key: hash
      end
    end

    # remove all data and create log dir
    def bootstrap
      FileUtils.rm_rf BASE_DATA_PATH
      FileUtils.mkdir_p LOG_FILE_PATH
    end

    # starting shard(mongo_d)
    def starting_shard
      (1..SHARDS_NUM).each do |i|
        FileUtils.mkdir_p "#{BASE_DATA_PATH}/shard#{i}"
        cmd = "mongod --shardsvr --port #{MONGO_D_PORT + i} --dbpath #{BASE_DATA_PATH}/shard#{i} --logpath #{LOG_FILE_PATH}/shard#{i}.log"
        @workers << spawn(cmd, :out => '/dev/null')
      end
    end

    # starting mongo_c
    def starting_mongo_c
      FileUtils.mkdir_p "#{BASE_DATA_PATH}/config"
      cmd = "mongod --configsvr --port #{MONGO_C_PORT} --dbpath #{BASE_DATA_PATH}/config --logpath #{LOG_FILE_PATH}/config.log"
      @workers << spawn(cmd, :out => '/dev/null')
    end

    # starting mongo_s
    def starting_mongo_s
      sleep 1
      cmd = "mongos --configdb localhost:#{MONGO_C_PORT} --port #{MONGO_S_PORT} --logpath #{LOG_FILE_PATH}/mongos.log --chunkSize #{CHUNK_SIZE}"
      @workers << spawn(cmd, :out => '/dev/null')
      sleep 1
    end

    # kill all mongo prosecc
    def shutdown
      puts ">> going to shutdown ..."
      @workers.each do |pid|
        Process.kill 'KILL', pid
      end
      exit 0
    end

    # trap signals
    def trap(&block)
      ['INT', 'TERM', 'KILL'].each do |signal|
        Signal.trap(signal) { shutdown }
      end
      block.call
    end

    # configure output format
    def output(subject, &blocks)
      puts "-- #{subject}"
      pp blocks.call
      puts ''
    end
  end
end

# if you want preconfigure
Sample::Mongo.configure do |mongo|
  # db_test = mongo.connection.db("sg")
  # coll = db_test.collection("sample.car")
  # (1..100000).each do |i|
  #   coll.insert({:_id => i, :maker => "toyota_#{i}", :name => "car_name_#{i}"})
  # end
end

# or 
Sample::Mongo.new