feat: Add new gcloud commands, API clients, and third-party libraries across various services.

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/__init__.py

@@ -0,0 +1,197 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2020 Google LLC
+#
+# 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 google.cloud.pubsublite import gapic_version as package_version
+
+__version__ = package_version.__version__
+
+
+from google.cloud.pubsublite_v1.services.admin_service.async_client import (
+    AdminServiceAsyncClient,
+)
+from google.cloud.pubsublite_v1.services.admin_service.client import AdminServiceClient
+from google.cloud.pubsublite_v1.services.cursor_service.async_client import (
+    CursorServiceAsyncClient,
+)
+from google.cloud.pubsublite_v1.services.cursor_service.client import (
+    CursorServiceClient,
+)
+from google.cloud.pubsublite_v1.services.partition_assignment_service.async_client import (
+    PartitionAssignmentServiceAsyncClient,
+)
+from google.cloud.pubsublite_v1.services.partition_assignment_service.client import (
+    PartitionAssignmentServiceClient,
+)
+from google.cloud.pubsublite_v1.services.publisher_service.async_client import (
+    PublisherServiceAsyncClient,
+)
+from google.cloud.pubsublite_v1.services.publisher_service.client import (
+    PublisherServiceClient,
+)
+from google.cloud.pubsublite_v1.services.subscriber_service.async_client import (
+    SubscriberServiceAsyncClient,
+)
+from google.cloud.pubsublite_v1.services.subscriber_service.client import (
+    SubscriberServiceClient,
+)
+from google.cloud.pubsublite_v1.services.topic_stats_service.async_client import (
+    TopicStatsServiceAsyncClient,
+)
+from google.cloud.pubsublite_v1.services.topic_stats_service.client import (
+    TopicStatsServiceClient,
+)
+from google.cloud.pubsublite_v1.types.admin import CreateSubscriptionRequest
+from google.cloud.pubsublite_v1.types.admin import CreateTopicRequest
+from google.cloud.pubsublite_v1.types.admin import DeleteSubscriptionRequest
+from google.cloud.pubsublite_v1.types.admin import DeleteTopicRequest
+from google.cloud.pubsublite_v1.types.admin import GetSubscriptionRequest
+from google.cloud.pubsublite_v1.types.admin import GetTopicPartitionsRequest
+from google.cloud.pubsublite_v1.types.admin import GetTopicRequest
+from google.cloud.pubsublite_v1.types.admin import ListSubscriptionsRequest
+from google.cloud.pubsublite_v1.types.admin import ListSubscriptionsResponse
+from google.cloud.pubsublite_v1.types.admin import ListTopicSubscriptionsRequest
+from google.cloud.pubsublite_v1.types.admin import ListTopicSubscriptionsResponse
+from google.cloud.pubsublite_v1.types.admin import ListTopicsRequest
+from google.cloud.pubsublite_v1.types.admin import ListTopicsResponse
+from google.cloud.pubsublite_v1.types.admin import OperationMetadata
+from google.cloud.pubsublite_v1.types.admin import SeekSubscriptionRequest
+from google.cloud.pubsublite_v1.types.admin import SeekSubscriptionResponse
+from google.cloud.pubsublite_v1.types.admin import TopicPartitions
+from google.cloud.pubsublite_v1.types.admin import UpdateSubscriptionRequest
+from google.cloud.pubsublite_v1.types.admin import UpdateTopicRequest
+from google.cloud.pubsublite_v1.types.common import AttributeValues
+from google.cloud.pubsublite_v1.types.common import Cursor
+from google.cloud.pubsublite_v1.types.common import ExportConfig
+from google.cloud.pubsublite_v1.types.common import PubSubMessage
+from google.cloud.pubsublite_v1.types.common import Reservation
+from google.cloud.pubsublite_v1.types.common import SequencedMessage
+from google.cloud.pubsublite_v1.types.common import Subscription
+from google.cloud.pubsublite_v1.types.common import TimeTarget
+from google.cloud.pubsublite_v1.types.common import Topic
+from google.cloud.pubsublite_v1.types.cursor import CommitCursorRequest
+from google.cloud.pubsublite_v1.types.cursor import CommitCursorResponse
+from google.cloud.pubsublite_v1.types.cursor import InitialCommitCursorRequest
+from google.cloud.pubsublite_v1.types.cursor import InitialCommitCursorResponse
+from google.cloud.pubsublite_v1.types.cursor import ListPartitionCursorsRequest
+from google.cloud.pubsublite_v1.types.cursor import ListPartitionCursorsResponse
+from google.cloud.pubsublite_v1.types.cursor import PartitionCursor
+from google.cloud.pubsublite_v1.types.cursor import SequencedCommitCursorRequest
+from google.cloud.pubsublite_v1.types.cursor import SequencedCommitCursorResponse
+from google.cloud.pubsublite_v1.types.cursor import StreamingCommitCursorRequest
+from google.cloud.pubsublite_v1.types.cursor import StreamingCommitCursorResponse
+from google.cloud.pubsublite_v1.types.publisher import InitialPublishRequest
+from google.cloud.pubsublite_v1.types.publisher import InitialPublishResponse
+from google.cloud.pubsublite_v1.types.publisher import MessagePublishRequest
+from google.cloud.pubsublite_v1.types.publisher import MessagePublishResponse
+from google.cloud.pubsublite_v1.types.publisher import PublishRequest
+from google.cloud.pubsublite_v1.types.publisher import PublishResponse
+from google.cloud.pubsublite_v1.types.subscriber import FlowControlRequest
+from google.cloud.pubsublite_v1.types.subscriber import (
+    InitialPartitionAssignmentRequest,
+)
+from google.cloud.pubsublite_v1.types.subscriber import InitialSubscribeRequest
+from google.cloud.pubsublite_v1.types.subscriber import InitialSubscribeResponse
+from google.cloud.pubsublite_v1.types.subscriber import MessageResponse
+from google.cloud.pubsublite_v1.types.subscriber import PartitionAssignment
+from google.cloud.pubsublite_v1.types.subscriber import PartitionAssignmentAck
+from google.cloud.pubsublite_v1.types.subscriber import PartitionAssignmentRequest
+from google.cloud.pubsublite_v1.types.subscriber import SeekRequest
+from google.cloud.pubsublite_v1.types.subscriber import SeekResponse
+from google.cloud.pubsublite_v1.types.subscriber import SubscribeRequest
+from google.cloud.pubsublite_v1.types.subscriber import SubscribeResponse
+from google.cloud.pubsublite_v1.types.topic_stats import ComputeMessageStatsRequest
+from google.cloud.pubsublite_v1.types.topic_stats import ComputeMessageStatsResponse
+from google.cloud.pubsublite.admin_client_interface import AdminClientInterface
+from google.cloud.pubsublite.admin_client import AdminClient
+
+__all__ = (
+    # Manual files
+    "AdminClient",
+    "AdminClientInterface",
+    # Generated files
+    "AdminServiceAsyncClient",
+    "AdminServiceClient",
+    "AttributeValues",
+    "CommitCursorRequest",
+    "CommitCursorResponse",
+    "ComputeMessageStatsRequest",
+    "ComputeMessageStatsResponse",
+    "CreateSubscriptionRequest",
+    "CreateTopicRequest",
+    "Cursor",
+    "CursorServiceAsyncClient",
+    "CursorServiceClient",
+    "DeleteSubscriptionRequest",
+    "DeleteTopicRequest",
+    "ExportConfig",
+    "FlowControlRequest",
+    "GetSubscriptionRequest",
+    "GetTopicPartitionsRequest",
+    "GetTopicRequest",
+    "InitialCommitCursorRequest",
+    "InitialCommitCursorResponse",
+    "InitialPartitionAssignmentRequest",
+    "InitialPublishRequest",
+    "InitialPublishResponse",
+    "InitialSubscribeRequest",
+    "InitialSubscribeResponse",
+    "ListPartitionCursorsRequest",
+    "ListPartitionCursorsResponse",
+    "ListSubscriptionsRequest",
+    "ListSubscriptionsResponse",
+    "ListTopicSubscriptionsRequest",
+    "ListTopicSubscriptionsResponse",
+    "ListTopicsRequest",
+    "ListTopicsResponse",
+    "MessagePublishRequest",
+    "MessagePublishResponse",
+    "MessageResponse",
+    "OperationMetadata",
+    "PartitionAssignment",
+    "PartitionAssignmentAck",
+    "PartitionAssignmentRequest",
+    "PartitionAssignmentServiceAsyncClient",
+    "PartitionAssignmentServiceClient",
+    "PartitionCursor",
+    "PubSubMessage",
+    "PublishRequest",
+    "PublishResponse",
+    "PublisherServiceAsyncClient",
+    "PublisherServiceClient",
+    "Reservation",
+    "SeekSubscriptionRequest",
+    "SeekSubscriptionResponse",
+    "SeekRequest",
+    "SeekResponse",
+    "SequencedCommitCursorRequest",
+    "SequencedCommitCursorResponse",
+    "SequencedMessage",
+    "StreamingCommitCursorRequest",
+    "StreamingCommitCursorResponse",
+    "SubscribeRequest",
+    "SubscribeResponse",
+    "SubscriberServiceAsyncClient",
+    "SubscriberServiceClient",
+    "Subscription",
+    "TimeTarget",
+    "Topic",
+    "TopicPartitions",
+    "TopicStatsServiceAsyncClient",
+    "TopicStatsServiceClient",
+    "UpdateSubscriptionRequest",
+    "UpdateTopicRequest",
+)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/admin_client.py

@@ -0,0 +1,153 @@
+# Copyright 2020 Google LLC
+#
+# 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 typing import Optional, List, Union
+
+from google.api_core.client_options import ClientOptions
+from google.api_core.operation import Operation
+from google.auth.credentials import Credentials
+from cloudsdk.google.protobuf.field_mask_pb2 import FieldMask  # pytype: disable=pyi-error
+
+from google.cloud.pubsublite.admin_client_interface import AdminClientInterface
+from google.cloud.pubsublite.internal.constructable_from_service_account import (
+    ConstructableFromServiceAccount,
+)
+from google.cloud.pubsublite.internal.endpoints import regional_endpoint
+from google.cloud.pubsublite.internal.wire.admin_client_impl import AdminClientImpl
+from google.cloud.pubsublite.types import (
+    CloudRegion,
+    SubscriptionPath,
+    LocationPath,
+    TopicPath,
+    BacklogLocation,
+    PublishTime,
+    EventTime,
+)
+from google.cloud.pubsublite.types.paths import ReservationPath
+from google.cloud.pubsublite_v1 import (
+    AdminServiceClient,
+    Subscription,
+    Topic,
+    Reservation,
+)
+
+
+class AdminClient(AdminClientInterface, ConstructableFromServiceAccount):
+    """
+    An admin client for Pub/Sub Lite. Only operates on a single region.
+    """
+
+    _impl: AdminClientInterface
+
+    def __init__(
+        self,
+        region: CloudRegion,
+        credentials: Optional[Credentials] = None,
+        transport: Optional[str] = None,
+        client_options: Optional[ClientOptions] = None,
+    ):
+        """
+        Create a new AdminClient.
+
+        Args:
+            region: The cloud region to connect to.
+            credentials: The credentials to use when connecting.
+            transport: The transport to use.
+            client_options: The client options to use when connecting. If used, must explicitly set `api_endpoint`.
+        """
+        if client_options is None:
+            client_options = ClientOptions(api_endpoint=regional_endpoint(region))
+        self._impl = AdminClientImpl(
+            AdminServiceClient(
+                client_options=client_options,
+                transport=transport,
+                credentials=credentials,
+            ),
+            region,
+        )
+
+    def region(self) -> CloudRegion:
+        return self._impl.region()
+
+    def create_topic(self, topic: Topic) -> Topic:
+        return self._impl.create_topic(topic)
+
+    def get_topic(self, topic_path: TopicPath) -> Topic:
+        return self._impl.get_topic(topic_path)
+
+    def get_topic_partition_count(self, topic_path: TopicPath) -> int:
+        return self._impl.get_topic_partition_count(topic_path)
+
+    def list_topics(self, location_path: LocationPath) -> List[Topic]:
+        return self._impl.list_topics(location_path)
+
+    def update_topic(self, topic: Topic, update_mask: FieldMask) -> Topic:
+        return self._impl.update_topic(topic, update_mask)
+
+    def delete_topic(self, topic_path: TopicPath):
+        return self._impl.delete_topic(topic_path)
+
+    def list_topic_subscriptions(self, topic_path: TopicPath) -> List[SubscriptionPath]:
+        return self._impl.list_topic_subscriptions(topic_path)
+
+    def create_subscription(
+        self,
+        subscription: Subscription,
+        target: Union[BacklogLocation, PublishTime, EventTime] = BacklogLocation.END,
+        starting_offset: Optional[BacklogLocation] = None,
+    ) -> Subscription:
+        return self._impl.create_subscription(subscription, target, starting_offset)
+
+    def get_subscription(self, subscription_path: SubscriptionPath) -> Subscription:
+        return self._impl.get_subscription(subscription_path)
+
+    def list_subscriptions(self, location_path: LocationPath) -> List[Subscription]:
+        return self._impl.list_subscriptions(location_path)
+
+    def update_subscription(
+        self, subscription: Subscription, update_mask: FieldMask
+    ) -> Subscription:
+        return self._impl.update_subscription(subscription, update_mask)
+
+    def seek_subscription(
+        self,
+        subscription_path: SubscriptionPath,
+        target: Union[BacklogLocation, PublishTime, EventTime],
+    ) -> Operation:
+        return self._impl.seek_subscription(subscription_path, target)
+
+    def delete_subscription(self, subscription_path: SubscriptionPath):
+        return self._impl.delete_subscription(subscription_path)
+
+    def create_reservation(self, reservation: Reservation) -> Reservation:
+        return self._impl.create_reservation(reservation)
+
+    def get_reservation(self, reservation_path: ReservationPath) -> Reservation:
+        return self._impl.get_reservation(reservation_path)
+
+    def list_reservations(self, location_path: LocationPath) -> List[Reservation]:
+        return self._impl.list_reservations(location_path)
+
+    def update_reservation(
+        self, reservation: Reservation, update_mask: FieldMask
+    ) -> Reservation:
+        return self._impl.update_reservation(reservation, update_mask)
+
+    def delete_reservation(self, reservation_path: ReservationPath):
+        return self._impl.delete_reservation(reservation_path)
+
+    def list_reservation_topics(
+        self, reservation_path: ReservationPath
+    ) -> List[TopicPath]:
+        return self._impl.list_reservation_topics(reservation_path)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/admin_client_interface.py

@@ -0,0 +1,151 @@
+# Copyright 2020 Google LLC
+#
+# 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 abc import ABC, abstractmethod
+from typing import List, Optional, Union
+
+from google.api_core.operation import Operation
+from google.cloud.pubsublite.types import (
+    CloudRegion,
+    TopicPath,
+    LocationPath,
+    SubscriptionPath,
+    BacklogLocation,
+    PublishTime,
+    EventTime,
+)
+from google.cloud.pubsublite.types.paths import ReservationPath
+from google.cloud.pubsublite_v1 import Topic, Subscription, Reservation
+from cloudsdk.google.protobuf.field_mask_pb2 import FieldMask  # pytype: disable=pyi-error
+
+
+class AdminClientInterface(ABC):
+    """
+    An admin client for Pub/Sub Lite. Only operates on a single region.
+    """
+
+    @abstractmethod
+    def region(self) -> CloudRegion:
+        """The region this client is for."""
+
+    @abstractmethod
+    def create_topic(self, topic: Topic) -> Topic:
+        """Create a topic, returns the created topic."""
+
+    @abstractmethod
+    def get_topic(self, topic_path: TopicPath) -> Topic:
+        """Get the topic object from the server."""
+
+    @abstractmethod
+    def get_topic_partition_count(self, topic_path: TopicPath) -> int:
+        """Get the number of partitions in the provided topic."""
+
+    @abstractmethod
+    def list_topics(self, location_path: LocationPath) -> List[Topic]:
+        """List the Pub/Sub lite topics that exist for a project in a given location."""
+
+    @abstractmethod
+    def update_topic(self, topic: Topic, update_mask: FieldMask) -> Topic:
+        """Update the masked fields of the provided topic."""
+
+    @abstractmethod
+    def delete_topic(self, topic_path: TopicPath):
+        """Delete a topic and all associated messages."""
+
+    @abstractmethod
+    def list_topic_subscriptions(self, topic_path: TopicPath) -> List[SubscriptionPath]:
+        """List the subscriptions that exist for a given topic."""
+
+    @abstractmethod
+    def create_subscription(
+        self,
+        subscription: Subscription,
+        target: Union[BacklogLocation, PublishTime, EventTime] = BacklogLocation.END,
+        starting_offset: Optional[BacklogLocation] = None,
+    ) -> Subscription:
+        """Create a subscription, returns the created subscription. By default
+        a subscription will only receive messages published after the
+        subscription was created.
+
+        `starting_offset` is deprecated. Use `target` to initialize the
+        subscription to a target location within the message backlog instead.
+        `starting_offset` has higher precedence if `target` is also set.
+
+        A seek is initiated if the target location is a publish or event time.
+        If the seek fails, the created subscription is not deleted.
+        """
+
+    @abstractmethod
+    def get_subscription(self, subscription_path: SubscriptionPath) -> Subscription:
+        """Get the subscription object from the server."""
+
+    @abstractmethod
+    def list_subscriptions(self, location_path: LocationPath) -> List[Subscription]:
+        """List the Pub/Sub lite subscriptions that exist for a project in a given location."""
+
+    @abstractmethod
+    def update_subscription(
+        self, subscription: Subscription, update_mask: FieldMask
+    ) -> Subscription:
+        """Update the masked fields of the provided subscription."""
+
+    @abstractmethod
+    def seek_subscription(
+        self,
+        subscription_path: SubscriptionPath,
+        target: Union[BacklogLocation, PublishTime, EventTime],
+    ) -> Operation:
+        """Initiate an out-of-band seek for a subscription to a specified target.
+
+        The seek target may be timestamps or named positions within the message
+        backlog See https://cloud.google.com/pubsub/lite/docs/seek for more
+        information.
+
+        Returns:
+            google.api_core.operation.Operation with:
+              result type: google.cloud.pubsublite.SeekSubscriptionResponse
+              metadata type: google.cloud.pubsublite.OperationMetadata
+        """
+
+    @abstractmethod
+    def delete_subscription(self, subscription_path: SubscriptionPath):
+        """Delete a subscription and all associated messages."""
+
+    @abstractmethod
+    def create_reservation(self, reservation: Reservation) -> Reservation:
+        """Create a reservation, returns the created reservation."""
+
+    @abstractmethod
+    def get_reservation(self, reservation_path: ReservationPath) -> Reservation:
+        """Get the reservation object from the server."""
+
+    @abstractmethod
+    def list_reservations(self, location_path: LocationPath) -> List[Reservation]:
+        """List the Pub/Sub lite reservations that exist for a project in a given location."""
+
+    @abstractmethod
+    def update_reservation(
+        self, reservation: Reservation, update_mask: FieldMask
+    ) -> Reservation:
+        """Update the masked fields of the provided reservation."""
+
+    @abstractmethod
+    def delete_reservation(self, reservation_path: ReservationPath):
+        """Delete a reservation and all associated messages."""
+
+    @abstractmethod
+    def list_reservation_topics(
+        self, reservation_path: ReservationPath
+    ) -> List[TopicPath]:
+        """List the subscriptions that exist for a given reservation."""

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/__init__.py

@@ -0,0 +1,40 @@
+# Copyright 2020 Google LLC
+#
+# 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.
+
+# flake8: noqa
+from .message_transformer import MessageTransformer
+from .nack_handler import NackHandler
+from .publisher_client import AsyncPublisherClient, PublisherClient
+from .publisher_client_interface import (
+    AsyncPublisherClientInterface,
+    PublisherClientInterface,
+)
+from .subscriber_client import AsyncSubscriberClient, SubscriberClient
+from .subscriber_client_interface import (
+    AsyncSubscriberClientInterface,
+    SubscriberClientInterface,
+)
+
+__all__ = (
+    "AsyncPublisherClient",
+    "AsyncPublisherClientInterface",
+    "AsyncSubscriberClient",
+    "AsyncSubscriberClientInterface",
+    "MessageTransformer",
+    "NackHandler",
+    "PublisherClient",
+    "PublisherClientInterface",
+    "SubscriberClient",
+    "SubscriberClientInterface",
+)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/ack_set_tracker.py

@@ -0,0 +1,56 @@
+# Copyright 2020 Google LLC
+#
+# 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 abc import abstractmethod, ABCMeta
+from typing import AsyncContextManager
+
+
+class AckSetTracker(AsyncContextManager, metaclass=ABCMeta):
+    """
+    An AckSetTracker tracks disjoint acknowledged messages and commits them when a contiguous prefix of tracked offsets
+    is aggregated.
+    """
+
+    @abstractmethod
+    def track(self, offset: int):
+        """
+        Track the provided offset.
+
+        Args:
+          offset: the offset to track.
+
+        Raises:
+          GoogleAPICallError: On an invalid offset to track.
+        """
+
+    @abstractmethod
+    def ack(self, offset: int):
+        """
+        Acknowledge the message with the provided offset. The offset must have previously been tracked.
+
+        Args:
+          offset: the offset to acknowledge.
+
+        Returns:
+          GoogleAPICallError: On a commit failure.
+        """
+
+    @abstractmethod
+    async def clear_and_commit(self):
+        """
+        Discard all outstanding acks and wait for the commit offset to be acknowledged by the server.
+
+        Raises:
+          GoogleAPICallError: If the committer has shut down due to a permanent error.
+        """

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/ack_set_tracker_impl.py

@@ -0,0 +1,75 @@
+# Copyright 2020 Google LLC
+#
+# 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 collections import deque
+from typing import Optional
+
+from google.api_core.exceptions import FailedPrecondition
+
+from google.cloud.pubsublite.cloudpubsub.internal.sorted_list import SortedList
+from google.cloud.pubsublite.cloudpubsub.internal.ack_set_tracker import AckSetTracker
+from google.cloud.pubsublite.internal.wire.committer import Committer
+from google.cloud.pubsublite_v1 import Cursor
+
+
+class AckSetTrackerImpl(AckSetTracker):
+    _committer: Committer
+
+    _receipts: "deque[int]"
+    _acks: SortedList[int]
+
+    def __init__(self, committer: Committer):
+        super().__init__()
+        self._committer = committer
+        self._receipts = deque()
+        self._acks = SortedList()
+
+    def track(self, offset: int):
+        if len(self._receipts) > 0:
+            last = self._receipts[0]
+            if last >= offset:
+                raise FailedPrecondition(
+                    f"Tried to track message {offset} which is before last tracked message {last}."
+                )
+        self._receipts.append(offset)
+
+    def ack(self, offset: int):
+        self._acks.push(offset)
+        prefix_acked_offset: Optional[int] = None
+        while len(self._receipts) != 0 and not self._acks.empty():
+            receipt = self._receipts.popleft()
+            ack = self._acks.peek()
+            if receipt == ack:
+                prefix_acked_offset = receipt
+                self._acks.pop()
+                continue
+            self._receipts.appendleft(receipt)
+            break
+        if prefix_acked_offset is None:
+            return
+        # Convert from last acked to first unacked.
+        cursor = Cursor()
+        cursor._pb.offset = prefix_acked_offset + 1
+        self._committer.commit(cursor)
+
+    async def clear_and_commit(self):
+        self._receipts.clear()
+        self._acks = SortedList()
+        await self._committer.wait_until_empty()
+
+    async def __aenter__(self):
+        await self._committer.__aenter__()
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        await self._committer.__aexit__(exc_type, exc_value, traceback)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/assigning_subscriber.py

