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

data/Rakefile

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+require "bundler"
+Bundler::GemHelper.install_tasks
+
+require "rake/testtask"
+
+Rake::TestTask.new(:test) do |test|
+  test.libs << "lib" << "test"
+  test.test_files = FileList["test/plugin/test_*.rb"]
+  test.verbose = true
+end
+
+task default: [:build]

data/fluent-plugin-gcloud-pubsub-custom.gemspec

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+$LOAD_PATH.push File.expand_path("lib", __dir__)
+
+Gem::Specification.new do |gem|
+  gem.name        = "fluent-plugin-gcloud-pubsub-custom-compress-batches"
+  gem.description = "Google Cloud Pub/Sub input/output plugin for Fluentd event collector - with payload compression. Forked from https://github.com/gocardless/fluent-plugin-gcloud-pubsub-custom"
+  gem.license     = "MIT"
+  gem.homepage    = "https://github.com/calvinaditya95/fluent-plugin-gcloud-pubsub-custom"
+  gem.summary     = "Google Cloud Pub/Sub input/output plugin for Fluentd event collector - with payload compression"
+  gem.version     = "1.3.4"
+  gem.authors     = ["Calvin Aditya"]
+  gem.email       = "calvin.aditya95@gmail.com"
+  gem.files       = `git ls-files`.split("\n")
+  gem.test_files  = `git ls-files -- {test,spec,features}/*`.split("\n")
+  gem.executables = `git ls-files -- bin/*`.split("\n").map { |f| File.basename(f) }
+  gem.require_paths = ["lib"]
+
+  gem.add_runtime_dependency "fluentd", [">= 0.14.15", "< 2"]
+  gem.add_runtime_dependency "google-cloud-pubsub", "~> 0.30.0"
+
+  # Use the same version constraint as fluent-plugin-prometheus currently specifies
+  gem.add_runtime_dependency "prometheus-client", "< 0.10"
+
+  gem.add_development_dependency "bundler"
+  gem.add_development_dependency "pry"
+  gem.add_development_dependency "pry-byebug"
+  gem.add_development_dependency "rake"
+  gem.add_development_dependency "rubocop", "~>0.83"
+  gem.add_development_dependency "test-unit"
+  gem.add_development_dependency "test-unit-rr"
+end

data/lib/fluent/plugin/gcloud_pubsub/client.rb

