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 => @host, :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-25,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 check_cluster
 152         puts "Performing Cluster Check (using node #{@nodes[0]})"
 153         show_nodes
 154     end
 155
 156     def alloc_slots
 157         slots_per_node = ClusterHashSlots/@nodes.length
 158         i = 0
 159         @nodes.each{|n|
 160             first = i*slots_per_node
 161             last = first+slots_per_node-1
 162             last = ClusterHashSlots-1 if i == @nodes.length-1
 163             n.add_slots first..last
 164             i += 1
 165         }
 166     end
 167
 168     def flush_nodes_config
 169         @nodes.each{|n|
 170             n.flush_node_config
 171         }
 172     end
 173
 174     def show_nodes
 175         @nodes.each{|n|
 176             puts n.info_string
 177         }
 178     end
 179
 180     def join_cluster
 181         # We use a brute force approach to make sure the node will meet
 182         # each other, that is, sending CLUSTER MEET messages to all the nodes
 183         # about the very same node.
 184         # Thanks to gossip this information should propagate across all the
 185         # cluster in a matter of seconds.
 186         first = false
 187         @nodes.each{|n|
 188             if !first then first = n.info; next; end # Skip the first node
 189             n.r.cluster("meet",first[:host],first[:port])
 190         }
 191     end
 192
 193     def yes_or_die(msg)
 194         print "#{msg} (type 'yes' to accept): "
 195         STDOUT.flush
 196         if !(STDIN.gets.chomp.downcase == "yes")
 197             puts "Aborting..."
 198             exit 1
 199         end
 200     end
 201
 202     # redis-trib subcommands implementations
 203
 204     def check_cluster_cmd
 205         node = ClusterNode.new(ARGV[1])
 206         node.connect(:abort => true)
 207         node.assert_cluster
 208         add_node(node)
 209         check_cluster
 210     end
 211
 212     def create_cluster_cmd
 213         puts "Creating cluster"
 214         ARGV[1..-1].each{|n|
 215             node = ClusterNode.new(n)
 216             node.connect(:abort => true)
 217             node.assert_cluster
 218             node.assert_empty
 219             add_node(node)
 220         }
 221         puts "Performing hash slots allocation on #{@nodes.length} nodes..."
 222         alloc_slots
 223         show_nodes
 224         yes_or_die "Can I set the above configuration?"
 225         flush_nodes_config
 226         puts "** Nodes configuration updated"
 227         puts "** Sending CLUSTER MEET messages to join the cluster"
 228         join_cluster
 229         check_cluster
 230     end
 231 end
 232
 233 COMMANDS={
 234     "create" => ["create_cluster_cmd", -2, "host1:port host2:port ... hostN:port"],
 235     "check" =>  ["check_cluster_cmd", 2, "host:port"]
 236 }
 237
 238 # Sanity check
 239 if ARGV.length == 0
 240     puts "Usage: redis-trib <command> <arguments ...>"
 241     puts
 242     COMMANDS.each{|k,v|
 243         puts "  #{k.ljust(20)} #{v[2]}"
 244     }
 245     puts
 246     exit 1
 247 end
 248
 249 rt = RedisTrib.new
 250 cmd_spec = COMMANDS[ARGV[0].downcase]
 251 if !cmd_spec
 252     puts "Unknown redis-trib subcommand '#{ARGV[0]}'"
 253     exit 1
 254 end
 255 rt.check_arity(cmd_spec[1],ARGV.length)
 256
 257 # Dispatch
 258 rt.send(cmd_spec[0])