@@ -0,0 +1,122 @@
+# Copyright 2020 Google LLC
+#
+# 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 asyncio import Future, Queue, ensure_future
+from typing import Callable, NamedTuple, Dict, List, Set, Optional
+
+from google.cloud.pubsub_v1.subscriber.message import Message
+
+from google.cloud.pubsublite.cloudpubsub.reassignment_handler import ReassignmentHandler
+from google.cloud.pubsublite.cloudpubsub.internal.single_subscriber import (
+    AsyncSingleSubscriber,
+)
+from google.cloud.pubsublite.internal.wait_ignore_cancelled import (
+    wait_ignore_cancelled,
+    wait_ignore_errors,
+)
+from google.cloud.pubsublite.internal.wire.assigner import Assigner
+from google.cloud.pubsublite.internal.wire.permanent_failable import PermanentFailable
+from google.cloud.pubsublite.types import Partition
+
+PartitionSubscriberFactory = Callable[[Partition], AsyncSingleSubscriber]
+
+
+class _RunningSubscriber(NamedTuple):
+    subscriber: AsyncSingleSubscriber
+    poller: Future
+
+
+class AssigningSingleSubscriber(AsyncSingleSubscriber, PermanentFailable):
+    _assigner_factory: Callable[[], Assigner]
+    _subscriber_factory: PartitionSubscriberFactory
+    _reassignment_handler: ReassignmentHandler
+
+    _subscribers: Dict[Partition, _RunningSubscriber]
+
+    # Lazily initialized to ensure they are initialized on the thread where __aenter__ is called.
+    _assigner: Optional[Assigner]
+    _batches: Optional["Queue[List[Message]]"]
+    _assign_poller: Future
+
+    def __init__(
+        self,
+        assigner_factory: Callable[[], Assigner],
+        subscriber_factory: PartitionSubscriberFactory,
+        reassignment_handler: ReassignmentHandler,
+    ):
+        """
+        Accepts a factory for an Assigner instead of an Assigner because GRPC asyncio uses the current thread's event
+        loop.
+        """
+        super().__init__()
+        self._assigner_factory = assigner_factory
+        self._subscriber_factory = subscriber_factory
+        self._reassignment_handler = reassignment_handler
+        self._assigner = None
+        self._subscribers = {}
+        self._batches = None
+
+    async def read(self) -> List[Message]:
+        return await self.await_unless_failed(self._batches.get())
+
+    async def _subscribe_action(self, subscriber: AsyncSingleSubscriber):
+        batch = await subscriber.read()
+        await self._batches.put(batch)
+
+    async def _start_subscriber(self, partition: Partition):
+        new_subscriber = self._subscriber_factory(partition)
+        await new_subscriber.__aenter__()
+        poller = ensure_future(
+            self.run_poller(lambda: self._subscribe_action(new_subscriber))
+        )
+        self._subscribers[partition] = _RunningSubscriber(new_subscriber, poller)
+
+    async def _stop_subscriber(self, running: _RunningSubscriber):
+        running.poller.cancel()
+        await wait_ignore_cancelled(running.poller)
+        await running.subscriber.__aexit__(None, None, None)
+
+    async def _assign_action(self):
+        assignment: Set[Partition] = await self._assigner.get_assignment()
+        old_assignment: Set[Partition] = set(self._subscribers.keys())
+        added_partitions = assignment - old_assignment
+        removed_partitions = old_assignment - assignment
+        for partition in added_partitions:
+            await self._start_subscriber(partition)
+        for partition in removed_partitions:
+            subscriber = self._subscribers[partition]
+            del self._subscribers[partition]
+            await self._stop_subscriber(subscriber)
+        maybe_awaitable = self._reassignment_handler.handle_reassignment(
+            old_assignment, assignment
+        )
+        if maybe_awaitable is not None:
+            await maybe_awaitable
+
+    async def __aenter__(self):
+        self._batches = Queue()
+        self._assigner = self._assigner_factory()
+        await self._assigner.__aenter__()
+        self._assign_poller = ensure_future(self.run_poller(self._assign_action))
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        self._assign_poller.cancel()
+        await wait_ignore_errors(self._assign_poller)
+        await wait_ignore_errors(
+            self._assigner.__aexit__(exc_type, exc_value, traceback)
+        )
+        for running in self._subscribers.values():
+            await wait_ignore_errors(self._stop_subscriber(running))
+        pass

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/async_publisher_impl.py

@@ -0,0 +1,56 @@
+# Copyright 2020 Google LLC
+#
+# 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 typing import Mapping, Callable, Optional
+
+from google.pubsub_v1 import PubsubMessage
+
+from google.cloud.pubsublite.cloudpubsub.message_transforms import (
+    from_cps_publish_message,
+)
+from google.cloud.pubsublite.cloudpubsub.internal.single_publisher import (
+    AsyncSinglePublisher,
+)
+from google.cloud.pubsublite.internal.wire.publisher import Publisher
+
+
+class AsyncSinglePublisherImpl(AsyncSinglePublisher):
+    _publisher_factory: Callable[[], Publisher]
+    _publisher: Optional[Publisher]
+
+    def __init__(self, publisher_factory: Callable[[], Publisher]):
+        """
+        Accepts a factory for a Publisher instead of a Publisher because GRPC asyncio uses the current thread's event
+        loop.
+        """
+        super().__init__()
+        self._publisher_factory = publisher_factory
+        self._publisher = None
+
+    async def publish(
+        self, data: bytes, ordering_key: str = "", **attrs: Mapping[str, str]
+    ) -> str:
+        cps_message = PubsubMessage(
+            data=data, ordering_key=ordering_key, attributes=attrs
+        )
+        psl_message = from_cps_publish_message(cps_message)
+        return (await self._publisher.publish(psl_message)).encode()
+
+    async def __aenter__(self):
+        self._publisher = self._publisher_factory()
+        await self._publisher.__aenter__()
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        await self._publisher.__aexit__(exc_type, exc_value, traceback)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/client_multiplexer.py

@@ -0,0 +1,122 @@
+# Copyright 2020 Google LLC
+#
+# 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 asyncio
+import threading
+from typing import Generic, TypeVar, Callable, Dict, Awaitable
+
+_Key = TypeVar("_Key")
+_Client = TypeVar("_Client")
+
+
+class ClientMultiplexer(Generic[_Key, _Client]):
+    _OpenedClientFactory = Callable[[_Key], _Client]
+    _ClientCloser = Callable[[_Client], None]
+
+    _factory: _OpenedClientFactory
+    _closer: _ClientCloser
+    _lock: threading.Lock
+    _live_clients: Dict[_Key, _Client]
+
+    def __init__(
+        self,
+        factory: _OpenedClientFactory,
+        closer: _ClientCloser = lambda client: client.__exit__(None, None, None),
+    ):
+        self._factory = factory
+        self._closer = closer
+        self._lock = threading.Lock()
+        self._live_clients = {}
+
+    def get_or_create(self, key: _Key) -> _Client:
+        with self._lock:
+            if key not in self._live_clients:
+                self._live_clients[key] = self._factory(key)
+            return self._live_clients[key]
+
+    def try_erase(self, key: _Key, client: _Client):
+        with self._lock:
+            if key not in self._live_clients:
+                return
+            current_client = self._live_clients[key]
+            if current_client is not client:
+                return
+            del self._live_clients[key]
+        self._closer(client)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        live_clients: Dict[_Key, _Client]
+        with self._lock:
+            live_clients = self._live_clients
+            self._live_clients = {}
+        for topic, client in live_clients.items():
+            self._closer(client)
+
+
+class AsyncClientMultiplexer(Generic[_Key, _Client]):
+    _OpenedClientFactory = Callable[[_Key], Awaitable[_Client]]
+    _ClientCloser = Callable[[_Client], Awaitable[None]]
+
+    _factory: _OpenedClientFactory
+    _closer: _ClientCloser
+    _live_clients: Dict[_Key, Awaitable[_Client]]
+
+    def __init__(
+        self,
+        factory: _OpenedClientFactory,
+        closer: _ClientCloser = lambda client: client.__aexit__(None, None, None),
+    ):
+        self._factory = factory
+        self._closer = closer
+        self._live_clients = {}
+
+    async def get_or_create(self, key: _Key) -> _Client:
+        if key not in self._live_clients:
+            self._live_clients[key] = asyncio.ensure_future(self._factory(key))
+        future = self._live_clients[key]
+        try:
+            return await future
+        except BaseException as e:
+            if key in self._live_clients and self._live_clients[key] is future:
+                del self._live_clients[key]
+            raise e
+
+    async def try_erase(self, key: _Key, client: _Client):
+        if key not in self._live_clients:
+            return
+        client_future = self._live_clients[key]
+        current_client = await client_future
+        if current_client is not client:
+            return
+        # duplicate check after await that no one raced with us
+        if (
+            key not in self._live_clients
+            or self._live_clients[key] is not client_future
+        ):
+            return
+        del self._live_clients[key]
+        await self._closer(client)
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        live_clients: Dict[_Key, Awaitable[_Client]]
+        live_clients = self._live_clients
+        self._live_clients = {}
+        for topic, client in live_clients.items():
+            await self._closer(await client)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/make_publisher.py

@@ -0,0 +1,124 @@
+# Copyright 2020 Google LLC
+#
+# 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 typing import Optional, Mapping
+
+from google.api_core.client_options import ClientOptions
+from google.auth.credentials import Credentials
+from google.cloud.pubsub_v1.types import BatchSettings
+
+from google.cloud.pubsublite.cloudpubsub.internal.async_publisher_impl import (
+    AsyncSinglePublisherImpl,
+)
+from google.cloud.pubsublite.cloudpubsub.internal.publisher_impl import (
+    SinglePublisherImpl,
+)
+from google.cloud.pubsublite.cloudpubsub.internal.single_publisher import (
+    AsyncSinglePublisher,
+    SinglePublisher,
+)
+from google.cloud.pubsublite.internal.publisher_client_id import PublisherClientId
+from google.cloud.pubsublite.internal.wire.make_publisher import (
+    make_publisher as make_wire_publisher,
+    DEFAULT_BATCHING_SETTINGS as WIRE_DEFAULT_BATCHING,
+)
+from google.cloud.pubsublite.internal.wire.merge_metadata import merge_metadata
+from google.cloud.pubsublite.internal.wire.pubsub_context import pubsub_context
+from google.cloud.pubsublite.types import TopicPath
+
+
+DEFAULT_BATCHING_SETTINGS = WIRE_DEFAULT_BATCHING
+
+
+def make_async_publisher(
+    topic: TopicPath,
+    transport: str,
+    per_partition_batching_settings: Optional[BatchSettings] = None,
+    credentials: Optional[Credentials] = None,
+    client_options: Optional[ClientOptions] = None,
+    metadata: Optional[Mapping[str, str]] = None,
+    client_id: Optional[PublisherClientId] = None,
+) -> AsyncSinglePublisher:
+    """
+    Make a new publisher for the given topic.
+
+    Args:
+      topic: The topic to publish to.
+      transport: The transport type to use.
+      per_partition_batching_settings: Settings for batching messages on each partition. The default is reasonable for most cases.
+      credentials: The credentials to use to connect. GOOGLE_DEFAULT_CREDENTIALS is used if None.
+      client_options: Other options to pass to the client. Note that if you pass any you must set api_endpoint.
+      metadata: Additional metadata to send with the RPC.
+      client_id: 128-bit unique client id. If set, enables publish idempotency for the session.
+
+    Returns:
+      A new AsyncPublisher.
+
+    Throws:
+      GoogleApiCallException on any error determining topic structure.
+    """
+    metadata = merge_metadata(pubsub_context(framework="CLOUD_PUBSUB_SHIM"), metadata)
+
+    def underlying_factory():
+        return make_wire_publisher(
+            topic=topic,
+            transport=transport,
+            per_partition_batching_settings=per_partition_batching_settings,
+            credentials=credentials,
+            client_options=client_options,
+            metadata=metadata,
+            client_id=client_id,
+        )
+
+    return AsyncSinglePublisherImpl(underlying_factory)
+
+
+def make_publisher(
+    topic: TopicPath,
+    transport: str,
+    per_partition_batching_settings: Optional[BatchSettings] = None,
+    credentials: Optional[Credentials] = None,
+    client_options: Optional[ClientOptions] = None,
+    metadata: Optional[Mapping[str, str]] = None,
+    client_id: Optional[PublisherClientId] = None,
+) -> SinglePublisher:
+    """
+    Make a new publisher for the given topic.
+
+    Args:
+      topic: The topic to publish to.
+      transport: The transport type to use.
+      per_partition_batching_settings: Settings for batching messages on each partition. The default is reasonable for most cases.
+      credentials: The credentials to use to connect. GOOGLE_DEFAULT_CREDENTIALS is used if None.
+      client_options: Other options to pass to the client. Note that if you pass any you must set api_endpoint.
+      metadata: Additional metadata to send with the RPC.
+      client_id: 128-bit unique client id. If set, enables publish idempotency for the session.
+
+    Returns:
+      A new Publisher.
+
+    Throws:
+      GoogleApiCallException on any error determining topic structure.
+    """
+    return SinglePublisherImpl(
+        make_async_publisher(
+            topic=topic,
+            transport=transport,
+            per_partition_batching_settings=per_partition_batching_settings,
+            credentials=credentials,
+            client_options=client_options,
+            metadata=metadata,
+            client_id=client_id,
+        )
+    )

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/make_subscriber.py

@@ -0,0 +1,246 @@
+# Copyright 2020 Google LLC
+#
+# 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 typing import Optional, Mapping, Set, AsyncIterator, Callable
+from uuid import uuid4
+
+from google.api_core.client_options import ClientOptions
+from google.auth.credentials import Credentials
+
+from google.cloud.pubsublite.cloudpubsub.reassignment_handler import (
+    ReassignmentHandler,
+    DefaultReassignmentHandler,
+)
+from google.cloud.pubsublite.cloudpubsub.message_transforms import (
+    to_cps_subscribe_message,
+    add_id_to_cps_subscribe_transformer,
+)
+from google.cloud.pubsublite.internal.wire.client_cache import ClientCache
+from google.cloud.pubsublite.types import FlowControlSettings
+from google.cloud.pubsublite.cloudpubsub.internal.ack_set_tracker_impl import (
+    AckSetTrackerImpl,
+)
+from google.cloud.pubsublite.cloudpubsub.internal.assigning_subscriber import (
+    PartitionSubscriberFactory,
+    AssigningSingleSubscriber,
+)
+from google.cloud.pubsublite.cloudpubsub.internal.single_partition_subscriber import (
+    SinglePartitionSingleSubscriber,
+)
+from google.cloud.pubsublite.cloudpubsub.message_transformer import MessageTransformer
+from google.cloud.pubsublite.cloudpubsub.nack_handler import (
+    NackHandler,
+    DefaultNackHandler,
+)
+from google.cloud.pubsublite.cloudpubsub.internal.single_subscriber import (
+    AsyncSingleSubscriber,
+)
+from google.cloud.pubsublite.internal.endpoints import regional_endpoint
+from google.cloud.pubsublite.internal.wire.assigner import Assigner
+from google.cloud.pubsublite.internal.wire.assigner_impl import AssignerImpl
+from google.cloud.pubsublite.internal.wire.committer_impl import CommitterImpl
+from google.cloud.pubsublite.internal.wire.fixed_set_assigner import FixedSetAssigner
+from google.cloud.pubsublite.internal.wire.gapic_connection import (
+    GapicConnectionFactory,
+)
+from google.cloud.pubsublite.internal.wire.merge_metadata import merge_metadata
+from google.cloud.pubsublite.internal.wire.pubsub_context import pubsub_context
+import google.cloud.pubsublite.internal.wire.subscriber_impl as wire_subscriber
+from google.cloud.pubsublite.internal.wire.subscriber_reset_handler import (
+    SubscriberResetHandler,
+)
+from google.cloud.pubsublite.types import Partition, SubscriptionPath
+from google.cloud.pubsublite.internal.routing_metadata import (
+    subscription_routing_metadata,
+)
+from google.cloud.pubsublite_v1 import (
+    SubscribeRequest,
+    InitialSubscribeRequest,
+    StreamingCommitCursorRequest,
+    PartitionAssignmentRequest,
+    InitialPartitionAssignmentRequest,
+    InitialCommitCursorRequest,
+)
+from google.cloud.pubsublite_v1.services.subscriber_service.async_client import (
+    SubscriberServiceAsyncClient,
+)
+from google.cloud.pubsublite_v1.services.partition_assignment_service.async_client import (
+    PartitionAssignmentServiceAsyncClient,
+)
+from google.cloud.pubsublite_v1.services.cursor_service.async_client import (
+    CursorServiceAsyncClient,
+)
+
+_DEFAULT_FLUSH_SECONDS = 0.1
+
+
+def _make_dynamic_assigner(
+    subscription: SubscriptionPath,
+    transport: str,
+    client_options: ClientOptions,
+    credentials: Optional[Credentials],
+    base_metadata: Optional[Mapping[str, str]],
+) -> Assigner:
+    if base_metadata is None:
+        base_metadata = {}
+
+    def assignment_connection_factory(
+        requests: AsyncIterator[PartitionAssignmentRequest],
+    ):
+        assignment_client = PartitionAssignmentServiceAsyncClient(credentials=credentials, transport=transport, client_options=client_options)  # type: ignore
+        return assignment_client.assign_partitions(
+            requests, metadata=list(base_metadata.items())
+        )
+
+    return AssignerImpl(
+        InitialPartitionAssignmentRequest(
+            subscription=str(subscription), client_id=uuid4().bytes
+        ),
+        GapicConnectionFactory(assignment_connection_factory),
+    )
+
+
+def _make_partition_subscriber_factory(
+    subscription: SubscriptionPath,
+    transport: str,
+    client_options: ClientOptions,
+    credentials: Optional[Credentials],
+    base_metadata: Optional[Mapping[str, str]],
+    flow_control_settings: FlowControlSettings,
+    nack_handler: NackHandler,
+    message_transformer: MessageTransformer,
+) -> PartitionSubscriberFactory:
+    subscribe_client_cache = ClientCache(
+        lambda: SubscriberServiceAsyncClient(
+            credentials=credentials, transport=transport, client_options=client_options
+        )
+    )
+    cursor_client_cache = ClientCache(
+        lambda: CursorServiceAsyncClient(
+            credentials=credentials, transport=transport, client_options=client_options
+        )
+    )
+
+    def factory(partition: Partition) -> AsyncSingleSubscriber:
+        final_metadata = merge_metadata(
+            base_metadata, subscription_routing_metadata(subscription, partition)
+        )
+
+        def subscribe_connection_factory(requests: AsyncIterator[SubscribeRequest]):
+            return subscribe_client_cache.get().subscribe(
+                requests, metadata=list(final_metadata.items())
+            )
+
+        def cursor_connection_factory(
+            requests: AsyncIterator[StreamingCommitCursorRequest],
+        ):
+            return cursor_client_cache.get().streaming_commit_cursor(
+                requests, metadata=list(final_metadata.items())
+            )
+
+        def subscriber_factory(reset_handler: SubscriberResetHandler):
+            return wire_subscriber.SubscriberImpl(
+                InitialSubscribeRequest(
+                    subscription=str(subscription), partition=partition.value
+                ),
+                _DEFAULT_FLUSH_SECONDS,
+                GapicConnectionFactory(subscribe_connection_factory),
+                reset_handler,
+            )
+
+        committer = CommitterImpl(
+            InitialCommitCursorRequest(
+                subscription=str(subscription), partition=partition.value
+            ),
+            _DEFAULT_FLUSH_SECONDS,
+            GapicConnectionFactory(cursor_connection_factory),
+        )
+        ack_set_tracker = AckSetTrackerImpl(committer)
+        return SinglePartitionSingleSubscriber(
+            subscriber_factory,
+            flow_control_settings,
+            ack_set_tracker,
+            nack_handler,
+            add_id_to_cps_subscribe_transformer(partition, message_transformer),
+        )
+
+    return factory
+
+
+def make_async_subscriber(
+    subscription: SubscriptionPath,
+    transport: str,
+    per_partition_flow_control_settings: FlowControlSettings,
+    nack_handler: Optional[NackHandler] = None,
+    reassignment_handler: Optional[ReassignmentHandler] = None,
+    message_transformer: Optional[MessageTransformer] = None,
+    fixed_partitions: Optional[Set[Partition]] = None,
+    credentials: Optional[Credentials] = None,
+    client_options: Optional[ClientOptions] = None,
+    metadata: Optional[Mapping[str, str]] = None,
+) -> AsyncSingleSubscriber:
+    """
+    Make a Pub/Sub Lite AsyncSubscriber.
+
+    Args:
+      subscription: The subscription to subscribe to.
+      transport: The transport type to use.
+      per_partition_flow_control_settings: The flow control settings for each partition subscribed to. Note that these
+        settings apply to each partition individually, not in aggregate.
+      nack_handler: An optional handler for when nack() is called on a Message. The default will fail the client.
+      message_transformer: An optional transformer from Pub/Sub Lite messages to Cloud Pub/Sub messages.
+      fixed_partitions: A fixed set of partitions to subscribe to. If not present, will instead use auto-assignment.
+      credentials: The credentials to use to connect. GOOGLE_DEFAULT_CREDENTIALS is used if None.
+      client_options: Other options to pass to the client. Note that if you pass any you must set api_endpoint.
+      metadata: Additional metadata to send with the RPC.
+
+    Returns:
+      A new AsyncSubscriber.
+    """
+    metadata = merge_metadata(pubsub_context(framework="CLOUD_PUBSUB_SHIM"), metadata)
+    if client_options is None:
+        client_options = ClientOptions(
+            api_endpoint=regional_endpoint(subscription.location.region)
+        )
+    assigner_factory: Callable[[], Assigner]
+    if fixed_partitions:
+        assigner_factory = lambda: FixedSetAssigner(fixed_partitions)  # noqa: E731
+    else:
+        assigner_factory = lambda: _make_dynamic_assigner(  # noqa: E731
+            subscription,
+            transport,
+            client_options,
+            credentials,
+            metadata,
+        )
+
+    if nack_handler is None:
+        nack_handler = DefaultNackHandler()
+    if reassignment_handler is None:
+        reassignment_handler = DefaultReassignmentHandler()
+    if message_transformer is None:
+        message_transformer = MessageTransformer.of_callable(to_cps_subscribe_message)
+    partition_subscriber_factory = _make_partition_subscriber_factory(
+        subscription,
+        transport,
+        client_options,
+        credentials,
+        metadata,
+        per_partition_flow_control_settings,
+        nack_handler,
+        message_transformer,
+    )
+    return AssigningSingleSubscriber(
+        assigner_factory, partition_subscriber_factory, reassignment_handler
+    )

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/managed_event_loop.py

@@ -0,0 +1,89 @@
+# Copyright 2020 Google LLC
+#
+# 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 asyncio import AbstractEventLoop, new_event_loop, run_coroutine_threadsafe
+from concurrent.futures import Future
+from threading import Thread, Lock
+from typing import ContextManager, Generic, TypeVar, Optional, Callable
+
+_T = TypeVar("_T")
+
+
+class _Lazy(Generic[_T]):
+    _Factory = Callable[[], _T]
+
+    _lock: Lock
+    _factory: _Factory
+    _impl: Optional[_T]
+
+    def __init__(self, factory: _Factory):
+        self._lock = Lock()
+        self._factory = factory
+        self._impl = None
+
+    def get(self) -> _T:
+        with self._lock:
+            if self._impl is None:
+                self._impl = self._factory()
+            return self._impl
+
+
+class _ManagedEventLoopImpl(ContextManager):
+    _loop: AbstractEventLoop
+    _thread: Thread
+
+    def __init__(self, name=None):
+        self._loop = new_event_loop()
+        self._thread = Thread(
+            target=lambda: self._loop.run_forever(), name=name, daemon=True
+        )
+
+    def __enter__(self):
+        self._thread.start()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self._loop.call_soon_threadsafe(self._loop.stop)
+        self._thread.join()
+
+    def submit(self, coro) -> Future:
+        return run_coroutine_threadsafe(coro, self._loop)
+
+
+# TODO(user): Remove when underlying issue is fixed.
+# This is a workaround for https://github.com/grpc/grpc/issues/25364, a grpc
+# issue which prevents grpc-asyncio working with multiple event loops in the
+# same process. This workaround enables multiple topic publishing as well as
+# publish/subscribe from the same process, but does not enable use with other
+# grpc-asyncio clients. Once this issue is fixed, roll back the PR which
+# introduced this to return to a single event loop per client for isolation.
+_global_event_loop: _Lazy[_ManagedEventLoopImpl] = _Lazy(
+    lambda: _ManagedEventLoopImpl(name="PubSubLiteEventLoopThread").__enter__()
+)
+
+
+class ManagedEventLoop(ContextManager):
+    _loop: _ManagedEventLoopImpl
+
+    def __init__(self, name=None):
+        self._loop = _global_event_loop.get()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        pass
+
+    def submit(self, coro) -> Future:
+        return self._loop.submit(coro)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/multiplexed_async_publisher_client.py

