google/cloud/pubsub_v1/publisher/client.py

# 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.

from __future__ import absolute_import

import copy
import logging
import os
import threading
import time
import typing
from typing import Any, Dict, Optional, Sequence, Tuple, Type, Union
import warnings

from google.api_core import gapic_v1
from google.auth.credentials import AnonymousCredentials  # type: ignore
from google.oauth2 import service_account  # type: ignore

from google.cloud.pubsub_v1 import types
from google.cloud.pubsub_v1.publisher import exceptions
from google.cloud.pubsub_v1.publisher import futures
from google.cloud.pubsub_v1.publisher._batch import thread
from google.cloud.pubsub_v1.publisher._sequencer import ordered_sequencer
from google.cloud.pubsub_v1.publisher._sequencer import unordered_sequencer
from google.cloud.pubsub_v1.publisher.flow_controller import FlowController
from google.pubsub_v1 import gapic_version as package_version
from google.pubsub_v1 import types as gapic_types
from google.pubsub_v1.services.publisher import client as publisher_client

__version__ = package_version.__version__

if typing.TYPE_CHECKING:  # pragma: NO COVER
    from google.cloud import pubsub_v1
    from google.cloud.pubsub_v1.publisher import _batch
    from google.pubsub_v1.services.publisher.client import OptionalRetry
    from google.pubsub_v1.types import pubsub as pubsub_types


_LOGGER = logging.getLogger(__name__)


_raw_proto_pubbsub_message = gapic_types.PubsubMessage.pb()

SequencerType = Union[
    ordered_sequencer.OrderedSequencer, unordered_sequencer.UnorderedSequencer
]


