google-cloud-pubsub 2.30.0__py3-none-any.whl

google/cloud/pubsub_v1/publisher/_sequencer/ordered_sequencer.py

@@ -0,0 +1,340 @@
+# Copyright 2019, Google LLC All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import enum
+import collections
+import threading
+import typing
+from typing import Deque, Iterable, Sequence
+
+from google.api_core import gapic_v1
+from google.cloud.pubsub_v1.publisher import futures
+from google.cloud.pubsub_v1.publisher import exceptions
+from google.cloud.pubsub_v1.publisher._sequencer import base as sequencer_base
+from google.cloud.pubsub_v1.publisher._batch import base as batch_base
+from google.cloud.pubsub_v1.open_telemetry.publish_message_wrapper import (
+    PublishMessageWrapper,
+)
+
+if typing.TYPE_CHECKING:  # pragma: NO COVER
+    from google.cloud.pubsub_v1 import types
+    from google.cloud.pubsub_v1.publisher import _batch
+    from google.cloud.pubsub_v1.publisher.client import Client as PublisherClient
+    from google.pubsub_v1.services.publisher.client import OptionalRetry
+
+
+class _OrderedSequencerStatus(str, enum.Enum):
+    """An enum-like class representing valid statuses for an OrderedSequencer.
+
+    Starting state: ACCEPTING_MESSAGES
+    Valid transitions:
+      ACCEPTING_MESSAGES -> PAUSED (on permanent error)
+      ACCEPTING_MESSAGES -> STOPPED  (when user calls stop() explicitly)
+      ACCEPTING_MESSAGES -> FINISHED  (all batch publishes finish normally)
+
+      PAUSED -> ACCEPTING_MESSAGES  (when user unpauses)
+      PAUSED -> STOPPED  (when user calls stop() explicitly)
+
+      STOPPED -> FINISHED (user stops client and the one remaining batch finishes
+                           publish)
+      STOPPED -> PAUSED (stop() commits one batch, which fails permanently)
+
+      FINISHED -> ACCEPTING_MESSAGES (publish happens while waiting for cleanup)
+      FINISHED -> STOPPED (when user calls stop() explicitly)
+    Illegal transitions:
+      PAUSED -> FINISHED (since all batches are cancelled on pause, there should
+                          not be any that finish normally. paused sequencers
+                          should not be cleaned up because their presence
+                          indicates that the ordering key needs to be resumed)
+      STOPPED -> ACCEPTING_MESSAGES (no way to make a user-stopped sequencer
+                                     accept messages again. this is okay since
+                                     stop() should only be called on shutdown.)
+      FINISHED -> PAUSED (no messages remain in flight, so they can't cause a
+                          permanent error and pause the sequencer)
+    """
+
+    # Accepting publishes and/or waiting for result of batch publish
+    ACCEPTING_MESSAGES = "accepting messages"
+    # Permanent error occurred. User must unpause this sequencer to resume
+    # publishing. This is done to maintain ordering.
+    PAUSED = "paused"
+    # No more publishes allowed. There may be an outstanding batch that will
+    # call the _batch_done_callback when it's done (success or error.)
+    STOPPED = "stopped"
+    # No more work to do. Waiting to be cleaned-up. A publish will transform
+    # this sequencer back into the normal accepting-messages state.
+    FINISHED = "finished"
+
+
+class OrderedSequencer(sequencer_base.Sequencer):
+    """Sequences messages into batches ordered by an ordering key for one topic.
+
+    A sequencer always has at least one batch in it, unless paused or stopped.
+    When no batches remain, the |publishes_done_callback| is called so the
+    client can perform cleanup.
+
+    Public methods are thread-safe.
+
+    Args:
+        client:
+            The publisher client used to create this sequencer.
+        topic:
+            The topic. The format for this is ``projects/{project}/topics/{topic}``.
+        ordering_key:
+            The ordering key for this sequencer.
+    """
+
+    def __init__(self, client: "PublisherClient", topic: str, ordering_key: str):
+        self._client = client
+        self._topic = topic
+        self._ordering_key = ordering_key
+        # Guards the variables below
+        self._state_lock = threading.Lock()
+        # Batches ordered from first (head/left) to last (right/tail).
+        # Invariant: always has at least one batch after the first publish,
+        # unless paused or stopped.
+        self._ordered_batches: Deque["_batch.thread.Batch"] = collections.deque()
+        # See _OrderedSequencerStatus for valid state transitions.
+        self._state = _OrderedSequencerStatus.ACCEPTING_MESSAGES
+
+    def is_finished(self) -> bool:
+        """Whether the sequencer is finished and should be cleaned up.
+
+        Returns:
+            Whether the sequencer is finished and should be cleaned up.
+        """
+        with self._state_lock:
+            return self._state == _OrderedSequencerStatus.FINISHED
+
+    def stop(self) -> None:
+        """Permanently stop this sequencer.
+
+        This differs from pausing, which may be resumed. Immediately commits
+        the first batch and cancels the rest.
+
+        Raises:
+            RuntimeError:
+                If called after stop() has already been called.
+        """
+        with self._state_lock:
+            if self._state == _OrderedSequencerStatus.STOPPED:
+                raise RuntimeError("Ordered sequencer already stopped.")
+
+            self._state = _OrderedSequencerStatus.STOPPED
+            if self._ordered_batches:
+                # Give only the first batch the chance to finish.
+                self._ordered_batches[0].commit()
+
+                # Cancel the rest of the batches and remove them from the deque
+                # of batches.
+                while len(self._ordered_batches) > 1:
+                    # Pops from the tail until it leaves only the head in the
+                    # deque.
+                    batch = self._ordered_batches.pop()
+                    batch.cancel(batch_base.BatchCancellationReason.CLIENT_STOPPED)
+
+    def commit(self) -> None:
+        """Commit the first batch, if unpaused.
+
+        If paused or no batches exist, this method does nothing.
+
+        Raises:
+            RuntimeError:
+                If called after stop() has already been called.
+        """
+        with self._state_lock:
+            if self._state == _OrderedSequencerStatus.STOPPED:
+                raise RuntimeError("Ordered sequencer already stopped.")
+
+            if self._state != _OrderedSequencerStatus.PAUSED and self._ordered_batches:
+                # It's okay to commit the same batch more than once. The
+                # operation is idempotent.
+                self._ordered_batches[0].commit()
+
+    def _batch_done_callback(self, success: bool) -> None:
+        """Deal with completion of a batch.
+
+        Called when a batch has finished publishing, with either a success
+        or a failure. (Temporary failures are retried infinitely when
+        ordering keys are enabled.)
+        """
+        ensure_cleanup_and_commit_timer_runs = False
+        with self._state_lock:
+            assert self._state != _OrderedSequencerStatus.PAUSED, (
+                "This method should not be called after pause() because "
+                "pause() should have cancelled all of the batches."
+            )
+            assert self._state != _OrderedSequencerStatus.FINISHED, (
+                "This method should not be called after all batches have been "
+                "finished."
+            )
+
+            # Message futures for the batch have been completed (either with a
+            # result or an exception) already, so remove the batch.
+            self._ordered_batches.popleft()
+
+            if success:
+                if len(self._ordered_batches) == 0:
+                    # Mark this sequencer as finished.
+                    # If new messages come in for this ordering key and this
+                    # sequencer hasn't been cleaned up yet, it will go back
+                    # into accepting-messages state. Otherwise, the client
+                    # must create a new OrderedSequencer.
+                    self._state = _OrderedSequencerStatus.FINISHED
+                    # Ensure cleanup thread runs at some point.
+                    ensure_cleanup_and_commit_timer_runs = True
+                elif len(self._ordered_batches) == 1:
+                    # Wait for messages and/or commit timeout
+                    # Ensure there's actually a commit timer thread that'll commit
+                    # after a delay.
+                    ensure_cleanup_and_commit_timer_runs = True
+                else:
+                    # If there is more than one batch, we know that the next batch
+                    # must be full and, therefore, ready to be committed.
+                    self._ordered_batches[0].commit()
+            else:
+                # Unrecoverable error detected
+                self._pause()
+
+        if ensure_cleanup_and_commit_timer_runs:
+            self._client.ensure_cleanup_and_commit_timer_runs()
+
+    def _pause(self) -> None:
+        """Pause this sequencer: set state to paused, cancel all batches, and
+        clear the list of ordered batches.
+
+        _state_lock must be taken before calling this method.
+        """
+        assert (
+            self._state != _OrderedSequencerStatus.FINISHED
+        ), "Pause should not be called after all batches have finished."
+        self._state = _OrderedSequencerStatus.PAUSED
+        for batch in self._ordered_batches:
+            batch.cancel(
+                batch_base.BatchCancellationReason.PRIOR_ORDERED_MESSAGE_FAILED
+            )
+        self._ordered_batches.clear()
+
+    def unpause(self) -> None:
+        """Unpause this sequencer.
+
+        Raises:
+            RuntimeError:
+                If called when the ordering key has not been paused.
+        """
+        with self._state_lock:
+            if self._state != _OrderedSequencerStatus.PAUSED:
+                raise RuntimeError("Ordering key is not paused.")
+            self._state = _OrderedSequencerStatus.ACCEPTING_MESSAGES
+
+    def _create_batch(
+        self,
+        commit_retry: "OptionalRetry" = gapic_v1.method.DEFAULT,
+        commit_timeout: "types.OptionalTimeout" = gapic_v1.method.DEFAULT,
+    ) -> "_batch.thread.Batch":
+        """Create a new batch using the client's batch class and other stored
+            settings.
+
+        Args:
+            commit_retry:
+                The retry settings to apply when publishing the batch.
+            commit_timeout:
+                The timeout to apply when publishing the batch.
+        """
+        return self._client._batch_class(
+            client=self._client,
+            topic=self._topic,
+            settings=self._client.batch_settings,
+            batch_done_callback=self._batch_done_callback,
+            commit_when_full=False,
+            commit_retry=commit_retry,
+            commit_timeout=commit_timeout,
+        )
+
+    def publish(
+        self,
+        wrapper: PublishMessageWrapper,
+        retry: "OptionalRetry" = gapic_v1.method.DEFAULT,
+        timeout: "types.OptionalTimeout" = gapic_v1.method.DEFAULT,
+    ) -> futures.Future:
+        """Publish message for this ordering key.
+
+        Args:
+            wrapper:
+                The Pub/Sub message wrapper.
+            retry:
+                The retry settings to apply when publishing the message.
+            timeout:
+                The timeout to apply when publishing the message.
+
+        Returns:
+            A class instance that conforms to Python Standard library's
+            :class:`~concurrent.futures.Future` interface (but not an
+            instance of that class). The future might return immediately with a
+            PublishToPausedOrderingKeyException if the ordering key is paused.
+            Otherwise, the future tracks the lifetime of the message publish.
+
+        Raises:
+            RuntimeError:
+                If called after this sequencer has been stopped, either by
+                a call to stop() or after all batches have been published.
+        """
+        with self._state_lock:
+            if self._state == _OrderedSequencerStatus.PAUSED:
+                errored_future = futures.Future()
+                exception = exceptions.PublishToPausedOrderingKeyException(
+                    self._ordering_key
+                )
+                errored_future.set_exception(exception)
+                return errored_future
+
+            # If waiting to be cleaned-up, convert to accepting messages to
+            # prevent this sequencer from being cleaned-up only to have another
+            # one with the same ordering key created immediately afterward.
+            if self._state == _OrderedSequencerStatus.FINISHED:
+                self._state = _OrderedSequencerStatus.ACCEPTING_MESSAGES
+
+            if self._state == _OrderedSequencerStatus.STOPPED:
+                raise RuntimeError("Cannot publish on a stopped sequencer.")
+
+            assert (
+                self._state == _OrderedSequencerStatus.ACCEPTING_MESSAGES
+            ), "Publish is only allowed in accepting-messages state."
+
+            if not self._ordered_batches:
+                new_batch = self._create_batch(
+                    commit_retry=retry, commit_timeout=timeout
+                )
+                self._ordered_batches.append(new_batch)
+
+            batch = self._ordered_batches[-1]
+            future = batch.publish(wrapper)
+            while future is None:
+                batch = self._create_batch(commit_retry=retry, commit_timeout=timeout)
+                self._ordered_batches.append(batch)
+                future = batch.publish(wrapper)
+
+            return future
+
+    # Used only for testing.
+    def _set_batch(self, batch: "_batch.thread.Batch") -> None:
+        self._ordered_batches = collections.deque([batch])
+
+    # Used only for testing.
+    def _set_batches(self, batches: Iterable["_batch.thread.Batch"]) -> None:
+        self._ordered_batches = collections.deque(batches)
+
+    # Used only for testing.
+    def _get_batches(self) -> Sequence["_batch.thread.Batch"]:
+        return self._ordered_batches