@@ -0,0 +1,73 @@
+# Copyright 2020 Google LLC
+#
+# 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 typing import Callable, Union, Mapping
+
+from google.api_core.exceptions import GoogleAPICallError
+
+from google.cloud.pubsublite.cloudpubsub.internal.client_multiplexer import (
+    AsyncClientMultiplexer,
+)
+from google.cloud.pubsublite.cloudpubsub.internal.single_publisher import (
+    AsyncSinglePublisher,
+)
+from google.cloud.pubsublite.cloudpubsub.publisher_client_interface import (
+    AsyncPublisherClientInterface,
+)
+from google.cloud.pubsublite.types import TopicPath
+
+
+AsyncPublisherFactory = Callable[[TopicPath], AsyncSinglePublisher]
+
+
+class MultiplexedAsyncPublisherClient(AsyncPublisherClientInterface):
+    _publisher_factory: AsyncPublisherFactory
+    _multiplexer: AsyncClientMultiplexer[TopicPath, AsyncSinglePublisher]
+
+    def __init__(self, publisher_factory: AsyncPublisherFactory):
+        self._publisher_factory = publisher_factory
+        self._multiplexer = AsyncClientMultiplexer(
+            lambda topic: self._create_and_open(topic)
+        )
+
+    async def _create_and_open(self, topic: TopicPath):
+        client = self._publisher_factory(topic)
+        await client.__aenter__()
+        return client
+
+    async def publish(
+        self,
+        topic: Union[TopicPath, str],
+        data: bytes,
+        ordering_key: str = "",
+        **attrs: Mapping[str, str]
+    ) -> str:
+        if isinstance(topic, str):
+            topic = TopicPath.parse(topic)
+
+        publisher = await self._multiplexer.get_or_create(topic)
+        try:
+            return await publisher.publish(
+                data=data, ordering_key=ordering_key, **attrs
+            )
+        except GoogleAPICallError as e:
+            await self._multiplexer.try_erase(topic, publisher)
+            raise e
+
+    async def __aenter__(self):
+        await self._multiplexer.__aenter__()
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        await self._multiplexer.__aexit__(exc_type, exc_value, traceback)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/multiplexed_async_subscriber_client.py

@@ -0,0 +1,92 @@
+# Copyright 2020 Google LLC
+#
+# 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 typing import (
+    Union,
+    AsyncIterator,
+    Awaitable,
+    Callable,
+    Optional,
+    Set,
+)
+
+from google.cloud.pubsub_v1.subscriber.message import Message
+
+from google.cloud.pubsublite.cloudpubsub.internal.single_subscriber import (
+    AsyncSubscriberFactory,
+    AsyncSingleSubscriber,
+)
+from google.cloud.pubsublite.cloudpubsub.subscriber_client_interface import (
+    AsyncSubscriberClientInterface,
+)
+from google.cloud.pubsublite.types import (
+    SubscriptionPath,
+    FlowControlSettings,
+    Partition,
+)
+
+
+async def _iterate_subscriber(
+    subscriber: AsyncSingleSubscriber, on_failure: Callable[[], Awaitable[None]]
+) -> AsyncIterator[Message]:
+    try:
+        while True:
+            batch = await subscriber.read()
+            for message in batch:
+                yield message
+    except:  # noqa: E722
+        await on_failure()
+        raise
+
+
+class MultiplexedAsyncSubscriberClient(AsyncSubscriberClientInterface):
+    _underlying_factory: AsyncSubscriberFactory
+    _live_clients: Set[AsyncSingleSubscriber]
+
+    def __init__(self, underlying_factory: AsyncSubscriberFactory):
+        self._underlying_factory = underlying_factory
+        self._live_clients = set()
+
+    async def subscribe(
+        self,
+        subscription: Union[SubscriptionPath, str],
+        per_partition_flow_control_settings: FlowControlSettings,
+        fixed_partitions: Optional[Set[Partition]] = None,
+    ) -> AsyncIterator[Message]:
+        if isinstance(subscription, str):
+            subscription = SubscriptionPath.parse(subscription)
+
+        subscriber = self._underlying_factory(
+            subscription, fixed_partitions, per_partition_flow_control_settings
+        )
+        await subscriber.__aenter__()
+        self._live_clients.add(subscriber)
+
+        return _iterate_subscriber(
+            subscriber, lambda: self._try_remove_client(subscriber)
+        )
+
+    async def __aenter__(self):
+        return self
+
+    async def _try_remove_client(self, client: AsyncSingleSubscriber):
+        if client in self._live_clients:
+            self._live_clients.remove(client)
+            await client.__aexit__(None, None, None)
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        live_clients = self._live_clients
+        self._live_clients = set()
+        for client in live_clients:
+            await client.__aexit__(None, None, None)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/multiplexed_publisher_client.py

@@ -0,0 +1,86 @@
+# Copyright 2020 Google LLC
+#
+# 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 concurrent.futures import Future
+from typing import Callable, Union, Mapping
+
+from google.api_core.exceptions import GoogleAPICallError
+
+from google.cloud.pubsublite.cloudpubsub.internal.client_multiplexer import (
+    ClientMultiplexer,
+)
+from google.cloud.pubsublite.cloudpubsub.internal.single_publisher import (
+    SinglePublisher,
+)
+from google.cloud.pubsublite.cloudpubsub.publisher_client_interface import (
+    PublisherClientInterface,
+)
+from google.cloud.pubsublite.types import TopicPath
+
+PublisherFactory = Callable[[TopicPath], SinglePublisher]
+
+
+class MultiplexedPublisherClient(PublisherClientInterface):
+    _publisher_factory: PublisherFactory
+    _multiplexer: ClientMultiplexer[TopicPath, SinglePublisher]
+
+    def __init__(self, publisher_factory: PublisherFactory):
+        self._publisher_factory = publisher_factory
+        self._multiplexer = ClientMultiplexer(
+            lambda topic: self._create_and_start_publisher(topic)
+        )
+
+    def publish(
+        self,
+        topic: Union[TopicPath, str],
+        data: bytes,
+        ordering_key: str = "",
+        **attrs: Mapping[str, str]
+    ) -> "Future[str]":
+        if isinstance(topic, str):
+            topic = TopicPath.parse(topic)
+        try:
+            publisher = self._multiplexer.get_or_create(topic)
+        except GoogleAPICallError as e:
+            failed = Future()
+            failed.set_exception(e)
+            return failed
+        future = publisher.publish(data=data, ordering_key=ordering_key, **attrs)
+        future.add_done_callback(
+            lambda fut: self._on_future_completion(topic, publisher, fut)
+        )
+        return future
+
+    def _create_and_start_publisher(self, topic: Union[TopicPath, str]):
+        publisher = self._publisher_factory(topic)
+        try:
+            return publisher.__enter__()
+        except GoogleAPICallError:
+            publisher.__exit__(None, None, None)
+            raise
+
+    def _on_future_completion(
+        self, topic: TopicPath, publisher: SinglePublisher, future: "Future[str]"
+    ):
+        try:
+            future.result()
+        except GoogleAPICallError:
+            self._multiplexer.try_erase(topic, publisher)
+
+    def __enter__(self):
+        self._multiplexer.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self._multiplexer.__exit__(exc_type, exc_value, traceback)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/multiplexed_subscriber_client.py

@@ -0,0 +1,95 @@
+# Copyright 2020 Google LLC
+#
+# 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 concurrent.futures.thread import ThreadPoolExecutor
+from typing import Union, Optional, Set
+from threading import Lock
+
+from google.cloud.pubsub_v1.subscriber.futures import StreamingPullFuture
+
+from google.cloud.pubsublite.cloudpubsub.internal.single_subscriber import (
+    AsyncSubscriberFactory,
+)
+from google.cloud.pubsublite.cloudpubsub.internal.subscriber_impl import SubscriberImpl
+from google.cloud.pubsublite.cloudpubsub.subscriber_client_interface import (
+    SubscriberClientInterface,
+    MessageCallback,
+)
+from google.cloud.pubsublite.types import (
+    SubscriptionPath,
+    FlowControlSettings,
+    Partition,
+)
+
+
+class MultiplexedSubscriberClient(SubscriberClientInterface):
+    _executor: ThreadPoolExecutor
+    _underlying_factory: AsyncSubscriberFactory
+
+    _lock: Lock
+    _live_clients: Set[StreamingPullFuture]
+
+    def __init__(
+        self, executor: ThreadPoolExecutor, underlying_factory: AsyncSubscriberFactory
+    ):
+        self._executor = executor
+        self._underlying_factory = underlying_factory
+        self._lock = Lock()
+        self._live_clients = set()
+
+    def subscribe(
+        self,
+        subscription: Union[SubscriptionPath, str],
+        callback: MessageCallback,
+        per_partition_flow_control_settings: FlowControlSettings,
+        fixed_partitions: Optional[Set[Partition]] = None,
+    ) -> StreamingPullFuture:
+        if isinstance(subscription, str):
+            subscription = SubscriptionPath.parse(subscription)
+
+        underlying = self._underlying_factory(
+            subscription, fixed_partitions, per_partition_flow_control_settings
+        )
+        subscriber = SubscriberImpl(underlying, callback, self._executor)
+        future = StreamingPullFuture(subscriber)
+        subscriber.__enter__()
+        future.add_done_callback(lambda fut: self._try_remove_client(future))
+        return future
+
+    @staticmethod
+    def _cancel_streaming_pull_future(fut: StreamingPullFuture):
+        try:
+            fut.cancel()
+            fut.result()
+        except:  # noqa: E722
+            pass
+
+    def _try_remove_client(self, future: StreamingPullFuture):
+        with self._lock:
+            if future not in self._live_clients:
+                return
+            self._live_clients.remove(future)
+        self._cancel_streaming_pull_future(future)
+
+    def __enter__(self):
+        self._executor.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        with self._lock:
+            live_clients = self._live_clients
+            self._live_clients = set()
+        for client in live_clients:
+            self._cancel_streaming_pull_future(client)
+        self._executor.__exit__(exc_type, exc_value, traceback)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/publisher_impl.py

@@ -0,0 +1,52 @@
+# Copyright 2020 Google LLC
+#
+# 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 concurrent.futures import Future
+from typing import Mapping
+
+from google.cloud.pubsublite.cloudpubsub.internal.managed_event_loop import (
+    ManagedEventLoop,
+)
+from google.cloud.pubsublite.cloudpubsub.internal.single_publisher import (
+    SinglePublisher,
+    AsyncSinglePublisher,
+)
+
+
+class SinglePublisherImpl(SinglePublisher):
+    _managed_loop: ManagedEventLoop
+    _underlying: AsyncSinglePublisher
+
+    def __init__(self, underlying: AsyncSinglePublisher):
+        super().__init__()
+        self._managed_loop = ManagedEventLoop("PublisherLoopThread")
+        self._underlying = underlying
+
+    def publish(
+        self, data: bytes, ordering_key: str = "", **attrs: Mapping[str, str]
+    ) -> "Future[str]":
+        return self._managed_loop.submit(
+            self._underlying.publish(data=data, ordering_key=ordering_key, **attrs)
+        )
+
+    def __enter__(self):
+        self._managed_loop.__enter__()
+        self._managed_loop.submit(self._underlying.__aenter__()).result()
+        return self
+
+    def __exit__(self, __exc_type, __exc_value, __traceback):
+        self._managed_loop.submit(
+            self._underlying.__aexit__(__exc_type, __exc_value, __traceback)
+        ).result()
+        self._managed_loop.__exit__(__exc_type, __exc_value, __traceback)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/single_partition_subscriber.py

@@ -0,0 +1,156 @@
+# Copyright 2020 Google LLC
+#
+# 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 asyncio
+from typing import Callable, List, Dict, NamedTuple
+
+from google.api_core.exceptions import GoogleAPICallError
+from google.cloud.pubsub_v1.subscriber.message import Message
+from google.pubsub_v1 import PubsubMessage
+
+from google.cloud.pubsublite.internal.wire.permanent_failable import adapt_error
+from google.cloud.pubsublite.types import FlowControlSettings
+from google.cloud.pubsublite.cloudpubsub.internal.ack_set_tracker import AckSetTracker
+from google.cloud.pubsublite.cloudpubsub.internal.wrapped_message import (
+    AckId,
+    WrappedMessage,
+)
+from google.cloud.pubsublite.cloudpubsub.message_transformer import MessageTransformer
+from google.cloud.pubsublite.cloudpubsub.nack_handler import NackHandler
+from google.cloud.pubsublite.cloudpubsub.internal.single_subscriber import (
+    AsyncSingleSubscriber,
+)
+from google.cloud.pubsublite.internal.wire.permanent_failable import PermanentFailable
+from google.cloud.pubsublite.internal.wire.subscriber import Subscriber
+from google.cloud.pubsublite.internal.wire.subscriber_reset_handler import (
+    SubscriberResetHandler,
+)
+from google.cloud.pubsublite_v1 import FlowControlRequest, SequencedMessage
+
+
+class _SizedMessage(NamedTuple):
+    message: PubsubMessage
+    size_bytes: int
+
+
+ResettableSubscriberFactory = Callable[[SubscriberResetHandler], Subscriber]
+
+
+class SinglePartitionSingleSubscriber(
+    PermanentFailable, AsyncSingleSubscriber, SubscriberResetHandler
+):
+    _underlying: Subscriber
+    _flow_control_settings: FlowControlSettings
+    _ack_set_tracker: AckSetTracker
+    _nack_handler: NackHandler
+    _transformer: MessageTransformer
+
+    _ack_generation_id: int
+    _messages_by_ack_id: Dict[AckId, _SizedMessage]
+
+    _loop: asyncio.AbstractEventLoop
+
+    def __init__(
+        self,
+        subscriber_factory: ResettableSubscriberFactory,
+        flow_control_settings: FlowControlSettings,
+        ack_set_tracker: AckSetTracker,
+        nack_handler: NackHandler,
+        transformer: MessageTransformer,
+    ):
+        super().__init__()
+        self._underlying = subscriber_factory(self)
+        self._flow_control_settings = flow_control_settings
+        self._ack_set_tracker = ack_set_tracker
+        self._nack_handler = nack_handler
+        self._transformer = transformer
+
+        self._ack_generation_id = 0
+        self._messages_by_ack_id = {}
+
+    async def handle_reset(self):
+        # Increment ack generation id to ignore unacked messages.
+        self._ack_generation_id += 1
+        await self._ack_set_tracker.clear_and_commit()
+
+    def _wrap_message(self, message: SequencedMessage.meta.pb) -> Message:
+        # Rewrap in the proto-plus-python wrapper for passing to the transform
+        rewrapped = SequencedMessage()
+        rewrapped._pb = message
+        cps_message = self._transformer.transform(rewrapped)
+        offset = message.cursor.offset
+        ack_id = AckId(self._ack_generation_id, offset)
+        self._ack_set_tracker.track(offset)
+        self._messages_by_ack_id[ack_id] = _SizedMessage(
+            cps_message, message.size_bytes
+        )
+        wrapped_message = WrappedMessage(
+            pb=cps_message._pb,
+            ack_id=ack_id,
+            ack_handler=lambda id, ack: self._on_ack_threadsafe(id, ack),
+        )
+        return wrapped_message
+
+    def _on_ack_threadsafe(self, ack_id: AckId, should_ack: bool) -> None:
+        """A function called when a message is acked, may happen from any thread."""
+        if should_ack:
+            self._loop.call_soon_threadsafe(lambda: self._handle_ack(ack_id))
+            return
+        try:
+            sized_message = self._messages_by_ack_id[ack_id]
+            # Call the threadsafe version on ack since the callback may be called from another thread.
+            self._nack_handler.on_nack(
+                sized_message.message, lambda: self._on_ack_threadsafe(ack_id, True)
+            )
+        except Exception as e:
+            e2 = adapt_error(e)
+            self._loop.call_soon_threadsafe(lambda: self.fail(e2))
+
+    async def read(self) -> List[Message]:
+        try:
+            latest_batch = await self.await_unless_failed(self._underlying.read())
+            return [self._wrap_message(message) for message in latest_batch]
+        except Exception as e:
+            e = adapt_error(e)  # This could be from user code
+            self.fail(e)
+            raise e
+
+    def _handle_ack(self, ack_id: AckId):
+        flow_control = FlowControlRequest()
+        flow_control._pb.allowed_messages = 1
+        flow_control._pb.allowed_bytes = self._messages_by_ack_id[ack_id].size_bytes
+        self._underlying.allow_flow(flow_control)
+        del self._messages_by_ack_id[ack_id]
+        # Always refill flow control tokens, but do not commit offsets from outdated generations.
+        if ack_id.generation == self._ack_generation_id:
+            try:
+                self._ack_set_tracker.ack(ack_id.offset)
+            except GoogleAPICallError as e:
+                self.fail(e)
+
+    async def __aenter__(self):
+        self._loop = asyncio.get_event_loop()
+        await self._ack_set_tracker.__aenter__()
+        await self._underlying.__aenter__()
+        self._underlying.allow_flow(
+            FlowControlRequest(
+                allowed_messages=self._flow_control_settings.messages_outstanding,
+                allowed_bytes=self._flow_control_settings.bytes_outstanding,
+            )
+        )
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        await self._underlying.__aexit__(exc_type, exc_value, traceback)
+        await self._ack_set_tracker.__aexit__(exc_type, exc_value, traceback)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/single_publisher.py

@@ -0,0 +1,74 @@
+# Copyright 2020 Google LLC
+#
+# 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 abc import abstractmethod, ABCMeta
+from typing import AsyncContextManager, Mapping, ContextManager
+from concurrent import futures
+
+
+class AsyncSinglePublisher(AsyncContextManager, metaclass=ABCMeta):
+    """
+    An AsyncPublisher publishes messages similar to Google Pub/Sub, but must be used in an
+    async context. Any publish failures are permanent.
+
+    Must be used in an `async with` block or have __aenter__() awaited before use.
+    """
+
+    @abstractmethod
+    async def publish(
+        self, data: bytes, ordering_key: str = "", **attrs: Mapping[str, str]
+    ) -> str:
+        """
+        Publish a message.
+
+        Args:
+          data: The bytestring payload of the message
+          ordering_key: The key to enforce ordering on, or "" for no ordering.
+          **attrs: Additional attributes to send.
+
+        Returns:
+          An ack id, which can be decoded using MessageMetadata.decode.
+
+        Raises:
+          GoogleApiCallError: On a permanent failure.
+        """
+        raise NotImplementedError()
+
+
+class SinglePublisher(ContextManager, metaclass=ABCMeta):
+    """
+    A Publisher publishes messages similar to Google Pub/Sub. Any publish failures are permanent.
+
+    Must be used in a `with` block or have __enter__() called before use.
+    """
+
+    @abstractmethod
+    def publish(
+        self, data: bytes, ordering_key: str = "", **attrs: Mapping[str, str]
+    ) -> "futures.Future[str]":
+        """
+        Publish a message.
+
+        Args:
+          data: The bytestring payload of the message
+          ordering_key: The key to enforce ordering on, or "" for no ordering.
+          **attrs: Additional attributes to send.
+
+        Returns:
+          A future completed with an ack id, which can be decoded using MessageMetadata.decode.
+
+        Raises:
+          GoogleApiCallError: On a permanent failure.
+        """
+        raise NotImplementedError()

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/single_subscriber.py

@@ -0,0 +1,55 @@
+# Copyright 2020 Google LLC
+#
+# 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 abc import abstractmethod, ABCMeta
+from typing import AsyncContextManager, Callable, List, Set, Optional
+
+from google.cloud.pubsub_v1.subscriber.message import Message
+
+from google.cloud.pubsublite.types import (
+    SubscriptionPath,
+    FlowControlSettings,
+    Partition,
+)
+
+
+class AsyncSingleSubscriber(AsyncContextManager, metaclass=ABCMeta):
+    """
+    A Cloud Pub/Sub asynchronous subscriber.
+
+    Must be used in an `async with` block or have __aenter__() awaited before use.
+    """
+
+    @abstractmethod
+    async def read(self) -> List[Message]:
+        """
+        Read the next batch off of the stream.
+
+        Returns:
+          The next batch of messages. ack() or nack() must eventually be called
+          exactly once on each message.
+
+          Pub/Sub Lite does not support nack() by default- if you do call nack(), it will immediately fail the client
+          unless you have a NackHandler installed.
+
+        Raises:
+          GoogleAPICallError: On a permanent error.
+        """
+        raise NotImplementedError()
+
+
+AsyncSubscriberFactory = Callable[
+    [SubscriptionPath, Optional[Set[Partition]], FlowControlSettings],
+    AsyncSingleSubscriber,
+]

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/sorted_list.py

@@ -0,0 +1,25 @@
+from typing import Generic, TypeVar, List, Optional
+import heapq
+
+_T = TypeVar("_T")
+
+
+class SortedList(Generic[_T]):
+    _vals: List[_T]
+
+    def __init__(self):
+        self._vals = []
+
+    def push(self, val: _T):
+        heapq.heappush(self._vals, val)
+
+    def peek(self) -> Optional[_T]:
+        if self.empty():
+            return None
+        return self._vals[0]
+
+    def pop(self):
+        heapq.heappop(self._vals)
+
+    def empty(self) -> bool:
+        return not bool(self._vals)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/streaming_pull_manager.py

@@ -0,0 +1,33 @@
+# Copyright 2020 Google LLC
+#
+# 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 abc import ABC, abstractmethod
+from typing import Optional, Callable
+
+from google.api_core.exceptions import GoogleAPICallError
+
+
+CloseCallback = Callable[["StreamingPullManager", Optional[GoogleAPICallError]], None]
+
+
+class StreamingPullManager(ABC):
+    """The API expected by StreamingPullFuture."""
+
+    @abstractmethod
+    def add_close_callback(self, close_callback: CloseCallback):
+        pass
+
+    @abstractmethod
+    def close(self):
+        pass

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/subscriber_impl.py

@@ -0,0 +1,115 @@
+# Copyright 2020 Google LLC
+#
+# 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 concurrent.futures
+import threading
+from concurrent.futures.thread import ThreadPoolExecutor
+from typing import ContextManager, Optional
+from google.api_core.exceptions import GoogleAPICallError
+from functools import partial
+
+from google.cloud.pubsublite.internal.wait_ignore_cancelled import wait_ignore_errors
+from google.cloud.pubsublite.cloudpubsub.internal.managed_event_loop import (
+    ManagedEventLoop,
+)
+from google.cloud.pubsublite.cloudpubsub.internal.streaming_pull_manager import (
+    StreamingPullManager,
+    CloseCallback,
+)
+from google.cloud.pubsublite.cloudpubsub.internal.single_subscriber import (
+    AsyncSingleSubscriber,
+)
+from google.cloud.pubsublite.cloudpubsub.subscriber_client_interface import (
+    MessageCallback,
+)
+
+
+class SubscriberImpl(ContextManager, StreamingPullManager):
+    _underlying: AsyncSingleSubscriber
+    _callback: MessageCallback
+    _unowned_executor: ThreadPoolExecutor
+
+    _event_loop: ManagedEventLoop
+
+    _poller_future: concurrent.futures.Future
+    _close_lock: threading.Lock
+    _failure: Optional[GoogleAPICallError]
+    _close_callback: Optional[CloseCallback]
+    _closed: bool
+
+    def __init__(
+        self,
+        underlying: AsyncSingleSubscriber,
+        callback: MessageCallback,
+        unowned_executor: ThreadPoolExecutor,
+    ):
+        self._underlying = underlying
+        self._callback = callback
+        self._unowned_executor = unowned_executor
+        self._event_loop = ManagedEventLoop("SubscriberLoopThread")
+        self._close_lock = threading.Lock()
+        self._failure = None
+        self._close_callback = None
+        self._closed = False
+
+    def add_close_callback(self, close_callback: CloseCallback):
+        """
+        A close callback must be set exactly once by the StreamingPullFuture managing this subscriber.
+
+        This two-phase init model is made necessary by the requirements of StreamingPullFuture.
+        """
+        with self._close_lock:
+            assert self._close_callback is None
+            self._close_callback = close_callback
+
+    def close(self):
+        with self._close_lock:
+            if self._closed:
+                return
+            self._closed = True
+        self.__exit__(None, None, None)
+
+    def _fail(self, error: GoogleAPICallError):
+        self._failure = error
+        self.close()
+
+    async def _poller(self):
+        try:
+            while True:
+                batch = await self._underlying.read()
+                self._unowned_executor.map(self._callback, batch)
+        except GoogleAPICallError as e:
+            self._unowned_executor.submit(partial(self._fail, e))
+
+    def __enter__(self):
+        assert self._close_callback is not None
+        self._event_loop.__enter__()
+        self._event_loop.submit(self._underlying.__aenter__()).result()
+        self._poller_future = self._event_loop.submit(self._poller())
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self._poller_future.cancel()
+        try:
+            self._poller_future.result()  # Ignore error.
+        except:  # noqa: E722
+            pass
+        self._event_loop.submit(
+            wait_ignore_errors(
+                self._underlying.__aexit__(exc_type, exc_value, traceback)
+            )
+        ).result()
+        self._event_loop.__exit__(exc_type, exc_value, traceback)
+        assert self._close_callback is not None
+        self._close_callback(self, self._failure)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/internal/wrapped_message.py

