Samples: Add concurrency control sample

samples/snippets/subscriber.py

@@ -1003,6 +1003,49 @@ def callback(message: pubsub_v1.subscriber.message.Message) -> None:
     # [END pubsub_dead_letter_delivery_attempt]
 
 
+def receive_messages_with_concurrency_control(
+    project_id: str, subscription_id: str, timeout: Optional[float] = None
+) -> None:
+    # [START pubsub_subscriber_concurrency_control]
+    from concurrent import futures
+    from google.cloud import pubsub_v1
+
+    # TODO(developer)
+    # project_id = "your-project-id"
+    # subscription_id = "your-subscription-id"
+    # Number of seconds the subscriber should listen for messages
+    # timeout = 5.0
+
+    # An optional executor to use. If not specified, a default one with maximum 10
+    # threads will be created.
+    executor = futures.ThreadPoolExecutor(max_workers=5)
+    # A thread pool-based scheduler. It must not be shared across SubscriberClients.
+    scheduler = pubsub_v1.subscriber.scheduler.ThreadScheduler(executor)
+
+    subscriber = pubsub_v1.SubscriberClient()
+    subscription_path = subscriber.subscription_path(project_id, subscription_id)
+
+    def callback(message: pubsub_v1.subscriber.message.Message) -> None:
+        print(f"Received {message.data!r}.")
+        message.ack()
+
+    streaming_pull_future = subscriber.subscribe(
+        subscription_path, callback=callback, scheduler=scheduler
+    )
+    print(f"Listening for messages on {subscription_path}..\n")
+
+    # Wrap subscriber in a 'with' block to automatically call close() when done.
+    with subscriber:
+        try:
+            # When `timeout` is not set, result() will block indefinitely,
+            # unless an exception is encountered first.
+            streaming_pull_future.result(timeout=timeout)
+        except TimeoutError:
+            streaming_pull_future.cancel()  # Trigger the shutdown.
+            streaming_pull_future.result()  # Block until the shutdown is complete.
+    # [END pubsub_subscriber_concurrency_control]
+
+
 if __name__ == "__main__":  # noqa
     parser = argparse.ArgumentParser(
         description=__doc__,
@@ -1183,6 +1226,15 @@ def callback(message: pubsub_v1.subscriber.message.Message) -> None:
         "timeout", default=None, type=float, nargs="?"
     )
 
+    receive_messages_with_concurrency_control_parser = subparsers.add_parser(
+        "receive-messages-with-concurrency-control",
+        help=receive_messages_with_concurrency_control.__doc__,
+    )
+    receive_messages_with_concurrency_control_parser.add_argument("subscription_id")
+    receive_messages_with_concurrency_control_parser.add_argument(
+        "timeout", default=None, type=float, nargs="?"
+    )
+
     args = parser.parse_args()
 
     if args.command == "list-in-topic":
@@ -1275,3 +1327,7 @@ def callback(message: pubsub_v1.subscriber.message.Message) -> None:
         receive_messages_with_delivery_attempts(
             args.project_id, args.subscription_id, args.timeout
         )
+    elif args.command == "receive-messages-with-concurrency-control":
+        receive_messages_with_concurrency_control(
+            args.project_id, args.subscription_id, args.timeout
+        )

samples/snippets/subscriber_test.py

@@ -1008,6 +1008,40 @@ def test_receive_synchronously(
     subscriber_client.delete_subscription(request={"subscription": subscription_path})
 
 
+def test_receive_messages_with_concurrency_control(
+    subscriber_client: pubsub_v1.SubscriberClient,
+    publisher_client: pubsub_v1.PublisherClient,
+    topic: str,
+    capsys: CaptureFixture[str],
+) -> None:
+    subscription_async_receive_messages_with_concurrency_control_name = f"subscription-test-subscription-async-receive-messages-with-concurrency-control-{PY_VERSION}-{UUID}"
+
+    subscription_path = subscriber_client.subscription_path(
+        PROJECT_ID, subscription_async_receive_messages_with_concurrency_control_name
+    )
+
+    try:
+        subscriber_client.get_subscription(request={"subscription": subscription_path})
+    except NotFound:
+        subscriber_client.create_subscription(
+            request={"name": subscription_path, "topic": topic}
+        )
+
+    _ = _publish_messages(publisher_client, topic)
+
+    subscriber.receive_messages_with_flow_control(
+        PROJECT_ID, subscription_async_receive_messages_with_concurrency_control_name, 5
+    )
+
+    out, _ = capsys.readouterr()
+    assert "Listening" in out
+    assert subscription_async_receive_messages_with_concurrency_control_name in out
+    assert "message" in out
+
+    # Clean up.
+    subscriber_client.delete_subscription(request={"subscription": subscription_path})
+
+
 @typed_flaky
 def test_receive_synchronously_with_lease(
     subscriber_client: pubsub_v1.SubscriberClient,