droonga-engine-absorb-data

#!/usr/bin/env ruby
#
# Copyright (C) 2014-2015 Droonga Project
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License version 2.1 as published by the Free Software Foundation.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

require "ostruct"
require "optparse"
require "socket"
require "coolio"

require "droonga/engine/version"
require "droonga/path"
require "droonga/node_name"
require "droonga/data_absorber_client"
require "droonga/serf"
require "droonga/client"
require "droonga/restarter"

module Droonga
  class AbsorbDataCommand
    def run
      @loop = Coolio::Loop.default

      parse_options
      assert_valid_options
      trap_signals

      puts "Start to absorb data from #{@options.source_dataset} at #{source_node.to_s}"
      puts "                       to #{@options.dataset} at #{destination_node.to_s}"
      puts "                      via #{@options.receiver_host} (this host)"
      puts ""
      puts "Absorbing..."

      succeeded = absorb

      if succeeded
        puts "Done."
      else
        do_cancel
      end
      succeeded
    end

    private
    def parse_options
      options = OpenStruct.new

      options.host    = Socket.gethostname
      options.port    = DataAbsorberClient::DEFAULT_PORT
      options.tag     = DataAbsorberClient::DEFAULT_TAG
      options.dataset = DataAbsorberClient::DEFAULT_DATASET

      options.source_host    = DataAbsorberClient::DEFAULT_HOST
      options.source_port    = DataAbsorberClient::DEFAULT_PORT
      options.source_tag     = DataAbsorberClient::DEFAULT_TAG
      options.source_dataset = DataAbsorberClient::DEFAULT_DATASET

      options.receiver_host = Socket.gethostname

      options.messages_per_second       = DataAbsorberClient::DEFAULT_MESSAGES_PER_SECOND
      options.progress_interval_seconds = DataAbsorberClient::DEFAULT_PROGRESS_INTERVAL_SECONDS

      options.verbose = false

      parser = OptionParser.new
      parser.version = Engine::VERSION

      parser.separator("")
      parser.separator("Destination node:")
      parser.on("--host=HOST",
                "Host name of the destination node.") do |host|
        options.host = host
      end
      parser.on("--port=PORT", Integer,
                "Port number of the destination node.",
                "(#{options.port})") do |port|
        options.port = port
      end
      parser.on("--tag=TAG", Integer,
                "Tag name of the destination node.",
                "(#{options.tag})") do |tag|
        options.tag = tag
      end
      parser.on("--dataset=DATASET",
                "Name of the destination dataset.",
                "(#{options.dataset})") do |dataset|
        options.dataset = dataset
      end

      parser.separator("")
      parser.separator("Source node:")
      parser.on("--source-host=HOST",
                "Host name of the source node.",
                "(#{options.source_host})") do |host|
        options.source_host = host
      end
      parser.on("--source-port=PORT", Integer,
                "Port number of the source node.",
                "(#{options.source_port})") do |host|
        options.source_host = host
      end
      parser.on("--source-tag=TAG",
                "Tag name of the source node.",
                "(#{options.source_tag})") do |tag|
        options.source_tag = tag
      end
      parser.on("--source-dataset=DATASET",
                "Name of the source dataset.",
                "(#{options.source_dataset})") do |dataset|
        options.source_dataset = dataset
      end

      parser.separator("")
      parser.separator("Connection:")
      parser.on("--receiver-host=HOST",
                "Host name of this computer.",
                "(#{options.receiver_host})") do |host|
        options.receiver_host = host
      end

      parser.separator("")
      parser.separator("Miscellaneous:")
      parser.on("--records-per-second=N", Integer,
                "Maximum number of records per second to be absorbed.",
                "'#{Client::RateLimiter::NO_LIMIT}' means no limit.",
                "(#{options.messages_per_second})") do |n|
        options.messages_per_second = n
      end
      parser.on("--progress-interval-seconds=N", Integer,
                "Interval seconds to report progress.",
                "(#{options.progress_interval_seconds})") do |n|
        options.progress_interval_seconds = n
      end
      parser.on("--[no-]verbose",
                "Output details for internal operations.",
                "(#{options.verbose})") do |verbose|
        options.verbose = verbose
      end

      parser.separator("")
      parser.separator("For backward compatibility:")
      parser.on("--destination-host=HOST",
                "Alias to \"--host\".") do |host|
        options.host = host
      end

      parser.parse!(ARGV)
      @options = options
    end

    def assert_valid_options
      unless @options.source_host
        raise "You must specify the source host via --source-host option."
      end
      unless @options.host
        raise "You must specify the destination host via --host option."
      end
    end

    def source_node
      @source_node ||= NodeName.new(:host => @options.source_host,
                                    :port => @options.source_port,
                                    :tag  => @options.source_tag)
    end

    def destination_node
      @destination_node ||= NodeName.new(:host => @options.host,
                                         :port => @options.port,
                                         :tag  => @options.tag)
    end

    def source_node_serf
      @source_node_serf ||= Serf.new(source_node.to_s,
                                     :verbose => @options.verbose)
    end

    def destination_node_serf
      @destination_node_serf ||= Serf.new(destination_node.to_s,
                                          :verbose => @options.verbose)
    end

    def absorber
      @absorber ||= prepare_absorber
    end

    def prepare_absorber
      absorber_options = {
        :host    => @options.host,
        :port    => @options.port,
        :tag     => @options.tag,
        :dataset => @options.dataset,

        :source_host    => @options.source_host,
        :source_port    => @options.source_port,
        :source_tag     => @options.source_tag,
        :source_dataset => @options.source_dataset,

        :receiver_host    => @options.receiver_host,

        :messages_per_second => @options.messages_per_second,
        :progress_interval_seconds => @options.progress_interval_seconds,

        :client_options   => {
          :backend => :coolio,
          :loop    => @loop,
        },
      }
      DataAbsorberClient.new(absorber_options)
    end

    def absorb
      last_progress = nil
      absorber.run do |progress|
        if last_progress
          printf("%s", "#{" " * last_progress[:message].size}\r")
        end
        printf("%s", "#{progress[:message]}\r")
        last_progress = progress
      end
      @loop.run

      if absorber.error_message
        puts(absorber.error_message)
        return false
      end

      puts ""

      timestamp = source_node_serf.last_processed_message_timestamp
      unless timestamp
        $stderr.puts("Couldn't get the time stamp of " +
                       "the last processed message from the source node.")
        return false
      end
      if timestamp and not timestamp.empty?
        puts "The timestamp of the last processed message in the source node: #{timestamp}"
        puts "Setting effective message timestamp for the destination node..."
        response = destination_node_serf.send_query("accept_messages_newer_than",
                                                    "node"      => destination_node.to_s,
                                                    "timestamp" => timestamp)
      end
      true
    end

    def trap_signals
      trap(:TERM) do
        trap(:TERM, "DEFAULT")
        do_cancel
      end

      trap(:INT) do
        trap(:INT, "DEFAULT")
        do_cancel
      end

      trap(:QUIT) do
        trap(:QUIT, "DEFAULT")
        do_cancel
      end
    end

    def do_cancel
      #XXX we have to write more codes to cancel remote processes!
    end
  end
end

exit(Droonga::AbsorbDataCommand.new.run)