@@ -0,0 +1,64 @@
+from concurrent import futures
+import logging
+from typing import NamedTuple, Callable
+
+from google.cloud.pubsub_v1.subscriber.message import Message
+from google.pubsub_v1 import PubsubMessage
+from google.cloud.pubsub_v1.subscriber.exceptions import AcknowledgeStatus
+
+
+class AckId(NamedTuple):
+    generation: int
+    offset: int
+
+    def encode(self) -> str:
+        return str(self.generation) + "," + str(self.offset)
+
+
+_SUCCESS_FUTURE = futures.Future()
+_SUCCESS_FUTURE.set_result(AcknowledgeStatus.SUCCESS)
+
+
+class WrappedMessage(Message):
+    _id: AckId
+    _ack_handler: Callable[[AckId, bool], None]
+
+    def __init__(
+        self,
+        pb: PubsubMessage.meta.pb,
+        ack_id: AckId,
+        ack_handler: Callable[[AckId, bool], None],
+    ):
+        super().__init__(pb, ack_id.encode(), 1, None)
+        self._id = ack_id
+        self._ack_handler = ack_handler
+
+    def ack(self):
+        self._ack_handler(self._id, True)
+
+    def ack_with_response(self) -> "futures.Future":
+        self._ack_handler(self._id, True)
+        return _SUCCESS_FUTURE
+
+    def nack(self):
+        self._ack_handler(self._id, False)
+
+    def nack_with_response(self) -> "futures.Future":
+        self._ack_handler(self._id, False)
+        return _SUCCESS_FUTURE
+
+    def drop(self):
+        logging.warning(
+            "Likely incorrect call to drop() on Pub/Sub Lite message. Pub/Sub Lite does not support redelivery in this way."
+        )
+
+    def modify_ack_deadline(self, seconds: int):
+        logging.warning(
+            "Likely incorrect call to modify_ack_deadline() on Pub/Sub Lite message. Pub/Sub Lite does not support redelivery in this way."
+        )
+
+    def modify_ack_deadline_with_response(self, seconds: int) -> "futures.Future":
+        logging.warning(
+            "Likely incorrect call to modify_ack_deadline_with_response() on Pub/Sub Lite message. Pub/Sub Lite does not support redelivery in this way."
+        )
+        return _SUCCESS_FUTURE

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/message_transformer.py

@@ -0,0 +1,46 @@
+# Copyright 2020 Google LLC
+#
+# 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 abc import ABC, abstractmethod
+from typing import Callable
+
+from google.pubsub_v1 import PubsubMessage
+
+from google.cloud.pubsublite_v1 import SequencedMessage
+
+
+class MessageTransformer(ABC):
+    """
+    A MessageTransformer turns Pub/Sub Lite message protos into Pub/Sub message protos.
+    """
+
+    @abstractmethod
+    def transform(self, source: SequencedMessage) -> PubsubMessage:
+        """Transform a SequencedMessage to a PubsubMessage.
+
+        Args:
+          source: The message to transform.
+
+        Raises:
+          GoogleAPICallError: To fail the client if raised inline.
+        """
+        pass
+
+    @staticmethod
+    def of_callable(transformer: Callable[[SequencedMessage], PubsubMessage]):
+        class CallableTransformer(MessageTransformer):
+            def transform(self, source: SequencedMessage) -> PubsubMessage:
+                return transformer(source)
+
+        return CallableTransformer()

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/message_transforms.py

@@ -0,0 +1,143 @@
+# Copyright 2020 Google LLC
+#
+# 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 datetime
+
+from google.api_core.exceptions import InvalidArgument
+from cloudsdk.google.protobuf.timestamp_pb2 import Timestamp  # pytype: disable=pyi-error
+from google.pubsub_v1 import PubsubMessage
+
+from google.cloud.pubsublite.cloudpubsub import MessageTransformer
+from google.cloud.pubsublite.internal import fast_serialize
+from google.cloud.pubsublite.types import Partition, MessageMetadata
+from google.cloud.pubsublite_v1 import AttributeValues, SequencedMessage, PubSubMessage
+
+PUBSUB_LITE_EVENT_TIME = "x-goog-pubsublite-event-time"
+
+
+def _encode_attribute_event_time_proto(ts: Timestamp) -> str:
+    return fast_serialize.dump([ts.seconds, ts.nanos])
+
+
+def _decode_attribute_event_time_proto(attr: str) -> Timestamp:
+    try:
+        ts = Timestamp()
+        loaded = fast_serialize.load(attr)
+        ts.seconds = loaded[0]
+        ts.nanos = loaded[1]
+        return ts
+    except Exception:  # noqa: E722
+        raise InvalidArgument("Invalid value for event time attribute.")
+
+
+def encode_attribute_event_time(dt: datetime.datetime) -> str:
+    ts = Timestamp()
+    ts.FromDatetime(dt.astimezone(datetime.timezone.utc))
+    return _encode_attribute_event_time_proto(ts)
+
+
+def decode_attribute_event_time(attr: str) -> datetime.datetime:
+    return (
+        _decode_attribute_event_time_proto(attr)
+        .ToDatetime()
+        .replace(tzinfo=datetime.timezone.utc)
+    )
+
+
+def _parse_attributes(values: AttributeValues) -> str:
+    if not len(values.values) == 1:
+        raise InvalidArgument(
+            "Received an unparseable message with multiple values for an attribute."
+        )
+    value: bytes = values.values[0]
+    try:
+        return value.decode("utf-8")
+    except UnicodeError:
+        raise InvalidArgument(
+            "Received an unparseable message with a non-utf8 attribute."
+        )
+
+
+def add_id_to_cps_subscribe_transformer(
+    partition: Partition, transformer: MessageTransformer
+) -> MessageTransformer:
+    def add_id_to_message(source: SequencedMessage):
+        source_pb = source._pb
+        message: PubsubMessage = transformer.transform(source)
+        message_pb = message._pb
+        if message_pb.message_id:
+            raise InvalidArgument(
+                "Message after transforming has the message_id field set."
+            )
+        message_pb.message_id = MessageMetadata._encode_parts(
+            partition.value, source_pb.cursor.offset
+        )
+        return message
+
+    return MessageTransformer.of_callable(add_id_to_message)
+
+
+def to_cps_subscribe_message(source: SequencedMessage) -> PubsubMessage:
+    source_pb = source._pb
+    out_pb = _to_cps_publish_message_proto(source_pb.message)
+    out_pb.publish_time.CopyFrom(source_pb.publish_time)
+    out = PubsubMessage()
+    out._pb = out_pb
+    return out
+
+
+def _to_cps_publish_message_proto(
+    source: PubSubMessage.meta.pb,
+) -> PubsubMessage.meta.pb:
+    out = PubsubMessage.meta.pb()
+    try:
+        out.ordering_key = source.key.decode("utf-8")
+    except UnicodeError:
+        raise InvalidArgument("Received an unparseable message with a non-utf8 key.")
+    if PUBSUB_LITE_EVENT_TIME in source.attributes:
+        raise InvalidArgument(
+            "Special timestamp attribute exists in wire message. Unable to parse message."
+        )
+    out.data = source.data
+    for key, values in source.attributes.items():
+        out.attributes[key] = _parse_attributes(values)
+    if source.HasField("event_time"):
+        out.attributes[PUBSUB_LITE_EVENT_TIME] = _encode_attribute_event_time_proto(
+            source.event_time
+        )
+    return out
+
+
+def to_cps_publish_message(source: PubSubMessage) -> PubsubMessage:
+    out = PubsubMessage()
+    out._pb = _to_cps_publish_message_proto(source._pb)
+    return out
+
+
+def from_cps_publish_message(source: PubsubMessage) -> PubSubMessage:
+    source_pb = source._pb
+    out = PubSubMessage()
+    out_pb = out._pb
+    if PUBSUB_LITE_EVENT_TIME in source_pb.attributes:
+        out_pb.event_time.CopyFrom(
+            _decode_attribute_event_time_proto(
+                source_pb.attributes[PUBSUB_LITE_EVENT_TIME]
+            )
+        )
+    out_pb.data = source_pb.data
+    out_pb.key = source_pb.ordering_key.encode("utf-8")
+    for key, value in source_pb.attributes.items():
+        if key != PUBSUB_LITE_EVENT_TIME:
+            out_pb.attributes[key].values.append(value.encode("utf-8"))
+    return out

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/nack_handler.py

@@ -0,0 +1,48 @@
+# Copyright 2020 Google LLC
+#
+# 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 abc import ABC, abstractmethod
+from typing import Callable
+
+from google.api_core.exceptions import FailedPrecondition
+from google.pubsub_v1 import PubsubMessage
+
+
+class NackHandler(ABC):
+    """
+    A NackHandler handles calls to the nack() method which is not expressible in Pub/Sub Lite.
+    """
+
+    @abstractmethod
+    def on_nack(self, message: PubsubMessage, ack: Callable[[], None]):
+        """Handle a negative acknowledgement. ack must eventually be called.
+
+        This method will be called on an event loop and should not block.
+
+        Args:
+          message: The nacked message.
+          ack: A callable to acknowledge the underlying message. This must eventually be called.
+
+        Raises:
+          GoogleAPICallError: To fail the client if raised inline.
+        """
+        pass
+
+
+class DefaultNackHandler(NackHandler):
+    def on_nack(self, message: PubsubMessage, ack: Callable[[], None]):
+        raise FailedPrecondition(
+            "You may not nack messages by default when using a PubSub Lite client. See NackHandler for how to customize"
+            " this."
+        )

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/publisher_client.py

@@ -0,0 +1,191 @@
+# Copyright 2020 Google LLC
+#
+# 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 concurrent.futures import Future
+from typing import Optional, Mapping, Union
+from uuid import uuid4
+
+from google.api_core.client_options import ClientOptions
+from google.auth.credentials import Credentials
+from google.cloud.pubsub_v1.types import BatchSettings
+
+from google.cloud.pubsublite.cloudpubsub.internal.make_publisher import (
+    make_publisher,
+    make_async_publisher,
+)
+from google.cloud.pubsublite.cloudpubsub.internal.multiplexed_async_publisher_client import (
+    MultiplexedAsyncPublisherClient,
+)
+from google.cloud.pubsublite.cloudpubsub.internal.multiplexed_publisher_client import (
+    MultiplexedPublisherClient,
+)
+from google.cloud.pubsublite.cloudpubsub.publisher_client_interface import (
+    PublisherClientInterface,
+    AsyncPublisherClientInterface,
+)
+from google.cloud.pubsublite.internal.constructable_from_service_account import (
+    ConstructableFromServiceAccount,
+)
+from google.cloud.pubsublite.internal.publisher_client_id import PublisherClientId
+from google.cloud.pubsublite.internal.require_started import RequireStarted
+from google.cloud.pubsublite.internal.wire.make_publisher import (
+    DEFAULT_BATCHING_SETTINGS as WIRE_DEFAULT_BATCHING,
+)
+from google.cloud.pubsublite.types import TopicPath
+
+
+def _get_client_id(enable_idempotence: bool):
+    return PublisherClientId(uuid4().bytes) if enable_idempotence else None
+
+
+class PublisherClient(PublisherClientInterface, ConstructableFromServiceAccount):
+    """
+    A PublisherClient publishes messages similar to Google Pub/Sub.
+    Any publish failures are unlikely to succeed if retried.
+
+    Must be used in a `with` block or have __enter__() called before use.
+    """
+
+    _impl: PublisherClientInterface
+    _require_started: RequireStarted
+
+    DEFAULT_BATCHING_SETTINGS = WIRE_DEFAULT_BATCHING
+    """
+    The default batching settings for a publisher client.
+    """
+
+    def __init__(
+        self,
+        *,
+        per_partition_batching_settings: Optional[BatchSettings] = None,
+        credentials: Optional[Credentials] = None,
+        transport: str = "grpc_asyncio",
+        client_options: Optional[ClientOptions] = None,
+        enable_idempotence: bool = False,
+    ):
+        """
+        Create a new PublisherClient.
+
+        Args:
+            per_partition_batching_settings: The settings for publish batching. Apply on a per-partition basis.
+            credentials: If provided, the credentials to use when connecting.
+            transport: The transport to use. Must correspond to an asyncio transport.
+            client_options: The client options to use when connecting. If used, must explicitly set `api_endpoint`.
+            enable_idempotence: Whether idempotence is enabled, where the server will ensure that unique messages within a single publisher session are stored only once.
+        """
+        client_id = _get_client_id(enable_idempotence)
+        self._impl = MultiplexedPublisherClient(
+            lambda topic: make_publisher(
+                topic=topic,
+                per_partition_batching_settings=per_partition_batching_settings,
+                credentials=credentials,
+                client_options=client_options,
+                transport=transport,
+                client_id=client_id,
+            )
+        )
+        self._require_started = RequireStarted()
+
+    def publish(
+        self,
+        topic: Union[TopicPath, str],
+        data: bytes,
+        ordering_key: str = "",
+        **attrs: Mapping[str, str],
+    ) -> "Future[str]":
+        self._require_started.require_started()
+        return self._impl.publish(
+            topic=topic, data=data, ordering_key=ordering_key, **attrs
+        )
+
+    def __enter__(self):
+        self._require_started.__enter__()
+        self._impl.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self._impl.__exit__(exc_type, exc_value, traceback)
+        self._require_started.__exit__(exc_type, exc_value, traceback)
+
+
+class AsyncPublisherClient(
+    AsyncPublisherClientInterface, ConstructableFromServiceAccount
+):
+    """
+    An AsyncPublisherClient publishes messages similar to Google Pub/Sub, but must be used in an
+    async context. Any publish failures are unlikely to succeed if retried.
+
+    Must be used in an `async with` block or have __aenter__() awaited before use.
+    """
+
+    _impl: AsyncPublisherClientInterface
+    _require_started: RequireStarted
+
+    DEFAULT_BATCHING_SETTINGS = WIRE_DEFAULT_BATCHING
+    """
+    The default batching settings for a publisher client.
+    """
+
+    def __init__(
+        self,
+        *,
+        per_partition_batching_settings: Optional[BatchSettings] = None,
+        credentials: Optional[Credentials] = None,
+        transport: str = "grpc_asyncio",
+        client_options: Optional[ClientOptions] = None,
+        enable_idempotence: bool = False,
+    ):
+        """
+        Create a new AsyncPublisherClient.
+
+        Args:
+            per_partition_batching_settings: The settings for publish batching. Apply on a per-partition basis.
+            credentials: If provided, the credentials to use when connecting.
+            transport: The transport to use. Must correspond to an asyncio transport.
+            client_options: The client options to use when connecting. If used, must explicitly set `api_endpoint`.
+            enable_idempotence: Whether idempotence is enabled, where the server will ensure that unique messages within a single publisher session are stored only once.
+        """
+        client_id = _get_client_id(enable_idempotence)
+        self._impl = MultiplexedAsyncPublisherClient(
+            lambda topic: make_async_publisher(
+                topic=topic,
+                per_partition_batching_settings=per_partition_batching_settings,
+                credentials=credentials,
+                client_options=client_options,
+                transport=transport,
+                client_id=client_id,
+            )
+        )
+        self._require_started = RequireStarted()
+
+    async def publish(
+        self,
+        topic: Union[TopicPath, str],
+        data: bytes,
+        ordering_key: str = "",
+        **attrs: Mapping[str, str],
+    ) -> str:
+        self._require_started.require_started()
+        return await self._impl.publish(
+            topic=topic, data=data, ordering_key=ordering_key, **attrs
+        )
+
+    async def __aenter__(self):
+        self._require_started.__enter__()
+        await self._impl.__aenter__()
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        await self._impl.__aexit__(exc_type, exc_value, traceback)
+        self._require_started.__exit__(exc_type, exc_value, traceback)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/publisher_client_interface.py

@@ -0,0 +1,88 @@
+# Copyright 2020 Google LLC
+#
+# 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 abc import abstractmethod, ABCMeta
+from concurrent.futures import Future
+from typing import ContextManager, Mapping, Union, AsyncContextManager
+
+from google.cloud.pubsublite.types import TopicPath
+
+
+class AsyncPublisherClientInterface(AsyncContextManager, metaclass=ABCMeta):
+    """
+    An AsyncPublisherClientInterface publishes messages similar to Google Pub/Sub, but must be used in an
+    async context. Any publish failures are unlikely to succeed if retried.
+
+    Must be used in an `async with` block or have __aenter__() awaited before use.
+    """
+
+    @abstractmethod
+    async def publish(
+        self,
+        topic: Union[TopicPath, str],
+        data: bytes,
+        ordering_key: str = "",
+        **attrs: Mapping[str, str],
+    ) -> str:
+        """
+        Publish a message.
+
+        Args:
+          topic: The topic to publish to. Publishes to new topics may have nontrivial startup latency.
+          data: The bytestring payload of the message
+          ordering_key: The key to enforce ordering on, or "" for no ordering.
+          **attrs: Additional attributes to send.
+
+        Returns:
+          An ack id, which can be decoded using MessageMetadata.decode.
+
+        Raises:
+          GoogleApiCallError: On a permanent failure.
+        """
+        raise NotImplementedError()
+
+
+class PublisherClientInterface(ContextManager, metaclass=ABCMeta):
+    """
+    A PublisherClientInterface publishes messages similar to Google Pub/Sub.
+    Any publish failures are unlikely to succeed if retried.
+
+    Must be used in a `with` block or have __enter__() called before use.
+    """
+
+    @abstractmethod
+    def publish(
+        self,
+        topic: Union[TopicPath, str],
+        data: bytes,
+        ordering_key: str = "",
+        **attrs: Mapping[str, str],
+    ) -> "Future[str]":
+        """
+        Publish a message.
+
+        Args:
+          topic: The topic to publish to. Publishes to new topics may have nontrivial startup latency.
+          data: The bytestring payload of the message
+          ordering_key: The key to enforce ordering on, or "" for no ordering.
+          **attrs: Additional attributes to send.
+
+        Returns:
+          A future completed with an ack id of type str, which can be decoded using
+          MessageMetadata.decode.
+
+        Raises:
+          GoogleApiCallError: On a permanent failure.
+        """
+        raise NotImplementedError()

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/reassignment_handler.py

@@ -0,0 +1,69 @@
+# Copyright 2020 Google LLC
+#
+# 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 abc import ABC, abstractmethod
+from typing import Set, Optional, Awaitable
+
+from google.cloud.pubsublite.types import Partition
+
+
+class ReassignmentHandler(ABC):
+    """
+    A ReassignmentHandler is called any time a new partition assignment is received from the server.
+    It will be called with both the previous and new assignments as decided by the backend.
+
+    The client library will not acknowledge the assignment until handleReassignment returns. The
+    assigning backend will not assign any of the partitions in `before` to another server unless the
+    assignment is acknowledged, or a client takes too long to acknowledged (currently 30 seconds from
+    the time the assignment is sent from server's point of view).
+
+    Because of the above, as long as reassignment handling is processed quickly, it can be used to
+    abort outstanding operations on partitions which are being assigned away from this client.
+    """
+
+    @abstractmethod
+    def handle_reassignment(
+        self, before: Set[Partition], after: Set[Partition]
+    ) -> Optional[Awaitable]:
+        """
+        Called with the previous and new assignment delivered to this client on an assignment change.
+        The assignment will not be acknowledged until this method returns, so it should complete
+        quickly, or the backend will assume it is non-responsive and assign all partitions away without
+        waiting for acknowledgement.
+
+        handle_reassignment will only be called after no new message deliveries will be started for the partition.
+        There may still be messages in flight on executors or in async callbacks.
+
+        Acks or nacks on messages from partitions being assigned away will have no effect.
+
+        This method will be called on an event loop and should not block.
+
+        Args:
+          before: The previous assignment.
+          after: The new assignment.
+
+        Returns:
+          Either None or an Awaitable to be waited on before acknowledging reassignment.
+
+        Raises:
+          GoogleAPICallError: To fail the client if raised.
+        """
+        pass
+
+
+class DefaultReassignmentHandler(ReassignmentHandler):
+    def handle_reassignment(
+        self, before: Set[Partition], after: Set[Partition]
+    ) -> Optional[Awaitable]:
+        return None

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/subscriber_client.py