@@ -0,0 +1,173 @@
+# frozen_string_literal: true
+
+require "google/cloud/pubsub"
+require "zlib"
+
+module Fluent
+  module GcloudPubSub
+    class Error < StandardError
+    end
+    class RetryableError < Error
+    end
+
+    COMPRESSION_ALGORITHM_ZLIB = "zlib"
+    # 30 is the ASCII record separator character
+    BATCHED_RECORD_SEPARATOR = 30.chr
+
+    class Message
+      attr_reader :message, :attributes
+
+      def initialize(message, attributes = {})
+        @message = message
+        @attributes = attributes
+      end
+
+      def bytesize
+        attr_size = 0
+        @attributes.each do |key, val|
+          attr_size += key.bytesize + val.bytesize
+        end
+        @message.bytesize + attr_size
+      end
+    end
+
+    class Publisher
+      def initialize(project, key, autocreate_topic, metric_prefix)
+        @pubsub = Google::Cloud::Pubsub.new project_id: project, credentials: key
+        @autocreate_topic = autocreate_topic
+        @topics = {}
+
+        # rubocop:disable Layout/LineLength
+        @compression_ratio =
+          Fluent::GcloudPubSub::Metrics.register_or_existing(:"#{metric_prefix}_messages_compressed_size_per_original_size_ratio") do
+            ::Prometheus::Client.registry.histogram(
+              :"#{metric_prefix}_messages_compressed_size_per_original_size_ratio",
+              "Compression ratio achieved on a batch of messages",
+              {},
+              # We expect compression for even a single message to be typically
+              # above 2x (0.5/50%), so bias the buckets towards the higher end
+              # of the range.
+              [0, 0.25, 0.5, 0.75, 0.85, 0.9, 0.95, 0.975, 1],
+            )
+          end
+
+        @compression_duration =
+          Fluent::GcloudPubSub::Metrics.register_or_existing(:"#{metric_prefix}_messages_compression_duration_seconds") do
+            ::Prometheus::Client.registry.histogram(
+              :"#{metric_prefix}_messages_compression_duration_seconds",
+              "Time taken to compress a batch of messages",
+              {},
+              [0, 0.0001, 0.0005, 0.001, 0.01, 0.05, 0.1, 0.25, 0.5, 1],
+            )
+          end
+        # rubocop:enable Layout/LineLength
+      end
+
+      def topic(topic_name)
+        return @topics[topic_name] if @topics.key? topic_name
+
+        client = @pubsub.topic topic_name
+        client = @pubsub.create_topic topic_name if client.nil? && @autocreate_topic
+        raise Error, "topic:#{topic_name} does not exist." if client.nil?
+
+        @topics[topic_name] = client
+        client
+      end
+
+      def publish(topic_name, messages, compress_batches = false)
+        if compress_batches
+          topic(topic_name).publish(*compress_messages_with_zlib(messages, topic_name))
+        else
+          topic(topic_name).publish do |batch|
+            messages.each do |m|
+              batch.publish m.message, m.attributes
+            end
+          end
+        end
+      rescue Google::Cloud::UnavailableError, Google::Cloud::DeadlineExceededError, Google::Cloud::InternalError => e
+        raise RetryableError, "Google api returns error:#{e.class} message:#{e}"
+      end
+
+      private
+
+      def compress_messages_with_zlib(messages, topic_name)
+        original_size = messages.sum(&:bytesize)
+        # This should never happen, only a programming error or major
+        # misconfiguration should lead to this situation. But checking against
+        # it here avoids a potential division by zero later on.
+        raise ArgumentError, "not compressing empty inputs" if original_size.zero?
+
+        # Here we're implicitly dropping the 'attributes' field of the messages
+        # that we're iterating over.
+        # This is fine, because the :attribute_keys config param is not
+        # supported when in compressed mode, so this field will always be
+        # empty.
+        packed_messages = messages.map(&:message).join(BATCHED_RECORD_SEPARATOR)
+
+        duration, compressed_messages = Fluent::GcloudPubSub::Metrics.measure_duration do
+          Zlib::Deflate.deflate(packed_messages)
+        end
+
+        @compression_duration.observe(
+          { topic: topic_name, algorithm: COMPRESSION_ALGORITHM_ZLIB },
+          duration,
+        )
+
+        compressed_size = compressed_messages.bytesize
+        @compression_ratio.observe(
+          { topic: topic_name, algorithm: COMPRESSION_ALGORITHM_ZLIB },
+          # If original = 1MiB and compressed = 256KiB; then metric value = 0.75 = 75% when plotted
+          1 - compressed_size.to_f / original_size,
+        )
+
+        [compressed_messages, { "compression_algorithm": COMPRESSION_ALGORITHM_ZLIB }]
+      end
+    end
+
+    class Subscriber
+      def initialize(project, key, topic_name, subscription_name)
+        pubsub = Google::Cloud::Pubsub.new project_id: project, credentials: key
+        if topic_name.nil?
+          @client = pubsub.subscription subscription_name
+        else
+          topic = pubsub.topic topic_name
+          @client = topic.subscription subscription_name
+        end
+        raise Error, "subscription:#{subscription_name} does not exist." if @client.nil?
+      end
+
+      def pull(immediate, max)
+        @client.pull immediate: immediate, max: max
+      rescue Google::Cloud::UnavailableError, Google::Cloud::DeadlineExceededError, Google::Cloud::InternalError => e
+        raise RetryableError, "Google pull api returns error:#{e.class} message:#{e}"
+      end
+
+      def acknowledge(messages)
+        @client.acknowledge messages
+      rescue Google::Cloud::UnavailableError, Google::Cloud::DeadlineExceededError, Google::Cloud::InternalError => e
+        raise RetryableError, "Google acknowledge api returns error:#{e.class} message:#{e}"
+      end
+    end
+
+    class MessageUnpacker
+      def self.unpack(message)
+        attributes = message.attributes
+        algorithm = attributes["compression_algorithm"]
+
+        case algorithm
+        when nil
+          # For an uncompressed message return the single line and attributes
+          [[message.message.data.chomp, message.attributes]]
+        when COMPRESSION_ALGORITHM_ZLIB
+          # Return all of the lines in the message, with empty attributes
+          Zlib::Inflate
+            .inflate(message.message.data)
+            .split(BATCHED_RECORD_SEPARATOR)
+            .map { |line| [line, {}] }
+        else
+          raise ArgumentError, "unknown compression algorithm: '#{algorithm}'"
+        end
+      end
+    end
+  end
+end

data/lib/fluent/plugin/gcloud_pubsub/metrics.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module Fluent
+  module GcloudPubSub
+    # Utilities for interacting with Prometheus metrics
+    module Metrics
+      def self.register_or_existing(metric_name)
+        return ::Prometheus::Client.registry.get(metric_name) if ::Prometheus::Client.registry.exist?(metric_name)
+
+        yield
+      end
+
+      # Time the elapsed execution of the provided block, return the duration
+      # as the first element followed by the result of the block.
+      def self.measure_duration
+        start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+        result = yield
+        finish = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+
+        [finish - start, *result]
+      end
+    end
+  end
+end

data/lib/fluent/plugin/in_gcloud_pubsub.rb

