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         @info = {}
  22         @info[:host] = s[0]
  23         @info[:port] = s[1]
  24         @info[:slots] = {}
  25         @dirty = false # True if we need to flush slots info into node.
  26         @friends = []
  27     end
  28
  29     def friends
  30         @friends
  31     end
  32
  33     def slots
  34         @info[:slots]
  35     end
  36
  37     def to_s
  38         "#{@info[:host]}:#{@info[: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 => @info[:host], :port => @info[: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.merge(info)
  89                 @info[: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                 } if slots
  98                 @dirty = false
  99                 @r.cluster("info").split("\n").each{|e|
 100                     k,v=e.split(":")
 101                     k = k.to_sym
 102                     v.chop!
 103                     if k != :cluster_state
 104                         @info[k] = v.to_i
 105                     else
 106                         @info[k] = v
 107                     end
 108                 }
 109             elsif o[:getfriends]
 110                 @friends << info
 111             end
 112         }
 113     end
 114
 115     def add_slots(slots)
 116         slots.each{|s|
 117             @info[:slots][s] = :new
 118         }
 119         @dirty = true
 120     end
 121
 122     def flush_node_config
 123         return if !@dirty
 124         new = []
 125         @info[:slots].each{|s,val|
 126             if val == :new
 127                 new << s
 128                 @info[:slots][s] = true
 129             end
 130         }
 131         @r.cluster("addslots",*new)
 132         @dirty = false
 133     end
 134
 135     def info_string
 136         # We want to display the hash slots assigned to this node
 137         # as ranges, like in: "1-5,8-9,20-25,30"
 138         #
 139         # Note: this could be easily written without side effects,
 140         # we use 'slots' just to split the computation into steps.
 141
 142         # First step: we want an increasing array of integers
 143         # for instance: [1,2,3,4,5,8,9,20,21,22,23,24,25,30]
 144         slots = @info[:slots].keys.sort
 145
 146         # As we want to aggregate adiacent slots we convert all the
 147         # slot integers into ranges (with just one element)
 148         # So we have something like [1..1,2..2, ... and so forth.
 149         slots.map!{|x| x..x}
 150
 151         # Finally we group ranges with adiacent elements.
 152         slots = slots.reduce([]) {|a,b|
 153             if !a.empty? && b.first == (a[-1].last)+1
 154                 a[0..-2] + [(a[-1].first)..(b.last)]
 155             else
 156                 a + [b]
 157             end
 158         }
 159
 160         # Now our task is easy, we just convert ranges with just one
 161         # element into a number, and a real range into a start-end format.
 162         # Finally we join the array using the comma as separator.
 163         slots = slots.map{|x|
 164             x.count == 1 ? x.first.to_s : "#{x.first}-#{x.last}"
 165         }.join(",")
 166
 167         "[#{@info[:cluster_state].upcase}] #{self.info[:name]} #{self.to_s.ljust(25)} slots:#{slots}"
 168     end
 169
 170     def info
 171         @info
 172     end
 173
 174     def is_dirty?
 175         @dirty
 176     end
 177
 178     def r
 179         @r
 180     end
 181 end
 182
 183 class RedisTrib
 184     def initialize
 185         @nodes = []
 186     end
 187
 188     def check_arity(req_args, num_args)
 189         if ((req_args > 0 and num_args != req_args) ||
 190            (req_args < 0 and num_args < req_args.abs))
 191            puts "Wrong number of arguments for specified sub command"
 192            exit 1
 193         end
 194     end
 195
 196     def add_node(node)
 197         @nodes << node
 198     end
 199
 200     def get_node_by_name(name)
 201         @nodes.each{|n|
 202             return n if n.info[:name] == name.downcase
 203         }
 204         return nil
 205     end
 206
 207     def check_cluster
 208         puts "Performing Cluster Check (using node #{@nodes[0]})"
 209         errors = []
 210         show_nodes
 211         # Check if all the slots are covered
 212         slots = {}
 213         @nodes.each{|n|
 214             slots = slots.merge(n.slots)
 215         }
 216         if slots.length == 4096
 217             puts "[OK] All 4096 slots covered."
 218         else
 219             errors << "[ERR] Not all 4096 slots are covered by nodes."
 220             puts errors[-1]
 221         end
 222         return errors
 223     end
 224
 225     def alloc_slots
 226         slots_per_node = ClusterHashSlots/@nodes.length
 227         i = 0
 228         @nodes.each{|n|
 229             first = i*slots_per_node
 230             last = first+slots_per_node-1
 231             last = ClusterHashSlots-1 if i == @nodes.length-1
 232             n.add_slots first..last
 233             i += 1
 234         }
 235     end
 236
 237     def flush_nodes_config
 238         @nodes.each{|n|
 239             n.flush_node_config
 240         }
 241     end
 242
 243     def show_nodes
 244         @nodes.each{|n|
 245             puts n.info_string
 246         }
 247     end
 248
 249     def join_cluster
 250         # We use a brute force approach to make sure the node will meet
 251         # each other, that is, sending CLUSTER MEET messages to all the nodes
 252         # about the very same node.
 253         # Thanks to gossip this information should propagate across all the
 254         # cluster in a matter of seconds.
 255         first = false
 256         @nodes.each{|n|
 257             if !first then first = n.info; next; end # Skip the first node
 258             n.r.cluster("meet",first[:host],first[:port])
 259         }
 260     end
 261
 262     def yes_or_die(msg)
 263         print "#{msg} (type 'yes' to accept): "
 264         STDOUT.flush
 265         if !(STDIN.gets.chomp.downcase == "yes")
 266             puts "Aborting..."
 267             exit 1
 268         end
 269     end
 270
 271     def load_cluster_info_from_node(nodeaddr)
 272         node = ClusterNode.new(ARGV[1])
 273         node.connect(:abort => true)
 274         node.assert_cluster
 275         node.load_info(:getfriends => true)
 276         add_node(node)
 277         node.friends.each{|f|
 278             fnode = ClusterNode.new(f[:addr])
 279             fnode.connect()
 280             fnode.load_info()
 281             add_node(fnode)
 282         }
 283     end
 284
 285     # Given a list of source nodes return a "resharding plan"
 286     # with what slots to move in order to move "numslots" slots to another
 287     # instance.
 288     def compute_reshard_table(sources,numslots)
 289         moved = []
 290         sources.each{|s|
 291             # Every node will provide a number of slots proportional to the
 292             # slots it has assigned.
 293             n = (numslots.to_f/4096*s.slots.length).ceil
 294             s.slots.keys.sort[(0...n)].each{|slot|
 295                 if moved.length < numslots
 296                     moved << {:source => s, :slot => slot}
 297                 end
 298             }
 299         }
 300         return moved
 301     end
 302
 303     def show_reshard_table(table)
 304         table.each{|e|
 305             puts "Moving slot #{e[:slot]} from #{e[:source].info[:name]}"
 306         }
 307     end
 308
 309     # redis-trib subcommands implementations
 310
 311     def check_cluster_cmd
 312         load_cluster_info_from_node(ARGV[1])
 313         check_cluster
 314     end
 315
 316     def reshard_cluster_cmd
 317         load_cluster_info_from_node(ARGV[1])
 318         errors = check_cluster
 319         if errors.length != 0
 320             puts "Please fix your cluster problems before resharding."
 321             exit 1
 322         end
 323         numslots = 0
 324         while numslots <= 0 or numslots > 4096
 325             print "How many slots do you want to move (from 1 to 4096)?"
 326             numslots = STDIN.gets.to_i
 327         end
 328         target = nil
 329         while not target
 330             print "What is the receiving node ID? "
 331             target = get_node_by_name(STDIN.gets.chop)
 332             if not target
 333                 puts "The specified node is not known, please retry."
 334             end
 335         end
 336         sources = []
 337         puts "Please enter all the source node IDs."
 338         puts "  Type 'all' to use all the nodes as source nodes for the hash slots."
 339         puts "  Type 'done' once you entered all the source nodes IDs."
 340         while true
 341             print "Source node ##{sources.length+1}:"
 342             line = STDIN.gets.chop
 343             src = get_node_by_name(line)
 344             if line == "done"
 345                 if sources.length == 0
 346                     puts "No source nodes given, operation aborted"
 347                     exit 1
 348                 else
 349                     break
 350                 end
 351             elsif line == "all"
 352                 @nodes.each{|n|
 353                     next if n.info[:name] == target.info[:name]
 354                     sources << n
 355                 }
 356                 break
 357             elsif not src
 358                 puts "The specified node is not known, please retry."
 359             elsif src.info[:name] == target.info[:name]
 360                 puts "It is not possible to use the target node as source node."
 361             else
 362                 sources << src
 363             end
 364         end
 365         puts "\nReady to move #{numslots} slots."
 366         puts "  Source nodes:"
 367         sources.each{|s| puts "    "+s.info_string}
 368         puts "  Destination node:"
 369         puts "    #{target.info_string}"
 370         reshard_table = compute_reshard_table(sources,numslots)
 371         show_reshard_table(reshard_table)
 372     end
 373
 374     def create_cluster_cmd
 375         puts "Creating cluster"
 376         ARGV[1..-1].each{|n|
 377             node = ClusterNode.new(n)
 378             node.connect(:abort => true)
 379             node.assert_cluster
 380             node.assert_empty
 381             add_node(node)
 382         }
 383         puts "Performing hash slots allocation on #{@nodes.length} nodes..."
 384         alloc_slots
 385         show_nodes
 386         yes_or_die "Can I set the above configuration?"
 387         flush_nodes_config
 388         puts "** Nodes configuration updated"
 389         puts "** Sending CLUSTER MEET messages to join the cluster"
 390         join_cluster
 391         check_cluster
 392     end
 393 end
 394
 395 COMMANDS={
 396     "create" => ["create_cluster_cmd", -2, "host1:port host2:port ... hostN:port"],
 397     "check" =>  ["check_cluster_cmd", 2, "host:port"],
 398     "reshard" =>  ["reshard_cluster_cmd", 2, "host:port"]
 399 }
 400
 401 # Sanity check
 402 if ARGV.length == 0
 403     puts "Usage: redis-trib <command> <arguments ...>"
 404     puts
 405     COMMANDS.each{|k,v|
 406         puts "  #{k.ljust(20)} #{v[2]}"
 407     }
 408     puts
 409     exit 1
 410 end
 411
 412 rt = RedisTrib.new
 413 cmd_spec = COMMANDS[ARGV[0].downcase]
 414 if !cmd_spec
 415     puts "Unknown redis-trib subcommand '#{ARGV[0]}'"
 416     exit 1
 417 end
 418 rt.check_arity(cmd_spec[1],ARGV.length)
 419
 420 # Dispatch
 421 rt.send(cmd_spec[0])