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