fluent-plugin-gcloud-pubsub-custom-subscriber 1.3.1

data/lib/fluent/plugin/in_gcloud_pubsub.rb

@@ -0,0 +1,253 @@
+require 'json'
+require 'webrick'
+
+require 'fluent/plugin/input'
+require 'fluent/plugin/parser'
+
+require 'fluent/plugin/gcloud_pubsub/client'
+
+module Fluent::Plugin
+  class GcloudPubSubInput < Input
+    Fluent::Plugin.register_input('gcloud_pubsub', self)
+
+    helpers :compat_parameters, :parser, :thread
+
+    DEFAULT_PARSER_TYPE = 'json'
+
+    class FailedParseError < StandardError
+    end
+
+    desc 'Set tag of messages.'
+    config_param :tag,                :string
+    desc 'Set key to be used as tag.'
+    config_param :tag_key,            :string,  default: nil
+    desc 'Set your GCP project.'
+    config_param :project,            :string,  default: nil
+    desc 'Set your credential file path.'
+    config_param :key,                :string,  default: nil
+    desc 'Set topic name to pull.'
+    config_param :topic,              :string
+    desc 'Set subscription name to pull.'
+    config_param :subscription,       :string
+    desc 'Pulling messages by intervals of specified seconds.'
+    config_param :pull_interval,      :float,   default: 5.0
+    desc 'Max messages pulling at once.'
+    config_param :max_messages,       :integer, default: 100
+    desc 'Setting `true`, keepalive connection to wait for new messages.'
+    config_param :return_immediately, :bool,    default: true
+    desc 'Set number of threads to pull messages.'
+    config_param :pull_threads,       :integer, default: 1
+    desc 'Specify the key of the attribute to be acquired as a record'
+    config_param :attribute_keys,     :array,   default: []
+    desc 'Set error type when parsing messages fails.'
+    config_param :parse_error_action, :enum,    default: :exception, list: [:exception, :warning]
+    # for HTTP RPC
+    desc 'If `true` is specified, HTTP RPC to stop or start pulling message is enabled.'
+    config_param :enable_rpc,         :bool,    default: false
+    desc 'Bind IP address for HTTP RPC.'
+    config_param :rpc_bind,           :string,  default: '0.0.0.0'
+    desc 'Port for HTTP RPC.'
+    config_param :rpc_port,           :integer, default: 24680
+
+    config_section :parse do
+      config_set_default :@type, DEFAULT_PARSER_TYPE
+    end
+
+    class RPCServlet < WEBrick::HTTPServlet::AbstractServlet
+      class Error < StandardError; end
+
+      def initialize(server, plugin)
+        super
+        @plugin = plugin
+      end
+
+      def do_GET(req, res)
+        begin
+          code, header, body = process(req, res)
+        rescue
+          code, header, body = render_json(500, {
+              'ok' => false,
+              'message' => 'Internal Server Error',
+              'error' => "#{$!}",
+              'backtrace'=> $!.backtrace
+          })
+        end
+
+        res.status = code
+        header.each_pair {|k,v|
+          res[k] = v
+        }
+        res.body = body
+      end
+
+      def render_json(code, obj)
+        [code, {'Content-Type' => 'application/json'}, obj.to_json]
+      end
+
+      def process(req, res)
+        ret = {'ok' => true}
+        case req.path_info
+        when '/stop'
+          @plugin.stop_pull
+        when '/start'
+          @plugin.start_pull
+        when '/status'
+          ret['status'] = @plugin.status_of_pull
+        else
+          raise Error.new "Invalid path_info: #{req.path_info}"
+        end
+        render_json(200, ret)
+      end
+    end
+
+    def configure(conf)
+      compat_parameters_convert(conf, :parser)
+      super
+      @rpc_srv = nil
+      @rpc_thread = nil
+      @stop_pull = false
+
+      @extract_tag = if @tag_key.nil?
+                       method(:static_tag)
+                     else
+                       method(:dynamic_tag)
+                     end
+
+      @parser = parser_create
+    end
+
+    def start
+      super
+      start_rpc if @enable_rpc
+
+      @subscriber = Fluent::GcloudPubSub::Subscriber.new @project, @key, @topic, @subscription
+      log.debug "connected subscription:#{@subscription} in project #{@project}"
+
+      @emit_guard = Mutex.new
+      @stop_subscribing = false
+      @subscribe_threads = []
+      @pull_threads.times do |idx|
+        @subscribe_threads.push thread_create("in_gcloud_pubsub_subscribe_#{idx}".to_sym, &method(:subscribe))
+      end
+    end
+
+    def shutdown
+      if @rpc_srv
+        @rpc_srv.shutdown
+        @rpc_srv = nil
+      end
+      if @rpc_thread
+        @rpc_thread = nil
+      end
+      @stop_subscribing = true
+      @subscribe_threads.each(&:join)
+      super
+    end
+
+    def stop_pull
+      @stop_pull = true
+      log.info "stop pull from subscription:#{@subscription}"
+    end
+
+    def start_pull
+      @stop_pull = false
+      log.info "start pull from subscription:#{@subscription}"
+    end
+
+    def status_of_pull
+      @stop_pull ? 'stopped' : 'started'
+    end
+
+    private
+
+    def static_tag(record)
+      @tag
+    end
+
+    def dynamic_tag(record)
+      record.delete(@tag_key) || @tag
+    end
+
+    def start_rpc
+      log.info "listening http rpc server on http://#{@rpc_bind}:#{@rpc_port}/"
+      @rpc_srv = WEBrick::HTTPServer.new(
+        {
+          BindAddress: @rpc_bind,
+          Port: @rpc_port,
+          Logger: WEBrick::Log.new(STDERR, WEBrick::Log::FATAL),
+          AccessLog: []
+        }
+      )
+      @rpc_srv.mount('/api/in_gcloud_pubsub/pull/', RPCServlet, self)
+      @rpc_thread = thread_create(:in_gcloud_pubsub_rpc_thread){
+        @rpc_srv.start
+      }
+    end
+
+    def subscribe
+      until @stop_subscribing
+        _subscribe unless @stop_pull
+
+        if @return_immediately || @stop_pull
+          sleep @pull_interval
+        end
+      end
+    rescue => ex
+      log.error "unexpected error", error_message: ex.to_s, error_class: ex.class.to_s
+      log.error_backtrace ex.backtrace
+    end
+
+    def _subscribe
+      messages = @subscriber.pull @return_immediately, @max_messages
+      if messages.length == 0
+        log.debug "no messages are pulled"
+        return
+      end
+
+      process messages
+      @subscriber.acknowledge messages
+
+      log.debug "#{messages.length} message(s) processed"
+    rescue Fluent::GcloudPubSub::RetryableError => ex
+      log.warn "Retryable error occurs. Fluentd will retry.", error_message: ex.to_s, error_class: ex.class.to_s
+    rescue => ex
+      log.error "unexpected error", error_message: ex.to_s, error_class: ex.class.to_s
+      log.error_backtrace ex.backtrace
+    end
+
+    def process(messages)
+      event_streams = Hash.new do |hsh, key|
+        hsh[key] = Fluent::MultiEventStream.new
+      end
+
+      messages.each do |m|
+        line = m.message.data.chomp
+        attributes = m.attributes
+        @parser.parse(line) do |time, record|
+          if time && record
+            @attribute_keys.each do |key|
+              record[key] = attributes[key]
+            end
+
+            event_streams[@extract_tag.call(record)].add(time, record)
+          else
+            case @parse_error_action
+            when :exception
+              raise FailedParseError.new "pattern not match: #{line}"
+            else
+              log.warn 'pattern not match', record: line
+            end
+          end
+        end
+      end
+
+      event_streams.each do |tag, es|
+        # There are some output plugins not to supposed to be called with multi-threading.
+        # Maybe remove in the future.
+        @emit_guard.synchronize do
+          router.emit_stream(tag, es)
+        end
+      end
+    end
+  end
+end

