gizzmo 0.1.4

data/lib/gizzard/thrift.rb

@@ -0,0 +1,160 @@
+#!/usr/bin/env ruby
+require 'vendor/thrift_client/simple'
+
+module Gizzard
+  module Thrift
+    T = ThriftClient::Simple
+
+    def self.struct(*args)
+      T::StructType.new(*args)
+    end
+
+    ShardException = T.make_exception(:ShardException,
+      T::Field.new(:description, T::STRING, 1)
+    )
+
+    ShardId = T.make_struct(:ShardId,
+      T::Field.new(:hostname, T::STRING, 1),
+      T::Field.new(:table_prefix, T::STRING, 2)
+    )
+
+    class ShardId
+      def inspect
+        "#{hostname}/#{table_prefix}"
+      end
+      
+      alias_method :to_unix, :inspect
+      
+      def self.parse(string)
+        new(*string.split("/"))
+      end
+    end
+
+    ShardInfo = T.make_struct(:ShardInfo,
+      T::Field.new(:id, struct(ShardId), 1),
+      T::Field.new(:class_name, T::STRING, 2),
+      T::Field.new(:source_type, T::STRING, 3),
+      T::Field.new(:destination_type, T::STRING, 4),
+      T::Field.new(:busy, T::I32, 5)
+    )
+
+    class ShardInfo
+      def busy?
+        busy && busy > 0
+      end
+
+      def inspect(short = false)
+        "#{id.inspect}" + (busy? ? " (BUSY)" : "")
+      end
+      
+      def to_unix
+        [id.to_unix, class_name, busy? ? "busy" : "unbusy"].join("\t")
+      end
+    end
+
+    LinkInfo = T.make_struct(:LinkInfo,
+      T::Field.new(:up_id, struct(ShardId), 1),
+      T::Field.new(:down_id, struct(ShardId), 2),
+      T::Field.new(:weight, T::I32, 3)
+    )
+
+    class LinkInfo
+      def inspect
+        "#{up_id.inspect} -> #{down_id.inspect}" + (weight == 1 ? "" : " <#{weight}>")
+      end
+      
+      def to_unix
+        [up_id.to_unix, down_id.to_unix, weight].join("\t")
+      end
+      
+    end
+
+    ShardMigration = T.make_struct(:ShardMigration,
+      T::Field.new(:source_id, struct(ShardId), 1),
+      T::Field.new(:destination_id, struct(ShardId), 2)
+    )
+
+    Forwarding = T.make_struct(:Forwarding,
+      T::Field.new(:table_id, T::I32, 1),
+      T::Field.new(:base_id, T::I64, 2),
+      T::Field.new(:shard_id, struct(ShardId), 3)
+    )
+
+    class Forwarding
+      #FIXME table_id is not human-readable
+      def inspect
+        "[#{table_id}] #{base_id.to_s(16)} -> #{shard_id.inspect}"
+      end
+    end
+
+    class ShardManager < T::ThriftService
+      def initialize(host, port, log_path, dry_run = false)
+        super(host, port)
+        @dry = dry_run
+        begin
+          @log = File.open(log_path, "a")
+        rescue
+          STDERR.puts "Error opening log file at #{log_path}.  Continuing..."
+        end
+      end
+      
+      def _proxy(method_name, *args)
+        cls = self.class.ancestors.find { |cls| cls.respond_to?(:_arg_structs) and cls._arg_structs[method_name.to_sym] }
+        arg_class, rv_class = cls._arg_structs[method_name.to_sym]
+        
+        # Writing methods return void. Methods should never both read and write. If this assumption
+        # is violated in the future, dry-run will fail!!
+        is_writing_method = rv_class._fields.first.type == ThriftClient::Simple::VOID
+        if @dry && is_writing_method
+          puts "Skipped writing: #{printable(method_name, args)}"
+        else
+          @log.puts printable(method_name, args, true)
+          super(method_name, *args)
+        end
+      rescue ThriftClient::Simple::ThriftException
+        if @dry
+          puts "Skipped reading: #{printable(method_name, args)}"
+        else
+          raise
+        end
+      end
+      
+      def printable(method_name, args, timestamp = false)
+        ts = timestamp ? "#{Time.now}\t" : ""
+        "#{ts}#{method_name}(#{args.map{|a| a.inspect}.join(', ')})"
+      end
+      
+      
+      thrift_method :create_shard, void, field(:shard, struct(ShardInfo), 1), :throws => exception(ShardException)
+      thrift_method :delete_shard, void, field(:id, struct(ShardId), 1)
+      thrift_method :get_shard, struct(ShardInfo), field(:id, struct(ShardId), 1)
+
+      thrift_method :add_link, void, field(:up_id, struct(ShardId), 1), field(:down_id, struct(ShardId), 2), field(:weight, i32, 3)
+      thrift_method :remove_link, void, field(:up_id, struct(ShardId), 1), field(:down_id, struct(ShardId), 2)
+
+      thrift_method :list_upward_links, list(struct(LinkInfo)), field(:id, struct(ShardId), 1)
+      thrift_method :list_downward_links, list(struct(LinkInfo)), field(:id, struct(ShardId), 1)
+
+      thrift_method :get_child_shards_of_class, list(struct(ShardInfo)), field(:parent_id, struct(ShardId), 1), field(:class_name, string, 2)
+
+      thrift_method :mark_shard_busy, void, field(:id, struct(ShardId), 1), field(:busy, i32, 2)
+      thrift_method :copy_shard, void, field(:source_id, struct(ShardId), 1), field(:destination_id, struct(ShardId), 2)
+
+      thrift_method :set_forwarding, void, field(:forwarding, struct(Forwarding), 1)
+      thrift_method :replace_forwarding, void, field(:old_id, struct(ShardId), 1), field(:new_id, struct(ShardId), 2)
+
+      thrift_method :get_forwarding, struct(Forwarding), field(:table_id, i32, 1), field(:base_id, i64, 2)
+      thrift_method :get_forwarding_for_shard, struct(Forwarding), field(:shard_id, struct(ShardId), 1)
+
+      thrift_method :get_forwardings, list(struct(Forwarding))
+      thrift_method :reload_forwardings, void
+
+      thrift_method :find_current_forwarding, struct(ShardInfo), field(:table_id, i32, 1), field(:id, i64, 2)
+
+      thrift_method :shards_for_hostname, list(struct(ShardInfo)), field(:hostname, string, 1)
+      thrift_method :get_busy_shards, list(struct(ShardInfo))
+
+      thrift_method :rebuild_schema, void
+    end
+  end
+end