class Client(publisher_client.PublisherClient):
    """A publisher client for Google Cloud Pub/Sub.

    This creates an object that is capable of publishing messages.
    Generally, you can instantiate this client with no arguments, and you
    get sensible defaults.

    Args:
        batch_settings:
            The settings for batch publishing.
        publisher_options:
            The options for the publisher client. Note that enabling message ordering
            will override the publish retry timeout to be infinite.
        kwargs:
            Any additional arguments provided are sent as keyword arguments to the
            underlying
            :class:`~google.cloud.pubsub_v1.gapic.publisher_client.PublisherClient`.
            Generally you should not need to set additional keyword
            arguments. Regional endpoints can be set via ``client_options`` that
            takes a single key-value pair that defines the endpoint.

    Example:

    .. code-block:: python

        from google.cloud import pubsub_v1

        publisher_client = pubsub_v1.PublisherClient(
            # Optional
            batch_settings = pubsub_v1.types.BatchSettings(
                max_bytes=1024,  # One kilobyte
                max_latency=1,   # One second
            ),

            # Optional
            publisher_options = pubsub_v1.types.PublisherOptions(
                enable_message_ordering=False,
                flow_control=pubsub_v1.types.PublishFlowControl(
                    message_limit=2000,
                    limit_exceeded_behavior=pubsub_v1.types.LimitExceededBehavior.BLOCK,
                ),
            ),

            # Optional
            client_options = {
                "api_endpoint": REGIONAL_ENDPOINT
            }
        )
    """

    def __init__(
        self,
        batch_settings: Union[types.BatchSettings, Sequence] = (),
        publisher_options: Union[types.PublisherOptions, Sequence] = (),
        **kwargs: Any,
    ):
        assert (
            type(batch_settings) is types.BatchSettings or len(batch_settings) == 0
        ), "batch_settings must be of type BatchSettings or an empty sequence."
        assert (
            type(publisher_options) is types.PublisherOptions
            or len(publisher_options) == 0
        ), "publisher_options must be of type PublisherOptions or an empty sequence."

        # Sanity check: Is our goal to use the emulator?
        # If so, create a grpc insecure channel with the emulator host
        # as the target.
        if os.environ.get("PUBSUB_EMULATOR_HOST"):
            kwargs["client_options"] = {
                "api_endpoint": os.environ.get("PUBSUB_EMULATOR_HOST")
            }
            kwargs["credentials"] = AnonymousCredentials()

        # For a transient failure, retry publishing the message infinitely.
        self.publisher_options = types.PublisherOptions(*publisher_options)
        self._enable_message_ordering = self.publisher_options[0]

        # Add the metrics headers, and instantiate the underlying GAPIC
        # client.
        super().__init__(**kwargs)
        self._target = self._transport._host
        self._batch_class = thread.Batch
        self.batch_settings = types.BatchSettings(*batch_settings)

        # The batches on the publisher client are responsible for holding
        # messages. One batch exists for each topic.
        self._batch_lock = self._batch_class.make_lock()
        # (topic, ordering_key) => sequencers object
        self._sequencers: Dict[Tuple[str, str], SequencerType] = {}
        self._is_stopped = False
        # Thread created to commit all sequencers after a timeout.
        self._commit_thread: Optional[threading.Thread] = None

        # The object controlling the message publishing flow
        self._flow_controller = FlowController(self.publisher_options.flow_control)

    @classmethod
    def from_service_account_file(  # type: ignore[override]
        cls,
        filename: str,
        batch_settings: Union[types.BatchSettings, Sequence] = (),
        **kwargs: Any,
    ) -> "Client":
        """Creates an instance of this client using the provided credentials
        file.

        Args:
            filename:
                The path to the service account private key JSON file.
            batch_settings:
                The settings for batch publishing.
            kwargs:
                Additional arguments to pass to the constructor.

        Returns:
            A Publisher instance that is the constructed client.
        """
        credentials = service_account.Credentials.from_service_account_file(filename)
        kwargs["credentials"] = credentials
        return cls(batch_settings, **kwargs)

    from_service_account_json = from_service_account_file  # type: ignore[assignment]

    @property
    def target(self) -> str:
        """Return the target (where the API is).

        Returns:
            The location of the API.
        """
        return self._target

    @property
    def api(self):
        """The underlying gapic API client.

        .. versionchanged:: 2.10.0
            Instead of a GAPIC ``PublisherClient`` client instance, this property is a
            proxy object to it with the same interface.

        .. deprecated:: 2.10.0
            Use the GAPIC methods and properties on the client instance directly
            instead of through the :attr:`api` attribute.
        """
        msg = (
            'The "api" property only exists for backward compatibility, access its '
            'attributes directly thorugh the client instance (e.g. "client.foo" '
            'instead of "client.api.foo").'
        )
        warnings.warn(msg, category=DeprecationWarning)
        return super()

    def _get_or_create_sequencer(self, topic: str, ordering_key: str) -> SequencerType:
        """Get an existing sequencer or create a new one given the (topic,
        ordering_key) pair.
        """
        sequencer_key = (topic, ordering_key)
        sequencer = self._sequencers.get(sequencer_key)
        if sequencer is None:
            if ordering_key == "":
                sequencer = unordered_sequencer.UnorderedSequencer(self, topic)
            else:
                sequencer = ordered_sequencer.OrderedSequencer(
                    self, topic, ordering_key
                )
            self._sequencers[sequencer_key] = sequencer

        return sequencer

    def resume_publish(self, topic: str, ordering_key: str) -> None:
        """Resume publish on an ordering key that has had unrecoverable errors.

        Args:
            topic: The topic to publish messages to.
            ordering_key: A string that identifies related messages for which
                publish order should be respected.

        Raises:
            RuntimeError:
                If called after publisher has been stopped by a `stop()` method
                call.
            ValueError:
                If the topic/ordering key combination has not been seen before
                by this client.
        """
        with self._batch_lock:
            if self._is_stopped:
                raise RuntimeError("Cannot resume publish on a stopped publisher.")

            if not self._enable_message_ordering:
                raise ValueError(
                    "Cannot resume publish on a topic/ordering key if ordering "
                    "is not enabled."
                )

            sequencer_key = (topic, ordering_key)
            sequencer = self._sequencers.get(sequencer_key)
            if sequencer is None:
                _LOGGER.debug(
                    "Error: The topic/ordering key combination has not "
                    "been seen before."
                )
            else:
                sequencer.unpause()

    def _gapic_publish(self, *args, **kwargs) -> "pubsub_types.PublishResponse":
        """Call the GAPIC public API directly."""
        return super().publish(*args, **kwargs)

    def publish(  # type: ignore[override]
        self,
        topic: str,
        data: bytes,
        ordering_key: str = "",
        retry: "OptionalRetry" = gapic_v1.method.DEFAULT,
        timeout: "types.OptionalTimeout" = gapic_v1.method.DEFAULT,
        **attrs: Union[bytes, str],
    ) -> "pubsub_v1.publisher.futures.Future":
        """Publish a single message.

        .. note::
            Messages in Pub/Sub are blobs of bytes. They are *binary* data,
            not text. You must send data as a bytestring
            (``bytes`` in Python 3; ``str`` in Python 2), and this library
            will raise an exception if you send a text string.

            The reason that this is so important (and why we do not try to
            coerce for you) is because Pub/Sub is also platform independent
            and there is no way to know how to decode messages properly on
            the other side; therefore, encoding and decoding is a required
            exercise for the developer.

        Add the given message to this object; this will cause it to be
        published once the batch either has enough messages or a sufficient
        period of time has elapsed.
        This method may block if LimitExceededBehavior.BLOCK is used in the
        flow control settings.

        Example:
            >>> from google.cloud import pubsub_v1
            >>> client = pubsub_v1.PublisherClient()
            >>> topic = client.topic_path('[PROJECT]', '[TOPIC]')
            >>> data = b'The rain in Wales falls mainly on the snails.'
            >>> response = client.publish(topic, data, username='guido')

        Args:
            topic: The topic to publish messages to.
            data: A bytestring representing the message body. This
                must be a bytestring.
            ordering_key: A string that identifies related messages for which
                publish order should be respected. Message ordering must be
                enabled for this client to use this feature.
            retry:
                Designation of what errors, if any, should be retried. If `ordering_key`
                is specified, the total retry deadline will be changed to "infinity".
                If given, it overides any retry passed into the client through
                the ``publisher_options`` argument.
            timeout:
                The timeout for the RPC request. Can be used to override any timeout
                passed in through ``publisher_options`` when instantiating the client.

            attrs: A dictionary of attributes to be
                sent as metadata. (These may be text strings or byte strings.)

        Returns:
            A :class:`~google.cloud.pubsub_v1.publisher.futures.Future`
            instance that conforms to Python Standard library's
            :class:`~concurrent.futures.Future` interface (but not an
            instance of that class).

        Raises:
            RuntimeError:
                If called after publisher has been stopped by a `stop()` method
                call.

            pubsub_v1.publisher.exceptions.MessageTooLargeError: If publishing
                the ``message`` would exceed the max size limit on the backend.
        """
        # Sanity check: Is the data being sent as a bytestring?
        # If it is literally anything else, complain loudly about it.
        if not isinstance(data, bytes):
            raise TypeError(
                "Data being published to Pub/Sub must be sent as a bytestring."
            )

        if not self._enable_message_ordering and ordering_key != "":
            raise ValueError(
                "Cannot publish a message with an ordering key when message "
                "ordering is not enabled."
            )

        # Coerce all attributes to text strings.
        for k, v in copy.copy(attrs).items():
            if isinstance(v, str):
                continue
            if isinstance(v, bytes):
                attrs[k] = v.decode("utf-8")
                continue
            raise TypeError(
                "All attributes being published to Pub/Sub must "
                "be sent as text strings."
            )

        # Create the Pub/Sub message object. For performance reasons, the message
        # should be constructed by directly using the raw protobuf class, and only
        # then wrapping it into the higher-level PubsubMessage class.
        vanilla_pb = _raw_proto_pubbsub_message(
            data=data, ordering_key=ordering_key, attributes=attrs
        )
        message = gapic_types.PubsubMessage.wrap(vanilla_pb)

        # Messages should go through flow control to prevent excessive
        # queuing on the client side (depending on the settings).
        try:
            self._flow_controller.add(message)
        except exceptions.FlowControlLimitError as exc:
            future = futures.Future()
            future.set_exception(exc)
            return future

        def on_publish_done(future):
            self._flow_controller.release(message)

        if retry is gapic_v1.method.DEFAULT:  # if custom retry not passed in
            retry = self.publisher_options.retry

        if timeout is gapic_v1.method.DEFAULT:  # if custom timeout not passed in
            timeout = self.publisher_options.timeout

        with self._batch_lock:
            if self._is_stopped:
                raise RuntimeError("Cannot publish on a stopped publisher.")

            # Set retry timeout to "infinite" when message ordering is enabled.
            # Note that this then also impacts messages added with an empty
            # ordering key.
            if self._enable_message_ordering:
                if retry is gapic_v1.method.DEFAULT:
                    # use the default retry for the publish GRPC method as a base
                    transport = self._transport
                    base_retry = transport._wrapped_methods[transport.publish]._retry
                    retry = base_retry.with_deadline(2.0**32)
                    # timeout needs to be overridden and set to infinite in
                    # addition to the retry deadline since both determine
                    # the duration for which retries are attempted.
                    timeout = 2.0**32
                elif retry is not None:
                    retry = retry.with_deadline(2.0**32)
                    timeout = 2.0**32

            # Delegate the publishing to the sequencer.
            sequencer = self._get_or_create_sequencer(topic, ordering_key)
            future = sequencer.publish(message, retry=retry, timeout=timeout)
            future.add_done_callback(on_publish_done)

            # Create a timer thread if necessary to enforce the batching
            # timeout.
            self._ensure_commit_timer_runs_no_lock()

            return future

    def ensure_cleanup_and_commit_timer_runs(self) -> None:
        """Ensure a cleanup/commit timer thread is running.

        If a cleanup/commit timer thread is already running, this does nothing.
        """
        with self._batch_lock:
            self._ensure_commit_timer_runs_no_lock()

    def _ensure_commit_timer_runs_no_lock(self) -> None:
        """Ensure a commit timer thread is running, without taking
        _batch_lock.

        _batch_lock must be held before calling this method.
        """
        if not self._commit_thread and self.batch_settings.max_latency < float("inf"):
            self._start_commit_thread()

    def _start_commit_thread(self) -> None:
        """Start a new thread to actually wait and commit the sequencers."""
        # NOTE: If the thread is *not* a daemon, a memory leak exists due to a CPython issue.
        # https://github.com/googleapis/python-pubsub/issues/395#issuecomment-829910303
        # https://github.com/googleapis/python-pubsub/issues/395#issuecomment-830092418
        self._commit_thread = threading.Thread(
            name="Thread-PubSubBatchCommitter",
            target=self._wait_and_commit_sequencers,
            daemon=True,
        )
        self._commit_thread.start()

    def _wait_and_commit_sequencers(self) -> None:
        """Wait up to the batching timeout, and commit all sequencers."""
        # Sleep for however long we should be waiting.
        time.sleep(self.batch_settings.max_latency)
        _LOGGER.debug("Commit thread is waking up")

        with self._batch_lock:
            if self._is_stopped:
                return
            self._commit_sequencers()
            self._commit_thread = None

    def _commit_sequencers(self) -> None:
        """Clean up finished sequencers and commit the rest."""
        finished_sequencer_keys = [
            key
            for key, sequencer in self._sequencers.items()
            if sequencer.is_finished()
        ]
        for sequencer_key in finished_sequencer_keys:
            del self._sequencers[sequencer_key]

        for sequencer in self._sequencers.values():
            sequencer.commit()

    def stop(self) -> None:
        """Immediately publish all outstanding messages.

        Asynchronously sends all outstanding messages and
        prevents future calls to `publish()`. Method should
        be invoked prior to deleting this `Client()` object
        in order to ensure that no pending messages are lost.

        .. note::

            This method is non-blocking. Use `Future()` objects
            returned by `publish()` to make sure all publish
            requests completed, either in success or error.

        Raises:
            RuntimeError:
                If called after publisher has been stopped by a `stop()` method
                call.
        """
        with self._batch_lock:
            if self._is_stopped:
                raise RuntimeError("Cannot stop a publisher already stopped.")

            self._is_stopped = True

            for sequencer in self._sequencers.values():
                sequencer.stop()

    # Used only for testing.
    def _set_batch(
        self, topic: str, batch: "_batch.thread.Batch", ordering_key: str = ""
    ) -> None:
        sequencer = self._get_or_create_sequencer(topic, ordering_key)
        sequencer._set_batch(batch)

    # Used only for testing.
    def _set_batch_class(self, batch_class: Type) -> None:
        self._batch_class = batch_class

    # Used only for testing.
    def _set_sequencer(
        self, topic: str, sequencer: SequencerType, ordering_key: str = ""
    ) -> None:
        sequencer_key = (topic, ordering_key)
        self._sequencers[sequencer_key] = sequencer