data/lib/fluent/plugin/out_gcloud_pubsub.rb

@@ -0,0 +1,110 @@
+require 'fluent/plugin/output'
+require 'fluent/plugin/gcloud_pubsub/client'
+require 'fluent/plugin_helper/inject'
+
+module Fluent::Plugin
+  class GcloudPubSubOutput < Output
+    include Fluent::PluginHelper::Inject
+
+    Fluent::Plugin.register_output('gcloud_pubsub', self)
+
+    helpers :compat_parameters, :formatter
+
+    DEFAULT_BUFFER_TYPE = "memory"
+    DEFAULT_FORMATTER_TYPE = "json"
+
+    desc 'Set your GCP project.'
+    config_param :project,            :string,  :default => nil
+    desc 'Set your credential file path.'
+    config_param :key,                :string,  :default => nil
+    desc 'Set topic name to publish.'
+    config_param :topic,              :string
+    desc "If set to `true`, specified topic will be created when it doesn't exist."
+    config_param :autocreate_topic,   :bool,    :default => false
+    desc 'Publishing messages count per request to Cloud Pub/Sub.'
+    config_param :max_messages,       :integer, :default => 1000
+    desc 'Publishing messages bytesize per request to Cloud Pub/Sub.'
+    config_param :max_total_size,     :integer, :default => 9800000  # 9.8MB
+    desc 'Limit bytesize per message.'
+    config_param :max_message_size,   :integer, :default => 4000000  # 4MB
+    desc 'Publishing the set field as an attribute'
+    config_param :attribute_keys,     :array,   :default => []
+
+    config_section :buffer do
+      config_set_default :@type, DEFAULT_BUFFER_TYPE
+    end
+
+    config_section :format do
+      config_set_default :@type, DEFAULT_FORMATTER_TYPE
+    end
+
+    def configure(conf)
+      compat_parameters_convert(conf, :buffer, :formatter)
+      super
+      placeholder_validate!(:topic, @topic)
+      @formatter = formatter_create
+    end
+
+    def start
+      super
+      @publisher = Fluent::GcloudPubSub::Publisher.new @project, @key, @autocreate_topic
+    end
+
+    def format(tag, time, record)
+      record = inject_values_to_record(tag, time, record)
+      attributes = {}
+      @attribute_keys.each do |key|
+        attributes[key] = record.delete(key)
+      end
+      [@formatter.format(tag, time, record), attributes].to_msgpack
+    end
+
+    def formatted_to_msgpack_binary?
+      true
+    end
+
+    def multi_workers_ready?
+      true
+    end
+
+    def write(chunk)
+      topic = extract_placeholders(@topic, chunk.metadata)
+
+      messages = []
+      size = 0
+
+      chunk.msgpack_each do |msg, attr|
+        msg = Fluent::GcloudPubSub::Message.new(msg, attr)
+        if msg.bytesize > @max_message_size
+          log.warn 'Drop a message because its size exceeds `max_message_size`', size: msg.bytesize
+          next
+        end
+        if messages.length + 1 > @max_messages || size + msg.bytesize > @max_total_size
+          publish(topic, messages)
+          messages = []
+          size = 0
+        end
+        messages << msg
+        size += msg.bytesize
+      end
+
+      if messages.length > 0
+        publish(topic, messages)
+      end
+    rescue Fluent::GcloudPubSub::RetryableError => ex
+      log.warn "Retryable error occurs. Fluentd will retry.", error_message: ex.to_s, error_class: ex.class.to_s
+      raise ex
+    rescue => ex
+      log.error "unexpected error", error_message: ex.to_s, error_class: ex.class.to_s
+      log.error_backtrace
+      raise ex
+    end
+
+    private
+
+    def publish(topic, messages)
+      log.debug "send message topic:#{topic} length:#{messages.length} size:#{messages.map(&:bytesize).inject(:+)}"
+      @publisher.publish(topic, messages)
+    end
+  end
+end