data/lib/gizzmo.rb

@@ -0,0 +1,172 @@
+#!/usr/bin/env ruby
+$: << File.dirname(__FILE__)
+class HelpNeededError < RuntimeError; end
+require "optparse"
+require "ostruct"
+require "gizzard"
+require "yaml"
+
+ORIGINAL_ARGV = ARGV.dup
+
+# Container for parsed options
+global_options     = OpenStruct.new
+subcommand_options = OpenStruct.new
+
+# Leftover arguments
+argv = nil
+
+begin
+  YAML.load_file(File.join(ENV["HOME"], ".gizzmorc")).each do |k, v|
+    global_options.send("#{k}=", v)
+  end
+rescue Errno::ENOENT
+  # Do nothing...
+rescue => e
+  abort "Unknown error loading ~/.gizzmorc: #{e.message}"
+end
+
+subcommands = {
+  'create' => OptionParser.new do |opts|
+    opts.banner = "Usage: #{$0} create [options] HOST TABLE_PREFIX CLASS_NAME"
+
+    opts.on("-s", "--source-type=TYPE") do |s|
+      subcommand_options.source_type = s
+    end
+
+    opts.on("-d", "--destination-type=TYPE") do |s|
+      subcommand_options.destination_type = s
+    end
+  end,
+  'wrap' => OptionParser.new do |opts|
+    opts.banner = "Usage: #{$0} wrap CLASS_NAME SHARD_ID_TO_WRAP [MORE SHARD_IDS...]"
+  end,
+  'subtree' => OptionParser.new do |opts|
+    opts.banner = "Usage: #{$0} subtree SHARD_ID"
+  end,
+  'delete' => OptionParser.new do |opts|
+    opts.banner = "Usage: #{$0} delete SHARD_ID_TO_DELETE [MORE SHARD_IDS]"
+  end,
+  'unwrap' => OptionParser.new do |opts|
+    opts.banner = "Usage: #{$0} unwrap SHARD_ID_TO_REMOVE [MORE SHARD_IDS]"
+  end,
+  'find' => OptionParser.new do |opts|
+    opts.banner = "Usage: #{$0} find [options]"
+
+    opts.on("-t", "--type=TYPE", "Return only shards of the specified TYPE") do |shard_type|
+      subcommand_options.shard_type = shard_type
+    end
+
+    opts.on("-H", "--host=HOST", "HOST of shard") do |shard_host|
+      subcommand_options.shard_host = shard_host
+    end
+  end,
+  'links' => OptionParser.new do |opts|
+    opts.banner = "Usage: #{$0} links SHARD_ID [MORE SHARD_IDS...]"
+  end,
+  'info' => OptionParser.new do |opts|
+    opts.banner = "Usage: #{$0} info SHARD_ID [MORE SHARD_IDS...]"
+  end,
+  'reload' => OptionParser.new do |opts|
+    opts.banner = "Usage: #{$0} reload"
+  end,
+  'addlink' => OptionParser.new do |opts|
+    opts.banner = "Usage: #{$0} link PARENT_SHARD_ID CHILD_SHARD_ID WEIGHT"
+  end,
+  'unlink' => OptionParser.new do |opts|
+    opts.banner = "Usage: #{$0} unlink PARENT_SHARD_ID CHILD_SHARD_ID"
+  end
+}
+
+global = OptionParser.new do |opts|
+  opts.banner = "Usage: #{$0} [global-options] SUBCOMMAND [subcommand-options]"
+  opts.separator ""
+  opts.separator "Subcommands:"
+  subcommands.keys.compact.sort.each do |sc|
+    opts.separator "  #{sc}"
+  end
+  opts.separator ""
+  opts.separator "You can type `#{$0} help SUBCOMMAND` for help on a specific subcommand."
+  opts.separator ""
+  opts.separator "Global options:"
+
+  opts.on("-H", "--host=HOSTNAME", "HOSTNAME of remote thrift service") do |host|
+    global_options.host = host
+  end
+
+  opts.on("-P", "--port=PORT", "PORT of remote thrift service") do |port|
+    global_options.port = port
+  end
+
+  opts.on("-D", "--dry-run", "") do |port|
+    global_options.dry = true
+  end
+
+  opts.on("-C", "--config=YAML_FILE", "YAML_FILE of option key/values") do |file|
+    YAML.load(File.open(file)).each do |k, v|
+      global_options.send("#{k}=", v)
+    end
+  end
+
+  opts.on("-L", "--log=LOG_FILE", "Path to LOG_FILE") do |file|
+    global_options.log = file
+  end
+end
+
+# Print banner if no args
+if ARGV.length == 0
+  STDERR.puts global
+  exit 1
+end
+
+# This
+def process_nested_parsers(global, subcommands)
+  begin
+    global.order!(ARGV) do |subcommand_name|
+      # puts args.inspect
+      subcommand = subcommands[subcommand_name]
+      argv = subcommand ? subcommand.parse!(ARGV) : ARGV
+      return subcommand_name, argv
+    end
+  rescue => e
+    STDERR.puts e.message
+    exit 1
+  end
+end
+
+
+subcommand_name, argv = process_nested_parsers(global, subcommands)
+
+# Print help sub-banners
+if subcommand_name == "help"
+  STDERR.puts subcommands[argv.shift] || global
+  exit 1
+end
+
+unless subcommands.include?(subcommand_name)
+  STDERR.puts "Subcommand not found: #{subcommand_name}"
+  exit 1
+end
+
+log = global_options.log || "/tmp/gizzmo.log"
+service = Gizzard::Thrift::ShardManager.new(global_options.host, global_options.port, log, global_options.dry)
+
+begin
+  Gizzard::Command.run(subcommand_name, service, global_options, argv, subcommand_options)
+rescue HelpNeededError => e
+  if e.class.name != e.message
+    STDERR.puts("=" * 80)
+    STDERR.puts e.message
+    STDERR.puts("=" * 80)
+  end
+  STDERR.puts subcommands[subcommand_name]
+  exit 1
+rescue ThriftClient::Simple::ThriftException => e
+  STDERR.puts e.message
+  exit 1
+rescue Errno::EPIPE
+  # This is just us trying to puts into a closed stdout.  For example, if you pipe into
+  # head -1, then this script will keep running after head closes.  We don't care, and
+  # seeing the backtrace is annoying!
+rescue Interrupt
+  exit 1
+end

