fix(sessions): resolve async deadlock in multiplexed session manager

Author: sinhasubham

google/cloud/spanner_v1/_async/database_sessions_manager.py

@@ -19,7 +19,6 @@
 from datetime import timedelta
 from enum import Enum
 from os import getenv
-import threading
 from threading import Thread
 from typing import Optional
 from weakref import ref
@@ -71,8 +70,7 @@ def __init__(self, database, pool):
         self._pool = pool
         self._multiplexed_session: Optional[Session] = None
         self._multiplexed_session_thread: Optional[CrossSync.Task] = None
-        # Use threading.Lock because this is accessed in a synchronous maintenance thread
-        self._multiplexed_session_lock: threading.Lock = threading.Lock()
+        self._multiplexed_session_lock: CrossSync.Lock = CrossSync.Lock()
         self._multiplexed_session_terminate_event: CrossSync.Event = CrossSync.Event()
 
     @CrossSync.convert
@@ -119,7 +117,7 @@ async def _get_multiplexed_session(self) -> Session:
 
         :rtype: :class:`~google.cloud.spanner_v1.session.Session`
         :returns: a multiplexed session."""
-        with CrossSync.rm_aio(self._multiplexed_session_lock):
+        async with self._multiplexed_session_lock:
             if self._multiplexed_session is None:
                 self._multiplexed_session = await self._build_multiplexed_session()
                 self._multiplexed_session_thread = self._build_maintenance_thread()
@@ -193,7 +191,7 @@ async def _maintain_multiplexed_session(session_manager_ref) -> None:
             if time() - session_created_time < refresh_interval_seconds:
                 await CrossSync.sleep(polling_interval_seconds)
                 continue
-            with manager._multiplexed_session_lock:
+            async with manager._multiplexed_session_lock:
                 await CrossSync.run_if_async(manager._multiplexed_session.delete)
                 manager._multiplexed_session = (
                     await manager._build_multiplexed_session()

google/cloud/spanner_v1/batch.py

@@ -243,7 +243,7 @@ def wrapped_method():
                     max_commit_delay=max_commit_delay,
                     request_options=request_options,
                 )
-                (call_metadata, error_augmenter) = database.with_error_augmentation(
+                call_metadata, error_augmenter = database.with_error_augmentation(
                     getattr(database, "_next_nth_request", 0), 1, metadata, span
                 )
                 commit_method = functools.partial(

google/cloud/spanner_v1/database.py

@@ -82,7 +82,6 @@
     trace_call,
 )
 from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
-
 from google.cloud.spanner_v1.table import Table
 
 SPANNER_DATA_SCOPE = "https://www.googleapis.com/auth/spanner.data"
@@ -211,11 +210,9 @@ def __init__(
     def _resource_info(self):
         """Resource information for metrics labels."""
         return {
-            "project": (
-                self._instance._client.project
-                if self._instance and self._instance._client
-                else None
-            ),
+            "project": self._instance._client.project
+            if self._instance and self._instance._client
+            else None,
             "instance": self._instance.instance_id if self._instance else None,
             "database": self.database_id,
         }
@@ -533,7 +530,7 @@ def with_error_augmentation(
             tuple: (metadata_list, context_manager)"""
         if span is None:
             span = get_current_span()