@@ -0,0 +1,194 @@
+# Copyright 2020 Google LLC
+#
+# 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 concurrent.futures.thread import ThreadPoolExecutor
+from typing import Optional, Union, Set, AsyncIterator
+
+from google.api_core.client_options import ClientOptions
+from google.auth.credentials import Credentials
+from google.cloud.pubsub_v1.subscriber.futures import StreamingPullFuture
+from google.cloud.pubsub_v1.subscriber.message import Message
+
+from google.cloud.pubsublite.cloudpubsub.reassignment_handler import ReassignmentHandler
+from google.cloud.pubsublite.cloudpubsub.internal.make_subscriber import (
+    make_async_subscriber,
+)
+from google.cloud.pubsublite.cloudpubsub.internal.multiplexed_async_subscriber_client import (
+    MultiplexedAsyncSubscriberClient,
+)
+from google.cloud.pubsublite.cloudpubsub.internal.multiplexed_subscriber_client import (
+    MultiplexedSubscriberClient,
+)
+from google.cloud.pubsublite.cloudpubsub.message_transformer import MessageTransformer
+from google.cloud.pubsublite.cloudpubsub.nack_handler import NackHandler
+from google.cloud.pubsublite.cloudpubsub.subscriber_client_interface import (
+    SubscriberClientInterface,
+    AsyncSubscriberClientInterface,
+    MessageCallback,
+)
+from google.cloud.pubsublite.internal.constructable_from_service_account import (
+    ConstructableFromServiceAccount,
+)
+from google.cloud.pubsublite.internal.require_started import RequireStarted
+from google.cloud.pubsublite.types import (
+    FlowControlSettings,
+    Partition,
+    SubscriptionPath,
+)
+
+
+class SubscriberClient(SubscriberClientInterface, ConstructableFromServiceAccount):
+    """
+    A SubscriberClient reads messages similar to Google Pub/Sub.
+    Any subscribe failures are unlikely to succeed if retried.
+
+    Must be used in a `with` block or have __enter__() called before use.
+    """
+
+    _impl: SubscriberClientInterface
+    _require_started: RequireStarted
+
+    def __init__(
+        self,
+        *,
+        executor: Optional[ThreadPoolExecutor] = None,
+        nack_handler: Optional[NackHandler] = None,
+        reassignment_handler: Optional[ReassignmentHandler] = None,
+        message_transformer: Optional[MessageTransformer] = None,
+        credentials: Optional[Credentials] = None,
+        transport: str = "grpc_asyncio",
+        client_options: Optional[ClientOptions] = None,
+    ):
+        """
+        Create a new SubscriberClient.
+
+        Args:
+            executor: A ThreadPoolExecutor to use. The client will shut it down on __exit__. If provided a single threaded executor, messages will be ordered per-partition, but take care that the callback does not block for too long as it will impede forward progress on all subscriptions.
+            nack_handler: A handler for when `nack()` is called. The default NackHandler raises an exception and fails the subscribe stream.
+            message_transformer: A transformer from Pub/Sub Lite messages to Cloud Pub/Sub messages. This may not return a message with "message_id" set.
+            credentials: If provided, the credentials to use when connecting.
+            transport: The transport to use. Must correspond to an asyncio transport.
+            client_options: The client options to use when connecting. If used, must explicitly set `api_endpoint`.
+        """
+        if executor is None:
+            executor = ThreadPoolExecutor()
+        self._impl = MultiplexedSubscriberClient(
+            executor,
+            lambda subscription, partitions, settings: make_async_subscriber(
+                subscription=subscription,
+                transport=transport,
+                per_partition_flow_control_settings=settings,
+                nack_handler=nack_handler,
+                reassignment_handler=reassignment_handler,
+                message_transformer=message_transformer,
+                fixed_partitions=partitions,
+                credentials=credentials,
+                client_options=client_options,
+            ),
+        )
+        self._require_started = RequireStarted()
+
+    def subscribe(
+        self,
+        subscription: Union[SubscriptionPath, str],
+        callback: MessageCallback,
+        per_partition_flow_control_settings: FlowControlSettings,
+        fixed_partitions: Optional[Set[Partition]] = None,
+    ) -> StreamingPullFuture:
+        self._require_started.require_started()
+        return self._impl.subscribe(
+            subscription,
+            callback,
+            per_partition_flow_control_settings,
+            fixed_partitions,
+        )
+
+    def __enter__(self):
+        self._require_started.__enter__()
+        self._impl.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self._impl.__exit__(exc_type, exc_value, traceback)
+        self._require_started.__exit__(exc_type, exc_value, traceback)
+
+
+class AsyncSubscriberClient(
+    AsyncSubscriberClientInterface, ConstructableFromServiceAccount
+):
+    """
+    An AsyncSubscriberClient reads messages similar to Google Pub/Sub, but must be used in an
+    async context.
+    Any subscribe failures are unlikely to succeed if retried.
+
+    Must be used in an `async with` block or have __aenter__() awaited before use.
+    """
+
+    _impl: AsyncSubscriberClientInterface
+    _require_started: RequireStarted
+
+    def __init__(
+        self,
+        *,
+        nack_handler: Optional[NackHandler] = None,
+        reassignment_handler: Optional[ReassignmentHandler] = None,
+        message_transformer: Optional[MessageTransformer] = None,
+        credentials: Optional[Credentials] = None,
+        transport: str = "grpc_asyncio",
+        client_options: Optional[ClientOptions] = None,
+    ):
+        """
+        Create a new AsyncSubscriberClient.
+
+        Args:
+            nack_handler: A handler for when `nack()` is called. The default NackHandler raises an exception and fails the subscribe stream.
+            message_transformer: A transformer from Pub/Sub Lite messages to Cloud Pub/Sub messages. This may not return a message with "message_id" set.
+            credentials: If provided, the credentials to use when connecting.
+            transport: The transport to use. Must correspond to an asyncio transport.
+            client_options: The client options to use when connecting. If used, must explicitly set `api_endpoint`.
+        """
+        self._impl = MultiplexedAsyncSubscriberClient(
+            lambda subscription, partitions, settings: make_async_subscriber(
+                subscription=subscription,
+                transport=transport,
+                per_partition_flow_control_settings=settings,
+                nack_handler=nack_handler,
+                reassignment_handler=reassignment_handler,
+                message_transformer=message_transformer,
+                fixed_partitions=partitions,
+                credentials=credentials,
+                client_options=client_options,
+            )
+        )
+        self._require_started = RequireStarted()
+
+    async def subscribe(
+        self,
+        subscription: Union[SubscriptionPath, str],
+        per_partition_flow_control_settings: FlowControlSettings,
+        fixed_partitions: Optional[Set[Partition]] = None,
+    ) -> AsyncIterator[Message]:
+        self._require_started.require_started()
+        return await self._impl.subscribe(
+            subscription, per_partition_flow_control_settings, fixed_partitions
+        )
+
+    async def __aenter__(self):
+        self._require_started.__enter__()
+        await self._impl.__aenter__()
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        await self._impl.__aexit__(exc_type, exc_value, traceback)
+        self._require_started.__exit__(exc_type, exc_value, traceback)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/cloudpubsub/subscriber_client_interface.py

@@ -0,0 +1,107 @@
+# Copyright 2020 Google LLC
+#
+# 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 abc import abstractmethod, ABCMeta
+from typing import (
+    ContextManager,
+    Union,
+    AsyncContextManager,
+    AsyncIterator,
+    Callable,
+    Optional,
+    Set,
+)
+
+from google.cloud.pubsub_v1.subscriber.futures import StreamingPullFuture
+from google.cloud.pubsub_v1.subscriber.message import Message
+
+from google.cloud.pubsublite.types import (
+    SubscriptionPath,
+    FlowControlSettings,
+    Partition,
+)
+
+
+class AsyncSubscriberClientInterface(AsyncContextManager, metaclass=ABCMeta):
+    """
+    An AsyncSubscriberClientInterface reads messages similar to Google Pub/Sub, but must be used in an
+    async context.
+    Any subscribe failures are unlikely to succeed if retried.
+
+    Must be used in an `async with` block or have __aenter__() awaited before use.
+    """
+
+    @abstractmethod
+    async def subscribe(
+        self,
+        subscription: Union[SubscriptionPath, str],
+        per_partition_flow_control_settings: FlowControlSettings,
+        fixed_partitions: Optional[Set[Partition]] = None,
+    ) -> AsyncIterator[Message]:
+        """
+        Read messages from a subscription.
+
+        Args:
+          subscription: The subscription to subscribe to.
+          per_partition_flow_control_settings: The flow control settings for each partition subscribed to. Note that these
+              settings apply to each partition individually, not in aggregate.
+          fixed_partitions: A fixed set of partitions to subscribe to. If not present, will instead use auto-assignment.
+
+        Returns:
+          An AsyncIterator with Messages that must have ack() called on each exactly once.
+
+        Raises:
+          GoogleApiCallError: On a permanent failure.
+        """
+        raise NotImplementedError()
+
+
+MessageCallback = Callable[[Message], None]
+
+
+class SubscriberClientInterface(ContextManager, metaclass=ABCMeta):
+    """
+    A SubscriberClientInterface reads messages similar to Google Pub/Sub.
+    Any subscribe failures are unlikely to succeed if retried.
+
+    Must be used in a `with` block or have __enter__() called before use.
+    """
+
+    @abstractmethod
+    def subscribe(
+        self,
+        subscription: Union[SubscriptionPath, str],
+        callback: MessageCallback,
+        per_partition_flow_control_settings: FlowControlSettings,
+        fixed_partitions: Optional[Set[Partition]] = None,
+    ) -> StreamingPullFuture:
+        """
+        This method starts a background thread to begin pulling messages from
+        a Pub/Sub Lite subscription and scheduling them to be processed using the
+        provided ``callback``.
+
+        Args:
+          subscription: The subscription to subscribe to.
+          callback: The callback function. This function receives the message as its only argument.
+          per_partition_flow_control_settings: The flow control settings for each partition subscribed to. Note that these
+              settings apply to each partition individually, not in aggregate.
+          fixed_partitions: A fixed set of partitions to subscribe to. If not present, will instead use auto-assignment.
+
+        Returns:
+          A StreamingPullFuture instance that can be used to manage the background stream.
+
+        Raises:
+          GoogleApiCallError: On a permanent failure.
+        """
+        raise NotImplementedError()

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/gapic_version.py

@@ -0,0 +1,16 @@
+# -*- coding: utf-8 -*-
+# Copyright 2022 Google LLC
+#
+# 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.
+#
+__version__ = "1.8.3"  # {x-release-please-version}

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/constructable_from_service_account.py

@@ -0,0 +1,32 @@
+# Copyright 2020 Google LLC
+#
+# 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 google.oauth2 import service_account
+
+
+class ConstructableFromServiceAccount:
+    @classmethod
+    def from_service_account_file(cls, filename, **kwargs):
+        f"""Creates an instance of this client using the provided credentials file.
+        Args:
+            filename (str): The path to the service account private key json
+                file.
+            kwargs: Additional arguments to pass to the constructor.
+        Returns:
+            A {cls.__name__}.
+        """
+        credentials = service_account.Credentials.from_service_account_file(filename)
+        return cls(credentials=credentials, **kwargs)
+
+    from_service_account_json = from_service_account_file

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/endpoints.py

@@ -0,0 +1,19 @@
+# Copyright 2020 Google LLC
+#
+# 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 google.cloud.pubsublite.types import CloudRegion
+
+
+def regional_endpoint(region: CloudRegion):
+    return f"dns:///{region}-pubsublite.googleapis.com"

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/fast_serialize.py

@@ -0,0 +1,13 @@
+"""
+A fast serialization method for lists of integers.
+"""
+
+from typing import List
+
+
+def dump(data: List[int]) -> str:
+    return ",".join(str(x) for x in data)
+
+
+def load(source: str) -> List[int]:
+    return [int(x) for x in source.split(",")]

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/publish_sequence_number.py

@@ -0,0 +1,22 @@
+# Copyright 2023 Google LLC
+#
+# 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 typing import NamedTuple
+
+
+class PublishSequenceNumber(NamedTuple):
+    value: int
+
+    def next(self) -> "PublishSequenceNumber":
+        return PublishSequenceNumber(self.value + 1)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/publisher_client_id.py

@@ -0,0 +1,19 @@
+# Copyright 2023 Google LLC
+#
+# 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 typing import NamedTuple
+
+
+class PublisherClientId(NamedTuple):
+    value: bytes

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/require_started.py

@@ -0,0 +1,35 @@
+# Copyright 2020 Google LLC
+#
+# 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 typing import ContextManager
+
+from google.api_core.exceptions import FailedPrecondition
+
+
+class RequireStarted(ContextManager):
+    def __init__(self):
+        self._started = False
+
+    def __enter__(self):
+        if self._started:
+            raise FailedPrecondition("__enter__ called twice.")
+        self._started = True
+        return self
+
+    def require_started(self):
+        if not self._started:
+            raise FailedPrecondition("__enter__ has never been called.")
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.require_started()

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/routing_metadata.py

@@ -0,0 +1,34 @@
+# Copyright 2020 Google LLC
+#
+# 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 typing import Mapping
+from urllib.parse import urlencode
+
+from google.cloud.pubsublite.types import Partition, TopicPath, SubscriptionPath
+
+_PARAMS_HEADER = "x-goog-request-params"
+
+
+def topic_routing_metadata(topic: TopicPath, partition: Partition) -> Mapping[str, str]:
+    encoded = urlencode({"partition": str(partition.value), "topic": str(topic)})
+    return {_PARAMS_HEADER: encoded}
+
+
+def subscription_routing_metadata(
+    subscription: SubscriptionPath, partition: Partition
+) -> Mapping[str, str]:
+    encoded = urlencode(
+        {"partition": str(partition.value), "subscription": str(subscription)}
+    )
+    return {_PARAMS_HEADER: encoded}

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/status_codes.py

@@ -0,0 +1,29 @@
+# Copyright 2020 Google LLC
+#
+# 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 grpc import StatusCode
+from google.api_core.exceptions import GoogleAPICallError
+
+retryable_codes = {
+    StatusCode.DEADLINE_EXCEEDED,
+    StatusCode.ABORTED,
+    StatusCode.INTERNAL,
+    StatusCode.UNAVAILABLE,
+    StatusCode.UNKNOWN,
+    StatusCode.CANCELLED,
+}
+
+
+def is_retryable(error: GoogleAPICallError) -> bool:
+    return error.grpc_status_code in retryable_codes

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wait_ignore_cancelled.py

@@ -0,0 +1,30 @@
+# Copyright 2020 Google LLC
+#
+# 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 asyncio import CancelledError
+from typing import Awaitable
+
+
+async def wait_ignore_cancelled(awaitable: Awaitable):
+    try:
+        await awaitable
+    except CancelledError:
+        pass
+
+
+async def wait_ignore_errors(awaitable: Awaitable):
+    try:
+        await awaitable
+    except:  # noqa: E722
+        pass

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/admin_client_impl.py

@@ -0,0 +1,210 @@
+# Copyright 2020 Google LLC
+#
+# 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 logging
+from typing import List, Optional, Union
+
+from google.api_core.exceptions import InvalidArgument
+from google.api_core.operation import Operation
+from cloudsdk.google.protobuf.field_mask_pb2 import FieldMask  # pytype: disable=pyi-error
+
+from google.cloud.pubsublite.admin_client_interface import AdminClientInterface
+from google.cloud.pubsublite.types import (
+    CloudRegion,
+    SubscriptionPath,
+    LocationPath,
+    TopicPath,
+    BacklogLocation,
+    PublishTime,
+    EventTime,
+)
+from google.cloud.pubsublite.types.paths import ReservationPath
+from google.cloud.pubsublite_v1 import (
+    Subscription,
+    Topic,
+    AdminServiceClient,
+    TopicPartitions,
+    Reservation,
+    TimeTarget,
+    SeekSubscriptionRequest,
+    CreateSubscriptionRequest,
+    ExportConfig,
+)
+
+log = logging.getLogger(__name__)
+
+
+class AdminClientImpl(AdminClientInterface):
+    _underlying: AdminServiceClient
+    _region: CloudRegion
+
+    def __init__(self, underlying: AdminServiceClient, region: CloudRegion):
+        self._underlying = underlying
+        self._region = region
+
+    def region(self) -> CloudRegion:
+        return self._region
+
+    def create_topic(self, topic: Topic) -> Topic:
+        path = TopicPath.parse(topic.name)
+        return self._underlying.create_topic(
+            parent=str(path.to_location_path()), topic=topic, topic_id=path.name
+        )
+
+    def get_topic(self, topic_path: TopicPath) -> Topic:
+        return self._underlying.get_topic(name=str(topic_path))
+
+    def get_topic_partition_count(self, topic_path: TopicPath) -> int:
+        partitions: TopicPartitions = self._underlying.get_topic_partitions(
+            name=str(topic_path)
+        )
+        return partitions.partition_count
+
+    def list_topics(self, location_path: LocationPath) -> List[Topic]:
+        return [x for x in self._underlying.list_topics(parent=str(location_path))]
+
+    def update_topic(self, topic: Topic, update_mask: FieldMask) -> Topic:
+        return self._underlying.update_topic(topic=topic, update_mask=update_mask)
+
+    def delete_topic(self, topic_path: TopicPath):
+        self._underlying.delete_topic(name=str(topic_path))
+
+    def list_topic_subscriptions(self, topic_path: TopicPath) -> List[SubscriptionPath]:
+        subscription_strings = [
+            x for x in self._underlying.list_topic_subscriptions(name=str(topic_path))
+        ]
+        return [SubscriptionPath.parse(x) for x in subscription_strings]
+
+    def create_subscription(
+        self,
+        subscription: Subscription,
+        target: Union[BacklogLocation, PublishTime, EventTime] = BacklogLocation.END,
+        starting_offset: Optional[BacklogLocation] = None,
+    ) -> Subscription:
+        if starting_offset:
+            log.warning("starting_offset is deprecated. Use target instead.")
+            target = starting_offset
+        path = SubscriptionPath.parse(subscription.name)
+        requires_seek = isinstance(target, PublishTime) or isinstance(target, EventTime)
+        requires_update = (
+            requires_seek
+            and subscription.export_config
+            and subscription.export_config.desired_state == ExportConfig.State.ACTIVE
+        )
+        if requires_update:
+            # Export subscriptions must be paused while seeking. The state is
+            # later updated to active.
+            subscription.export_config.desired_state = ExportConfig.State.PAUSED
+
+        # Request 1 - Create the subscription.
+        skip_backlog = False
+        if isinstance(target, BacklogLocation):
+            skip_backlog = target == BacklogLocation.END
+        response = self._underlying.create_subscription(
+            request=CreateSubscriptionRequest(
+                parent=str(path.to_location_path()),
+                subscription=subscription,
+                subscription_id=path.name,
+                skip_backlog=skip_backlog,
+            )
+        )
+        # Request 2 (optional) - seek the subscription.
+        if requires_seek:
+            self.seek_subscription(subscription_path=path, target=target)
+        # Request 3 (optional) - make the export subscription active.
+        if requires_update:
+            response = self.update_subscription(
+                subscription=Subscription(
+                    name=response.name,
+                    export_config=ExportConfig(
+                        desired_state=ExportConfig.State.ACTIVE,
+                    ),
+                ),
+                update_mask=FieldMask(paths=["export_config.desired_state"]),
+            )
+        return response
+
+    def get_subscription(self, subscription_path: SubscriptionPath) -> Subscription:
+        return self._underlying.get_subscription(name=str(subscription_path))
+
+    def list_subscriptions(self, location_path: LocationPath) -> List[Subscription]:
+        return [
+            x for x in self._underlying.list_subscriptions(parent=str(location_path))
+        ]
+
+    def update_subscription(
+        self, subscription: Subscription, update_mask: FieldMask
+    ) -> Subscription:
+        return self._underlying.update_subscription(
+            subscription=subscription, update_mask=update_mask
+        )
+
+    def seek_subscription(
+        self,
+        subscription_path: SubscriptionPath,
+        target: Union[BacklogLocation, PublishTime, EventTime],
+    ) -> Operation:
+        request = SeekSubscriptionRequest(name=str(subscription_path))
+        if isinstance(target, PublishTime):
+            request.time_target = TimeTarget(publish_time=target.value)
+        elif isinstance(target, EventTime):
+            request.time_target = TimeTarget(event_time=target.value)
+        elif isinstance(target, BacklogLocation):
+            if target == BacklogLocation.END:
+                request.named_target = SeekSubscriptionRequest.NamedTarget.HEAD
+            else:
+                request.named_target = SeekSubscriptionRequest.NamedTarget.TAIL
+        else:
+            raise InvalidArgument("A valid seek target must be specified.")
+        return self._underlying.seek_subscription(request=request)
+
+    def delete_subscription(self, subscription_path: SubscriptionPath):
+        self._underlying.delete_subscription(name=str(subscription_path))
+
+    def create_reservation(self, reservation: Reservation) -> Reservation:
+        path = ReservationPath.parse(reservation.name)
+        return self._underlying.create_reservation(
+            parent=str(path.to_location_path()),
+            reservation=reservation,
+            reservation_id=path.name,
+        )
+
+    def get_reservation(self, reservation_path: ReservationPath) -> Reservation:
+        return self._underlying.get_reservation(name=str(reservation_path))
+
+    def list_reservations(self, location_path: LocationPath) -> List[Reservation]:
+        return [
+            x for x in self._underlying.list_reservations(parent=str(location_path))
+        ]
+
+    def update_reservation(
+        self, reservation: Reservation, update_mask: FieldMask
+    ) -> Reservation:
+        return self._underlying.update_reservation(
+            reservation=reservation, update_mask=update_mask
+        )
+
+    def delete_reservation(self, reservation_path: ReservationPath):
+        self._underlying.delete_reservation(name=str(reservation_path))
+
+    def list_reservation_topics(
+        self, reservation_path: ReservationPath
+    ) -> List[TopicPath]:
+        subscription_strings = [
+            x
+            for x in self._underlying.list_reservation_topics(
+                name=str(reservation_path)
+            )
+        ]
+        return [TopicPath.parse(x) for x in subscription_strings]

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/assigner.py

@@ -0,0 +1,29 @@
+# Copyright 2020 Google LLC
+#
+# 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 abc import abstractmethod, ABCMeta
+from typing import AsyncContextManager, Set
+
+from google.cloud.pubsublite.types.partition import Partition
+
+
+class Assigner(AsyncContextManager, metaclass=ABCMeta):
+    """
+    An assigner will deliver a continuous stream of assignments when called into. Perform all necessary work with the
+    assignment before attempting to get the next one.
+    """
+
+    @abstractmethod
+    async def get_assignment(self) -> Set[Partition]:
+        raise NotImplementedError()

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/assigner_impl.py

@@ -0,0 +1,132 @@
+# Copyright 2020 Google LLC
+#
+# 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 asyncio
+from typing import Optional, Set
+
+import logging
+
+from google.cloud.pubsublite.internal.wait_ignore_cancelled import wait_ignore_errors
+from google.cloud.pubsublite.internal.wire.assigner import Assigner
+from google.cloud.pubsublite.internal.wire.retrying_connection import (
+    RetryingConnection,
+    ConnectionFactory,
+)
+from google.api_core.exceptions import FailedPrecondition, GoogleAPICallError
+from google.cloud.pubsublite.internal.wire.connection_reinitializer import (
+    ConnectionReinitializer,
+)
+from google.cloud.pubsublite.internal.wire.connection import Connection
+from google.cloud.pubsublite.types.partition import Partition
+from google.cloud.pubsublite_v1.types import (
+    PartitionAssignmentRequest,
+    PartitionAssignment,
+    InitialPartitionAssignmentRequest,
+    PartitionAssignmentAck,
+)
+
+_LOGGER = logging.getLogger(__name__)
+
+# Maximum bytes per batch at 3.5 MiB to avoid GRPC limit of 4 MiB
+_MAX_BYTES = int(3.5 * 1024 * 1024)
+
+# Maximum messages per batch at 1000
+_MAX_MESSAGES = 1000
+
+
+class AssignerImpl(
+    Assigner, ConnectionReinitializer[PartitionAssignmentRequest, PartitionAssignment]
+):
+    _initial: InitialPartitionAssignmentRequest
+    _connection: RetryingConnection[PartitionAssignmentRequest, PartitionAssignment]
+
+    _outstanding_assignment: bool
+
+    _receiver: Optional[asyncio.Future]
+
+    # A queue that may only hold one element with the next assignment.
+    _new_assignment: "asyncio.Queue[Set[Partition]]"
+
+    def __init__(
+        self,
+        initial: InitialPartitionAssignmentRequest,
+        factory: ConnectionFactory[PartitionAssignmentRequest, PartitionAssignment],
+    ):
+        self._initial = initial
+        self._connection = RetryingConnection(factory, self)
+        self._outstanding_assignment = False
+        self._receiver = None
+        self._new_assignment = asyncio.Queue(maxsize=1)
+
+    async def __aenter__(self):
+        await self._connection.__aenter__()
+        return self
+
+    def _start_receiver(self):
+        assert self._receiver is None
+        self._receiver = asyncio.ensure_future(self._receive_loop())
+
+    async def _stop_receiver(self):
+        if self._receiver:
+            self._receiver.cancel()
+            await wait_ignore_errors(self._receiver)
+            self._receiver = None
+
+    async def _receive_loop(self):
+        while True:
+            response = await self._connection.read()
+            if self._outstanding_assignment or not self._new_assignment.empty():
+                self._connection.fail(
+                    FailedPrecondition(
+                        "Received a duplicate assignment on the stream while one was outstanding."
+                    )
+                )
+                return
+            self._outstanding_assignment = True
+            partitions = set()
+            for partition in response.partitions:
+                partitions.add(Partition(partition))
+            _LOGGER.info(f"Received new assignment with partitions: {partitions}.")
+            self._new_assignment.put_nowait(partitions)
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self._stop_receiver()
+        await self._connection.__aexit__(exc_type, exc_val, exc_tb)
+
+    async def stop_processing(self, error: GoogleAPICallError):
+        await self._stop_receiver()
+        self._outstanding_assignment = False
+        while not self._new_assignment.empty():
+            self._new_assignment.get_nowait()
+
+    async def reinitialize(
+        self,
+        connection: Connection[PartitionAssignmentRequest, PartitionAssignment],
+    ):
+        await connection.write(PartitionAssignmentRequest(initial=self._initial))
+        self._start_receiver()
+
+    async def get_assignment(self) -> Set[Partition]:
+        if self._outstanding_assignment:
+            try:
+                await self._connection.write(
+                    PartitionAssignmentRequest(ack=PartitionAssignmentAck())
+                )
+                self._outstanding_assignment = False
+            except GoogleAPICallError as e:
+                # If there is a failure to ack, keep going. The stream likely restarted.
+                _LOGGER.debug(
+                    f"Assignment ack attempt failed due to stream failure: {e}"
+                )
+        return await self._connection.await_unless_failed(self._new_assignment.get())

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/client_cache.py