data/lib/vendor/thrift_client/simple.rb

@@ -0,0 +1,334 @@
+require 'socket'
+require 'getoptlong'
+
+class ThriftClient
+
+  # This is a simplified form of thrift, useful for clients only, and not
+  # making any attempt to have good performance. It's intended to be used by
+  # small command-line tools that don't want to install a dozen ruby files.
+  module Simple
+    VERSION_1 = 0x8001
+
+    # message types
+    CALL, REPLY, EXCEPTION = (1..3).to_a
+
+    # value types
+    STOP, VOID, BOOL, BYTE, DOUBLE, _, I16, _, I32, _, I64, STRING, STRUCT, MAP, SET, LIST = (0..15).to_a
+
+    FORMATS = {
+      BYTE => "c",
+      DOUBLE => "G",
+      I16 => "n",
+      I32 => "N",
+    }
+
+    SIZES = {
+      BYTE => 1,
+      DOUBLE => 8,
+      I16 => 2,
+      I32 => 4,
+    }
+
+    module ComplexType
+      module Extends
+        def type_id=(n)
+          @type_id = n
+        end
+
+        def type_id
+          @type_id
+        end
+      end
+
+      module Includes
+        def to_i
+          self.class.type_id
+        end
+
+        def to_s
+          args = self.values.map { |v| self.class.type_id == STRUCT ? v.name : v.to_s }.join(", ")
+          "#{self.class.name}.new(#{args})"
+        end
+      end
+    end
+
+    def self.make_type(type_id, name, *args)
+      klass = Struct.new("STT_#{name}", *args)
+      klass.send(:extend, ComplexType::Extends)
+      klass.send(:include, ComplexType::Includes)
+      klass.type_id = type_id
+      klass
+    end
+
+    ListType = make_type(LIST, "ListType", :element_type)
+    MapType = make_type(MAP, "MapType", :key_type, :value_type)
+    SetType = make_type(SET, "SetType", :element_type)
+    StructType = make_type(STRUCT, "StructType", :struct_class)
+
+    class << self
+      def pack_value(type, value)
+        case type
+        when BOOL
+          [ value ? 1 : 0 ].pack("c")
+        when STRING
+          [ value.size, value ].pack("Na*")
+        when I64
+          [ value >> 32, value & 0xffffffff ].pack("NN")
+        when ListType
+          [ type.element_type.to_i, value.size ].pack("cN") + value.map { |item| pack_value(type.element_type, item) }.join("")
+        when MapType
+          [ type.key_type.to_i, type.value_type.to_i, value.size ].pack("ccN") + value.map { |k, v| pack_value(type.key_type, k) + pack_value(type.value_type, v) }.join("")
+        when SetType
+          [ type.element_type.to_i, value.size ].pack("cN") + value.map { |item| pack_value(type.element_type, item) }.join("")
+        when StructType
+          value._pack
+        else
+          [ value ].pack(FORMATS[type])
+        end
+      end
+
+      def pack_request(method_name, arg_struct, request_id=0)
+        [ VERSION_1, CALL, method_name.to_s.size, method_name.to_s, request_id, arg_struct._pack ].pack("nnNa*Na*")
+      end
+
+      def read_value(s, type)
+        case type
+        when BOOL
+          s.read(1).unpack("c").first != 0
+        when STRING
+          len = s.read(4).unpack("N").first
+          s.read(len)
+        when I64
+          hi, lo = s.read(8).unpack("NN")
+          rv = (hi << 32) | lo
+          (rv >= (1 << 63)) ? (rv - (1 << 64)) : rv
+        when LIST
+          read_list(s)
+        when MAP
+          read_map(s)
+        when STRUCT
+          read_struct(s, UnknownStruct)
+        when ListType
+          read_list(s, type.element_type)
+        when MapType
+          read_map(s, type.key_type, type.value_type)
+        when StructType
+          read_struct(s, type.struct_class)
+        else
+          rv = s.read(SIZES[type]).unpack(FORMATS[type]).first
+          case type
+          when I16
+            (rv >= (1 << 15)) ? (rv - (1 << 16)) : rv
+          when I32
+            (rv >= (1 << 31)) ? (rv - (1 << 32)) : rv
+          else
+            rv
+          end
+        end
+      end
+
+      def read_list(s, element_type=nil)
+        etype, len = s.read(5).unpack("cN")
+        expected_type = (element_type and element_type.to_i == etype.to_i) ? element_type : etype
+        rv = []
+        len.times do
+          rv << read_value(s, expected_type)
+        end
+        rv
+      end
+
+      def read_map(s, key_type=nil, value_type=nil)
+        ktype, vtype, len = s.read(6).unpack("ccN")
+        rv = {}
+        expected_key_type, expected_value_type = if key_type and value_type and key_type.to_i == ktype and value_type.to_i == vtype
+          [ key_type, value_type ]
+        else
+          [ ktype, vtype ]
+        end
+        len.times do
+          key = read_value(s, expected_key_type)
+          value = read_value(s, expected_value_type)
+          rv[key] = value
+        end
+        rv
+      end
+
+      def read_struct(s, struct_class)
+        struct = struct_class.new
+        while true
+          ftype = s.read(1).unpack("c").first
+          return struct if ftype == STOP
+          fid = s.read(2).unpack("n").first
+
+          if field = struct_class._fields.find { |f| (f.fid == fid) and (f.type.to_i == ftype) }
+            struct[field.name] = read_value(s, field.type)
+          else
+            $stderr.puts "Warning: Unknown struct field encountered. (recieved id: #{fid})"
+            raise "Warning: Unknown struct field encountered. (recieved id: #{fid})"
+            read_value(s, ftype)
+          end
+        end
+      end
+      
+      def read_response(s, rv_class)
+        version, message_type, method_name_len = s.read(8).unpack("nnN")
+        method_name = s.read(method_name_len)
+        seq_id = s.read(4).unpack("N").first
+        if message_type == EXCEPTION
+          exception = read_struct(s, ExceptionStruct)
+          raise ThriftException, exception.message
+        end
+        response = read_struct(s, rv_class)
+        raise response.ex if response.respond_to?(:ex) and response.ex
+        [ method_name, seq_id, response.rv ]
+      end
+    end
+
+    ## ----------------------------------------
+
+    class Field
+      attr_accessor :name, :type, :fid
+
+      def initialize(name, type, fid)
+        @name = name
+        @type = type
+        @fid = fid
+      end
+
+      def pack(value)
+        value.nil? ? "" : [ type.to_i, fid, ThriftClient::Simple.pack_value(type, value) ].pack("cna*")
+      end
+    end
+
+    class ThriftException < RuntimeError
+      def initialize(reason)
+        @reason = reason
+      end
+
+      def to_s
+        "ThriftException(#{@reason.inspect})"
+      end
+    end
+
+    module ThriftStruct
+      module Include
+        def _pack
+          self.class._fields.map { |f| f.pack(self[f.name]) }.join + [ STOP ].pack("c")
+        end
+      end
+
+      module Extend
+        def _fields
+          @fields
+        end
+
+        def _fields=(f)
+          @fields = f
+        end
+      end
+    end
+
+    def self.make_struct(name, *fields)
+      st_name = "ST_#{name.to_s.tr(':', '_')}"
+      if Struct.constants.include?(st_name)
+        warn "#{caller[0]}: Struct::#{st_name} is already defined; returning original class."
+        Struct.const_get(st_name)
+      else
+        names = fields.map { |f| f.name.to_sym }
+        klass = Struct.new(st_name, *names)
+        klass.send(:include, ThriftStruct::Include)
+        klass.send(:extend, ThriftStruct::Extend)
+        klass._fields = fields
+        klass
+      end
+    end
+
+    def self.make_exception(name, *fields)
+      struct_class = self.make_struct(name, *fields)
+      ex_class = Class.new(StandardError)
+
+      (class << struct_class; self end).send(:define_method, :exception_class) { ex_class }
+      (class << ex_class; self end).send(:define_method, :struct_class) { struct_class }
+
+      ex_class.class_eval do
+        attr_reader :struct
+
+        def initialize
+          @struct = self.class.struct_class.new
+        end
+
+        def self._fields
+          struct_class._fields
+        end
+
+        def to_s
+          method = [:message, :description].find {|m| struct.respond_to? m }
+          struct.send method || :to_s
+        end
+
+        alias message to_s
+
+        def method_missing(method, *args)
+          struct.send(method, *args)
+        end
+      end
+
+      ex_class
+    end
+
+    ExceptionStruct = make_struct(:ProtocolException, Field.new(:message, STRING, 1), Field.new(:type, I32, 2))
+    UnknownStruct = make_struct(:Unknown)
+
+    class ThriftService
+      def initialize(host, port)
+        @host = host
+        @port = port
+      end
+
+      def self._arg_structs
+        @_arg_structs = {} if @_arg_structs.nil?
+        @_arg_structs
+      end
+
+      def self.thrift_method(name, rtype, *args)
+        options = args.last.is_a?(Hash) ? args.pop : {}
+        fields = [ ThriftClient::Simple::Field.new(:rv, rtype, 0),
+                   (options[:throws] ? ThriftClient::Simple::Field.new(:ex, options[:throws], 1) : nil)
+                 ].compact
+
+        arg_struct = ThriftClient::Simple.make_struct("Args__#{self.name}__#{name}", *args)
+        rv_struct = ThriftClient::Simple.make_struct("Retval__#{self.name}__#{name}", *fields)
+
+        _arg_structs[name.to_sym] = [ arg_struct, rv_struct ]
+
+        arg_names = args.map { |a| a.name.to_s }.join(", ")
+        class_eval "def #{name}(#{arg_names}); _proxy(:#{name}#{args.size > 0 ? ', ' : ''}#{arg_names}); end"
+      end
+
+      def _proxy(method_name, *args)
+        cls = self.class.ancestors.find { |cls| cls.respond_to?(:_arg_structs) and cls._arg_structs[method_name.to_sym] }
+        arg_class, rv_class = cls._arg_structs[method_name.to_sym]
+        arg_struct = arg_class.new(*args)
+        sock = TCPSocket.new(@host, @port)
+        sock.write(ThriftClient::Simple.pack_request(method_name, arg_struct))
+        rv = ThriftClient::Simple.read_response(sock, rv_class)
+        sock.close
+        rv[2]
+      end
+
+      # convenience. robey is lazy.
+      { :field => "Field.new",
+        :struct => "StructType.new",
+        :exception => "StructType.new",
+        :list => "ListType.new",
+        :map => "MapType.new",
+      }.each do |new_name, old_name|
+        class_eval "def self.#{new_name}(*args); ThriftClient::Simple::#{old_name}(*args); end"
+      end
+
+#      alias exception struct
+
+      [ :void, :bool, :byte, :double, :i16, :i32, :i64, :string ].each { |sym| class_eval "def self.#{sym}; ThriftClient::Simple::#{sym.to_s.upcase}; end" }
+    end
+  end
+end