-        (metadata, request_id) = _metadata_with_request_id_and_req_id(
+        metadata, request_id = _metadata_with_request_id_and_req_id(
             self._nth_client_id,
             self._channel_id,
             nth_request,
@@ -810,7 +807,7 @@ def execute_pdml():
                 session = self._sessions_manager.get_session(transaction_type)
                 try:
                     add_span_event(span, "Starting BeginTransaction")
-                    (call_metadata, error_augmenter) = self.with_error_augmentation(
+                    call_metadata, error_augmenter = self.with_error_augmentation(
                         self._next_nth_request, 1, metadata, span
                     )
                     with error_augmenter:

google/cloud/spanner_v1/database_sessions_manager.py

@@ -19,7 +19,6 @@
 from datetime import timedelta
 from enum import Enum
 from os import getenv
-import threading
 from threading import Thread
 from typing import Optional
 from weakref import ref
@@ -69,7 +68,9 @@ def __init__(self, database, pool):
         self._pool = pool
         self._multiplexed_session: Optional[Session] = None
         self._multiplexed_session_thread: Optional[CrossSync._Sync_Impl.Task] = None
-        self._multiplexed_session_lock: threading.Lock = threading.Lock()
+        self._multiplexed_session_lock: CrossSync._Sync_Impl.Lock = (
+            CrossSync._Sync_Impl.Lock()
+        )
         self._multiplexed_session_terminate_event: CrossSync._Sync_Impl.Event = (
             CrossSync._Sync_Impl.Event()
         )

google/cloud/spanner_v1/instance.py

@@ -479,7 +479,9 @@ def database(
                 database_role=database_role,
                 enable_drop_protection=enable_drop_protection,
             )
-        db._pool.bind(db)
+        res = db._pool.bind(db)
+        if res is not None:
+            res
         return db
 
     def list_databases(self, page_size=None):

google/cloud/spanner_v1/pool.py

@@ -304,7 +304,7 @@ def _fill_pool(self):
                     f"Creating {request.session_count} sessions",
                     span_event_attributes,
                 )
-                (call_metadata, error_augmenter) = database.with_error_augmentation(
+                call_metadata, error_augmenter = database.with_error_augmentation(
                     database._next_nth_request, 1, metadata, span
                 )
                 with error_augmenter:
@@ -612,7 +612,7 @@ def bind(self, database):
         ) as span, MetricsCapture(self._resource_info):
             returned_session_count = 0
             while returned_session_count < self.size:
-                (call_metadata, error_augmenter) = database.with_error_augmentation(
+                call_metadata, error_augmenter = database.with_error_augmentation(
                     database._next_nth_request, 1, metadata, span
                 )
                 with error_augmenter:
@@ -654,7 +654,7 @@ def get(self, timeout=None):
         ping_after = None
         session = None
         try:
-            (ping_after, session) = CrossSync._Sync_Impl.queue_get(
+            ping_after, session = CrossSync._Sync_Impl.queue_get(
                 self._sessions, block=True, timeout=timeout
             )
         except CrossSync._Sync_Impl.QueueEmpty as e:
@@ -698,9 +698,7 @@ def clear(self):
         """Delete all sessions in the pool."""
         while True:
             try:
-                (_, session) = CrossSync._Sync_Impl.queue_get(
-                    self._sessions, block=False
-                )
+                _, session = CrossSync._Sync_Impl.queue_get(self._sessions, block=False)
             except CrossSync._Sync_Impl.QueueEmpty:
                 break
             else:
@@ -713,7 +711,7 @@ def ping(self):
         or during the "idle" phase of an event loop."""
         while True:
             try:
-                (ping_after, session) = CrossSync._Sync_Impl.queue_get(
+                ping_after, session = CrossSync._Sync_Impl.queue_get(
                     self._sessions, block=False
                 )
             except CrossSync._Sync_Impl.QueueEmpty:

google/cloud/spanner_v1/session.py

@@ -188,7 +188,7 @@ def create(self):
             observability_options=observability_options,
             metadata=metadata,
         ) as span, MetricsCapture(self._resource_info):
-            (call_metadata, error_augmenter) = database.with_error_augmentation(
+            call_metadata, error_augmenter = database.with_error_augmentation(
                 nth_request, 1, metadata, span
             )
             with error_augmenter:
@@ -232,7 +232,7 @@ def exists(self):
             observability_options=observability_options,
             metadata=metadata,
         ) as span, MetricsCapture(self._resource_info):
-            (call_metadata, error_augmenter) = database.with_error_augmentation(
+            call_metadata, error_augmenter = database.with_error_augmentation(
                 nth_request, 1, metadata, span
             )
             with error_augmenter:
@@ -283,7 +283,7 @@ def delete(self):
             observability_options=observability_options,
             metadata=metadata,
         ) as span, MetricsCapture(self._resource_info):
-            (call_metadata, error_augmenter) = database.with_error_augmentation(
+            call_metadata, error_augmenter = database.with_error_augmentation(
                 nth_request, 1, metadata, span
             )
             with error_augmenter:
@@ -300,7 +300,7 @@ def ping(self):
         metadata = _metadata_with_prefix(database.name)
         nth_request = database._next_nth_request
         with trace_call("CloudSpanner.Session.ping", self) as span:
-            (call_metadata, error_augmenter) = database.with_error_augmentation(
+            call_metadata, error_augmenter = database.with_error_augmentation(
                 nth_request, 1, metadata, span
             )
             with error_augmenter:

google/cloud/spanner_v1/snapshot.py

@@ -322,7 +322,7 @@ def execute_sql(
                 raise ValueError("Transaction has not begun.")
         if params is not None:
             params_pb = Struct(
-                fields={key: _make_value_pb(value) for (key, value) in params.items()}
+                fields={key: _make_value_pb(value) for key, value in params.items()}
             )
         else:
             params_pb = {}
@@ -513,7 +513,7 @@ def partition_query(
             raise ValueError("Cannot partition a single-use transaction.")
         if params is not None:
             params_pb = Struct(
-                fields={key: _make_value_pb(value) for (key, value) in params.items()}
+                fields={key: _make_value_pb(value) for key, value in params.items()}
             )
         else:
             params_pb = Struct()
@@ -614,7 +614,7 @@ def wrapped_method():
                 begin_transaction_request = BeginTransactionRequest(
                     **begin_request_kwargs
                 )
-                (call_metadata, error_augmenter) = database.with_error_augmentation(
+                call_metadata, error_augmenter = database.with_error_augmentation(
                     nth_request, attempt.increment(), metadata, span
                 )
                 begin_transaction_method = functools.partial(

google/cloud/spanner_v1/streamed.py

@@ -147,7 +147,7 @@ def _consume_next(self):
 
     def __iter__(self):
         while True:
-            (iter_rows, self._rows[:]) = (self._rows[:], ())
+            iter_rows, self._rows[:] = (self._rows[:], ())
             while iter_rows:
                 yield iter_rows.pop(0)
             if self._done:
@@ -230,7 +230,7 @@ def to_dict_list(self):
             rows.append(
                 {
                     column: value
-                    for (column, value) in zip(
+                    for column, value in zip(
                         [column.name for column in self._metadata.row_type.fields], row
                     )
                 }
@@ -291,7 +291,7 @@ def _merge_array(lhs, rhs, type_):
     if element_type.code in _UNMERGEABLE_TYPES:
         lhs.list_value.values.extend(rhs.list_value.values)
         return lhs
-    (lhs, rhs) = (list(lhs.list_value.values), list(rhs.list_value.values))
+    lhs, rhs = (list(lhs.list_value.values), list(rhs.list_value.values))
     if not len(lhs) or not len(rhs):
         return Value(list_value=ListValue(values=lhs + rhs))
     first = rhs.pop(0)
@@ -316,7 +316,7 @@ def _merge_array(lhs, rhs, type_):
 def _merge_struct(lhs, rhs, type_):
     """Helper for '_merge_by_type'."""
     fields = type_.struct_type.fields
-    (lhs, rhs) = (list(lhs.list_value.values), list(rhs.list_value.values))
+    lhs, rhs = (list(lhs.list_value.values), list(rhs.list_value.values))
     if not len(lhs) or not len(rhs):
         return Value(list_value=ListValue(values=lhs + rhs))
     candidate_type = fields[len(lhs) - 1].type_

google/cloud/spanner_v1/transaction.py

@@ -162,7 +162,7 @@ def rollback(self) -> None:
 
                 def wrapped_method(*args, **kwargs):
                     attempt.increment()
-                    (call_metadata, error_augmenter) = database.with_error_augmentation(
+                    call_metadata, error_augmenter = database.with_error_augmentation(
                         nth_request, attempt.value, metadata, span
                     )
                     rollback_method = functools.partial(
@@ -269,7 +269,7 @@ def wrapped_method(*args, **kwargs):
                 is_multiplexed = getattr(self._session, "is_multiplexed", False)
                 if is_multiplexed and self._precommit_token is not None:
                     commit_request_args["precommit_token"] = self._precommit_token
-                (call_metadata, error_augmenter) = database.with_error_augmentation(
+                call_metadata, error_augmenter = database.with_error_augmentation(
                     nth_request, attempt.value, metadata, span
                 )
                 commit_method = functools.partial(
@@ -300,7 +300,7 @@ def before_next_retry(nth_retry, delay_in_seconds):
             if commit_response_pb._pb.HasField("precommit_token"):
                 add_span_event(span, commit_retry_event_name)
                 nth_request = database._next_nth_request
-                (call_metadata, error_augmenter) = database.with_error_augmentation(
+                call_metadata, error_augmenter = database.with_error_augmentation(
                     nth_request, 1, metadata, span
                 )
                 with error_augmenter:
@@ -338,7 +338,7 @@ def _make_params_pb(params, param_types):
             If ``params`` is None but ``param_types`` is not None."""
         if params:
             return Struct(
-                fields={key: _make_value_pb(value) for (key, value) in params.items()}
+                fields={key: _make_value_pb(value) for key, value in params.items()}
             )
         return {}
 
@@ -417,7 +417,7 @@ def execute_update(
             metadata.append(
                 _metadata_with_leader_aware_routing(database._route_to_leader_enabled)
             )
-        (seqno, self._execute_sql_request_count) = (
+        seqno, self._execute_sql_request_count = (
             self._execute_sql_request_count,
             self._execute_sql_request_count + 1,
         )
@@ -454,7 +454,7 @@ def execute_update(
 
         def wrapped_method(*args, **kwargs):
             attempt.increment()
-            (call_metadata, error_augmenter) = database.with_error_augmentation(
+            call_metadata, error_augmenter = database.with_error_augmentation(
                 nth_request, attempt.value, metadata
             )
             execute_sql_method = functools.partial(
@@ -544,7 +544,7 @@ def batch_update(
             if isinstance(statement, str):
                 parsed.append(ExecuteBatchDmlRequest.Statement(sql=statement))
             else:
-                (dml, params, param_types) = statement
+                dml, params, param_types = statement
                 params_pb = self._make_params_pb(params, param_types)
                 parsed.append(
                     ExecuteBatchDmlRequest.Statement(
@@ -556,7 +556,7 @@ def batch_update(
             metadata.append(
                 _metadata_with_leader_aware_routing(database._route_to_leader_enabled)
             )
-        (seqno, self._execute_sql_request_count) = (
+        seqno, self._execute_sql_request_count = (
             self._execute_sql_request_count,
             self._execute_sql_request_count + 1,
         )
@@ -590,7 +590,7 @@ def batch_update(
 
         def wrapped_method(*args, **kwargs):
             attempt.increment()
-            (call_metadata, error_augmenter) = database.with_error_augmentation(
+            call_metadata, error_augmenter = database.with_error_augmentation(
                 nth_request, attempt.value, metadata
             )
             execute_batch_dml_method = functools.partial(