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