waterdrop 2.6.6 → 2.6.8

data/lib/waterdrop/clients/buffered.rb

@@ -19,17 +19,73 @@ module WaterDrop
         super
         @messages = []
         @topics = Hash.new { |k, v| k[v] = [] }
+
+        @transaction_mutex = Mutex.new
+        @transaction_active = false
+        @transaction_messages = []
+        @transaction_topics = Hash.new { |k, v| k[v] = [] }
+        @transaction_level = 0
       end
 
       # "Produces" message to Kafka: it acknowledges it locally, adds it to the internal buffer
       # @param message [Hash] `WaterDrop::Producer#produce_sync` message hash
       def produce(message)
-        # We pre-validate the message payload, so topic is ensured to be present
-        @topics[message.fetch(:topic)] << message
-        @messages << message
+        if @transaction_active
+          @transaction_topics[message.fetch(:topic)] << message
+          @transaction_messages << message
+        else
+          # We pre-validate the message payload, so topic is ensured to be present
+          @topics[message.fetch(:topic)] << message
+          @messages << message
+        end
+
         SyncResponse.new
       end
 
+      # Yields the code pretending it is in a transaction
+      # Supports our aborting transaction flow
+      # Moves messages the appropriate buffers only if transaction is successful
+      def transaction
+        @transaction_level += 1
+
+        return yield if @transaction_mutex.owned?
+
+        @transaction_mutex.lock
+        @transaction_active = true
+
+        result = nil
+        commit = false
+
+        catch(:abort) do
+          result = yield
+          commit = true
+        end
+
+        commit || raise(WaterDrop::Errors::AbortTransaction)
+
+        # Transfer transactional data on success
+        @transaction_topics.each do |topic, messages|
+          @topics[topic] += messages
+        end
+
+        @messages += @transaction_messages
+
+        result
+      rescue StandardError => e
+        return if e.is_a?(WaterDrop::Errors::AbortTransaction)
+
+        raise
+      ensure
+        @transaction_level -= 1
+
+        if @transaction_level.zero? && @transaction_mutex.owned?
+          @transaction_topics.clear
+          @transaction_messages.clear
+          @transaction_active = false
+          @transaction_mutex.unlock
+        end
+      end
+
       # Returns messages produced to a given topic
       # @param topic [String]
       def messages_for(topic)

data/lib/waterdrop/clients/dummy.rb

@@ -25,6 +25,26 @@ module WaterDrop
         true
       end
 
+      # Yields the code pretending it is in a transaction
+      # Supports our aborting transaction flow
+      def transaction
+        result = nil
+        commit = false
+
+        catch(:abort) do
+          result = yield
+          commit = true
+        end
+
+        commit || raise(WaterDrop::Errors::AbortTransaction)
+
+        result
+      rescue StandardError => e
+        return if e.is_a?(WaterDrop::Errors::AbortTransaction)
+
+        raise
+      end
+
       # @param _args [Object] anything really, this dummy is suppose to support anything
       # @return [self] returns self for chaining cases
       def method_missing(*_args)

data/lib/waterdrop/clients/rdkafka.rb

@@ -11,7 +11,9 @@ module WaterDrop
         # @param producer [WaterDrop::Producer] producer instance with its config, etc
         # @note We overwrite this that way, because we do not care
         def new(producer)
-          client = ::Rdkafka::Config.new(producer.config.kafka.to_h).producer
+          config = producer.config.kafka.to_h
+
+          client = ::Rdkafka::Config.new(config).producer
 
           # This callback is not global and is per client, thus we do not have to wrap it with a
           # callbacks manager to make it work
@@ -20,6 +22,9 @@ module WaterDrop
             producer.config.monitor
           )
 
+          # Switch to the transactional mode if user provided the transactional id
+          client.init_transactions if config.key?(:'transactional.id')
+
           client
         end
       end

data/lib/waterdrop/config.rb

@@ -64,6 +64,11 @@ module WaterDrop
     # option [Numeric] how many seconds should we wait with the backoff on queue having space for
     # more messages before re-raising the error.
     setting :wait_timeout_on_queue_full, default: 10
+
+    setting :wait_backoff_on_transaction_command, default: 0.5
+
+    setting :max_attempts_on_transaction_command, default: 5
+
     # option [Boolean] should we send messages. Setting this to false can be really useful when
     #   testing and or developing because when set to false, won't actually ping Kafka but will
     #   run all the validations, etc

data/lib/waterdrop/contracts/message.rb

@@ -26,7 +26,10 @@ module WaterDrop
         @max_payload_size = max_payload_size
       end
 
-      required(:topic) { |val| val.is_a?(String) && TOPIC_REGEXP.match?(val) }
+      required(:topic) do |val|
+        (val.is_a?(String) || val.is_a?(Symbol)) && TOPIC_REGEXP.match?(val.to_s)
+      end
+
       required(:payload) { |val| val.nil? || val.is_a?(String) }
       optional(:key) { |val| val.nil? || (val.is_a?(String) && !val.empty?) }
       optional(:partition) { |val| val.is_a?(Integer) && val >= -1 }

data/lib/waterdrop/errors.rb

@@ -32,6 +32,9 @@ module WaterDrop
     # Raised when there is an inline error during single message produce operations
     ProduceError = Class.new(BaseError)
 
+    # Raise it within a transaction to abort it
+    AbortTransaction = Class.new(BaseError)
+
     # Raised when during messages producing something bad happened inline
     class ProduceManyError < ProduceError
       attr_reader :dispatched

data/lib/waterdrop/instrumentation/callbacks/delivery.rb

@@ -17,10 +17,10 @@ module WaterDrop
         # Emits delivery details to the monitor
         # @param delivery_report [Rdkafka::Producer::DeliveryReport] delivery report
         def call(delivery_report)
-          if delivery_report.error.to_i.positive?
-            instrument_error(delivery_report)
-          else
+          if delivery_report.error.to_i.zero?
             instrument_acknowledged(delivery_report)
+          else
+            instrument_error(delivery_report)
           end
         end
 
@@ -36,6 +36,7 @@ module WaterDrop
             offset: delivery_report.offset,
             partition: delivery_report.partition,
             topic: delivery_report.topic_name,
+            delivery_report: delivery_report,
             type: 'librdkafka.dispatch_error'
           )
         end
@@ -47,7 +48,8 @@ module WaterDrop
             producer_id: @producer_id,
             offset: delivery_report.offset,
             partition: delivery_report.partition,
-            topic: delivery_report.topic_name
+            topic: delivery_report.topic_name,
+            delivery_report: delivery_report
           )
         end
       end

data/lib/waterdrop/instrumentation/logger_listener.rb

@@ -112,9 +112,14 @@ module WaterDrop
         debug(event, messages)
       end
 
+      # @param event [Dry::Events::Event] event that happened with the details
+      def on_buffer_purged(event)
+        info(event, 'Successfully purging buffer')
+      end
+
       # @param event [Dry::Events::Event] event that happened with the details
       def on_producer_closed(event)
-        info event, 'Closing producer'
+        info(event, 'Closing producer')
       end
 
       # @param event [Dry::Events::Event] event that happened with the error details
@@ -125,6 +130,21 @@ module WaterDrop
         error(event, "Error occurred: #{error} - #{type}")
       end
 
+      # @param event [Dry::Events::Event] event that happened with the details
+      def on_transaction_started(event)
+        info(event, 'Starting transaction')
+      end
+
+      # @param event [Dry::Events::Event] event that happened with the details
+      def on_transaction_aborted(event)
+        info(event, 'Aborting transaction')
+      end
+
+      # @param event [Dry::Events::Event] event that happened with the details
+      def on_transaction_committed(event)
+        info(event, 'Committing transaction')
+      end
+
       private
 
       # @return [Boolean] should we report the messages details in the debug mode.

data/lib/waterdrop/instrumentation/notifications.rb

@@ -18,8 +18,13 @@ module WaterDrop
         messages.produced_sync
         messages.buffered
 
+        transaction.started
+        transaction.committed
+        transaction.aborted
+
         buffer.flushed_async
         buffer.flushed_sync
+        buffer.purged
 
         statistics.emitted
 

data/lib/waterdrop/producer/async.rb

@@ -60,8 +60,10 @@ module WaterDrop
           producer_id: id,
           messages: messages
         ) do
-          messages.each do |message|
-            dispatched << produce(message)
+          with_transaction_if_transactional do
+            messages.each do |message|
+              dispatched << produce(message)
+            end
           end
 
           dispatched

data/lib/waterdrop/producer/sync.rb

@@ -63,8 +63,10 @@ module WaterDrop
         dispatched = []
 
         @monitor.instrument('messages.produced_sync', producer_id: id, messages: messages) do
-          messages.each do |message|
-            dispatched << produce(message)
+          with_transaction_if_transactional do
+            messages.each do |message|
+              dispatched << produce(message)
+            end
           end
 
           dispatched.map! do |handler|

data/lib/waterdrop/producer/transactions.rb

@@ -0,0 +1,165 @@
+# frozen_string_literal: true
+
+module WaterDrop
+  class Producer
+    # Transactions related producer functionalities
+    module Transactions
+      # Creates a transaction.
+      #
+      # Karafka transactions work in a similar manner to SQL db transactions though there are some
+      # crucial differences. When you start a transaction, all messages produced during it will
+      # be delivered together or will fail together. The difference is, that messages from within
+      # a single transaction can be delivered and will have a delivery handle but will be then
+      # compacted prior to moving the LSO forward. This means, that not every delivery handle for
+      # async dispatches will emit a queue purge error. None for sync as the delivery has happened
+      # but they will never be visible by the transactional consumers.
+      #
+      # Transactions **are** thread-safe however they lock a mutex. This means, that for
+      # high-throughput transactional messages production in multiple threads
+      # (for example in Karafka), it may be much better to use few instances that can work in
+      # parallel.
+      #
+      # Please note, that if a producer is configured as transactional, it **cannot** produce
+      # messages outside of transactions, that is why by default all dispatches will be wrapped
+      # with a transaction. One transaction per single dispatch and for `produce_many` it will be
+      # a single transaction wrapping all messages dispatches (not one per message).
+      #
+      # @return Block result
+      #
+      # @example Simple transaction
+      #   producer.transaction do
+      #     producer.produce_async(topic: 'topic', payload: 'data')
+      #   end
+      #
+      # @example Aborted transaction - messages producer won't be visible by consumers
+      #   producer.transaction do
+      #     producer.produce_sync(topic: 'topic', payload: 'data')
+      #     throw(:abort)
+      #   end
+      #
+      # @example Use block result last handler to wait on all messages ack
+      #   handler = producer.transaction do
+      #               producer.produce_async(topic: 'topic', payload: 'data')
+      #             end
+      #
+      #   handler.wait
+      def transaction
+        # This will safely allow us to support one operation transactions so a transactional
+        # producer can work without the transactional block if needed
+        return yield if @transaction_mutex.owned?
+
+        @transaction_mutex.synchronize do
+          transactional_instrument(:committed) do
+            with_transactional_error_handling(:begin) do
+              transactional_instrument(:started) { client.begin_transaction }
+            end
+
+            result = nil
+            commit = false
+
+            catch(:abort) do
+              result = yield
+              commit = true
+            end
+
+            commit || raise(WaterDrop::Errors::AbortTransaction)
+
+            with_transactional_error_handling(:commit) do
+              client.commit_transaction
+            end
+
+            result
+          rescue StandardError => e
+            with_transactional_error_handling(:abort) do
+              transactional_instrument(:aborted) { client.abort_transaction }
+            end
+
+            raise unless e.is_a?(WaterDrop::Errors::AbortTransaction)
+          end
+        end
+      end
+
+      # @return [Boolean] Is this producer a transactional one
+      def transactional?
+        return @transactional if instance_variable_defined?(:'@transactional')
+
+        @transactional = config.kafka.to_h.key?(:'transactional.id')
+      end
+
+      private
+
+      # Runs provided code with a transaction wrapper if transactions are enabled.
+      # This allows us to simplify the async and sync batch dispatchers because we can ensure that
+      # their internal dispatches will be wrapped only with a single transaction and not
+      # a transaction per message
+      # @param block [Proc] code we want to run
+      def with_transaction_if_transactional(&block)
+        transactional? ? transaction(&block) : yield
+      end
+
+      # Instruments the transactional operation with producer id
+      #
+      # @param key [Symbol] transaction operation key
+      # @param block [Proc] block to run inside the instrumentation or nothing if not given
+      def transactional_instrument(key, &block)
+        @monitor.instrument("transaction.#{key}", producer_id: id, &block)
+      end
+
+      # Error handling for transactional operations is a bit special. There are three types of
+      # errors coming from librdkafka:
+      #   - retryable - indicates that a given operation (like offset commit) can be retried after
+      #     a backoff and that is should be operating later as expected. We try to retry those
+      #     few times before finally failing.
+      #   - fatal - errors that will not recover no matter what (for example being fenced out)
+      #   - abortable - error from which we cannot recover but for which we should abort the
+      #     current transaction.
+      #
+      # The code below handles this logic also publishing the appropriate notifications via our
+      # notifications pipeline.
+      #
+      # @param action [Symbol] action type
+      # @param allow_abortable [Boolean] should we allow for the abortable flow. This is set to
+      #   false internally to prevent attempts to abort from failed abort operations
+      def with_transactional_error_handling(action, allow_abortable: true)
+        attempt ||= 0
+        attempt += 1
+
+        yield
+      rescue ::Rdkafka::RdkafkaError => e
+        # Decide if there is a chance to retry given error
+        do_retry = e.retryable? && attempt < config.max_attempts_on_transaction_command
+
+        @monitor.instrument(
+          'error.occurred',
+          producer_id: id,
+          caller: self,
+          error: e,
+          type: "transaction.#{action}",
+          retry: do_retry,
+          attempt: attempt
+        )
+
+        raise if e.fatal?
+
+        if do_retry
+          # Backoff more and more before retries
+          sleep(config.wait_backoff_on_transaction_command * attempt)
+
+          retry
+        end
+
+        if e.abortable? && allow_abortable
+          # Always attempt to abort but if aborting fails with an abortable error, do not attempt
+          # to abort from abort as this could create an infinite loop
+          with_transactional_error_handling(:abort, allow_abortable: false) do
+            transactional_instrument(:aborted) { @client.abort_transaction }
+          end
+
+          raise
+        end
+
+        raise
+      end
+    end
+  end
+end

data/lib/waterdrop/producer.rb

@@ -7,6 +7,7 @@ module WaterDrop
     include Sync
     include Async
     include Buffer
+    include Transactions
     include ::Karafka::Core::Helpers::Time
 
     # Which of the inline flow errors do we want to intercept and re-bind
@@ -38,6 +39,7 @@ module WaterDrop
       @buffer_mutex = Mutex.new
       @connecting_mutex = Mutex.new
       @operating_mutex = Mutex.new
+      @transaction_mutex = Mutex.new
 
       @status = Status.new
       @messages = Concurrent::Array.new
@@ -117,8 +119,25 @@ module WaterDrop
       @client
     end
 
+    # Purges data from both the buffer queue as well as the librdkafka queue.
+    #
+    # @note This is an operation that can cause data loss. Keep that in mind. It will not only
+    #   purge the internal WaterDrop buffer but will also purge the librdkafka queue as well as
+    #   will cancel any outgoing messages dispatches.
+    def purge
+      @monitor.instrument('buffer.purged', producer_id: id) do
+        @buffer_mutex.synchronize do
+          @messages = Concurrent::Array.new
+        end
+
+        @client.purge
+      end
+    end
+
     # Flushes the buffers in a sync way and closes the producer
-    def close
+    # @param force [Boolean] should we force closing even with outstanding messages after the
+    #   max wait timeout
+    def close(force: false)
       @operating_mutex.synchronize do
         return unless @status.active?
 
