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