@@ -0,0 +1,42 @@
+# Copyright 2020 Google LLC
+#
+# 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 threading
+from typing import Generic, TypeVar, Callable, Optional
+
+_Client = TypeVar("_Client")
+_MAX_CLIENT_USES = 75  # GRPC channels are limited to 100 concurrent streams.
+
+
+class ClientCache(Generic[_Client]):
+    _ClientFactory = Callable[[], _Client]
+
+    _factory: _ClientFactory
+    _latest: Optional[_Client]
+    _remaining_uses: int
+    _lock: threading.Lock
+
+    def __init__(self, factory: _ClientFactory):
+        self._factory = factory
+        self._latest = None
+        self._remaining_uses = 0
+        self._lock = threading.Lock()
+
+    def get(self) -> _Client:
+        with self._lock:
+            if self._remaining_uses <= 0:
+                self._remaining_uses = _MAX_CLIENT_USES
+                self._latest = self._factory()
+            self._remaining_uses -= 1
+            return self._latest

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/committer.py

@@ -0,0 +1,44 @@
+# Copyright 2020 Google LLC
+#
+# 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 abc import abstractmethod, ABCMeta
+from typing import AsyncContextManager
+
+from google.cloud.pubsublite_v1 import Cursor
+
+
+class Committer(AsyncContextManager, metaclass=ABCMeta):
+    """
+    A Committer is able to commit subscribers' completed offsets.
+    """
+
+    @abstractmethod
+    def commit(self, cursor: Cursor) -> None:
+        """
+        Start the commit for a cursor.
+
+        Raises:
+          GoogleAPICallError: When the committer terminates in failure.
+        """
+        pass
+
+    @abstractmethod
+    async def wait_until_empty(self):
+        """
+        Flushes pending commits and waits for all outstanding commit responses from the server.
+
+        Raises:
+          GoogleAPICallError: When the committer terminates in failure.
+        """
+        pass

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/committer_impl.py

@@ -0,0 +1,176 @@
+# Copyright 2020 Google LLC
+#
+# 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 asyncio
+from typing import Optional, List
+
+import logging
+
+from google.cloud.pubsublite.internal.wait_ignore_cancelled import wait_ignore_errors
+from google.cloud.pubsublite.internal.wire.committer import Committer
+from google.cloud.pubsublite.internal.wire.retrying_connection import (
+    RetryingConnection,
+    ConnectionFactory,
+)
+from google.api_core.exceptions import FailedPrecondition, GoogleAPICallError
+from google.cloud.pubsublite.internal.wire.connection_reinitializer import (
+    ConnectionReinitializer,
+)
+from google.cloud.pubsublite.internal.wire.connection import Connection
+from google.cloud.pubsublite_v1 import Cursor
+from google.cloud.pubsublite_v1.types import (
+    StreamingCommitCursorRequest,
+    StreamingCommitCursorResponse,
+    InitialCommitCursorRequest,
+)
+
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class CommitterImpl(
+    Committer,
+    ConnectionReinitializer[
+        StreamingCommitCursorRequest, StreamingCommitCursorResponse
+    ],
+):
+    _initial: InitialCommitCursorRequest
+    _flush_seconds: float
+    _connection: RetryingConnection[
+        StreamingCommitCursorRequest, StreamingCommitCursorResponse
+    ]
+
+    _next_to_commit: Optional[Cursor]
+    _outstanding_commits: List[Cursor]
+
+    _receiver: Optional[asyncio.Future]
+    _flusher: Optional[asyncio.Future]
+    _empty: asyncio.Event
+
+    def __init__(
+        self,
+        initial: InitialCommitCursorRequest,
+        flush_seconds: float,
+        factory: ConnectionFactory[
+            StreamingCommitCursorRequest, StreamingCommitCursorResponse
+        ],
+    ):
+        self._initial = initial
+        self._flush_seconds = flush_seconds
+        self._connection = RetryingConnection(factory, self)
+        self._next_to_commit = None
+        self._outstanding_commits = []
+        self._receiver = None
+        self._flusher = None
+        self._empty = asyncio.Event()
+        self._empty.set()
+
+    async def __aenter__(self):
+        await self._connection.__aenter__()
+        return self
+
+    def _start_loopers(self):
+        assert self._receiver is None
+        assert self._flusher is None
+        self._receiver = asyncio.ensure_future(self._receive_loop())
+        self._flusher = asyncio.ensure_future(self._flush_loop())
+
+    async def _stop_loopers(self):
+        if self._receiver:
+            self._receiver.cancel()
+            await wait_ignore_errors(self._receiver)
+            self._receiver = None
+        if self._flusher:
+            self._flusher.cancel()
+            await wait_ignore_errors(self._flusher)
+            self._flusher = None
+
+    def _handle_response(self, response: StreamingCommitCursorResponse):
+        if "commit" not in response:
+            self._connection.fail(
+                FailedPrecondition(
+                    "Received an invalid subsequent response on the commit stream."
+                )
+            )
+        if response.commit.acknowledged_commits > len(self._outstanding_commits):
+            self._connection.fail(
+                FailedPrecondition(
+                    "Received a commit response on the stream with no outstanding commits."
+                )
+            )
+        for _ in range(response.commit.acknowledged_commits):
+            self._outstanding_commits.pop(0)
+        if len(self._outstanding_commits) == 0:
+            self._empty.set()
+
+    async def _receive_loop(self):
+        while True:
+            response = await self._connection.read()
+            self._handle_response(response)
+
+    async def _flush_loop(self):
+        while True:
+            await asyncio.sleep(self._flush_seconds)
+            await self._flush()
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self._stop_loopers()
+        if not self._connection.error():
+            await self._flush()
+        await self._connection.__aexit__(exc_type, exc_val, exc_tb)
+
+    async def _flush(self):
+        if self._next_to_commit is None:
+            return
+        req = StreamingCommitCursorRequest()
+        req.commit.cursor = self._next_to_commit
+        self._outstanding_commits.append(self._next_to_commit)
+        self._next_to_commit = None
+        self._empty.clear()
+        try:
+            await self._connection.write(req)
+        except GoogleAPICallError as e:
+            _LOGGER.debug(f"Failed commit on stream: {e}")
+
+    async def wait_until_empty(self):
+        await self._flush()
+        await self._connection.await_unless_failed(self._empty.wait())
+
+    def commit(self, cursor: Cursor) -> None:
+        if self._connection.error():
+            raise self._connection.error()
+        self._next_to_commit = cursor
+
+    async def stop_processing(self, error: GoogleAPICallError):
+        await self._stop_loopers()
+
+    async def reinitialize(
+        self,
+        connection: Connection[
+            StreamingCommitCursorRequest, StreamingCommitCursorResponse
+        ],
+    ):
+        await connection.write(StreamingCommitCursorRequest(initial=self._initial))
+        response = await connection.read()
+        if "initial" not in response:
+            self._connection.fail(
+                FailedPrecondition(
+                    "Received an invalid initial response on the publish stream."
+                )
+            )
+        if self._next_to_commit is None:
+            if self._outstanding_commits:
+                self._next_to_commit = self._outstanding_commits[-1]
+        self._outstanding_commits = []
+        self._start_loopers()

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/connection.py

@@ -0,0 +1,55 @@
+# Copyright 2020 Google LLC
+#
+# 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 typing import Generic, TypeVar, AsyncContextManager
+from abc import abstractmethod, ABCMeta
+
+Request = TypeVar("Request")
+Response = TypeVar("Response")
+
+
+class Connection(
+    AsyncContextManager["Connection"], Generic[Request, Response], metaclass=ABCMeta
+):
+    """
+    A connection to an underlying stream. Only one call to 'read' may be outstanding at a time.
+    """
+
+    @abstractmethod
+    async def write(self, request: Request) -> None:
+        """
+        Write a message to the stream.
+
+        Raises:
+          GoogleAPICallError: When the connection terminates in failure.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    async def read(self) -> Response:
+        """
+        Read a message off of the stream.
+
+        Raises:
+          GoogleAPICallError: When the connection terminates in failure.
+        """
+        raise NotImplementedError()
+
+
+class ConnectionFactory(Generic[Request, Response], metaclass=ABCMeta):
+    """A factory for producing Connections."""
+
+    @abstractmethod
+    async def new(self) -> Connection[Request, Response]:
+        raise NotImplementedError()

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/connection_reinitializer.py

@@ -0,0 +1,52 @@
+# Copyright 2020 Google LLC
+#
+# 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 typing import Generic
+from abc import ABCMeta, abstractmethod
+from google.api_core.exceptions import GoogleAPICallError
+from google.cloud.pubsublite.internal.wire.connection import (
+    Connection,
+    Request,
+    Response,
+)
+
+
+class ConnectionReinitializer(Generic[Request, Response], metaclass=ABCMeta):
+    """A class capable of reinitializing a connection after a new one has been created."""
+
+    @abstractmethod
+    async def stop_processing(self, error: GoogleAPICallError):
+        """Tear down internal state processing the current connection in
+        response to a stream error.
+
+        Args:
+            error: The error that caused the stream to break
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    async def reinitialize(
+        self,
+        connection: Connection[Request, Response],
+    ):
+        """Reinitialize a connection. Must ensure no calls to the associated RetryingConnection
+        occur until this completes.
+
+        Args:
+            connection: The connection to reinitialize
+
+        Raises:
+            GoogleAPICallError: If it fails to reinitialize.
+        """
+        raise NotImplementedError()

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/default_routing_policy.py

@@ -0,0 +1,47 @@
+# Copyright 2020 Google LLC
+#
+# 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 hashlib
+import random
+
+from google.cloud.pubsublite.internal.wire.routing_policy import RoutingPolicy
+from google.cloud.pubsublite.types.partition import Partition
+from google.cloud.pubsublite_v1.types import PubSubMessage
+
+
+class DefaultRoutingPolicy(RoutingPolicy):
+    """
+    The default routing policy which routes based on sha256 % num_partitions using the key if set or round robin if
+    unset.
+    """
+
+    _num_partitions: int
+    _current_round_robin: Partition
+
+    def __init__(self, num_partitions: int):
+        self._num_partitions = num_partitions
+        self._current_round_robin = Partition(random.randint(0, num_partitions - 1))
+
+    def route(self, message: PubSubMessage) -> Partition:
+        """Route the message using the key if set or round robin if unset."""
+        if not message.key:
+            result = Partition(self._current_round_robin.value)
+            self._current_round_robin = Partition(
+                (self._current_round_robin.value + 1) % self._num_partitions
+            )
+            return result
+        sha = hashlib.sha256()
+        sha.update(message.key)
+        as_int = int.from_bytes(sha.digest(), byteorder="big")
+        return Partition(as_int % self._num_partitions)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/fixed_set_assigner.py

@@ -0,0 +1,42 @@
+# Copyright 2020 Google LLC
+#
+# 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 asyncio
+from typing import Set
+
+from google.cloud.pubsublite.internal.wire.assigner import Assigner
+from google.cloud.pubsublite.types.partition import Partition
+
+
+class FixedSetAssigner(Assigner):
+    _partitions: Set[Partition]
+    _returned_set: bool
+
+    def __init__(self, partitions: Set[Partition]):
+        self._partitions = partitions
+        self._returned_set = False
+
+    async def get_assignment(self) -> Set[Partition]:
+        """Only returns an assignment the first iteration."""
+        if self._returned_set:
+            await asyncio.sleep(float("inf"))
+            raise RuntimeError("Should never happen.")
+        self._returned_set = True
+        return self._partitions
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        pass

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/flow_control_batcher.py

@@ -0,0 +1,93 @@
+# Copyright 2020 Google LLC
+#
+# 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 typing import List, Optional
+
+from google.cloud.pubsublite_v1 import FlowControlRequest, SequencedMessage
+
+_EXPEDITE_BATCH_REQUEST_RATIO = 0.5
+_MAX_INT64 = 0x7FFFFFFFFFFFFFFF
+
+
+def _clamp(val: int):
+    if val > _MAX_INT64:
+        return _MAX_INT64
+    if val < 0:
+        return 0
+    return val
+
+
+class _AggregateRequest:
+    _request: FlowControlRequest.meta.pb
+
+    def __init__(self):
+        self._request = FlowControlRequest.meta.pb()
+
+    def __add__(self, other: FlowControlRequest):
+        other_pb = other._pb
+        self._request.allowed_bytes = (
+            self._request.allowed_bytes + other_pb.allowed_bytes
+        )
+        self._request.allowed_bytes = min(self._request.allowed_bytes, _MAX_INT64)
+        self._request.allowed_messages = (
+            self._request.allowed_messages + other_pb.allowed_messages
+        )
+        self._request.allowed_messages = min(self._request.allowed_messages, _MAX_INT64)
+        return self
+
+    def to_optional(self) -> Optional[FlowControlRequest]:
+        allowed_messages = _clamp(self._request.allowed_messages)
+        allowed_bytes = _clamp(self._request.allowed_bytes)
+        if allowed_messages == 0 and allowed_bytes == 0:
+            return None
+        request = FlowControlRequest()
+        request._pb.allowed_messages = allowed_messages
+        request._pb.allowed_bytes = allowed_bytes
+        return request
+
+
+def _exceeds_expedite_ratio(pending: int, client: int):
+    if client <= 0:
+        return False
+    return (pending / client) >= _EXPEDITE_BATCH_REQUEST_RATIO
+
+
+class FlowControlBatcher:
+    _client_tokens: _AggregateRequest
+    _pending_tokens: _AggregateRequest
+
+    def __init__(self):
+        self._client_tokens = _AggregateRequest()
+        self._pending_tokens = _AggregateRequest()
+
+    def add(self, request: FlowControlRequest):
+        self._client_tokens += request
+        self._pending_tokens += request
+
+    def on_messages(self, messages: List[SequencedMessage]):
+        byte_size = 0
+        for message in messages:
+            byte_size += message.size_bytes
+        self._client_tokens += FlowControlRequest(
+            allowed_bytes=-byte_size, allowed_messages=-len(messages)
+        )
+
+    def request_for_restart(self) -> Optional[FlowControlRequest]:
+        self._pending_tokens = _AggregateRequest()
+        return self._client_tokens.to_optional()
+
+    def release_pending_request(self) -> Optional[FlowControlRequest]:
+        request = self._pending_tokens
+        self._pending_tokens = _AggregateRequest()
+        return request.to_optional()

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/gapic_connection.py

@@ -0,0 +1,108 @@
+# Copyright 2020 Google LLC
+#
+# 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 typing import (
+    cast,
+    AsyncIterator,
+    TypeVar,
+    Optional,
+    Callable,
+    AsyncIterable,
+    Awaitable,
+)
+import asyncio
+
+from google.api_core.exceptions import GoogleAPICallError, FailedPrecondition
+
+from google.cloud.pubsublite.internal.wire.connection import (
+    Connection,
+    Request,
+    Response,
+    ConnectionFactory,
+)
+from google.cloud.pubsublite.internal.wire.work_item import WorkItem
+from google.cloud.pubsublite.internal.wire.permanent_failable import PermanentFailable
+
+T = TypeVar("T")
+
+
+class GapicConnection(
+    Connection[Request, Response], AsyncIterator[Request], PermanentFailable
+):
+    """A Connection wrapping a gapic AsyncIterator[Request/Response] pair."""
+
+    _write_queue: "asyncio.Queue[WorkItem[Request, None]]"
+    _response_it: Optional[AsyncIterator[Response]]
+
+    def __init__(self):
+        super().__init__()
+        self._write_queue = asyncio.Queue(maxsize=1)
+
+    def set_response_it(self, response_it: AsyncIterator[Response]):
+        self._response_it = response_it
+
+    async def write(self, request: Request) -> None:
+        item = WorkItem(request)
+        await self.await_unless_failed(self._write_queue.put(item))
+        await self.await_unless_failed(item.response_future)
+
+    async def read(self) -> Response:
+        if self._response_it is None:
+            self.fail(FailedPrecondition("GapicConnection not initialized."))
+            raise self.error()
+        try:
+            response_it = cast(AsyncIterator[Response], self._response_it)
+            return await self.await_unless_failed(response_it.__anext__())
+        except StopAsyncIteration:
+            self.fail(FailedPrecondition("Server sent unprompted half close."))
+        except GoogleAPICallError as e:
+            self.fail(e)
+        raise self.error()
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
+        pass
+
+    async def __anext__(self) -> Request:
+        item: WorkItem[Request, None] = await self.await_unless_failed(
+            self._write_queue.get()
+        )
+        item.response_future.set_result(None)
+        return item.request
+
+    def __aiter__(self) -> AsyncIterator[Response]:
+        return self
+
+
+class GapicConnectionFactory(ConnectionFactory[Request, Response]):
+    """A ConnectionFactory that produces GapicConnections."""
+
+    _producer = Callable[[AsyncIterator[Request]], Awaitable[AsyncIterable[Response]]]
+
+    def __init__(
+        self,
+        producer: Callable[
+            [AsyncIterator[Request]], Awaitable[AsyncIterable[Response]]
+        ],
+    ):
+        self._producer = producer
+
+    async def new(self) -> Connection[Request, Response]:
+        conn = GapicConnection[Request, Response]()
+        response_fut = self._producer(conn)
+        response_iterable = await response_fut
+        conn.set_response_it(response_iterable.__aiter__())
+        return conn

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/make_publisher.py

@@ -0,0 +1,132 @@
+# Copyright 2020 Google LLC
+#
+# 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 typing import AsyncIterator, Mapping, Optional
+
+from google.cloud.pubsub_v1.types import BatchSettings
+
+from google.cloud.pubsublite.admin_client import AdminClient
+from google.cloud.pubsublite.internal.endpoints import regional_endpoint
+from google.cloud.pubsublite.internal.publisher_client_id import PublisherClientId
+from google.cloud.pubsublite.internal.publish_sequence_number import (
+    PublishSequenceNumber,
+)
+from google.cloud.pubsublite.internal.wire.client_cache import ClientCache
+from google.cloud.pubsublite.internal.wire.default_routing_policy import (
+    DefaultRoutingPolicy,
+)
+from google.cloud.pubsublite.internal.wire.gapic_connection import (
+    GapicConnectionFactory,
+)
+from google.cloud.pubsublite.internal.wire.merge_metadata import merge_metadata
+from google.cloud.pubsublite.internal.wire.partition_count_watcher_impl import (
+    PartitionCountWatcherImpl,
+)
+from google.cloud.pubsublite.internal.wire.partition_count_watching_publisher import (
+    PartitionCountWatchingPublisher,
+)
+from google.cloud.pubsublite.internal.wire.publisher import Publisher
+from google.cloud.pubsublite.internal.wire.single_partition_publisher import (
+    SinglePartitionPublisher,
+)
+from google.cloud.pubsublite.types import Partition, TopicPath
+from google.cloud.pubsublite.internal.routing_metadata import topic_routing_metadata
+from google.cloud.pubsublite_v1 import InitialPublishRequest, PublishRequest
+from google.cloud.pubsublite_v1.services.publisher_service import async_client
+from google.api_core.client_options import ClientOptions
+from google.auth.credentials import Credentials
+
+DEFAULT_BATCHING_SETTINGS = BatchSettings(
+    max_bytes=(
+        3 * 1024 * 1024
+    ),  # 3 MiB to stay 1 MiB below GRPC's 4 MiB per-message limit.
+    max_messages=1000,
+    max_latency=0.05,  # 50 ms
+)
+DEFAULT_PARTITION_POLL_PERIOD = 600  # ten minutes
+
+
+def make_publisher(
+    topic: TopicPath,
+    transport: str,
+    per_partition_batching_settings: Optional[BatchSettings] = None,
+    credentials: Optional[Credentials] = None,
+    client_options: Optional[ClientOptions] = None,
+    metadata: Optional[Mapping[str, str]] = None,
+    client_id: Optional[PublisherClientId] = None,
+) -> Publisher:
+    """
+    Make a new publisher for the given topic.
+
+    Args:
+      topic: The topic to publish to.
+      transport: The transport type to use.
+      per_partition_batching_settings: Settings for batching messages on each partition. The default is reasonable for most cases.
+      credentials: The credentials to use to connect. GOOGLE_DEFAULT_CREDENTIALS is used if None.
+      client_options: Other options to pass to the client. Note that if you pass any you must set api_endpoint.
+      metadata: Additional metadata to send with the RPC.
+      client_id: 128-bit unique client id. If set, enables publish idempotency for the session.
+
+    Returns:
+      A new Publisher.
+
+    Throws:
+      GoogleApiCallException on any error determining topic structure.
+    """
+    if per_partition_batching_settings is None:
+        per_partition_batching_settings = DEFAULT_BATCHING_SETTINGS
+    admin_client = AdminClient(
+        region=topic.location.region,
+        credentials=credentials,
+        client_options=client_options,
+    )
+    if client_options is None:
+        client_options = ClientOptions(
+            api_endpoint=regional_endpoint(topic.location.region)
+        )
+    client_cache = ClientCache(
+        lambda: async_client.PublisherServiceAsyncClient(
+            credentials=credentials, transport=transport, client_options=client_options
+        )
+    )
+
+    def publisher_factory(partition: Partition):
+        def connection_factory(requests: AsyncIterator[PublishRequest]):
+            final_metadata = merge_metadata(
+                metadata, topic_routing_metadata(topic, partition)
+            )
+            return client_cache.get().publish(
+                requests, metadata=list(final_metadata.items())
+            )
+
+        initial_request = InitialPublishRequest(
+            topic=str(topic), partition=partition.value
+        )
+        if client_id:
+            initial_request.client_id = client_id.value
+        return SinglePartitionPublisher(
+            initial_request,
+            per_partition_batching_settings,
+            GapicConnectionFactory(connection_factory),
+            PublishSequenceNumber(0),
+        )
+
+    def policy_factory(partition_count: int):
+        return DefaultRoutingPolicy(partition_count)
+
+    return PartitionCountWatchingPublisher(
+        PartitionCountWatcherImpl(admin_client, topic, DEFAULT_PARTITION_POLL_PERIOD),
+        publisher_factory,
+        policy_factory,
+    )

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/merge_metadata.py

@@ -0,0 +1,31 @@
+# Copyright 2020 Google LLC
+#
+# 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 typing import Mapping, Optional
+
+
+def merge_metadata(
+    a: Optional[Mapping[str, str]], b: Optional[Mapping[str, str]]
+) -> Mapping[str, str]:
+    """
+    Merge the two sets of metadata if either exists. The second map overwrites the first.
+    """
+    result = {}
+    if a:
+        for k, v in a.items():
+            result[k] = v
+    if b:
+        for k, v in b.items():
+            result[k] = v
+    return result

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/partition_count_watcher.py

@@ -0,0 +1,22 @@
+# Copyright 2020 Google LLC
+#
+# 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 abc import abstractmethod, ABCMeta
+from typing import AsyncContextManager
+
+
+class PartitionCountWatcher(AsyncContextManager, metaclass=ABCMeta):
+    @abstractmethod
+    async def get_partition_count(self) -> int:
+        raise NotImplementedError()

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/partition_count_watcher_impl.py

@@ -0,0 +1,75 @@
+# Copyright 2020 Google LLC
+#
+# 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 logging
+from concurrent.futures.thread import ThreadPoolExecutor
+import asyncio
+
+from google.cloud.pubsublite import AdminClientInterface
+from google.cloud.pubsublite.internal.wait_ignore_cancelled import wait_ignore_cancelled
+from google.cloud.pubsublite.internal.wire.partition_count_watcher import (
+    PartitionCountWatcher,
+)
+from google.cloud.pubsublite.internal.wire.permanent_failable import PermanentFailable
+from google.cloud.pubsublite.types import TopicPath
+from google.api_core.exceptions import GoogleAPICallError
+
+
+class PartitionCountWatcherImpl(PartitionCountWatcher, PermanentFailable):
+    _admin: AdminClientInterface
+    _topic_path: TopicPath
+    _duration: float
+    _any_success: bool
+    _thread: ThreadPoolExecutor
+    _queue: asyncio.Queue
+    _partition_loop_poller: asyncio.Future
+
+    def __init__(
+        self, admin: AdminClientInterface, topic_path: TopicPath, duration: float
+    ):
+        super().__init__()
+        self._admin = admin
+        self._topic_path = topic_path
+        self._duration = duration
+        self._any_success = False
+
+    async def __aenter__(self):
+        self._thread = ThreadPoolExecutor(max_workers=1)
+        self._queue = asyncio.Queue(maxsize=1)
+        self._partition_loop_poller = asyncio.ensure_future(
+            self.run_poller(self._poll_partition_loop)
+        )
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self._partition_loop_poller.cancel()
+        await wait_ignore_cancelled(self._partition_loop_poller)
+        self._thread.shutdown(wait=False)
+
+    def _get_partition_count_sync(self) -> int:
+        return self._admin.get_topic_partition_count(self._topic_path)
+
+    async def _poll_partition_loop(self):
+        try:
+            partition_count = await asyncio.get_event_loop().run_in_executor(
+                self._thread, self._get_partition_count_sync
+            )
+            self._any_success = True
+            await self._queue.put(partition_count)
+        except GoogleAPICallError as e:
+            if not self._any_success:
+                raise e
+            logging.exception("Failed to retrieve partition count")
+        await asyncio.sleep(self._duration)
+
+    async def get_partition_count(self) -> int:
+        return await self.await_unless_failed(self._queue.get())

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/partition_count_watching_publisher.py