google/cloud/pubsub_v1/publisher/_sequencer/unordered_sequencer.py

@@ -0,0 +1,167 @@
+# Copyright 2019, Google LLC All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import typing
+from typing import Optional
+
+from google.api_core import gapic_v1
+
+from google.cloud.pubsub_v1.publisher._sequencer import base
+from google.cloud.pubsub_v1.open_telemetry.publish_message_wrapper import (
+    PublishMessageWrapper,
+)
+
+if typing.TYPE_CHECKING:  # pragma: NO COVER
+    from google.cloud.pubsub_v1.publisher import _batch
+    from google.cloud.pubsub_v1.publisher import futures
+    from google.cloud.pubsub_v1.publisher.client import Client as PublisherClient
+    from google.pubsub_v1.services.publisher.client import OptionalRetry
+
+    from google.cloud.pubsub_v1 import types
+
+
+class UnorderedSequencer(base.Sequencer):
+    """Sequences messages into batches for one topic without any ordering.
+
+    Public methods are NOT thread-safe.
+    """
+
+    def __init__(self, client: "PublisherClient", topic: str):
+        self._client = client
+        self._topic = topic
+        self._current_batch: Optional["_batch.thread.Batch"] = None
+        self._stopped = False
+
+    def is_finished(self) -> bool:
+        """Whether the sequencer is finished and should be cleaned up.
+
+        Returns:
+            Whether the sequencer is finished and should be cleaned up.
+        """
+        # TODO: Implement. Not implementing yet because of possible performance
+        # impact due to extra locking required. This does mean that
+        # UnorderedSequencers don't get cleaned up, but this is the same as
+        # previously existing behavior.
+        return False
+
+    def stop(self) -> None:
+        """Stop the sequencer.
+
+        Subsequent publishes will fail.
+
+        Raises:
+            RuntimeError:
+                If called after stop() has already been called.
+        """
+        if self._stopped:
+            raise RuntimeError("Unordered sequencer already stopped.")
+        self.commit()
+        self._stopped = True
+
+    def commit(self) -> None:
+        """Commit the batch.
+
+        Raises:
+            RuntimeError:
+                If called after stop() has already been called.
+        """
+        if self._stopped:
+            raise RuntimeError("Unordered sequencer already stopped.")
+        if self._current_batch:
+            self._current_batch.commit()
+
+            # At this point, we lose track of the old batch, but we don't
+            # care since we just committed it.
+            # Setting this to None guarantees the next publish() creates a new
+            # batch.
+            self._current_batch = None
+
+    def unpause(self) -> typing.NoReturn:
+        """Not relevant for this class."""
+        raise NotImplementedError
+
+    def _create_batch(
+        self,
+        commit_retry: "OptionalRetry" = gapic_v1.method.DEFAULT,
+        commit_timeout: "types.OptionalTimeout" = gapic_v1.method.DEFAULT,
+    ) -> "_batch.thread.Batch":
+        """Create a new batch using the client's batch class and other stored
+            settings.
+
+        Args:
+            commit_retry:
+                The retry settings to apply when publishing the batch.
+            commit_timeout:
+                The timeout to apply when publishing the batch.
+        """
+        return self._client._batch_class(
+            client=self._client,
+            topic=self._topic,
+            settings=self._client.batch_settings,
+            batch_done_callback=None,
+            commit_when_full=True,
+            commit_retry=commit_retry,
+            commit_timeout=commit_timeout,
+        )
+
+    def publish(
+        self,
+        wrapper: PublishMessageWrapper,
+        retry: "OptionalRetry" = gapic_v1.method.DEFAULT,
+        timeout: "types.OptionalTimeout" = gapic_v1.method.DEFAULT,
+    ) -> "futures.Future":
+        """Batch message into existing or new batch.
+
+        Args:
+            wrapper:
+                The Pub/Sub message wrapper.
+            retry:
+                The retry settings to apply when publishing the message.
+            timeout:
+                The timeout to apply when publishing the message.
+
+        Returns:
+            An object conforming to the :class:`~concurrent.futures.Future` interface.
+            The future tracks the publishing status of the message.
+
+        Raises:
+            RuntimeError:
+                If called after stop() has already been called.
+
+            pubsub_v1.publisher.exceptions.MessageTooLargeError: If publishing
+                the ``message`` would exceed the max size limit on the backend.
+        """
+        if self._stopped:
+            raise RuntimeError("Unordered sequencer already stopped.")
+
+        if not self._current_batch:
+            newbatch = self._create_batch(commit_retry=retry, commit_timeout=timeout)
+            self._current_batch = newbatch
+
+        batch = self._current_batch
+        future = None
+        while future is None:
+            # Might throw MessageTooLargeError
+            future = batch.publish(wrapper)
+            # batch is full, triggering commit_when_full
+            if future is None:
+                batch = self._create_batch(commit_retry=retry, commit_timeout=timeout)
+                # At this point, we lose track of the old batch, but we don't
+                # care since it's already committed (because it was full.)
+                self._current_batch = batch
+        return future
+
+    # Used only for testing.
+    def _set_batch(self, batch: "_batch.thread.Batch") -> None:
+        self._current_batch = batch