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 #{addr}"
  18             exit 1
  19         end
  20         @r = nil
  21         @host = s[0]
  22         @port = s[1]
  23         @slots = {}
  24         @dirty = false
  25     end
  26
  27     def to_s
  28         "#{@host}:#{@port}"
  29     end
  30
  31     def connect(o={})
  32         xputs "Connecting to node #{self}: "
  33         begin
  34             @r = Redis.new(:host => @ost, :port => @port)
  35             @r.ping
  36         rescue
  37             puts "ERROR"
  38             puts "Sorry, can't connect to node #{self}"
  39             exit 1 if o[:abort]
  40             @r = nil
  41         end
  42         puts "OK"
  43     end
  44
  45     def assert_cluster
  46         info = @r.info
  47         if !info["cluster_enabled"] || info["cluster_enabled"].to_i == 0
  48             puts "Error: Node #{self} is not configured as a cluster node."
  49             exit 1
  50         end
  51     end
  52
  53     def assert_empty
  54         if !(@r.cluster("info").split("\r\n").index("cluster_known_nodes:1")) ||
  55             (@r.info['db0'])
  56             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."
  57             exit 1
  58         end
  59     end
  60
  61     def add_slots(slots)
  62         slots.each{|s|
  63             @slots[s] = :new
  64         }
  65         @dirty = true
  66     end
  67
  68     def flush_node_config
  69         return if !@dirty
  70         new = []
  71         @slots.each{|s,val|
  72             if val == :new
  73                 new << s
  74                 @slots[s] = true
  75             end
  76         }
  77         @r.cluster("addslots",*new)
  78         @dirty = false
  79     end
  80
  81     def info_string
  82         # We want to display the hash slots assigned to this node
  83         # as ranges, like in: "1-5,8,9,20-35,30"
  84         #
  85         # Note: this could be easily written without side effects,
  86         # we use 'slots' just to split the computation into steps.
  87
  88         # First step: we want an increasing array of integers
  89         # for instance: [1,2,3,4,5,8,9,20,21,22,23,24,25,30]
  90         slots = @slots.keys.sort
  91
  92         # As we want to aggregate adiacent slots we convert all the
  93         # slot integers into ranges (with just one element)
  94         # So we have something like [1..1,2..2, ... and so forth.
  95         slots.map!{|x| x..x}
  96
  97         # Finally we group ranges with adiacent elements.
  98         slots = slots.reduce([]) {|a,b|
  99             if !a.empty? && b.first == (a[-1].last)+1
 100                 a[0..-2] + [(a[-1].first)..(b.last)]
 101             else
 102                 a + [b]
 103             end
 104         }
 105
 106         # Now our task is easy, we just convert ranges with just one
 107         # element into a number, and a real range into a start-end format.
 108         # Finally we join the array using the comma as separator.
 109         slots = slots.map{|x|
 110             x.count == 1 ? x.first.to_s : "#{x.first}-#{x.last}"
 111         }.join(",")
 112
 113         "#{self.to_s.ljust(25)} slots:#{slots}"
 114     end
 115
 116     def info
 117         {
 118             :host => @host,
 119             :port => @port,
 120             :slots => @slots,
 121             :dirty => @dirty
 122         }
 123     end
 124
 125     def is_dirty?
 126         @dirty
 127     end
 128
 129     def r
 130         @r
 131     end
 132 end
 133
 134 class RedisTrib
 135     def initialize
 136         @nodes = []
 137     end
 138
 139     def check_arity(req_args, num_args)
 140         if ((req_args > 0 and num_args != req_args) ||
 141            (req_args < 0 and num_args < req_args.abs))
 142            puts "Wrong number of arguments for specified sub command"
 143            exit 1
 144         end
 145     end
 146
 147     def add_node(node)
 148         @nodes << node
 149     end
 150
 151     def create_cluster
 152         puts "Creating cluster"
 153         ARGV[1..-1].each{|n|
 154             node = ClusterNode.new(n)
 155             node.connect(:abort => true)
 156             node.assert_cluster
 157             node.assert_empty
 158             add_node(node)
 159         }
 160         puts "Performing hash slots allocation on #{@nodes.length} nodes..."
 161         alloc_slots
 162         show_nodes
 163         yes_or_die "Can I set the above configuration?"
 164         flush_nodes_config
 165         puts "** Nodes configuration updated"
 166         puts "** Sending CLUSTER MEET messages to join the cluster"
 167         join_cluster
 168         check_cluster
 169     end
 170
 171     def check_cluster
 172         puts "Performing Cluster Check (node #{ARGV[1]})"
 173         node = ClusterNode.new(ARGV[1])
 174         node.connect(:abort => true)
 175         node.assert_cluster
 176         node.add_slots(10..15)
 177         node.add_slots(30..30)
 178         node.add_slots(5..5)
 179         add_node(node)
 180         show_nodes
 181     end
 182
 183     def alloc_slots
 184         slots_per_node = ClusterHashSlots/@nodes.length
 185         i = 0
 186         @nodes.each{|n|
 187             first = i*slots_per_node
 188             last = first+slots_per_node-1
 189             last = ClusterHashSlots-1 if i == @nodes.length-1
 190             n.add_slots first..last
 191             i += 1
 192         }
 193     end
 194
 195     def flush_nodes_config
 196         @nodes.each{|n|
 197             n.flush_node_config
 198         }
 199     end
 200
 201     def show_nodes
 202         @nodes.each{|n|
 203             puts n.info_string
 204         }
 205     end
 206
 207     def join_cluster
 208         # We use a brute force approach to make sure the node will meet
 209         # each other, that is, sending CLUSTER MEET messages to all the nodes
 210         # about the very same node.
 211         # Thanks to gossip this information should propagate across all the
 212         # cluster in a matter of seconds.
 213         first = false
 214         @nodes.each{|n|
 215             if !first then first = n.info; next; end # Skip the first node
 216             n.r.cluster("meet",first[:host],first[:port])
 217         }
 218     end
 219
 220     def yes_or_die(msg)
 221         print "#{msg} (type 'yes' to accept): "
 222         STDOUT.flush
 223         if !(STDIN.gets.chomp.downcase == "yes")
 224             puts "Aborting..."
 225             exit 1
 226         end
 227     end
 228 end
 229
 230 COMMANDS={
 231     "create" => ["create_cluster", -2, "host1:port host2:port ... hostN:port"],
 232     "check" =>  ["check_cluster", 2, "host:port"]
 233 }
 234
 235 # Sanity check
 236 if ARGV.length == 0
 237     puts "Usage: redis-trib <command> <arguments ...>"
 238     puts
 239     COMMANDS.each{|k,v|
 240         puts "  #{k.ljust(20)} #{v[2]}"
 241     }
 242     puts
 243     exit 1
 244 end
 245
 246 rt = RedisTrib.new
 247 cmd_spec = COMMANDS[ARGV[0].downcase]
 248 if !cmd_spec
 249     puts "Unknown redis-trib subcommand '#{ARGV[0]}'"
 250     exit 1
 251 end
 252 rt.check_arity(cmd_spec[1],ARGV.length)
 253
 254 # Dispatch
 255 rt.send(cmd_spec[0])