@@ -0,0 +1,96 @@
+# Copyright 2020 Google LLC
+#
+# 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 asyncio
+import sys
+from typing import Callable, Dict
+
+from google.cloud.pubsublite.internal.wait_ignore_cancelled import wait_ignore_cancelled
+from google.cloud.pubsublite.internal.wire.partition_count_watcher import (
+    PartitionCountWatcher,
+)
+from google.cloud.pubsublite.internal.wire.publisher import Publisher
+from google.cloud.pubsublite.internal.wire.routing_policy import RoutingPolicy
+from google.cloud.pubsublite.types import MessageMetadata, Partition
+from google.cloud.pubsublite_v1 import PubSubMessage
+
+
+class PartitionCountWatchingPublisher(Publisher):
+    _publishers: Dict[Partition, Publisher]
+    _publisher_factory: Callable[[Partition], Publisher]
+    _policy_factory: Callable[[int], RoutingPolicy]
+    _watcher: PartitionCountWatcher
+    _partition_count_poller: asyncio.Future
+
+    def __init__(
+        self,
+        watcher: PartitionCountWatcher,
+        publisher_factory: Callable[[Partition], Publisher],
+        policy_factory: Callable[[int], RoutingPolicy],
+    ):
+        self._publishers = {}
+        self._publisher_factory = publisher_factory
+        self._policy_factory = policy_factory
+        self._watcher = watcher
+        self._partition_count_poller = None
+
+    async def __aenter__(self):
+        try:
+            await self._watcher.__aenter__()
+            await self._poll_partition_count_action()
+        except Exception:
+            await self._watcher.__aexit__(*sys.exc_info())
+            raise
+        self._partition_count_poller = asyncio.ensure_future(
+            self._watch_partition_count()
+        )
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if self._partition_count_poller is not None:
+            self._partition_count_poller.cancel()
+            await wait_ignore_cancelled(self._partition_count_poller)
+            await self._watcher.__aexit__(exc_type, exc_val, exc_tb)
+        for publisher in self._publishers.values():
+            await publisher.__aexit__(exc_type, exc_val, exc_tb)
+
+    async def _poll_partition_count_action(self):
+        partition_count = await self._watcher.get_partition_count()
+        await self._handle_partition_count_update(partition_count)
+
+    async def _watch_partition_count(self):
+        while True:
+            await self._poll_partition_count_action()
+
+    async def _handle_partition_count_update(self, partition_count: int):
+        current_count = len(self._publishers)
+        if current_count == partition_count:
+            return
+        if current_count > partition_count:
+            return
+
+        new_publishers = {
+            Partition(index): self._publisher_factory(Partition(index))
+            for index in range(current_count, partition_count)
+        }
+        await asyncio.gather(*[p.__aenter__() for p in new_publishers.values()])
+        routing_policy = self._policy_factory(partition_count)
+
+        self._publishers.update(new_publishers)
+        self._routing_policy = routing_policy
+
+    async def publish(self, message: PubSubMessage) -> MessageMetadata:
+        partition = self._routing_policy.route(message)
+        assert partition in self._publishers
+        publisher = self._publishers[partition]
+        return await publisher.publish(message)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/permanent_failable.py

@@ -0,0 +1,103 @@
+# Copyright 2020 Google LLC
+#
+# 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 asyncio
+from typing import Awaitable, TypeVar, Optional, Callable
+
+from google.api_core.exceptions import GoogleAPICallError, Unknown
+
+from google.cloud.pubsublite.internal.wait_ignore_cancelled import wait_ignore_errors
+
+T = TypeVar("T")
+
+
+def adapt_error(e: Exception) -> GoogleAPICallError:
+    if isinstance(e, GoogleAPICallError):
+        return e
+    return Unknown("Had an unknown error", errors=[e])
+
+
+class _TaskWithCleanup:
+    def __init__(self, a: Awaitable):
+        self._task = asyncio.ensure_future(a)
+
+    async def __aenter__(self):
+        return self._task
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if not self._task.done():
+            self._task.cancel()
+            await wait_ignore_errors(self._task)
+
+
+class PermanentFailable:
+    """A class that can experience permanent failures, with helpers for forwarding these to client actions."""
+
+    _maybe_failure_task: Optional[asyncio.Future]
+
+    def __init__(self):
+        self._maybe_failure_task = None
+
+    @property
+    def _failure_task(self) -> asyncio.Future:
+        """Get the failure task, initializing it lazily, since it needs to be initialized in the event loop."""
+        if self._maybe_failure_task is None:
+            self._maybe_failure_task = asyncio.Future()
+        return self._maybe_failure_task
+
+    async def await_unless_failed(self, awaitable: Awaitable[T]) -> T:
+        """
+        Await the awaitable, unless fail() is called first.
+        Args:
+          awaitable: An awaitable
+
+        Returns: The result of the awaitable
+        Raises: The permanent error if fail() is called or the awaitable raises one.
+        """
+        async with _TaskWithCleanup(awaitable) as task:
+            if self._failure_task.done():
+                raise self._failure_task.exception()
+            done, _ = await asyncio.wait(
+                [task, self._failure_task], return_when=asyncio.FIRST_COMPLETED
+            )
+            if task in done:
+                return await task
+            raise self._failure_task.exception()
+
+    async def run_poller(self, poll_action: Callable[[], Awaitable[None]]):
+        """
+        Run a polling loop, which runs poll_action forever unless this is failed.
+        Args:
+          poll_action: A callable returning an awaitable to run in a loop. Note that async functions which return once
+          satisfy this.
+        """
+        try:
+            while True:
+                await self.await_unless_failed(poll_action())
+        except asyncio.CancelledError:
+            pass
+        except Exception as e:
+            self.fail(adapt_error(e))
+
+    def fail(self, err: GoogleAPICallError):
+        if not self._failure_task.done():
+            self._failure_task.set_exception(err)
+            # Ensure that even if _failure_task is never used, the exception is
+            # retrieved and the asyncio runtime doesn't print an error.
+            self._failure_task.exception()
+
+    def error(self) -> Optional[GoogleAPICallError]:
+        if not self._failure_task.done():
+            return None
+        return self._failure_task.exception()

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/publisher.py

@@ -0,0 +1,40 @@
+# Copyright 2020 Google LLC
+#
+# 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 abc import abstractmethod, ABCMeta
+from typing import AsyncContextManager
+from google.cloud.pubsublite_v1.types import PubSubMessage
+from google.cloud.pubsublite.types import MessageMetadata
+
+
+class Publisher(AsyncContextManager, metaclass=ABCMeta):
+    """
+    A Pub/Sub Lite asynchronous wire protocol publisher.
+    """
+
+    @abstractmethod
+    async def publish(self, message: PubSubMessage) -> MessageMetadata:
+        """
+        Publish the provided message.
+
+        Args:
+          message: The message to be published.
+
+        Returns:
+          Metadata about the published message.
+
+        Raises:
+          GoogleAPICallError: On a permanent error.
+        """
+        raise NotImplementedError()

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/pubsub_context.py

@@ -0,0 +1,56 @@
+# Copyright 2020 Google LLC
+#
+# 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 base64 import b64encode
+from typing import Mapping, Optional, NamedTuple
+
+import logging
+import pkg_resources
+from cloudsdk.google.protobuf import struct_pb2  # pytype: disable=pyi-error
+
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class _Semver(NamedTuple):
+    major: int
+    minor: int
+
+
+def _version() -> _Semver:
+    try:
+        version = pkg_resources.get_distribution("google-cloud-pubsublite").version
+    except pkg_resources.DistributionNotFound:
+        _LOGGER.info(
+            "Failed to extract the google-cloud-pubsublite semver version. DistributionNotFound."
+        )
+        return _Semver(0, 0)
+    splits = version.split(".")
+    if len(splits) != 3:
+        _LOGGER.info(f"Failed to extract semver from {version}.")
+        return _Semver(0, 0)
+    return _Semver(int(splits[0]), int(splits[1]))
+
+
+def pubsub_context(framework: Optional[str] = None) -> Mapping[str, str]:
+    """Construct the pubsub context mapping for the given framework."""
+    context = struct_pb2.Struct()
+    context.fields["language"].string_value = "PYTHON"
+    if framework:
+        context.fields["framework"].string_value = framework
+    version = _version()
+    context.fields["major_version"].number_value = version.major
+    context.fields["minor_version"].number_value = version.minor
+    encoded = b64encode(context.SerializeToString()).decode("utf-8")
+    return {"x-goog-pubsub-context": encoded}

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/reset_signal.py

@@ -0,0 +1,50 @@
+# Copyright 2021 Google LLC
+#
+# 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 google.api_core.exceptions import GoogleAPICallError
+from google.cloud.pubsublite.internal.status_codes import is_retryable
+from google.rpc.error_details_pb2 import ErrorInfo
+from google.rpc import status_pb2
+
+
+def is_reset_signal(error: GoogleAPICallError) -> bool:
+    """
+    Determines whether the given error contains the stream RESET signal, sent by
+    the server to instruct clients to reset stream state.
+
+    Returns: True if the error contains the RESET signal.
+    """
+    if not is_retryable(error) or not error.response:
+        return False
+    call = error.response
+    if call.trailing_metadata() is None:
+        return False
+    for key, value in call.trailing_metadata():
+        if key == "grpc-status-details-bin":
+            rich_status = status_pb2.Status.FromString(value)
+            if (
+                call.code().value[0] != rich_status.code
+                or call.details() != rich_status.message
+            ):
+                return False
+            for detail in rich_status.details:
+                if detail.Is(ErrorInfo.DESCRIPTOR):
+                    info = ErrorInfo()
+                    if (
+                        detail.Unpack(info)
+                        and info.reason == "RESET"
+                        and info.domain == "pubsublite.googleapis.com"
+                    ):
+                        return True
+    return False

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/retrying_connection.py

@@ -0,0 +1,172 @@
+# Copyright 2020 Google LLC
+#
+# 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 asyncio
+from asyncio import Future
+import logging
+import traceback
+
+from google.api_core.exceptions import Cancelled
+from google.cloud.pubsublite.internal.wire.permanent_failable import adapt_error
+from google.cloud.pubsublite.internal.status_codes import is_retryable
+from google.cloud.pubsublite.internal.wait_ignore_cancelled import (
+    wait_ignore_errors,
+    wait_ignore_cancelled,
+)
+from google.cloud.pubsublite.internal.wire.connection_reinitializer import (
+    ConnectionReinitializer,
+)
+from google.cloud.pubsublite.internal.wire.connection import (
+    Connection,
+    Request,
+    Response,
+    ConnectionFactory,
+)
+from google.cloud.pubsublite.internal.wire.work_item import WorkItem
+from google.cloud.pubsublite.internal.wire.permanent_failable import PermanentFailable
+
+_MIN_BACKOFF_SECS = 0.01
+_MAX_BACKOFF_SECS = 10
+
+
+class RetryingConnection(Connection[Request, Response], PermanentFailable):
+    """A connection which performs retries on an underlying stream when experiencing retryable errors."""
+
+    _connection_factory: ConnectionFactory[Request, Response]
+    _reinitializer: ConnectionReinitializer[Request, Response]
+    _initialized_once: asyncio.Event
+
+    _loop_task: asyncio.Future
+
+    _write_queue: "asyncio.Queue[WorkItem[Request, None]]"
+    _read_queue: "asyncio.Queue[Response]"
+
+    def __init__(
+        self,
+        connection_factory: ConnectionFactory[Request, Response],
+        reinitializer: ConnectionReinitializer[Request, Response],
+    ):
+        super().__init__()
+        self._connection_factory = connection_factory
+        self._reinitializer = reinitializer
+        self._initialized_once = asyncio.Event()
+        self._write_queue = asyncio.Queue(maxsize=1)
+        self._read_queue = asyncio.Queue(maxsize=1)
+
+    async def __aenter__(self):
+        self._loop_task = asyncio.ensure_future(self._run_loop())
+        await self.await_unless_failed(self._initialized_once.wait())
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self.fail(Cancelled("Connection shutting down."))
+        self._loop_task.cancel()
+        await wait_ignore_errors(self._loop_task)
+
+    async def write(self, request: Request) -> None:
+        item = WorkItem(request)
+        await self.await_unless_failed(self._write_queue.put(item))
+        return await self.await_unless_failed(item.response_future)
+
+    async def read(self) -> Response:
+        return await self.await_unless_failed(self._read_queue.get())
+
+    async def _run_loop(self):
+        """
+        Processes actions on this connection and handles retries until cancelled.
+        """
+        try:
+            bad_retries = 0
+            while not self.error():
+                try:
+                    conn_fut = self._connection_factory.new()
+                    async with (await conn_fut) as connection:
+                        await self._reinitializer.reinitialize(
+                            connection  # pytype: disable=name-error
+                        )
+                        self._initialized_once.set()
+                        bad_retries = 0
+                        await self._loop_connection(
+                            connection  # pytype: disable=name-error
+                        )
+                except Exception as e:
+                    if self.error():
+                        return
+                    e = adapt_error(e)
+                    logging.debug(
+                        "Saw a stream failure. Cause: \n%s", traceback.format_exc()
+                    )
+                    if not is_retryable(e):
+                        self.fail(e)
+                        return
+                    try:
+                        await self._reinitializer.stop_processing(e)
+                    except Exception as stop_error:
+                        self.fail(adapt_error(stop_error))
+                        return
+                    while not self._write_queue.empty():
+                        response_future = self._write_queue.get_nowait().response_future
+                        if not response_future.cancelled():
+                            response_future.set_exception(e)
+                    self._read_queue = asyncio.Queue(maxsize=1)
+                    self._write_queue = asyncio.Queue(maxsize=1)
+                    await wait_ignore_cancelled(
+                        asyncio.sleep(
+                            min(
+                                _MAX_BACKOFF_SECS,
+                                _MIN_BACKOFF_SECS * (2**bad_retries),
+                            )
+                        )
+                    )
+                    bad_retries += 1
+        except Exception as e:
+            logging.error(
+                "Saw a stream failure which was unhandled. Cause: \n%s",
+                traceback.format_exc(),
+            )
+            self.fail(adapt_error(e))
+
+    async def _loop_connection(self, connection: Connection[Request, Response]):
+        read_task: "Future[Response]" = asyncio.ensure_future(connection.read())
+        write_task: "Future[WorkItem[Request, None]]" = asyncio.ensure_future(
+            self._write_queue.get()
+        )
+        try:
+            while True:
+                done, _ = await asyncio.wait(
+                    [write_task, read_task], return_when=asyncio.FIRST_COMPLETED
+                )
+                if write_task in done:
+                    await self._handle_write(connection, await write_task)
+                    write_task = asyncio.ensure_future(self._write_queue.get())
+                if read_task in done:
+                    await self._read_queue.put(await read_task)
+                    read_task = asyncio.ensure_future(connection.read())
+        finally:
+            read_task.cancel()
+            write_task.cancel()
+            await wait_ignore_errors(read_task)
+            await wait_ignore_errors(write_task)
+
+    @staticmethod
+    async def _handle_write(
+        connection: Connection[Request, Response], to_write: WorkItem[Request, Response]
+    ):
+        try:
+            await connection.write(to_write.request)
+            to_write.response_future.set_result(None)
+        except Exception as e:
+            e = adapt_error(e)
+            to_write.response_future.set_exception(e)
+            raise e

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/routing_policy.py

@@ -0,0 +1,34 @@
+# Copyright 2020 Google LLC
+#
+# 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 abc import ABC, abstractmethod
+
+from google.cloud.pubsublite.types.partition import Partition
+from google.cloud.pubsublite_v1.types.common import PubSubMessage
+
+
+class RoutingPolicy(ABC):
+    """A policy for how to route messages."""
+
+    @abstractmethod
+    def route(self, message: PubSubMessage) -> Partition:
+        """
+        Route a message to a given partition.
+        Args:
+          message: The message to route
+
+        Returns: The partition to route to
+
+        """
+        raise NotImplementedError()

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/routing_publisher.py

@@ -0,0 +1,48 @@
+# Copyright 2020 Google LLC
+#
+# 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 asyncio
+from typing import Mapping
+
+from google.cloud.pubsublite.internal.wire.publisher import Publisher
+from google.cloud.pubsublite.internal.wire.routing_policy import RoutingPolicy
+from google.cloud.pubsublite.types import Partition, MessageMetadata
+from google.cloud.pubsublite_v1 import PubSubMessage
+
+
+class RoutingPublisher(Publisher):
+    _routing_policy: RoutingPolicy
+    _publishers: Mapping[Partition, Publisher]
+
+    def __init__(
+        self, routing_policy: RoutingPolicy, publishers: Mapping[Partition, Publisher]
+    ):
+        self._routing_policy = routing_policy
+        self._publishers = publishers
+
+    async def __aenter__(self):
+        enter_futures = [
+            publisher.__aenter__() for publisher in self._publishers.values()
+        ]
+        await asyncio.gather(*enter_futures)
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        for publisher in self._publishers.values():
+            await publisher.__aexit__(exc_type, exc_val, exc_tb)
+
+    async def publish(self, message: PubSubMessage) -> MessageMetadata:
+        partition = self._routing_policy.route(message)
+        assert partition in self._publishers
+        return await self._publishers[partition].publish(message)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/serial_batcher.py

@@ -0,0 +1,83 @@
+# Copyright 2020 Google LLC
+#
+# 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 abc import abstractmethod, ABCMeta
+from typing import Generic, List, NamedTuple
+import asyncio
+
+from google.cloud.pubsublite.internal.wire.connection import Request, Response
+from google.cloud.pubsublite.internal.wire.work_item import WorkItem
+
+
+class BatchSize(NamedTuple):
+    element_count: int
+    byte_count: int
+
+    def __add__(self, other: "BatchSize") -> "BatchSize":
+        return BatchSize(
+            self.element_count + other.element_count, self.byte_count + other.byte_count
+        )
+
+
+class RequestSizer(Generic[Request], metaclass=ABCMeta):
+    """A RequestSizer determines the size of a request."""
+
+    @abstractmethod
+    def get_size(self, request: Request) -> BatchSize:
+        """
+        Args:
+          request: A single request.
+
+        Returns: The BatchSize of this request
+        """
+        raise NotImplementedError()
+
+
+class IgnoredRequestSizer(RequestSizer[Request]):
+    def get_size(self, request) -> BatchSize:
+        return BatchSize(0, 0)
+
+
+class SerialBatcher(Generic[Request, Response]):
+    _sizer: RequestSizer[Request]
+    _requests: List[WorkItem[Request, Response]]  # A list of outstanding requests
+    _batch_size: BatchSize
+
+    def __init__(self, sizer: RequestSizer[Request] = IgnoredRequestSizer()):
+        self._sizer = sizer
+        self._requests = []
+        self._batch_size = BatchSize(0, 0)
+
+    def add(self, request: Request) -> "asyncio.Future[Response]":
+        """Add a new request to this batcher.
+
+        Args:
+          request: The request to send.
+
+        Returns:
+          A future that will resolve to the response or a GoogleAPICallError.
+        """
+        item = WorkItem[Request, Response](request)
+        self._requests.append(item)
+        self._batch_size += self._sizer.get_size(request)
+        return item.response_future
+
+    def size(self) -> BatchSize:
+        return self._batch_size
+
+    def flush(self) -> List[WorkItem[Request, Response]]:
+        requests = self._requests
+        self._requests = []
+        self._batch_size = BatchSize(0, 0)
+        return requests

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/single_partition_publisher.py

@@ -0,0 +1,232 @@
+# Copyright 2020 Google LLC
+#
+# 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 asyncio
+from typing import Optional, List, NamedTuple
+
+import logging
+from google.cloud.pubsub_v1.types import BatchSettings
+
+from google.cloud.pubsublite.internal.publish_sequence_number import (
+    PublishSequenceNumber,
+)
+from google.cloud.pubsublite.internal.wait_ignore_cancelled import wait_ignore_errors
+from google.cloud.pubsublite.internal.wire.publisher import Publisher
+from google.cloud.pubsublite.internal.wire.retrying_connection import (
+    RetryingConnection,
+    ConnectionFactory,
+)
+from google.api_core.exceptions import FailedPrecondition, GoogleAPICallError
+from google.cloud.pubsublite.internal.wire.connection_reinitializer import (
+    ConnectionReinitializer,
+)
+from google.cloud.pubsublite.internal.wire.connection import Connection
+from google.cloud.pubsublite.internal.wire.serial_batcher import (
+    SerialBatcher,
+    RequestSizer,
+    BatchSize,
+)
+from google.cloud.pubsublite.types import Partition, MessageMetadata
+from google.cloud.pubsublite_v1.types import (
+    PubSubMessage,
+    Cursor,
+    PublishRequest,
+    PublishResponse,
+    InitialPublishRequest,
+)
+from google.cloud.pubsublite.internal.wire.work_item import WorkItem
+
+_LOGGER = logging.getLogger(__name__)
+
+# Maximum bytes per batch at 3.5 MiB to avoid GRPC limit of 4 MiB
+_MAX_BYTES = int(3.5 * 1024 * 1024)
+
+# Maximum messages per batch at 1000
+_MAX_MESSAGES = 1000
+
+
+class _MessageWithSequence(NamedTuple):
+    message: PubSubMessage
+    sequence_number: PublishSequenceNumber
+
+
+class SinglePartitionPublisher(
+    Publisher,
+    ConnectionReinitializer[PublishRequest, PublishResponse],
+    RequestSizer[_MessageWithSequence],
+):
+    _initial: InitialPublishRequest
+    _batching_settings: BatchSettings
+    _connection: RetryingConnection[PublishRequest, PublishResponse]
+
+    _next_sequence: PublishSequenceNumber
+    _batcher: SerialBatcher[_MessageWithSequence, Cursor]
+    _outstanding_writes: List[List[WorkItem[_MessageWithSequence, Cursor]]]
+
+    _receiver: Optional[asyncio.Future]
+    _flusher: Optional[asyncio.Future]
+
+    def __init__(
+        self,
+        initial: InitialPublishRequest,
+        batching_settings: BatchSettings,
+        factory: ConnectionFactory[PublishRequest, PublishResponse],
+        initial_sequence: PublishSequenceNumber,
+    ):
+        self._initial = initial
+        self._batching_settings = batching_settings
+        self._connection = RetryingConnection(factory, self)
+        self._next_sequence = initial_sequence
+        self._batcher = SerialBatcher(self)
+        self._outstanding_writes = []
+        self._receiver = None
+        self._flusher = None
+
+    @property
+    def _partition(self) -> Partition:
+        return Partition(self._initial.partition)
+
+    async def __aenter__(self):
+        await self._connection.__aenter__()
+        return self
+
+    def _start_loopers(self):
+        assert self._receiver is None
+        assert self._flusher is None
+        self._receiver = asyncio.ensure_future(self._receive_loop())
+        self._flusher = asyncio.ensure_future(self._flush_loop())
+
+    async def _stop_loopers(self):
+        if self._receiver:
+            self._receiver.cancel()
+            await wait_ignore_errors(self._receiver)
+            self._receiver = None
+        if self._flusher:
+            self._flusher.cancel()
+            await wait_ignore_errors(self._flusher)
+            self._flusher = None
+
+    def _handle_response(self, response: PublishResponse):
+        if "message_response" not in response:
+            self._connection.fail(
+                FailedPrecondition(
+                    "Received an invalid subsequent response on the publish stream."
+                )
+            )
+        if not self._outstanding_writes:
+            self._connection.fail(
+                FailedPrecondition(
+                    "Received an publish response on the stream with no outstanding publishes."
+                )
+            )
+        ranges = response.message_response.cursor_ranges
+        ranges.sort(key=lambda r: r.start_index)
+        batch: List[WorkItem[_MessageWithSequence]] = self._outstanding_writes.pop(0)
+        range_idx = 0
+        for msg_idx, item in enumerate(batch):
+            if range_idx < len(ranges) and ranges[range_idx].end_index <= msg_idx:
+                range_idx += 1
+            offset = -1
+            if (
+                range_idx < len(ranges)
+                and msg_idx >= ranges[range_idx].start_index
+                and msg_idx < ranges[range_idx].end_index
+            ):
+                offset_in_range = msg_idx - ranges[range_idx].start_index
+                offset = ranges[range_idx].start_cursor.offset + offset_in_range
+            item.response_future.set_result(Cursor(offset=offset))
+
+    async def _receive_loop(self):
+        while True:
+            response = await self._connection.read()
+            self._handle_response(response)
+
+    async def _flush_loop(self):
+        while True:
+            await asyncio.sleep(self._batching_settings.max_latency)
+            await self._flush()
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if self._connection.error():
+            self._fail_if_retrying_failed()
+        else:
+            await self._flush()
+        await self._stop_loopers()
+        await self._connection.__aexit__(exc_type, exc_val, exc_tb)
+
+    def _fail_if_retrying_failed(self):
+        if self._connection.error():
+            for batch in self._outstanding_writes:
+                for item in batch:
+                    item.response_future.set_exception(self._connection.error())
+
+    async def _flush(self):
+        batch = self._batcher.flush()
+        if not batch:
+            return
+        self._outstanding_writes.append(batch)
+        aggregate = PublishRequest()
+        aggregate.message_publish_request.messages = [
+            item.request.message for item in batch
+        ]
+        if self._initial.client_id:
+            aggregate.message_publish_request.first_sequence_number = batch[
+                0
+            ].request.sequence_number.value
+        try:
+            await self._connection.write(aggregate)
+        except GoogleAPICallError as e:
+            _LOGGER.debug(f"Failed publish on stream: {e}")
+            self._fail_if_retrying_failed()
+
+    async def publish(self, message: PubSubMessage) -> MessageMetadata:
+        future = self._batcher.add(_MessageWithSequence(message, self._next_sequence))
+        self._next_sequence = self._next_sequence.next()
+        if self._should_flush():
+            await self._flush()
+        return MessageMetadata(self._partition, await future)
+
+    async def stop_processing(self, error: GoogleAPICallError):
+        await self._stop_loopers()
+
+    async def reinitialize(
+        self,
+        connection: Connection[PublishRequest, PublishResponse],
+    ):
+        await connection.write(PublishRequest(initial_request=self._initial))
+        response = await connection.read()
+        if "initial_response" not in response:
+            self._connection.fail(
+                FailedPrecondition(
+                    "Received an invalid initial response on the publish stream."
+                )
+            )
+        for batch in self._outstanding_writes:
+            aggregate = PublishRequest()
+            aggregate.message_publish_request.messages = [
+                item.request.message for item in batch
+            ]
+            await connection.write(aggregate)
+        self._start_loopers()
+
+    def get_size(self, request: _MessageWithSequence) -> BatchSize:
+        return BatchSize(
+            element_count=1, byte_count=PubSubMessage.pb(request.message).ByteSize()
+        )
+
+    def _should_flush(self) -> bool:
+        size = self._batcher.size()
+        return (size.element_count >= self._batching_settings.max_messages) or (
+            size.byte_count >= self._batching_settings.max_bytes
+        )

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/subscriber.py

