src/redis-trib.rb

   1 #!/usr/bin/env ruby
   2
   3 require 'rubygems'
   4 require 'redis'
   5
   6 ClusterHashSlots = 4096
   7
   8 def xputs(s)
   9     printf s
  10     STDOUT.flush
  11 end
  12
  13 class ClusterNode
  14     def initialize(addr)
  15         s = addr.split(":")
  16         if s.length != 2
  17             puts "Invalid node name #{node}"
  18             exit 1
  19         end
  20         @host = s[0]
  21         @port = s[1]
  22         @slots = {}
  23         @dirty = false
  24     end
  25
  26     def to_s
  27         "#{@host}:#{@port}"
  28     end
  29
  30     def connect
  31         xputs "Connecting to node #{self}: "
  32         begin
  33             @r = Redis.new(:host => @ost, :port => @port)
  34             @r.ping
  35         rescue
  36             puts "ERROR"
  37             puts "Sorry, can't connect to node #{self}"
  38         end
  39         puts "OK"
  40     end
  41
  42     def assert_cluster
  43         info = @r.info
  44         if !info["cluster_enabled"] || info["cluster_enabled"].to_i == 0
  45             puts "Error: Node #{self} is not configured as a cluster node."
  46             exit 1
  47         end
  48     end
  49
  50     def assert_empty
  51         if !(@r.cluster("info").split("\r\n").index("cluster_known_nodes:1")) ||
  52             (@r.info['db0'])
  53             puts "Error: Node #{self} is not empty. Either the node already knows other nodes (check with nodes-info) or contains some key in database 0."
  54             exit 1
  55         end
  56     end
  57
  58     def add_slots(slots)
  59         slots.each{|s|
  60             @slots[s] = :new
  61         }
  62         @dirty = true
  63     end
  64
  65     def flush_node_config
  66         return if !@dirty
  67         new = []
  68         @slots.each{|s,val|
  69             if val == :new
  70                 new << s
  71                 @slots[s] = true
  72             end
  73         }
  74         @r.cluster("addslots",*new)
  75         @dirty = false
  76     end
  77
  78     def info
  79         slots = @slots.map{|k,v| k}.reduce{|a,b|
  80             a = [(a..a)] if !a.is_a?(Array)
  81             if b == (a[-1].last)+1
  82                 a[-1] = (a[-1].first)..b
  83                 a
  84             else
  85                 a << (b..b)
  86             end
  87         }.map{|x|
  88             (x.first == x.last) ? x.first.to_s : "#{x.first}-#{x.last}"
  89         }.join(",")
  90         "#{self.to_s.ljust(25)} slots:#{slots}"
  91     end
  92
  93     def is_dirty?
  94         @dirty
  95     end
  96
  97     def r
  98         @r
  99     end
 100 end
 101
 102 class RedisTrib
 103     def initialize
 104         @nodes = []
 105     end
 106
 107     def check_arity(req_args, num_args)
 108         if ((req_args > 0 and num_args != req_args) ||
 109            (req_args < 0 and num_args < req_args.abs))
 110            puts "Wrong number of arguments for specified sub command"
 111            exit 1
 112         end
 113     end
 114
 115     def create_cluster
 116         puts "Creating cluster"
 117         ARGV[1..-1].each{|n|
 118             node = ClusterNode.new(n)
 119             node.connect
 120             node.assert_cluster
 121             node.assert_empty
 122             @nodes << node
 123         }
 124         puts "Performing hash slots allocation on #{@nodes.length} nodes..."
 125         alloc_slots
 126         show_nodes
 127         yes_or_die "Can I set the above configuration?"
 128         flush_nodes_config
 129         puts "** Nodes configuration updated"
 130         puts "Sending CLUSTER MEET messages to join the cluster"
 131         join_cluster
 132     end
 133
 134     def alloc_slots
 135         slots_per_node = ClusterHashSlots/@nodes.length
 136         i = 0
 137         @nodes.each{|n|
 138             first = i*slots_per_node
 139             last = first+slots_per_node-1
 140             last = ClusterHashSlots-1 if i == @nodes.length-1
 141             n.add_slots first..last
 142             i += 1
 143         }
 144     end
 145
 146     def flush_nodes_config
 147         @nodes.each{|n|
 148             n.flush_node_config
 149         }
 150     end
 151
 152     def show_nodes
 153         @nodes.each{|n|
 154             puts n.info
 155         }
 156     end
 157
 158     def join_cluster
 159     end
 160
 161     def yes_or_die(msg)
 162         print "#{msg} (type 'yes' to accept): "
 163         STDOUT.flush
 164         if !(STDIN.gets.chomp.downcase == "yes")
 165             puts "Aborting..."
 166             exit 1
 167         end
 168     end
 169 end
 170
 171 COMMANDS={
 172     "create-cluster" => ["create_cluster", -2]
 173 }
 174
 175 # Sanity check
 176 if ARGV.length == 0
 177     puts "Usage: redis-trib <command> <arguments ...>"
 178     exit 1
 179 end
 180
 181 rt = RedisTrib.new
 182 cmd_spec = COMMANDS[ARGV[0].downcase]
 183 if !cmd_spec
 184     puts "Unknown redis-trib subcommand '#{ARGV[0]}'"
 185     exit 1
 186 end
 187 rt.check_arity(cmd_spec[1],ARGV.length)
 188
 189 # Dispatch
 190 rt.send(cmd_spec[0])