data/test/plugin/test_in_gcloud_pubsub.rb

@@ -0,0 +1,387 @@
+require 'net/http'
+require 'json'
+
+require_relative "../test_helper"
+require "fluent/test/driver/input"
+
+class GcloudPubSubInputTest < Test::Unit::TestCase
+  CONFIG = %[
+      tag test
+      project project-test
+      topic topic-test
+      subscription subscription-test
+      key key-test
+  ]
+
+  DEFAULT_HOST = '127.0.0.1'
+  DEFAULT_PORT = 24680
+
+  class DummyInvalidMsgData
+    def data
+      'foo:bar'
+    end
+  end
+  class DummyInvalidMessage
+    def message
+      DummyInvalidMsgData.new
+    end
+    def attributes
+      {"attr_1" => "a", "attr_2" => "b"}
+    end
+  end
+
+  def create_driver(conf=CONFIG)
+    Fluent::Test::Driver::Input.new(Fluent::Plugin::GcloudPubSubInput).configure(conf)
+  end
+
+  def http_get(path)
+    http = Net::HTTP.new(DEFAULT_HOST, DEFAULT_PORT)
+    req = Net::HTTP::Get.new(path, {'Content-Type' => 'application/x-www-form-urlencoded'})
+    http.request(req)
+  end
+
+  setup do
+    Fluent::Test.setup
+  end
+
+  sub_test_case 'configure' do
+    test 'all params are configured' do
+      d = create_driver(%[
+        tag test
+        project project-test
+        topic topic-test
+        subscription subscription-test
+        key key-test
+        max_messages 1000
+        return_immediately true
+        pull_interval 2
+        pull_threads 3
+        attribute_keys attr-test
+        enable_rpc true
+        rpc_bind 127.0.0.1
+        rpc_port 24681
+      ])
+
+      assert_equal('test', d.instance.tag)
+      assert_equal('project-test', d.instance.project)
+      assert_equal('topic-test', d.instance.topic)
+      assert_equal('subscription-test', d.instance.subscription)
+      assert_equal('key-test', d.instance.key)
+      assert_equal(2.0, d.instance.pull_interval)
+      assert_equal(1000, d.instance.max_messages)
+      assert_equal(true, d.instance.return_immediately)
+      assert_equal(3, d.instance.pull_threads)
+      assert_equal(['attr-test'], d.instance.attribute_keys)
+      assert_equal(true, d.instance.enable_rpc)
+      assert_equal('127.0.0.1', d.instance.rpc_bind)
+      assert_equal(24681, d.instance.rpc_port)
+    end
+
+    test 'default values are configured' do
+      d = create_driver
+      assert_equal(5.0, d.instance.pull_interval)
+      assert_equal(100, d.instance.max_messages)
+      assert_equal(true, d.instance.return_immediately)
+      assert_equal(1, d.instance.pull_threads)
+      assert_equal([], d.instance.attribute_keys)
+      assert_equal(false, d.instance.enable_rpc)
+      assert_equal('0.0.0.0', d.instance.rpc_bind)
+      assert_equal(24680, d.instance.rpc_port)
+    end
+  end
+
+  sub_test_case 'start' do
+    setup do
+      @topic_mock = mock!
+      @pubsub_mock = mock!.topic('topic-test').at_least(1) { @topic_mock }
+      stub(Google::Cloud::Pubsub).new { @pubsub_mock }
+    end
+
+    test '40x error occurred on connecting to Pub/Sub' do
+      @topic_mock.subscription('subscription-test').once do
+        raise Google::Cloud::NotFoundError.new('TEST')
+      end
+
+      d = create_driver
+      assert_raise Google::Cloud::NotFoundError do
+        d.run {}
+      end
+    end
+
+    test '50x error occurred on connecting to Pub/Sub' do
+      @topic_mock.subscription('subscription-test').once do
+        raise Google::Cloud::UnavailableError.new('TEST')
+      end
+
+      d = create_driver
+      assert_raise Google::Cloud::UnavailableError do
+        d.run {}
+      end
+    end
+
+    test 'subscription is nil' do
+      @topic_mock.subscription('subscription-test').once { nil }
+
+      d = create_driver
+      assert_raise Fluent::GcloudPubSub::Error do
+        d.run {}
+      end
+    end
+  end
+
+  sub_test_case 'emit' do
+    class DummyMsgData
+      def data
+        '{"foo": "bar"}'
+      end
+    end
+    class DummyMessage
+      def message
+        DummyMsgData.new
+      end
+      def attributes
+        {"attr_1" => "a", "attr_2" => "b"}
+      end
+    end
+
+    class DummyMsgDataWithTagKey
+      def initialize(tag)
+        @tag = tag
+      end
+      def data
+        '{"foo": "bar", "test_tag_key": "' + @tag + '"}'
+      end
+    end
+    class DummyMessageWithTagKey
+      def initialize(tag)
+        @tag = tag
+      end
+      def message
+        DummyMsgDataWithTagKey.new @tag
+      end
+      def attributes
+        {"attr_1" => "a", "attr_2" => "b"}
+      end
+    end
+
+    setup do
+      @subscriber = mock!
+      @topic_mock = mock!.subscription('subscription-test') { @subscriber }
+      @pubsub_mock = mock!.topic('topic-test') { @topic_mock }
+      stub(Google::Cloud::Pubsub).new { @pubsub_mock }
+    end
+
+    test 'empty' do
+      @subscriber.pull(immediate: true, max: 100).at_least(1) { [] }
+      @subscriber.acknowledge.times(0)
+
+      d = create_driver
+      d.run(expect_emits: 1, timeout: 3)
+
+      assert_true d.events.empty?
+    end
+
+    test 'simple' do
+      messages = Array.new(1, DummyMessage.new)
+      @subscriber.pull(immediate: true, max: 100).at_least(1) { messages }
+      @subscriber.acknowledge(messages).at_least(1)
+
+      d = create_driver
+      d.run(expect_emits: 1, timeout: 3)
+      emits = d.events
+
+      assert(1 <= emits.length)
+      emits.each do |tag, time, record|
+        assert_equal("test", tag)
+        assert_equal({"foo" => "bar"}, record)
+      end
+    end
+
+    test 'multithread' do
+      messages = Array.new(1, DummyMessage.new)
+      @subscriber.pull(immediate: true, max: 100).at_least(2) { messages }
+      @subscriber.acknowledge(messages).at_least(2)
+
+      d = create_driver("#{CONFIG}\npull_threads 2")
+      d.run(expect_emits: 2, timeout: 1)
+      emits = d.events
+
+      assert(2 <= emits.length)
+      emits.each do |tag, time, record|
+        assert_equal("test", tag)
+        assert_equal({"foo" => "bar"}, record)
+      end
+    end
+
+    test 'with tag_key' do
+      messages = [
+        DummyMessageWithTagKey.new('tag1'),
+        DummyMessageWithTagKey.new('tag2'),
+        DummyMessage.new
+      ]
+      @subscriber.pull(immediate: true, max: 100).at_least(1) { messages }
+      @subscriber.acknowledge(messages).at_least(1)
+
+      d = create_driver("#{CONFIG}\ntag_key test_tag_key")
+      d.run(expect_emits: 1, timeout: 3)
+      emits = d.events
+
+      assert(3 <= emits.length)
+      # test tag
+      assert_equal("tag1", emits[0][0])
+      assert_equal("tag2", emits[1][0])
+      assert_equal("test", emits[2][0])
+      # test record
+      emits.each do |tag, time, record|
+        assert_equal({"foo" => "bar"}, record)
+      end
+    end
+
+    test 'invalid messages with parse_error_action exception ' do
+      messages = Array.new(1, DummyInvalidMessage.new)
+      @subscriber.pull(immediate: true, max: 100).at_least(1) { messages }
+      @subscriber.acknowledge.times(0)
+
+      d = create_driver
+      d.run(expect_emits: 1, timeout: 3)
+      assert_true d.events.empty?
+    end
+
+    test 'with attributes' do
+      messages = Array.new(1, DummyMessage.new)
+      @subscriber.pull(immediate: true, max: 100).at_least(1) { messages }
+      @subscriber.acknowledge(messages).at_least(1)
+
+      d = create_driver("#{CONFIG}\nattribute_keys attr_1")
+      d.run(expect_emits: 1, timeout: 3)
+      emits = d.events
+
+      assert(1 <= emits.length)
+      emits.each do |tag, time, record|
+        assert_equal("test", tag)
+        assert_equal({"foo" => "bar", "attr_1" => "a"}, record)
+      end
+    end
+
+    test 'invalid messages with parse_error_action warning' do
+      messages = Array.new(1, DummyInvalidMessage.new)
+      @subscriber.pull(immediate: true, max: 100).at_least(1) { messages }
+      @subscriber.acknowledge(messages).at_least(1)
+
+      d = create_driver("#{CONFIG}\nparse_error_action warning")
+      d.run(expect_emits: 1, timeout: 3)
+      assert_true d.events.empty?
+    end
+
+    test 'retry if raised error' do
+      class UnknownError < StandardError
+      end
+      @subscriber.pull(immediate: true, max: 100).at_least(2) { raise UnknownError.new('test') }
+      @subscriber.acknowledge.times(0)
+
+      d = create_driver(CONFIG + 'pull_interval 0.5')
+      d.run(expect_emits: 1, timeout: 0.8)
+
+      assert_equal(0.5, d.instance.pull_interval)
+      assert_true d.events.empty?
+    end
+
+    test 'retry if raised RetryableError on pull' do
+      @subscriber.pull(immediate: true, max: 100).at_least(2) { raise Google::Cloud::UnavailableError.new('TEST') }
+      @subscriber.acknowledge.times(0)
+
+      d = create_driver("#{CONFIG}\npull_interval 0.5")
+      d.run(expect_emits: 1, timeout: 0.8)
+
+      assert_equal(0.5, d.instance.pull_interval)
+      assert_true d.events.empty?
+    end
+
+    test 'retry if raised RetryableError on acknowledge' do
+      messages = Array.new(1, DummyMessage.new)
+      @subscriber.pull(immediate: true, max: 100).at_least(2) { messages }
+      @subscriber.acknowledge(messages).at_least(2) { raise Google::Cloud::UnavailableError.new('TEST') }
+
+      d = create_driver("#{CONFIG}\npull_interval 0.5")
+      d.run(expect_emits: 2, timeout: 3)
+      emits = d.events
+
+      # not acknowledged, but already emitted to engine.
+      assert(2 <= emits.length)
+      emits.each do |tag, time, record|
+        assert_equal("test", tag)
+        assert_equal({"foo" => "bar"}, record)
+      end
+    end
+
+    test 'stop by http rpc' do
+      messages = Array.new(1, DummyMessage.new)
+      @subscriber.pull(immediate: true, max: 100).once { messages }
+      @subscriber.acknowledge(messages).once
+
+      d = create_driver("#{CONFIG}\npull_interval 1.0\nenable_rpc true")
+      assert_equal(false, d.instance.instance_variable_get(:@stop_pull))
+
+      d.run {
+        http_get('/api/in_gcloud_pubsub/pull/stop')
+        sleep 0.75
+        # d.run sleeps 0.5 sec
+      }
+      emits = d.events
+
+      assert_equal(1, emits.length)
+      assert_true d.instance.instance_variable_get(:@stop_pull)
+
+      emits.each do |tag, time, record|
+        assert_equal("test", tag)
+        assert_equal({"foo" => "bar"}, record)
+      end
+    end
+
+    test 'start by http rpc' do
+      messages = Array.new(1, DummyMessage.new)
+      @subscriber.pull(immediate: true, max: 100).at_least(1) { messages }
+      @subscriber.acknowledge(messages).at_least(1)
+
+      d = create_driver("#{CONFIG}\npull_interval 1.0\nenable_rpc true")
+      d.instance.stop_pull
+      assert_equal(true, d.instance.instance_variable_get(:@stop_pull))
+
+      d.run(expect_emits: 1, timeout: 3) {
+        http_get('/api/in_gcloud_pubsub/pull/start')
+        sleep 0.75
+        # d.run sleeps 0.5 sec
+      }
+      emits = d.events
+
+      assert_equal(true, emits.length > 0)
+      assert_false d.instance.instance_variable_get(:@stop_pull)
+
+      emits.each do |tag, time, record|
+        assert_equal("test", tag)
+        assert_equal({"foo" => "bar"}, record)
+      end
+    end
+
+    test 'get status by http rpc when started' do
+      d = create_driver("#{CONFIG}\npull_interval 1.0\nenable_rpc true")
+      assert_false d.instance.instance_variable_get(:@stop_pull)
+
+      d.run {
+        res = http_get('/api/in_gcloud_pubsub/pull/status')
+        assert_equal({"ok" => true, "status" => "started"}, JSON.parse(res.body))
+      }
+    end
+
+    test 'get status by http rpc when stopped' do
+      d = create_driver("#{CONFIG}\npull_interval 1.0\nenable_rpc true")
+      d.instance.stop_pull
+      assert_true d.instance.instance_variable_get(:@stop_pull)
+
+      d.run {
+        res = http_get('/api/in_gcloud_pubsub/pull/status')
+        assert_equal({"ok" => true, "status" => "stopped"}, JSON.parse(res.body))
+      }
+    end
+  end
+end