Create a gist now

Instantly share code, notes, and snippets.

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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment