google-cloud-pubsub 0.33.1 → 2.15.1

data/lib/google/cloud/pubsub/async_publisher.rb

@@ -15,20 +15,24 @@
 
 require "monitor"
 require "concurrent"
+require "google/cloud/pubsub/errors"
+require "google/cloud/pubsub/flow_controller"
+require "google/cloud/pubsub/async_publisher/batch"
 require "google/cloud/pubsub/publish_result"
 require "google/cloud/pubsub/service"
+require "google/cloud/pubsub/convert"
 
 module Google
   module Cloud
-    module Pubsub
+    module PubSub
       ##
       # Used to publish multiple messages in batches to a topic. See
-      # {Google::Cloud::Pubsub::Topic#async_publisher}
+      # {Google::Cloud::PubSub::Topic#async_publisher}
       #
       # @example
       #   require "google/cloud/pubsub"
       #
-      #   pubsub = Google::Cloud::Pubsub.new
+      #   pubsub = Google::Cloud::PubSub.new
       #
       #   topic = pubsub.topic "my-topic"
       #   topic.publish_async "task completed" do |result|
@@ -39,79 +43,128 @@ module Google
       #     end
       #   end
       #
-      #   topic.async_publisher.stop.wait!
+      #   topic.async_publisher.stop!
       #
-      # @attr_reader [String] topic_name The name of the topic the messages
-      #   are published to. In the form of
-      #   "/projects/project-identifier/topics/topic-name".
-      # @attr_reader [Integer] max_bytes The maximum size of messages to be
-      #   collected before the batch is published. Default is 10,000,000
-      #   (10MB).
-      # @attr_reader [Integer] max_messages The maximum number of messages to
-      #   be collected before the batch is published. Default is 1,000.
-      # @attr_reader [Numeric] interval The number of seconds to collect
-      #   messages before the batch is published. Default is 0.25.
-      # @attr_reader [Numeric] publish_threads The number of threads used to
-      #   publish messages. Default is 4.
-      # @attr_reader [Numeric] callback_threads The number of threads to
-      #   handle the published messages' callbacks. Default is 8.
+      # @attr_reader [String] topic_name The name of the topic the messages are published to. The value is a
+      #   fully-qualified topic name in the form `projects/{project_id}/topics/{topic_id}`.
+      # @attr_reader [Integer] max_bytes The maximum size of messages to be collected before the batch is published.
+      #   Default is 1,000,000 (1MB).
+      # @attr_reader [Integer] max_messages The maximum number of messages to be collected before the batch is
+      #   published. Default is 100.
+      # @attr_reader [Numeric] interval The number of seconds to collect messages before the batch is published. Default
+      #   is 0.01.
+      # @attr_reader [Numeric] publish_threads The number of threads used to publish messages. Default is 2.
+      # @attr_reader [Numeric] callback_threads The number of threads to handle the published messages' callbacks.
+      #   Default is 4.
       #
       class AsyncPublisher
         include MonitorMixin
 
-        attr_reader :topic_name, :max_bytes, :max_messages, :interval,
-                    :publish_threads, :callback_threads
+        attr_reader :topic_name
+        attr_reader :max_bytes
+        attr_reader :max_messages
+        attr_reader :interval
+        attr_reader :publish_threads
+        attr_reader :callback_threads
+        attr_reader :flow_control
         ##
         # @private Implementation accessors
         attr_reader :service, :batch, :publish_thread_pool,
-                    :callback_thread_pool
+                    :callback_thread_pool, :flow_controller,
+                    :compress, :compression_bytes_threshold
 
         ##
         # @private Create a new instance of the object.
-        def initialize topic_name, service, max_bytes: 10000000,
-                       max_messages: 1000, interval: 0.25, threads: {}
+        def initialize topic_name,
+                       service,
+                       max_bytes: 1_000_000,
+                       max_messages: 100,
+                       interval: 0.01,
+                       threads: {},
+                       flow_control: {},
+                       compress: nil,
+                       compression_bytes_threshold: nil
+          # init MonitorMixin
+          super()
           @topic_name = service.topic_path topic_name
           @service    = service
 
           @max_bytes        = max_bytes
           @max_messages     = max_messages
           @interval         = interval
-          @publish_threads  = (threads[:publish] || 4).to_i
-          @callback_threads = (threads[:callback] || 8).to_i
-
+          @publish_threads  = (threads[:publish] || 2).to_i
+          @callback_threads = (threads[:callback] || 4).to_i
+          @flow_control = {
+            message_limit: 10 * @max_messages,
+            byte_limit: 10 * @max_bytes
+          }.merge(flow_control).freeze
+
+          @published_at = nil
+          @publish_thread_pool = Concurrent::ThreadPoolExecutor.new max_threads: @publish_threads
+          @callback_thread_pool = Concurrent::ThreadPoolExecutor.new max_threads: @callback_threads
+
+          @ordered = false
+          @batches = {}
           @cond = new_cond
-
-          # init MonitorMixin
-          super()
+          @flow_controller = FlowController.new(**@flow_control)
+          @thread = Thread.new { run_background }
+          @compress = compress || Google::Cloud::PubSub::DEFAULT_COMPRESS
+          @compression_bytes_threshold = compression_bytes_threshold ||
+                                         Google::Cloud::PubSub::DEFAULT_COMPRESSION_BYTES_THRESHOLD
         end
 
         ##
         # Add a message to the async publisher to be published to the topic.
         # Messages will be collected in batches and published together.
-        # See {Google::Cloud::Pubsub::Topic#publish_async}
-        def publish data = nil, attributes = {}, &block
-          msg = create_pubsub_message data, attributes
+        # See {Google::Cloud::PubSub::Topic#publish_async}
+        #
+        # @param [String, File] data The message payload. This will be converted
+        #   to bytes encoded as ASCII-8BIT.
+        # @param [Hash] attributes Optional attributes for the message.
+        # @param [String] ordering_key Identifies related messages for which
+        #   publish order should be respected.
+        # @yield [result] the callback for when the message has been published
+        # @yieldparam [PublishResult] result the result of the asynchronous
+        #   publish
+        # @raise [Google::Cloud::PubSub::AsyncPublisherStopped] when the
+        #   publisher is stopped. (See {#stop} and {#stopped?}.)
+        # @raise [Google::Cloud::PubSub::OrderedMessagesDisabled] when
+        #   publishing a message with an `ordering_key` but ordered messages are
+        #   not enabled. (See {#message_ordering?} and
+        #   {#enable_message_ordering!}.)
+        # @raise [Google::Cloud::PubSub::OrderingKeyError] when publishing a
+        #   message with an `ordering_key` that has already failed when
+        #   publishing. Use {#resume_publish} to allow this `ordering_key` to be
+        #   published again.
+        #
+        def publish data = nil, attributes = nil, ordering_key: nil, **extra_attrs, &callback
+          msg = Convert.pubsub_message data, attributes, ordering_key, extra_attrs
+          begin
+            @flow_controller.acquire msg.to_proto.bytesize
+          rescue FlowControlLimitError => e
+            stop_publish ordering_key, e if ordering_key
+            raise
+          end
 
           synchronize do
-            raise "Can't publish when stopped." if @stopped
-
-            if @batch.nil?
-              @batch ||= Batch.new self
-              @batch.add msg, block
-            else
-              unless @batch.try_add msg, block
-                publish_batch!
-                @batch = Batch.new self
-                @batch.add msg, block
-              end
-            end
-
-            init_resources!
-
-            publish_batch! if @batch.ready?
+            raise AsyncPublisherStopped if @stopped
+            raise OrderedMessagesDisabled if !@ordered && !msg.ordering_key.empty? # default is empty string
 
+            batch = resolve_batch_for_message msg
+            if batch.canceled?
+              @flow_controller.release msg.to_proto.bytesize
+              raise OrderingKeyError, batch.ordering_key
+            end
+            batch_action = batch.add msg, callback
+            if batch_action == :full
+              publish_batches!
+            elsif @published_at.nil?
+              # Set initial time to now to start the background counter
+              @published_at = Time.now
+            end
             @cond.signal
           end
+
           nil
         end
 
@@ -127,36 +180,53 @@ module Google
             break if @stopped
 
             @stopped = true
-            publish_batch!
+            publish_batches! stop: true
             @cond.signal
-            @publish_thread_pool.shutdown if @publish_thread_pool
+            @publish_thread_pool.shutdown
           end
 
           self
         end
 
         ##
-        # Blocks until the publisher is fully stopped, all pending messages
-        # have been published, and all callbacks have completed. Does not stop
-        # the publisher. To stop the publisher, first call {#stop} and then
-        # call {#wait!} to block until the publisher is stopped.
+        # Blocks until the publisher is fully stopped, all pending messages have
+        # been published, and all callbacks have completed, or until `timeout`
+        # seconds have passed.
+        #
+        # Does not stop the publisher. To stop the publisher, first call {#stop}
+        # and then call {#wait!} to block until the publisher is stopped
+        #
+        # @param [Number, nil] timeout The number of seconds to block until the
+        #   publisher is fully stopped. Default will block indefinitely.
         #
         # @return [AsyncPublisher] returns self so calls can be chained.
         def wait! timeout = nil
           synchronize do
-            if @publish_thread_pool
-              @publish_thread_pool.wait_for_termination timeout
-            end
+            @publish_thread_pool.wait_for_termination timeout
 
-            if @callback_thread_pool
-              @callback_thread_pool.shutdown
-              @callback_thread_pool.wait_for_termination timeout
-            end
+            @callback_thread_pool.shutdown
+            @callback_thread_pool.wait_for_termination timeout
           end
 
           self
         end
 
+        ##
+        # Stop this publisher and block until the publisher is fully stopped,
+        # all pending messages have been published, and all callbacks have
+        # completed, or until `timeout` seconds have passed.
+        #
+        # The same as calling {#stop} and {#wait!}.
+        #
+        # @param [Number, nil] timeout The number of seconds to block until the
+        #   publisher is fully stopped. Default will block indefinitely.
+        #
+        # @return [AsyncPublisher] returns self so calls can be chained.
+        def stop! timeout = nil
+          stop
+          wait! timeout
+        end
+
         ##
         # Forces all messages in the current batch to be published
         # immediately.
@@ -164,7 +234,7 @@ module Google
         # @return [AsyncPublisher] returns self so calls can be chained.
         def flush
           synchronize do
-            publish_batch!
+            publish_batches!
             @cond.signal
           end
 
@@ -187,151 +257,190 @@ module Google
           synchronize { @stopped }
         end
 
-        protected
+        ##
+        # Enables message ordering for messages with ordering keys. When
+        # enabled, messages published with the same `ordering_key` will be
+        # delivered in the order they were published.
+        #
+        # See {#message_ordering?}. See {Topic#publish_async},
+        # {Subscription#listen}, and {Message#ordering_key}.
+        #
+        def enable_message_ordering!
+          synchronize { @ordered = true }
+        end
+
+        ##
+        # Whether message ordering for messages with ordering keys has been
+        # enabled. When enabled, messages published with the same `ordering_key`
+        # will be delivered in the order they were published. When disabled,
+        # messages may be delivered in any order.
+        #
+        # See {#enable_message_ordering!}. See {Topic#publish_async},
+        # {Subscription#listen}, and {Message#ordering_key}.
+        #
+        # @return [Boolean]
+        #
+        def message_ordering?
+          synchronize { @ordered }
+        end
 
-        def init_resources!
-          @first_published_at   ||= Time.now
-          @publish_thread_pool  ||= Concurrent::FixedThreadPool.new \
-            @publish_threads
-          @callback_thread_pool ||= Concurrent::FixedThreadPool.new \
-            @callback_threads
-          @thread ||= Thread.new { run_background }
+        ##
+        # Resume publishing ordered messages for the provided ordering key.
+        #
+        # @param [String] ordering_key Identifies related messages for which
+        #   publish order should be respected.
+        #
+        # @return [boolean] `true` when resumed, `false` otherwise.
+        #
+        def resume_publish ordering_key
+          synchronize do
+            batch = resolve_batch_for_ordering_key ordering_key
+            return if batch.nil?
+            batch.resume!
+          end
         end
 
+        protected
+
         def run_background
           synchronize do
             until @stopped
-              if @batch.nil?
+              if @published_at.nil?
                 @cond.wait
                 next
               end
 
-              time_since_first_publish = Time.now - @first_published_at
+              time_since_first_publish = Time.now - @published_at
               if time_since_first_publish > @interval
-                # interval met, publish the batch...
-                publish_batch!
+                # interval met, flush the batches...
+                publish_batches!
                 @cond.wait
               else
                 # still waiting for the interval to publish the batch...
-                @cond.wait(@interval - time_since_first_publish)
+                timeout = @interval - time_since_first_publish
+                @cond.wait timeout
               end
             end
           end
         end
 
-        def publish_batch!
-          return unless @batch
+        def resolve_batch_for_message msg
+          @batches[msg.ordering_key] ||= Batch.new self, msg.ordering_key
+        end
+
+        def resolve_batch_for_ordering_key ordering_key
+          @batches[ordering_key]
+        end
 
-          publish_batch_async @topic_name, @batch
-          @batch = nil
-          @first_published_at = nil
+        def stop_publish ordering_key, err
+          synchronize do
+            batch = resolve_batch_for_ordering_key ordering_key
+            return if batch.nil?
+            items = batch.cancel!
+            items.each do |item|
+              @flow_controller.release item.bytesize
+              next unless item.callback
+
+              publish_result = PublishResult.from_error item.msg, err
+              execute_callback_async item.callback, publish_result
+            end
+          end
+        end
+
+        def publish_batches! stop: nil
+          @batches.reject! { |_ordering_key, batch| batch.empty? }
+          @batches.each_value do |batch|
+            ready = batch.publish! stop: stop
+            publish_batch_async @topic_name, batch if ready
+          end
+          # Set published_at to nil to wait indefinitely
+          @published_at = nil
         end
 
         def publish_batch_async topic_name, batch
+          # TODO: raise unless @publish_thread_pool.running?
           return unless @publish_thread_pool.running?
 
-          Concurrent::Future.new(executor: @publish_thread_pool) do
-            begin
-              grpc = @service.publish topic_name, batch.messages
-              batch.items.zip(Array(grpc.message_ids)) do |item, id|
-                next unless item.callback
+          Concurrent::Promises.future_on(
+            @publish_thread_pool, topic_name, batch
+          ) { |t, b| publish_batch_sync t, b }
+        end
 
-                item.msg.message_id = id
-                publish_result = PublishResult.from_grpc(item.msg)
-                execute_callback_async item.callback, publish_result
-              end
-            rescue StandardError => e
-              batch.items.each do |item|
+        # rubocop:disable Metrics/AbcSize
+
+        def publish_batch_sync topic_name, batch
+          # The only batch methods that are safe to call from the loop are
+          # rebalance! and reset! because they are the only methods that are
+          # synchronized.
+          loop do
+            items = batch.rebalance!
+
+            unless items.empty?
+              grpc = @service.publish topic_name,
+                                      items.map(&:msg),
+                                      compress: compress && batch.total_message_bytes >= compression_bytes_threshold
+              items.zip Array(grpc.message_ids) do |item, id|
+                @flow_controller.release item.bytesize
                 next unless item.callback
 
-                publish_result = PublishResult.from_error(item.msg, e)
+                item.msg.message_id = id
+                publish_result = PublishResult.from_grpc item.msg
                 execute_callback_async item.callback, publish_result
               end
             end
-          end.execute
-        end
-
-        def execute_callback_async callback, publish_result
-          return unless @callback_thread_pool.running?
 
-          Concurrent::Future.new(executor: @callback_thread_pool) do
-            callback.call publish_result
-          end.execute
-        end
-
-        def create_pubsub_message data, attributes
-          attributes ||= {}
-          if data.is_a?(::Hash) && attributes.empty?
-            attributes = data
-            data = nil
+            break unless batch.reset!
           end
-          # Convert IO-ish objects to strings
-          if data.respond_to?(:read) && data.respond_to?(:rewind)
-            data.rewind
-            data = data.read
+        rescue StandardError => e
+          items = batch.items
+
+          unless batch.ordering_key.empty?
+            retry if publish_batch_error_retryable? e
+            # Cancel the batch if the error is not to be retried.
+            begin
+              raise OrderingKeyError, batch.ordering_key
+            rescue OrderingKeyError => e
+              # The existing e variable is not set to OrderingKeyError
+              # Get all unsent messages for the callback
+              items = batch.cancel!
+            end
           end
-          # Convert data to encoded byte array to match the protobuf defn
-          data_bytes = \
-            String(data).dup.force_encoding(Encoding::ASCII_8BIT).freeze
 
-          # Convert attributes to strings to match the protobuf definition
-          attributes = Hash[attributes.map { |k, v| [String(k), String(v)] }]
+          items.each do |item|
+            @flow_controller.release item.bytesize
+            next unless item.callback
 
-          Google::Pubsub::V1::PubsubMessage.new data: data_bytes,
-                                                attributes: attributes
-        end
-
-        ##
-        # @private
-        class Batch
-          attr_reader :messages, :callbacks
-
-          def initialize publisher
-            @publisher = publisher
-            @messages = []
-            @callbacks = []
-            @total_message_bytes = publisher.topic_name.bytesize + 2
+            publish_result = PublishResult.from_error item.msg, e
+            execute_callback_async item.callback, publish_result
           end
 
-          def add msg, callback
-            @messages << msg
-            @callbacks << callback
-            @total_message_bytes += msg.to_proto.bytesize + 2
-          end
+          # publish will retry indefinitely, as long as there are unsent items.
+          retry if batch.reset!
+        end
 
-          def try_add msg, callback
-            new_message_count = total_message_count + 1
-            new_message_bytes = total_message_bytes + msg.to_proto.bytesize + 2
-            if new_message_count > @publisher.max_messages ||
-               new_message_bytes >= @publisher.max_bytes
-              return false
-            end
-            add msg, callback
-            true
-          end
+        # rubocop:enable Metrics/AbcSize
 
-          def ready?
-            total_message_count >= @publisher.max_messages ||
-              total_message_bytes >= @publisher.max_bytes
-          end
+        PUBLISH_RETRY_ERRORS = [
+          GRPC::Cancelled, GRPC::DeadlineExceeded, GRPC::Internal,
+          GRPC::ResourceExhausted, GRPC::Unauthenticated, GRPC::Unavailable
+        ].freeze
 
-          def total_message_count
-            @messages.count
-          end
+        def publish_batch_error_retryable? error
+          PUBLISH_RETRY_ERRORS.any? { |klass| error.is_a? klass }
+        end
 
-          def total_message_bytes
-            @total_message_bytes
-          end
+        def execute_callback_async callback, publish_result
+          return unless @callback_thread_pool.running?
 
-          def items
-            @messages.zip(@callbacks).map do |msg, callback|
-              Item.new msg, callback
-            end
+          Concurrent::Promises.future_on(
+            @callback_thread_pool, callback, publish_result
+          ) do |cback, p_result|
+            cback.call p_result
           end
-
-          Item = Struct.new :msg, :callback
         end
       end
     end
+
+    Pubsub = PubSub unless const_defined? :Pubsub
   end
 end