@@ -0,0 +1,303 @@
+# frozen_string_literal: true
+
+require "json"
+require "webrick"
+
+require "fluent/plugin/input"
+require "fluent/plugin/parser"
+
+require "fluent/plugin/gcloud_pubsub/client"
+require "fluent/plugin/gcloud_pubsub/metrics"
+
+require "prometheus/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, default: nil
+    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 "Acquire these fields from attributes on the Pub/Sub message and merge them into the record"
+    config_param :attribute_keys,     :array, default: []
+    desc "Set error type when parsing messages fails."
+    config_param :parse_error_action, :enum, default: :exception, list: %i[exception warning]
+    desc "The prefix for Prometheus metric names"
+    config_param :metric_prefix, :string, default: "fluentd_input_gcloud_pubsub"
+    # 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: 24_680
+
+    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
+
+      # rubocop:disable Naming/MethodName
+      def do_GET(req, res)
+        begin
+          code, header, body = process(req, res)
+        rescue StandardError
+          code, header, body = render_json(500, {
+                                             "ok" => false,
+                                             "message" => "Internal Server Error",
+                                             "error" => $ERROR_INFO.to_s,
+                                             "backtrace" => $ERROR_INFO.backtrace,
+                                           })
+        end
+
+        res.status = code
+        header.each_pair do |k, v|
+          res[k] = v
+        end
+        res.body = body
+      end
+      # rubocop:enable Naming/MethodName
+
+      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, "Invalid path_info: #{req.path_info}"
+        end
+        render_json(200, ret)
+      end
+    end
+
+    # rubocop:disable Metrics/MethodLength
+    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
+
+      @messages_pulled =
+        Fluent::GcloudPubSub::Metrics.register_or_existing(:"#{@metric_prefix}_messages_pulled") do
+          ::Prometheus::Client.registry.histogram(
+            :"#{@metric_prefix}_messages_pulled",
+            "Number of Pub/Sub messages pulled by the subscriber on each invocation",
+            {},
+            [0, 1, 10, 50, 100, 250, 500, 1000],
+          )
+        end
+
+      @messages_pulled_bytes =
+        Fluent::GcloudPubSub::Metrics.register_or_existing(:"#{@metric_prefix}_messages_pulled_bytes") do
+          ::Prometheus::Client.registry.histogram(
+            :"#{@metric_prefix}_messages_pulled_bytes",
+            "Total size in bytes of the Pub/Sub messages pulled by the subscriber on each invocation",
+            {},
+            [100, 1000, 10_000, 100_000, 1_000_000, 5_000_000, 10_000_000],
+          )
+        end
+
+      @pull_errors =
+        Fluent::GcloudPubSub::Metrics.register_or_existing(:"#{@metric_prefix}_pull_errors_total") do
+          ::Prometheus::Client.registry.counter(
+            :"#{@metric_prefix}_pull_errors_total",
+            "Errors encountered while pulling or processing messages",
+            {},
+          )
+        end
+    end
+    # rubocop:enable Metrics/MethodLength
+
+    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
+      @rpc_thread = nil if @rpc_thread
+      @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) do
+        @rpc_srv.start
+      end
+    end
+
+    def subscribe
+      until @stop_subscribing
+        _subscribe unless @stop_pull
+
+        sleep @pull_interval if @return_immediately || @stop_pull
+      end
+    rescue StandardError => e
+      log.error "unexpected error", error_message: e.to_s, error_class: e.class.to_s
+      log.error_backtrace e.backtrace
+    end
+
+    def _subscribe
+      messages = @subscriber.pull @return_immediately, @max_messages
+      @messages_pulled.observe(common_labels, messages.size)
+      if messages.empty?
+        log.debug "no messages are pulled"
+        return
+      end
+
+      messages_size = messages.sum do |message|
+        message.data.bytesize + message.attributes.sum { |k, v| k.bytesize + v.bytesize }
+      end
+      @messages_pulled_bytes.observe(common_labels, messages_size)
+
+      process messages
+      @subscriber.acknowledge messages
+
+      log.debug "#{messages.length} message(s) processed"
+    rescue Fluent::GcloudPubSub::RetryableError => e
+      @pull_errors.increment(common_labels.merge({ retryable: true }))
+      log.warn "Retryable error occurs. Fluentd will retry.", error_message: e.to_s, error_class: e.class.to_s
+    rescue StandardError => e
+      @pull_errors.increment(common_labels.merge({ retryable: false }))
+      log.error "unexpected error", error_message: e.to_s, error_class: e.class.to_s
+      log.error_backtrace e.backtrace
+    end
+
+    def process(messages)
+      event_streams = Hash.new do |hsh, key|
+        hsh[key] = Fluent::MultiEventStream.new
+      end
+
+      messages.each do |m|
+        lines_attributes = Fluent::GcloudPubSub::MessageUnpacker.unpack(m)
+
+        lines_attributes.each do |line, 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, "pattern not match: #{line}"
+              else
+                log.warn "pattern not match", record: line
+              end
+            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
+
+    def common_labels
+      { subscription: @subscription }
+    end
+  end
+end