@@ -149,7 +168,26 @@ module WaterDrop
           # We also mark it as closed only if it was connected, if not, it would trigger a new
           # connection that anyhow would be immediately closed
           if @client
-            client.close
+            # Why do we trigger it early instead of just having `#close` do it?
+            # The linger.ms time will be ignored for the duration of the call,
+            # queued messages will be sent to the broker as soon as possible.
+            begin
+              # `max_wait_timeout` is in seconds at the moment
+              @client.flush(@config.max_wait_timeout * 1_000) unless @client.closed?
+            # We can safely ignore timeouts here because any left outstanding requests
+            # will anyhow force wait on close if not forced.
+            # If forced, we will purge the queue and just close
+            rescue ::Rdkafka::RdkafkaError, Rdkafka::AbstractHandle::WaitTimeoutError
+              nil
+            ensure
+              # Purge fully the local queue in case of a forceful shutdown just to be sure, that
+              # there are no dangling messages. In case flush was successful, there should be
+              # none but we do it just in case it timed out
+              purge if force
+            end
+
+            @client.close
+
             @client = nil
           end
 
@@ -162,6 +200,11 @@ module WaterDrop
       end
     end
 
+    # Closes the producer with forced close after timeout, purging any outgoing data
+    def close!
+      close(force: true)
+    end
+
     private
 
     # Ensures that we don't run any operations when the producer is not configured or when it
@@ -211,7 +254,15 @@ module WaterDrop
         ensure_active!
       end
 
-      client.produce(**message)
+      # In case someone defines topic as a symbol, we need to convert it into a string as
+      # librdkafka does not accept symbols
+      message = message.merge(topic: message[:topic].to_s) if message[:topic].is_a?(Symbol)
+
+      if transactional?
+        transaction { client.produce(**message) }
+      else
+        client.produce(**message)
+      end
     rescue SUPPORTED_FLOW_ERRORS.first => e
       # Unless we want to wait and retry and it's a full queue, we raise normally
       raise unless @config.wait_on_queue_full

data/lib/waterdrop/version.rb

@@ -3,5 +3,5 @@
 # WaterDrop library
 module WaterDrop
   # Current WaterDrop version
-  VERSION = '2.6.6'
+  VERSION = '2.6.8'
 end

data/waterdrop.gemspec

@@ -16,7 +16,7 @@ Gem::Specification.new do |spec|
   spec.description   = spec.summary
   spec.license       = 'MIT'
 
-  spec.add_dependency 'karafka-core', '>= 2.1.1', '< 3.0.0'
+  spec.add_dependency 'karafka-core', '>= 2.2.3', '< 3.0.0'
   spec.add_dependency 'zeitwerk', '~> 2.3'
 
   if $PROGRAM_NAME.end_with?('gem')
@@ -31,10 +31,10 @@ Gem::Specification.new do |spec|
   spec.metadata = {
     'funding_uri' => 'https://karafka.io/#become-pro',
     'homepage_uri' => 'https://karafka.io',
-    'changelog_uri' => 'https://github.com/karafka/waterdrop/blob/master/CHANGELOG.md',
+    'changelog_uri' => 'https://karafka.io/docs/Changelog-WaterDrop',
     'bug_tracker_uri' => 'https://github.com/karafka/waterdrop/issues',
     'source_code_uri' => 'https://github.com/karafka/waterdrop',
-    'documentation_uri' => 'https://github.com/karafka/waterdrop#readme',
+    'documentation_uri' => 'https://karafka.io/docs/#waterdrop',
     'rubygems_mfa_required' => 'true'
   }
 end