perf(spanner): support PartialResultSet.last with background stream draining

packages/google-cloud-spanner/google/cloud/spanner_v1/_async/_helpers.py

@@ -1,5 +1,6 @@
 import asyncio
 import inspect
+import os
 import time
 
 from google.api_core.exceptions import Aborted
@@ -153,3 +154,44 @@ def _create_experimental_host_transport(
         client_key,
         interceptors=interceptors,
     )
+
+
+_PENDING_DRAIN_TASKS = set()
+if hasattr(os, "register_at_fork"):
+    os.register_at_fork(after_in_child=_PENDING_DRAIN_TASKS.clear)
+
+
+def _drain_stream(iterator):
+    """Drain an async stream iterator to EOF in the background.
+
+    Called when PartialResultSet.last is True to allow the caller to return immediately
+    while consuming trailing gRPC metadata so the stream terminates cleanly with status OK.
+    """
+    if iterator is None:
+        return
+
+    async def _drain():
+        try:
+            async for _ in iterator:
+                pass
+        except asyncio.CancelledError:
+            if hasattr(iterator, "cancel"):
+                try:
+                    iterator.cancel()
+                except Exception:
+                    pass
+            raise
+        except Exception:
+            pass
+
+    try:
+        task = asyncio.create_task(_drain())
+        _PENDING_DRAIN_TASKS.add(task)
+        task.add_done_callback(_PENDING_DRAIN_TASKS.discard)
+    except RuntimeError:
+        # Event loop may be closed or not running.
+        if hasattr(iterator, "cancel"):
+            try:
+                iterator.cancel()
+            except Exception:
+                pass

packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py

@@ -28,7 +28,7 @@
 from google.protobuf.struct_pb2 import Struct
 
 from google.cloud.aio._cross_sync import CrossSync
-from google.cloud.spanner_v1._async._helpers import _retry
+from google.cloud.spanner_v1._async._helpers import _drain_stream, _retry
 from google.cloud.spanner_v1._async.streamed import StreamedResultSet
 from google.cloud.spanner_v1._helpers import (
     AtomicCounter,
@@ -114,94 +114,117 @@ async def _restart_on_unavailable(
     attempt = 1
     nth_request = getattr(request_id_manager, "_next_nth_request", 0)
     current_request_id = None
-
-    while True:
-        try:
-            # Get results iterator.
-            if iterator is None:
-                with (
-                    trace_call(
-                        trace_name,
-                        session,
-                        attributes,
-                        observability_options=observability_options,
-                        metadata=metadata,
-                    ) as span,
-                    MetricsCapture(resource_info),
-                ):
-                    (
-                        call_metadata,
-                        current_request_id,
-                    ) = request_id_manager.metadata_and_request_id(
-                        nth_request,
-                        attempt,
-                        metadata,
-                        span,
-                    )
-                    iterator = await CrossSync.run_if_async(
-                        method,
-                        request=request,
-                        metadata=call_metadata,
-                    )
-
-            # Add items from iterator to buffer.
-            item: PartialResultSet
-            async for item in iterator:
-                item_buffer.append(item)
-
-                # Update the transaction from the response.
+    stream_finished = False
+
+    try:
+        while True:
+            try:
+                # Get results iterator.
+                if iterator is None:
+                    with (
+                        trace_call(
+                            trace_name,
+                            session,
+                            attributes,
+                            observability_options=observability_options,
+                            metadata=metadata,
+                        ) as span,
+                        MetricsCapture(resource_info),
+                    ):
+                        (
+                            call_metadata,
+                            current_request_id,
+                        ) = request_id_manager.metadata_and_request_id(
+                            nth_request,
+                            attempt,
+                            metadata,
+                            span,
+                        )
+                        iterator = await CrossSync.run_if_async(
+                            method,
+                            request=request,
+                            metadata=call_metadata,
+                        )
+
+                # Add items from iterator to buffer.
+                item: PartialResultSet
+                async for item in iterator:
+                    item_buffer.append(item)
+
+                    # Update the transaction from the response.
+                    if transaction is not None:
+                        transaction._update_for_result_set_pb(item)
+                    if (
+                        item._pb is not None
+                        and item._pb.HasField("precommit_token")
+                        and transaction is not None
+                    ):
+                        await transaction._update_for_precommit_token_pb(
+                            item.precommit_token
+                        )
+
+                    try:
+                        item_is_last = item.last
+                    except AttributeError:
+                        item_is_last = False
+
+                    if item_is_last:
+                        stream_finished = True
+                        _drain_stream(iterator)
+                        iterator = None
+                        break
+
+                    if item.resume_token:
+                        resume_token = item.resume_token
+                        break
+
+            except ServiceUnavailable:
+                del item_buffer[:]
+                request.resume_token = resume_token
                 if transaction is not None:
-                    transaction._update_for_result_set_pb(item)
-                if (
-                    item._pb is not None
-                    and item._pb.HasField("precommit_token")
-                    and transaction is not None
-                ):
-                    await transaction._update_for_precommit_token_pb(
-                        item.precommit_token
-                    )
-
-                if item.resume_token:
-                    resume_token = item.resume_token
-                    break
-
-        except ServiceUnavailable:
-            del item_buffer[:]
-            request.resume_token = resume_token
-            if transaction is not None:
-                transaction_selector = transaction._build_transaction_selector_pb()
-            request.transaction = transaction_selector
-            attempt += 1
-            iterator = None
-            continue
-
-        except InternalServerError as exc:
-            resumable_error = any(
-                resumable_message in exc.message
-                for resumable_message in _STREAM_RESUMPTION_INTERNAL_ERROR_MESSAGES
-            )
-            if not resumable_error:
+                    transaction_selector = transaction._build_transaction_selector_pb()
+                request.transaction = transaction_selector
+                attempt += 1
+                iterator = None
+                continue
+
+            except InternalServerError as exc:
+                resumable_error = any(
+                    resumable_message in exc.message
+                    for resumable_message in _STREAM_RESUMPTION_INTERNAL_ERROR_MESSAGES
+                )
+                if not resumable_error:
+                    raise _augment_error_with_request_id(exc, current_request_id)
+                del item_buffer[:]
+                request.resume_token = resume_token
+                if transaction is not None:
+                    transaction_selector = transaction._build_transaction_selector_pb()
+                attempt += 1
+                request.transaction = transaction_selector
+                iterator = None
+                continue
+
+            except Exception as exc:
+                # Augment any other exception with the request ID
                 raise _augment_error_with_request_id(exc, current_request_id)
-            del item_buffer[:]
-            request.resume_token = resume_token
-            if transaction is not None:
-                transaction_selector = transaction._build_transaction_selector_pb()
-            attempt += 1
-            request.transaction = transaction_selector
-            iterator = None
-            continue
 
-        except Exception as exc:
-            # Augment any other exception with the request ID
-            raise _augment_error_with_request_id(exc, current_request_id)
+            if len(item_buffer) == 0:
+                iterator = None
+                break
 
-        if len(item_buffer) == 0:
-            break
+            for item in item_buffer:
+                yield item
 
-        for item in item_buffer:
-            yield item
+            del item_buffer[:]
 
-        del item_buffer[:]
+            if stream_finished:
+                break
+    finally:
+        if iterator is not None and hasattr(iterator, "cancel"):
+            try:
+                iterator.cancel()
+            except Exception:
+                pass
 
 
 class _SnapshotBase(_SessionWrapper):