# -*- encoding: binary -*-
require 'thread'

# This class communicates with the MogileFS trackers.
# You should not have to use this directly unless you are developing
# support for new commands or plugins for MogileFS
class MogileFS::Backend

  # Adds MogileFS commands +names+.
  def self.add_command(*names)
    names.each do |name|
      define_method name do |*args|
        do_request(name, args[0] || {}, false)
      end
    end
  end

  # adds idempotent MogileFS commands +names+, these commands may be retried
  # transparently on a different tracker if there is a network/server error.
  def self.add_idempotent_command(*names)
    names.each do |name|
      define_method name do |*args|
        do_request(name, args[0] || {}, true)
      end
    end
  end

  BACKEND_ERRORS = {} # :nodoc:

  # this converts an error code from a mogilefsd tracker to an exception:
  #
  # Examples of some exceptions that get created:
  #   class AfterMismatchError < MogileFS::Error; end
  #   class DomainNotFoundError < MogileFS::Error; end
  #   class InvalidCharsError < MogileFS::Error; end
  def self.add_error(err_snake)
    err_camel = err_snake.gsub(/(?:^|_)([a-z])/) { $1.upcase }
    err_camel << 'Error' unless /Error\z/ =~ err_camel
    unless const_defined?(err_camel)
      const_set(err_camel, Class.new(MogileFS::Error))
    end
    BACKEND_ERRORS[err_snake] = const_get(err_camel)
  end

  def self.const_missing(name) # :nodoc:
    if /Error\z/ =~ name.to_s
      const_set(name, Class.new(MogileFS::Error))
    else
      super name
    end
  end

  ##
  # The last error

  attr_reader :lasterr

  ##
  # The string attached to the last error

  attr_reader :lasterrstr

  ##
  # Creates a new MogileFS::Backend.
  #
  # :hosts is a required argument and must be an Array containing one or more
  # 'hostname:port' pairs as Strings.
  #
  # :timeout adjusts the request timeout before an error is returned.

  def initialize(args)
    @hosts = args[:hosts]
    @fail_timeout = args[:fail_timeout] || 5
    raise ArgumentError, "must specify at least one host" unless @hosts
    raise ArgumentError, "must specify at least one host" if @hosts.empty?
    unless @hosts == @hosts.select { |h| h =~ /:\d+$/ } then
      raise ArgumentError, ":hosts must be in 'host:port' form"
    end

    @mutex = Mutex.new
    @timeout = args[:timeout] || 3
    @socket = nil
    @lasterr = nil
    @lasterrstr = nil
    @pending = []

    @dead = {}
  end

  ##
  # Closes this backend's socket.

  def shutdown
    @mutex.synchronize { shutdown_unlocked }
  end

  # MogileFS::MogileFS commands

  add_command :create_open
  add_command :create_close
  add_idempotent_command :get_paths
  add_idempotent_command :noop
  add_command :delete
  add_idempotent_command :sleep
  add_command :rename
  add_idempotent_command :list_keys
  add_idempotent_command :file_info
  add_idempotent_command :file_debug

  # MogileFS::Backend commands

  add_idempotent_command :get_hosts
  add_idempotent_command :get_devices
  add_idempotent_command :list_fids
  add_idempotent_command :stats
  add_idempotent_command :get_domains
  add_command :create_device
  add_command :create_domain
  add_command :delete_domain
  add_command :create_class
  add_command :update_class
  add_command :updateclass
  add_command :delete_class
  add_command :create_host
  add_command :update_host
  add_command :delete_host
  add_command :set_state
  add_command :set_weight
  add_command :replicate_now

  def shutdown_unlocked(do_raise = false) # :nodoc:
    @pending = []
    if @socket
      @socket.close rescue nil # ignore errors
      @socket = nil
    end
    raise if do_raise
  end

  def dispatch_unlocked(request, timeout = @timeout) # :nodoc:
    begin
      io = socket
      io.timed_write(request, timeout)
      io
    rescue SystemCallError, MogileFS::RequestTruncatedError  => err
      @dead[@active_host] = [ Time.now, err ]
      shutdown_unlocked
      retry
    end
  end

  def pipeline_gets_unlocked(io, timeout) # :nodoc:
    line = io.timed_gets(timeout) or
      raise MogileFS::PipelineError,
            "EOF with #{@pending.size} requests in-flight"
    ready = @pending.shift
    ready[1].call(parse_response(line, ready[0]))
  end

  def timeout_update(timeout, t0) # :nodoc:
    timeout -= (Time.now - t0)
    timeout < 0 ? 0 : timeout
  end

  # try to read any responses we have pending already before filling
  # the pipeline more requests.  This usually takes very little time,
  # but trackers may return huge responses and we could be on a slow
  # network.
  def pipeline_drain_unlocked(io, timeout) # :nodoc:
    set = [ io ]
    while @pending.size > 0
      t0 = Time.now
      r = IO.select(set, set, nil, timeout)
      timeout = timeout_update(timeout, t0)

      if r && r[0][0]
        t0 = Time.now
        pipeline_gets_unlocked(io, timeout)
        timeout = timeout_update(timeout, t0)
      else
        return timeout
      end
    end
    timeout
  end

  # dispatch a request like do_request, but queue +block+ for execution
  # upon receiving a response.  It is the users' responsibility to ensure
  # &block is executed in the correct order.  Trackers with multiple
  # queryworkers are not guaranteed to return responses in the same
  # order they were requested.
  def pipeline_dispatch(cmd, args, &block) # :nodoc:
    request = make_request(cmd, args)
    timeout = @timeout

    @mutex.synchronize do
      io = socket
      timeout = pipeline_drain_unlocked(io, timeout)

      # send the request out...
      begin
        io.timed_write(request, timeout)
        @pending << [ request, block ]
      rescue SystemCallError, MogileFS::RequestTruncatedError => err
        @dead[@active_host] = [ Time.now, err ]
        shutdown_unlocked(@pending[0])
        io = socket
        retry
      end

      @pending.size
    end
  end

  def pipeline_wait(count = nil) # :nodoc:
    @mutex.synchronize do
      io = socket
      count ||= @pending.size
      @pending.size < count and
        raise MogileFS::Error,
              "pending=#{@pending.size} < expected=#{count} failed"
      begin
        count.times { pipeline_gets_unlocked(io, @timeout) }
      rescue
        shutdown_unlocked(true)
      end
    end
  end

  # Performs the +cmd+ request with +args+.
  def do_request(cmd, args, idempotent = false)
    no_raise = args.delete(:ruby_no_raise)
    request = make_request(cmd, args)
    line = nil
    failed = false
    @mutex.synchronize do
      begin
        io = dispatch_unlocked(request)
        line = io.timed_gets(@timeout)
        break if /\r?\n\z/ =~ line

        line and raise MogileFS::InvalidResponseError,
                       "Invalid response from server: #{line.inspect}"

        idempotent or
          raise EOFError, "end of file reached after: #{request.inspect}"
        # fall through to retry in loop
      rescue SystemCallError,
             MogileFS::InvalidResponseError # truncated response
        # we got a successful timed_write, but not a timed_gets
        if idempotent
          failed = true
          shutdown_unlocked(false)
          retry
        end
        shutdown_unlocked(true)
      rescue MogileFS::UnreadableSocketError, MogileFS::Timeout
        shutdown_unlocked(true)
      rescue
        # we DO NOT want the response we timed out waiting for, to crop up later
        # on, on the same socket, intersperesed with a subsequent request!  we
        # close the socket if there's any error.
        shutdown_unlocked(true)
      end while idempotent
      shutdown_unlocked if failed
    end # @mutex.synchronize
    parse_response(line, no_raise ? request : nil)
  end

  # Makes a new request string for +cmd+ and +args+.
  def make_request(cmd, args)
    "#{cmd} #{url_encode args}\r\n"
  end

  # this converts an error code from a mogilefsd tracker to an exception
  # Most of these exceptions should already be defined, but since the
  # MogileFS server code is liable to change and we may not always be
  # able to keep up with the changes
  def error(err_snake)
    BACKEND_ERRORS[err_snake] || self.class.add_error(err_snake)
  end

  # Turns the +line+ response from the server into a Hash of options, an
  # error, or raises, as appropriate.
  def parse_response(line, request = nil)
    case line
    when /\AOK\s+\d*\s*(\S*)\r?\n\z/
      url_decode($1)
    when /\AERR\s+(\w+)\s*([^\r\n]*)/
      @lasterr = $1
      @lasterrstr = $2 ? url_unescape($2) : nil
      if request
        request = " request=#{request.strip}"
        @lasterrstr = @lasterrstr ? (@lasterrstr << request) : request
        return error(@lasterr).new(@lasterrstr)
      end
      raise error(@lasterr).new(@lasterrstr)
    else
      raise MogileFS::InvalidResponseError,
            "Invalid response from server: #{line.inspect}"
    end
  end

  # this command is special since the cache is per-tracker, so we connect
  # to all backends and not just one
  def clear_cache(types = %w(all))
    opts = {}
    types.each { |type| opts[type] = 1 }

    sockets = @hosts.map do |host|
      MogileFS::Socket.start(*(host.split(/:/))) rescue nil
    end
    sockets.compact!

    wpending = sockets
    rpending = []
    request = make_request("clear_cache", opts)
    while wpending[0] || rpending[0]
      r = IO.select(rpending, wpending, nil, @timeout) or return
      rpending -= r[0]
      wpending -= r[1]
      r[0].each { |io| io.timed_gets(0) rescue nil }
      r[1].each do |io|
        begin
          io.timed_write(request, 0)
          rpending << io
        rescue
        end
      end
    end
    nil
    ensure
      sockets.each { |io| io.close }
  end

  # Returns a socket connected to a MogileFS tracker.
  def socket
    return @socket if @socket and not @socket.closed?

    @hosts.shuffle.each do |host|
      next if dead = @dead[host] and dead[0] > (Time.now - @fail_timeout)

      begin
        addr, port = host.split(/:/)
        @socket = MogileFS::Socket.tcp(addr, port, @timeout)
        @active_host = host
      rescue SystemCallError, MogileFS::Timeout => err
        @dead[host] = [ Time.now, err ]
        next
      end

      return @socket
    end

    errors = @dead.map { |host,(_,e)| "#{host} - #{e.message} (#{e.class})" }
    raise MogileFS::UnreachableBackendError,
          "couldn't connect to any tracker: #{errors.join(', ')}"
  end

  # Turns a url params string into a Hash.
  def url_decode(str) # :nodoc:
    rv = {}
    str.split(/&/).each do |pair|
      k, v = pair.split(/=/, 2).map! { |x| url_unescape(x) }
      rv[k.freeze] = v
    end
    rv
  end

  # :stopdoc:
  # TODO: see if we can use existing URL-escape/unescaping routines
  # in the Ruby standard library, Perl MogileFS seems to NIH these
  #  routines, too
  # :startdoc:

  # Turns a Hash (or Array of pairs) into a url params string.
  def url_encode(params) # :nodoc:
    params.map do |k,v|
      "#{url_escape k.to_s}=#{url_escape v.to_s}"
    end.join("&")
  end

  # Escapes naughty URL characters.
  if ''.respond_to?(:ord) # Ruby 1.9
    def url_escape(str) # :nodoc:
      str.gsub(/([^\w\,\-.\/\\\: ])/) { "%%%02x" % $1.ord }.tr(' ', '+')
    end
  else # Ruby 1.8
    def url_escape(str) # :nodoc:
      str.gsub(/([^\w\,\-.\/\\\: ])/) { "%%%02x" % $1[0] }.tr(' ', '+')
    end
  end

  # Unescapes naughty URL characters.
  def url_unescape(str) # :nodoc:
    str.tr('+', ' ').gsub(/%([a-f0-9][a-f0-9])/i) { [$1.to_i(16)].pack 'C' }
  end
end