fluent-plugin-gcloud-pubsub-custom-compress-batches 1.3.4

data/lib/fluent/plugin/out_gcloud_pubsub.rb

@@ -0,0 +1,167 @@
+# frozen_string_literal: true
+
+require "fluent/plugin/output"
+require "fluent/plugin/gcloud_pubsub/client"
+require "fluent/plugin/gcloud_pubsub/metrics"
+require "fluent/plugin_helper/inject"
+require "prometheus/client"
+
+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: 9_800_000 # 9.8MB
+    desc "Limit bytesize per message."
+    config_param :max_message_size, :integer, default: 4_000_000 # 4MB
+    desc "Extract these fields from the record and send them as attributes on the Pub/Sub message. " \
+         "Cannot be set if compress_batches is enabled."
+    config_param :attribute_keys, :array, default: []
+    desc "The prefix for Prometheus metric names"
+    config_param :metric_prefix, :string, default: "fluentd_output_gcloud_pubsub"
+    desc "If set to `true`, messages will be batched and compressed before publication"
+    config_param :compress_batches, :bool, default: false
+
+    config_section :buffer do
+      config_set_default :@type, DEFAULT_BUFFER_TYPE
+    end
+
+    config_section :format do
+      config_set_default :@type, DEFAULT_FORMATTER_TYPE
+    end
+
+    # rubocop:disable Metrics/MethodLength
+    def configure(conf)
+      compat_parameters_convert(conf, :buffer, :formatter)
+      super
+      placeholder_validate!(:topic, @topic)
+      @formatter = formatter_create
+
+      if @compress_batches && !@attribute_keys.empty?
+        # The attribute_keys option is implemented by extracting keys from the
+        # record and setting them on the Pub/Sub message.
+        # This is not possible in compressed mode, because we're sending just a
+        # single Pub/Sub message that comprises many records, therefore the
+        # attribute keys would clash.
+        raise Fluent::ConfigError, ":attribute_keys cannot be used when compression is enabled"
+      end
+
+      @messages_published =
+        Fluent::GcloudPubSub::Metrics.register_or_existing(:"#{@metric_prefix}_messages_published_per_batch") do
+          ::Prometheus::Client.registry.histogram(
+            :"#{@metric_prefix}_messages_published_per_batch",
+            "Number of records published to Pub/Sub per buffer flush",
+            {},
+            [1, 10, 50, 100, 250, 500, 1000],
+          )
+        end
+
+      @bytes_published =
+        Fluent::GcloudPubSub::Metrics.register_or_existing(:"#{@metric_prefix}_messages_published_bytes") do
+          ::Prometheus::Client.registry.histogram(
+            :"#{@metric_prefix}_messages_published_bytes",
+            "Total size in bytes of the records published to Pub/Sub",
+            {},
+            [100, 1000, 10_000, 100_000, 1_000_000, 5_000_000, 10_000_000],
+          )
+        end
+
+      @compression_enabled =
+        Fluent::GcloudPubSub::Metrics.register_or_existing(:"#{@metric_prefix}_compression_enabled") do
+          ::Prometheus::Client.registry.gauge(
+            :"#{@metric_prefix}_compression_enabled",
+            "Whether compression/batching is enabled",
+            {},
+          )
+        end
+      @compression_enabled.set(common_labels, @compress_batches ? 1 : 0)
+    end
+    # rubocop:enable Metrics/MethodLength
+
+    def start
+      super
+      @publisher = Fluent::GcloudPubSub::Publisher.new @project, @key, @autocreate_topic, @metric_prefix
+    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
+
+      publish(topic, messages) unless messages.empty?
+    rescue Fluent::GcloudPubSub::RetryableError => e
+      log.warn "Retryable error occurs. Fluentd will retry.", error_message: e.to_s, error_class: e.class.to_s
+      raise e
+    rescue StandardError => e
+      log.error "unexpected error", error_message: e.to_s, error_class: e.class.to_s
+      log.error_backtrace
+      raise e
+    end
+
+    private
+
+    def publish(topic, messages)
+      size = messages.map(&:bytesize).inject(:+)
+      log.debug "send message topic:#{topic} length:#{messages.length} size:#{size}"
+
+      @messages_published.observe(common_labels, messages.length)
+      @bytes_published.observe(common_labels, size)
+
+      @publisher.publish(topic, messages, @compress_batches)
+    end
+
+    def common_labels
+      { topic: @topic }
+    end
+  end
+end

data/test/plugin/test_in_gcloud_pubsub.rb

@@ -0,0 +1,455 @@
+# frozen_string_literal: true
+
+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 = 24_680
+
+  class DummyInvalidMsgData
+    def data
+      "foo:bar"
+    end
+  end
+  class DummyInvalidMessage
+    def message
+      DummyInvalidMsgData.new
+    end
+
+    def data
+      message.data
+    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(24_681, 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(24_680, 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, "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, "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 data
+        message.data
+      end
+
+      def attributes
+        { "attr_1" => "a", "attr_2" => "b" }
+      end
+    end
+
+    class DummyCompressedMessageData
+      attr_reader :data
+
+      def initialize(messages)
+        @data = Zlib::Deflate.deflate(messages.join(30.chr))
+      end
+    end
+
+    class DummyCompressedMessage
+      attr_reader :message
+
+      def initialize(messages)
+        @message = DummyCompressedMessageData.new(messages)
+      end
+
+      def data
+        message.data
+      end
+
+      def attributes
+        { "compression_algorithm" => "zlib" }
+      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 data
+        message.data
+      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(emits.length >= 1)
+      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(emits.length >= 2)
+      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(emits.length >= 3)
+      # 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(emits.length >= 1)
+      emits.each do |tag, _time, record|
+        assert_equal("test", tag)
+        assert_equal({ "foo" => "bar", "attr_1" => "a" }, record)
+      end
+    end
+
+    test "compressed batch of messages" do
+      original_messages = [
+        { foo: "bar" },
+        { baz: "qux" },
+      ]
+      messages = Array.new(1, DummyCompressedMessage.new(original_messages.map(&:to_json)))
+
+      @subscriber.pull(immediate: true, max: 100).once { messages }
+      @subscriber.acknowledge(messages).at_least(1)
+
+      d = create_driver
+      d.run(expect_emits: 1, timeout: 3)
+      emits = d.events
+
+      output_records = emits.map do |e|
+        # Pick out only the record element, i.e. ignore the time and tag
+        record = e[2]
+        # Convert the keys from strings to symbols, to allow for strict comparison
+        record.map { |k, v| [k.to_sym, v] }.to_h
+      end
+
+      assert_equal(original_messages, output_records)
+    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, "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, "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, "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(emits.length >= 2)
+      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 do
+        http_get("/api/in_gcloud_pubsub/pull/stop")
+        sleep 0.75
+        # d.run sleeps 0.5 sec
+      end
+      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) do
+        http_get("/api/in_gcloud_pubsub/pull/start")
+        sleep 0.75
+        # d.run sleeps 0.5 sec
+      end
+      emits = d.events
+
+      assert_equal(true, !emits.empty?)
+      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 do
+        res = http_get("/api/in_gcloud_pubsub/pull/status")
+        assert_equal({ "ok" => true, "status" => "started" }, JSON.parse(res.body))
+      end
+    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 do
+        res = http_get("/api/in_gcloud_pubsub/pull/status")
+        assert_equal({ "ok" => true, "status" => "stopped" }, JSON.parse(res.body))
+      end
+    end
+  end
+end