@@ -0,0 +1,43 @@
+# Copyright 2020 Google LLC
+#
+# 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 abc import abstractmethod, ABCMeta
+from typing import AsyncContextManager, List
+from google.cloud.pubsublite_v1.types import SequencedMessage, FlowControlRequest
+
+
+class Subscriber(AsyncContextManager, metaclass=ABCMeta):
+    """
+    A Pub/Sub Lite asynchronous wire protocol subscriber.
+    """
+
+    @abstractmethod
+    async def read(self) -> List[SequencedMessage.meta.pb]:
+        """
+        Read a batch of messages off of the stream.
+
+        Returns:
+          The next batch of messages.
+
+        Raises:
+          GoogleAPICallError: On a permanent error.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def allow_flow(self, request: FlowControlRequest):
+        """
+        Allow an additional amount of messages and bytes to be sent to this client.
+        """
+        raise NotImplementedError()

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/subscriber_impl.py

@@ -0,0 +1,205 @@
+# Copyright 2020 Google LLC
+#
+# 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 asyncio
+from copy import deepcopy
+from typing import Optional, List
+
+from google.api_core.exceptions import GoogleAPICallError, FailedPrecondition
+
+from google.cloud.pubsublite.internal.wait_ignore_cancelled import wait_ignore_errors
+from google.cloud.pubsublite.internal.wire.connection import (
+    Connection,
+    ConnectionFactory,
+)
+from google.cloud.pubsublite.internal.wire.connection_reinitializer import (
+    ConnectionReinitializer,
+)
+from google.cloud.pubsublite.internal.wire.flow_control_batcher import (
+    FlowControlBatcher,
+)
+from google.cloud.pubsublite.internal.wire.reset_signal import is_reset_signal
+from google.cloud.pubsublite.internal.wire.retrying_connection import RetryingConnection
+from google.cloud.pubsublite.internal.wire.subscriber import Subscriber
+from google.cloud.pubsublite_v1 import (
+    SubscribeRequest,
+    SubscribeResponse,
+    FlowControlRequest,
+    SequencedMessage,
+    InitialSubscribeRequest,
+    SeekRequest,
+    Cursor,
+)
+from google.cloud.pubsublite.internal.wire.subscriber_reset_handler import (
+    SubscriberResetHandler,
+)
+
+
+class SubscriberImpl(
+    Subscriber, ConnectionReinitializer[SubscribeRequest, SubscribeResponse]
+):
+    _base_initial: InitialSubscribeRequest
+    _token_flush_seconds: float
+    _connection: RetryingConnection[SubscribeRequest, SubscribeResponse]
+    _reset_handler: SubscriberResetHandler
+
+    _outstanding_flow_control: FlowControlBatcher
+
+    _last_received_offset: Optional[int]
+
+    _message_queue: "asyncio.Queue[List[SequencedMessage.meta.pb]]"
+
+    _receiver: Optional[asyncio.Future]
+    _flusher: Optional[asyncio.Future]
+
+    def __init__(
+        self,
+        base_initial: InitialSubscribeRequest,
+        token_flush_seconds: float,
+        factory: ConnectionFactory[SubscribeRequest, SubscribeResponse],
+        reset_handler: SubscriberResetHandler,
+    ):
+        self._base_initial = base_initial
+        self._token_flush_seconds = token_flush_seconds
+        self._connection = RetryingConnection(factory, self)
+        self._reset_handler = reset_handler
+        self._outstanding_flow_control = FlowControlBatcher()
+        self._reinitializing = False
+        self._last_received_offset = None
+        self._message_queue = asyncio.Queue()
+        self._receiver = None
+        self._flusher = None
+
+    async def __aenter__(self):
+        await self._connection.__aenter__()
+        return self
+
+    def _start_loopers(self):
+        assert self._receiver is None
+        assert self._flusher is None
+        self._receiver = asyncio.ensure_future(self._receive_loop())
+        self._flusher = asyncio.ensure_future(self._flush_loop())
+
+    async def _stop_loopers(self):
+        if self._receiver:
+            self._receiver.cancel()
+            await wait_ignore_errors(self._receiver)
+            self._receiver = None
+        if self._flusher:
+            self._flusher.cancel()
+            await wait_ignore_errors(self._flusher)
+            self._flusher = None
+
+    def _handle_response(self, response: SubscribeResponse):
+        if "messages" not in response:
+            self._connection.fail(
+                FailedPrecondition(
+                    "Received an invalid subsequent response on the subscribe stream."
+                )
+            )
+            return
+        # Workaround for incredibly slow proto-plus-python accesses
+        messages = list(
+            response.messages.messages._pb  # pytype: disable=attribute-error
+        )
+        self._outstanding_flow_control.on_messages(messages)
+        for message in messages:
+            if (
+                self._last_received_offset is not None
+                and message.cursor.offset <= self._last_received_offset
+            ):
+                self._connection.fail(
+                    FailedPrecondition(
+                        "Received an invalid out of order message from the server. Message is {}, previous last received is {}.".format(
+                            message.cursor.offset, self._last_received_offset
+                        )
+                    )
+                )
+                return
+            self._last_received_offset = message.cursor.offset
+        # queue is unbounded.
+        self._message_queue.put_nowait(messages)
+
+    async def _receive_loop(self):
+        while True:
+            response = await self._connection.read()
+            self._handle_response(response)
+
+    async def _try_send_tokens(self):
+        req = self._outstanding_flow_control.release_pending_request()
+        if req is None:
+            return
+        try:
+            await self._connection.write(SubscribeRequest(flow_control=req))
+        except GoogleAPICallError:
+            # May be transient, in which case these tokens will be resent.
+            pass
+
+    async def _flush_loop(self):
+        while True:
+            await asyncio.sleep(self._token_flush_seconds)
+            await self._try_send_tokens()
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self._stop_loopers()
+        await self._connection.__aexit__(exc_type, exc_val, exc_tb)
+
+    async def stop_processing(self, error: GoogleAPICallError):
+        await self._stop_loopers()
+        if is_reset_signal(error):
+            # Discard undelivered messages and refill flow control tokens.
+            while not self._message_queue.empty():
+                batch: List[SequencedMessage.meta.pb] = self._message_queue.get_nowait()
+                allowed_bytes = sum(message.size_bytes for message in batch)
+                self._outstanding_flow_control.add(
+                    FlowControlRequest(
+                        allowed_messages=len(batch),
+                        allowed_bytes=allowed_bytes,
+                    )
+                )
+
+            await self._reset_handler.handle_reset()
+            self._last_received_offset = None
+
+    async def reinitialize(
+        self, connection: Connection[SubscribeRequest, SubscribeResponse]
+    ):
+        initial = deepcopy(self._base_initial)
+        if self._last_received_offset is not None:
+            initial.initial_location = SeekRequest(
+                cursor=Cursor(offset=self._last_received_offset + 1)
+            )
+        else:
+            initial.initial_location = SeekRequest(
+                named_target=SeekRequest.NamedTarget.COMMITTED_CURSOR
+            )
+        await connection.write(SubscribeRequest(initial=initial))
+        response = await connection.read()
+        if "initial" not in response:
+            self._connection.fail(
+                FailedPrecondition(
+                    "Received an invalid initial response on the subscribe stream."
+                )
+            )
+            return
+        tokens = self._outstanding_flow_control.request_for_restart()
+        if tokens is not None:
+            await connection.write(SubscribeRequest(flow_control=tokens))
+        self._start_loopers()
+
+    async def read(self) -> List[SequencedMessage.meta.pb]:
+        return await self._connection.await_unless_failed(self._message_queue.get())
+
+    def allow_flow(self, request: FlowControlRequest):
+        self._outstanding_flow_control.add(request)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/subscriber_reset_handler.py

@@ -0,0 +1,28 @@
+# Copyright 2021 Google LLC
+#
+# 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 abc import ABCMeta, abstractmethod
+
+
+class SubscriberResetHandler(metaclass=ABCMeta):
+    """Helps to reset subscriber state when the `RESET` signal is received from the server."""
+
+    @abstractmethod
+    async def handle_reset(self):
+        """Reset subscriber state.
+
+        Raises:
+            GoogleAPICallError: If reset handling fails. The subscriber will shut down.
+        """
+        raise NotImplementedError()

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/internal/wire/work_item.py

@@ -0,0 +1,29 @@
+# Copyright 2020 Google LLC
+#
+# 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 asyncio
+from typing import Generic
+
+from google.cloud.pubsublite.internal.wire.connection import Request, Response
+
+
+class WorkItem(Generic[Request, Response]):
+    """An item of work and a future to complete when it is finished."""
+
+    request: Request
+    response_future: "asyncio.Future[Response]"
+
+    def __init__(self, request: Request):
+        self.request = request
+        self.response_future = asyncio.Future()

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/testing/test_reset_signal.py

@@ -0,0 +1,44 @@
+# Copyright 2021 Google LLC
+#
+# 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 unittest.mock import MagicMock
+from google.api_core.exceptions import Aborted, GoogleAPICallError
+from cloudsdk.google.protobuf.any_pb2 import Any  # pytype: disable=pyi-error
+from google.rpc.error_details_pb2 import ErrorInfo
+from google.rpc.status_pb2 import Status
+import grpc
+from grpc_status import rpc_status
+
+
+def make_call(status_pb: Status) -> grpc.Call:
+    status = rpc_status.to_status(status_pb)
+    mock_call = MagicMock(spec=grpc.Call)
+    mock_call.details.return_value = status.details
+    mock_call.code.return_value = status.code
+    mock_call.trailing_metadata.return_value = status.trailing_metadata
+    return mock_call
+
+
+def make_call_without_metadata(status_pb: Status) -> grpc.Call:
+    mock_call = make_call(status_pb)
+    # Causes rpc_status.from_call to return None.
+    mock_call.trailing_metadata.return_value = None
+    return mock_call
+
+
+def make_reset_signal() -> GoogleAPICallError:
+    any = Any()
+    any.Pack(ErrorInfo(reason="RESET", domain="pubsublite.googleapis.com"))
+    status_pb = Status(code=10, details=[any])
+    return Aborted("", response=make_call(status_pb))

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/testing/test_utils.py

@@ -0,0 +1,78 @@
+# Copyright 2020 Google LLC
+#
+# 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 asyncio
+import threading
+from typing import List, Union, Any, TypeVar, Generic, Optional, Callable, Awaitable
+
+from unittest.mock import AsyncMock
+
+T = TypeVar("T")
+
+
+async def async_iterable(elts: List[Union[Any, Exception]]):
+    for elt in elts:
+        if isinstance(elt, Exception):
+            raise elt
+        yield elt
+
+
+def make_queue_waiter(
+    started_q: "asyncio.Queue[None]", result_q: "asyncio.Queue[Union[T, Exception]]"
+) -> Callable[[], Awaitable[T]]:
+    """
+    Given a queue to notify when started and a queue to get results from, return a waiter which
+    notifies started_q when started and returns from result_q when done.
+    """
+
+    async def waiter(*args, **kwargs):
+        await started_q.put(None)
+        result = await result_q.get()
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    return waiter
+
+
+class QueuePair:
+    called: asyncio.Queue
+    results: asyncio.Queue
+
+    def __init__(self):
+        self.called = asyncio.Queue()
+        self.results = asyncio.Queue()
+
+
+def wire_queues(mock: AsyncMock) -> QueuePair:
+    queues = QueuePair()
+    mock.side_effect = make_queue_waiter(queues.called, queues.results)
+    return queues
+
+
+class Box(Generic[T]):
+    val: Optional[T]
+
+
+def run_on_thread(func: Callable[[], T]) -> T:
+    box = Box()
+
+    def set_box():
+        box.val = func()
+
+    # Initialize watcher on another thread with a different event loop.
+    thread = threading.Thread(target=set_box)
+    thread.start()
+    thread.join()
+    return box.val

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/types/__init__.py

@@ -0,0 +1,37 @@
+# Copyright 2020 Google LLC
+#
+# 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 .location import CloudRegion, CloudZone
+from .partition import Partition
+from .paths import LocationPath, ReservationPath, TopicPath, SubscriptionPath
+from .message_metadata import MessageMetadata
+from .flow_control_settings import FlowControlSettings, DISABLED_FLOW_CONTROL
+from .backlog_location import BacklogLocation, PublishTime, EventTime
+
+__all__ = (
+    "BacklogLocation",
+    "CloudRegion",
+    "CloudZone",
+    "EventTime",
+    "FlowControlSettings",
+    "LocationPath",
+    "MessageMetadata",
+    "Partition",
+    "PublishTime",
+    "Reservation",
+    "ReservationPath",
+    "SubscriptionPath",
+    "TopicPath",
+)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/types/backlog_location.py

@@ -0,0 +1,38 @@
+# Copyright 2020 Google LLC
+#
+# 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
+from datetime import datetime
+from typing import NamedTuple
+
+
+class BacklogLocation(enum.Enum):
+    """A location with respect to the message backlog. BEGINNING refers to the
+    location of the oldest retained message. END refers to the location past
+    all currently published messages, skipping the entire message backlog."""
+
+    BEGINNING = 0
+    END = 1
+
+
+class PublishTime(NamedTuple):
+    """The publish timestamp of a message."""
+
+    value: datetime
+
+
+class EventTime(NamedTuple):
+    """A user-defined event timestamp of a message."""
+
+    value: datetime

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/types/flow_control_settings.py

@@ -0,0 +1,25 @@
+# Copyright 2020 Google LLC
+#
+# 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 typing import NamedTuple
+
+
+class FlowControlSettings(NamedTuple):
+    messages_outstanding: int
+    bytes_outstanding: int
+
+
+_MAX_INT64 = 0x7FFFFFFFFFFFFFFF
+
+DISABLED_FLOW_CONTROL = FlowControlSettings(_MAX_INT64, _MAX_INT64)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/types/location.py

@@ -0,0 +1,51 @@
+# Copyright 2020 Google LLC
+#
+# 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 typing import NamedTuple
+
+from google.api_core.exceptions import InvalidArgument
+
+
+class CloudRegion(NamedTuple):
+    name: str
+
+    def __str__(self):
+        return self.name
+
+    @staticmethod
+    def parse(to_parse: str):
+        splits = to_parse.split("-")
+        if len(splits) != 2:
+            raise InvalidArgument("Invalid region name: " + to_parse)
+        return CloudRegion(name=splits[0] + "-" + splits[1])
+
+    @property
+    def region(self):
+        return self
+
+
+class CloudZone(NamedTuple):
+    region: CloudRegion
+    zone_id: str
+
+    def __str__(self):
+        return f"{self.region.name}-{self.zone_id}"
+
+    @staticmethod
+    def parse(to_parse: str):
+        splits = to_parse.split("-")
+        if len(splits) != 3 or len(splits[2]) != 1:
+            raise InvalidArgument("Invalid zone name: " + to_parse)
+        region = CloudRegion(name=splits[0] + "-" + splits[1])
+        return CloudZone(region, zone_id=splits[2])

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/types/message_metadata.py

@@ -0,0 +1,53 @@
+# Copyright 2020 Google LLC
+#
+# 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 typing import NamedTuple
+
+from google.cloud.pubsublite.internal import fast_serialize
+from google.cloud.pubsublite_v1.types.common import Cursor
+from google.cloud.pubsublite.types.partition import Partition
+
+
+class MessageMetadata(NamedTuple):
+    """Information about a message in Pub/Sub Lite.
+
+    Attributes:
+        partition (Partition):
+            The partition of the topic that the message was published to.
+        cursor (Cursor):
+            A cursor containing the offset that the message was assigned.
+            If this MessageMetadata was returned for a publish result and
+            publish idempotence was enabled, the offset may be -1 when the
+            message was identified as a duplicate of an already successfully
+            published message, but the server did not have sufficient
+            information to return the message's offset at publish time. Messages
+            received by subscribers will always have the correct offset.
+    """
+
+    partition: Partition
+    cursor: Cursor
+
+    def encode(self) -> str:
+        return self._encode_parts(self.partition.value, self.cursor._pb.offset)
+
+    @staticmethod
+    def _encode_parts(partition: int, offset: int) -> str:
+        return fast_serialize.dump([partition, offset])
+
+    @staticmethod
+    def decode(source: str) -> "MessageMetadata":
+        loaded = fast_serialize.load(source)
+        cursor = Cursor()
+        cursor._pb.offset = loaded[1]
+        return MessageMetadata(partition=Partition(loaded[0]), cursor=cursor)

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/types/partition.py

@@ -0,0 +1,19 @@
+# Copyright 2020 Google LLC
+#
+# 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 typing import NamedTuple
+
+
+class Partition(NamedTuple):
+    value: int

gcloud auth application-default login/google-cloud-sdk/lib/third_party/google/cloud/pubsublite/types/paths.py

@@ -0,0 +1,130 @@
+# Copyright 2020 Google LLC
+#
+# 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 typing import NamedTuple, Union
+
+from google.api_core.exceptions import InvalidArgument
+
+from google.cloud.pubsublite.types.location import CloudZone, CloudRegion
+
+
+def _parse_location(to_parse: str) -> Union[CloudRegion, CloudZone]:
+    try:
+        return CloudZone.parse(to_parse)
+    except InvalidArgument:
+        pass
+    try:
+        return CloudRegion.parse(to_parse)
+    except InvalidArgument:
+        pass
+    raise InvalidArgument("Invalid location name: " + to_parse)
+
+
+class LocationPath(NamedTuple):
+    project: Union[int, str]
+    location: Union[CloudRegion, CloudZone]
+
+    def __str__(self):
+        return f"projects/{self.project}/locations/{self.location}"
+
+    @staticmethod
+    def parse(to_parse: str) -> "LocationPath":
+        splits = to_parse.split("/")
+        if len(splits) != 6 or splits[0] != "projects" or splits[2] != "locations":
+            raise InvalidArgument(
+                "Location path must be formatted like projects/{project_number}/locations/{location} but was instead "
+                + to_parse
+            )
+        return LocationPath(splits[1], _parse_location(splits[3]))
+
+
+class TopicPath(NamedTuple):
+    project: Union[int, str]
+    location: Union[CloudRegion, CloudZone]
+    name: str
+
+    def __str__(self):
+        return f"projects/{self.project}/locations/{self.location}/topics/{self.name}"
+
+    def to_location_path(self):
+        return LocationPath(self.project, self.location)
+
+    @staticmethod
+    def parse(to_parse: str) -> "TopicPath":
+        splits = to_parse.split("/")
+        if (
+            len(splits) != 6
+            or splits[0] != "projects"
+            or splits[2] != "locations"
+            or splits[4] != "topics"
+        ):
+            raise InvalidArgument(
+                "Topic path must be formatted like projects/{project_number}/locations/{location}/topics/{name} but was instead "
+                + to_parse
+            )
+        return TopicPath(splits[1], _parse_location(splits[3]), splits[5])
+
+
+class SubscriptionPath(NamedTuple):
+    project: Union[int, str]
+    location: Union[CloudRegion, CloudZone]
+    name: str
+
+    def __str__(self):
+        return f"projects/{self.project}/locations/{self.location}/subscriptions/{self.name}"
+
+    def to_location_path(self):
+        return LocationPath(self.project, self.location)
+
+    @staticmethod
+    def parse(to_parse: str) -> "SubscriptionPath":
+        splits = to_parse.split("/")
+        if (
+            len(splits) != 6
+            or splits[0] != "projects"
+            or splits[2] != "locations"
+            or splits[4] != "subscriptions"
+        ):
+            raise InvalidArgument(
+                "Subscription path must be formatted like projects/{project_number}/locations/{location}/subscriptions/{name} but was instead "
+                + to_parse
+            )
+        return SubscriptionPath(splits[1], _parse_location(splits[3]), splits[5])
+
+
+class ReservationPath(NamedTuple):
+    project: Union[int, str]
+    location: CloudRegion
+    name: str
+
+    def __str__(self):
+        return f"projects/{self.project}/locations/{self.location}/reservations/{self.name}"
+
+    def to_location_path(self):
+        return LocationPath(self.project, self.location)
+
+    @staticmethod
+    def parse(to_parse: str) -> "ReservationPath":
+        splits = to_parse.split("/")
+        if (
+            len(splits) != 6
+            or splits[0] != "projects"
+            or splits[2] != "locations"
+            or splits[4] != "reservations"
+        ):
+            raise InvalidArgument(
+                "Reservation path must be formatted like projects/{project_number}/locations/{location}/reservations/{name} but was instead "
+                + to_parse
+            )
+        return ReservationPath(splits[1], CloudRegion.parse(splits[3]), splits[5])