diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/database_sessions_manager.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/database_sessions_manager.py index 12a114ab4e09..30c70d797694 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/database_sessions_manager.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/database_sessions_manager.py @@ -18,6 +18,7 @@ import asyncio import threading +import time from datetime import timedelta from enum import Enum from os import getenv @@ -128,6 +129,10 @@ async def _get_multiplexed_session(self) -> Session: :rtype: :class:`~google.cloud.spanner_v1.session.Session` :returns: a multiplexed session.""" + session = self._multiplexed_session + if session is not None: + return session + with self._init_lock: if self._multiplexed_session_lock is None: self._multiplexed_session_lock = CrossSync.Lock() @@ -136,10 +141,12 @@ async def _get_multiplexed_session(self) -> Session: 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() + session = await self._build_multiplexed_session() + maintenance_thread = self._build_maintenance_thread(session) if not CrossSync.is_async: - self._multiplexed_session_thread.start() + maintenance_thread.start() + self._multiplexed_session_thread = maintenance_thread + self._multiplexed_session = session return self._multiplexed_session @CrossSync.convert @@ -156,26 +163,63 @@ async def _build_multiplexed_session(self) -> Session: await session.create() return session - def _build_maintenance_thread(self) -> CrossSync.Task: + def _build_maintenance_thread( + self, session: Optional[Session] = None + ) -> CrossSync.Task: """Builds and returns a multiplexed session maintenance thread for the database session manager. This thread will periodically delete and recreate the multiplexed session to ensure that it is always valid. + :type session: :class:`~google.cloud.spanner_v1.session.Session` + :param session: (Optional) The multiplexed session to maintain. + :rtype: :class:`CrossSync.Task` :returns: a multiplexed session maintenance thread.""" + session_to_maintain = ( + session if session is not None else self._multiplexed_session + ) session_manager_ref = ref(self) if CrossSync.is_async: return CrossSync.create_task( self._maintain_multiplexed_session, session_manager_ref ) else: + session_id = ( + session_to_maintain.session_id + if session_to_maintain is not None + else "" + ) return Thread( target=self._maintain_multiplexed_session, - name=f"maintenance-multiplexed-session-{self._multiplexed_session.session_id}", + name=f"maintenance-multiplexed-session-{session_id}", args=[session_manager_ref], daemon=True, ) + @CrossSync.convert + async def _rotate_multiplexed_session(self) -> bool: + """Rotates the multiplexed session by building and swapping in a new session. + + :rtype: bool + :returns: True if the session was successfully refreshed, False otherwise. + """ + try: + new_session = await self._build_multiplexed_session() + except Exception: + return False + + async with self._multiplexed_session_lock: + old_session = self._multiplexed_session + self._multiplexed_session = new_session + + if old_session is not None: + try: + await CrossSync.run_if_async(old_session.delete) + except Exception: + pass + + return True + @staticmethod @CrossSync.convert async def _maintain_multiplexed_session(session_manager_ref) -> None: @@ -196,24 +240,26 @@ async def _maintain_multiplexed_session(session_manager_ref) -> None: refresh_interval_seconds = ( manager._MAINTENANCE_THREAD_REFRESH_INTERVAL.total_seconds() ) - from time import time - - session_created_time = time() + session_created_time = time.monotonic() while True: manager = session_manager_ref() if manager is None: return - if manager._multiplexed_session_terminate_event.is_set(): + terminate_event = manager._multiplexed_session_terminate_event + if terminate_event.is_set(): return - if time() - session_created_time < refresh_interval_seconds: - await CrossSync.sleep(polling_interval_seconds) - continue - async with manager._multiplexed_session_lock: - await CrossSync.run_if_async(manager._multiplexed_session.delete) - manager._multiplexed_session = ( - await manager._build_multiplexed_session() - ) - session_created_time = time() + + if time.monotonic() - session_created_time >= refresh_interval_seconds: + if await manager._rotate_multiplexed_session(): + session_created_time = time.monotonic() + manager = None + continue + + manager = None + await CrossSync.event_wait( + terminate_event, + timeout=polling_interval_seconds, + ) @classmethod def _use_multiplexed(cls, transaction_type: TransactionType) -> bool: @@ -247,5 +293,6 @@ async def close(self) -> None: else: self._multiplexed_session_thread.join() if self._multiplexed_session is not None: - await self._multiplexed_session.delete() + session_to_delete = self._multiplexed_session self._multiplexed_session = None + await session_to_delete.delete() diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/database_sessions_manager.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/database_sessions_manager.py index 021e844a7406..1b2c6231f46e 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/database_sessions_manager.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/database_sessions_manager.py @@ -18,6 +18,7 @@ """Manage sessions for a database.""" import threading +import time from datetime import timedelta from enum import Enum from os import getenv @@ -126,6 +127,10 @@ def _get_multiplexed_session(self) -> Session: :rtype: :class:`~google.cloud.spanner_v1.session.Session` :returns: a multiplexed session.""" + session = self._multiplexed_session + if session is not None: + return session + with self._init_lock: if self._multiplexed_session_lock is None: self._multiplexed_session_lock = CrossSync._Sync_Impl.Lock() @@ -133,9 +138,11 @@ def _get_multiplexed_session(self) -> Session: self._multiplexed_session_terminate_event = CrossSync._Sync_Impl.Event() with self._multiplexed_session_lock: if self._multiplexed_session is None: - self._multiplexed_session = self._build_multiplexed_session() - self._multiplexed_session_thread = self._build_maintenance_thread() - self._multiplexed_session_thread.start() + session = self._build_multiplexed_session() + maintenance_thread = self._build_maintenance_thread(session) + maintenance_thread.start() + self._multiplexed_session_thread = maintenance_thread + self._multiplexed_session = session return self._multiplexed_session def _build_multiplexed_session(self) -> Session: @@ -151,21 +158,55 @@ def _build_multiplexed_session(self) -> Session: session.create() return session - def _build_maintenance_thread(self) -> CrossSync._Sync_Impl.Task: + def _build_maintenance_thread( + self, session: Optional[Session] = None + ) -> CrossSync._Sync_Impl.Task: """Builds and returns a multiplexed session maintenance thread for the database session manager. This thread will periodically delete and recreate the multiplexed session to ensure that it is always valid. + :type session: :class:`~google.cloud.spanner_v1.session.Session` + :param session: (Optional) The multiplexed session to maintain. + :rtype: :class:`CrossSync._Sync_Impl.Task` :returns: a multiplexed session maintenance thread.""" + session_to_maintain = ( + session if session is not None else self._multiplexed_session + ) session_manager_ref = ref(self) + session_id = ( + session_to_maintain.session_id if session_to_maintain is not None else "" + ) return Thread( target=self._maintain_multiplexed_session, - name=f"maintenance-multiplexed-session-{self._multiplexed_session.session_id}", + name=f"maintenance-multiplexed-session-{session_id}", args=[session_manager_ref], daemon=True, ) + def _rotate_multiplexed_session(self) -> bool: + """Rotates the multiplexed session by building and swapping in a new session. + + :rtype: bool + :returns: True if the session was successfully refreshed, False otherwise. + """ + try: + new_session = self._build_multiplexed_session() + except Exception: + return False + + with self._multiplexed_session_lock: + old_session = self._multiplexed_session + self._multiplexed_session = new_session + + if old_session is not None: + try: + CrossSync._Sync_Impl.run_if_async(old_session.delete) + except Exception: + pass + + return True + @staticmethod def _maintain_multiplexed_session(session_manager_ref) -> None: """Maintains the multiplexed session for the database session manager. @@ -185,22 +226,25 @@ def _maintain_multiplexed_session(session_manager_ref) -> None: refresh_interval_seconds = ( manager._MAINTENANCE_THREAD_REFRESH_INTERVAL.total_seconds() ) - from time import time - - session_created_time = time() + session_created_time = time.monotonic() while True: manager = session_manager_ref() if manager is None: return - if manager._multiplexed_session_terminate_event.is_set(): + terminate_event = manager._multiplexed_session_terminate_event + if terminate_event.is_set(): return - if time() - session_created_time < refresh_interval_seconds: - CrossSync._Sync_Impl.sleep(polling_interval_seconds) - continue - with manager._multiplexed_session_lock: - CrossSync._Sync_Impl.run_if_async(manager._multiplexed_session.delete) - manager._multiplexed_session = manager._build_multiplexed_session() - session_created_time = time() + if time.monotonic() - session_created_time >= refresh_interval_seconds: + if manager._rotate_multiplexed_session(): + session_created_time = time.monotonic() + manager = None + continue + + manager = None + CrossSync._Sync_Impl.event_wait( + terminate_event, + timeout=polling_interval_seconds, + ) @classmethod def _use_multiplexed(cls, transaction_type: TransactionType) -> bool: @@ -226,5 +270,6 @@ def close(self) -> None: if self._multiplexed_session_thread is not None: self._multiplexed_session_thread.join() if self._multiplexed_session is not None: - self._multiplexed_session.delete() + session_to_delete = self._multiplexed_session self._multiplexed_session = None + session_to_delete.delete() diff --git a/packages/google-cloud-spanner/tests/unit/_async/test_sessions_manager_extra.py b/packages/google-cloud-spanner/tests/unit/_async/test_sessions_manager_extra.py index 56e07f970ba1..8a40dd942801 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_sessions_manager_extra.py +++ b/packages/google-cloud-spanner/tests/unit/_async/test_sessions_manager_extra.py @@ -106,6 +106,7 @@ async def fake_coro(): task = asyncio.create_task(fake_coro()) manager._multiplexed_session_thread = task manager._multiplexed_session = mock.AsyncMock() + manager._multiplexed_session_terminate_event = mock.Mock() with mock.patch( "google.cloud.spanner_v1._async.database_sessions_manager.CrossSync.is_async", @@ -114,16 +115,19 @@ async def fake_coro(): await manager.close() # task is cancelled and awaited in close() self.assertTrue(task.done()) + manager._multiplexed_session_terminate_event.set.assert_called_once() # Sync branch of close manager._multiplexed_session_thread = mock.Mock() manager._multiplexed_session = mock.AsyncMock() + manager._multiplexed_session_terminate_event = mock.Mock() with mock.patch( "google.cloud.spanner_v1._async.database_sessions_manager.CrossSync.is_async", False, ): await manager.close() self.assertTrue(manager._multiplexed_session_thread.join.called) + manager._multiplexed_session_terminate_event.set.assert_called_once() async def test_maintain_multiplexed_session_refresh(self): # coverage for line 196-202 @@ -135,16 +139,10 @@ async def test_maintain_multiplexed_session_refresh(self): # We need to simulate time passing and then terminating refresh_interval = manager._MAINTENANCE_THREAD_REFRESH_INTERVAL.total_seconds() - from time import time + import time from weakref import ref - start_time = time() - - # Mock time to skip forward - # First call in while loop: manager._multiplexed_session_terminate_event.is_set() -> False - # Second call in while loop: time() - session_created_time < refresh_interval -> False (after mock) - # Inside the 'if' block: - # We want it to run once then terminate + start_time = time.monotonic() call_count = 0 @@ -158,7 +156,10 @@ def mock_time(): return start_time + refresh_interval + 10 return start_time - with mock.patch("time.time", side_effect=mock_time): + with mock.patch( + "google.cloud.spanner_v1._async.database_sessions_manager.time.monotonic", + side_effect=mock_time, + ): async def mock_build(): manager._multiplexed_session_terminate_event.set() @@ -172,7 +173,6 @@ async def mock_build(): self.assertTrue(manager._multiplexed_session_terminate_event.is_set()) async def test_maintain_multiplexed_session_manager_gone_in_loop(self): - # coverage for line 191 # Mock session_manager_ref() within the loop manager = DatabaseSessionsManager(self.database, self.pool) call_count = 0 @@ -184,14 +184,15 @@ def mock_ref(): return manager return None - with mock.patch("time.time", return_value=0): - # Wait, it's a static method, so we pass the ref + with mock.patch( + "google.cloud.spanner_v1._async.database_sessions_manager.time.monotonic", + return_value=0, + ): r = mock.Mock(side_effect=mock_ref) await DatabaseSessionsManager._maintain_multiplexed_session(r) self.assertEqual(call_count, 2) async def test_maintain_multiplexed_session_loop_sleep(self): - # coverage for line 196 manager = DatabaseSessionsManager(self.database, self.pool) manager._multiplexed_session_lock = asyncio.Lock() manager._multiplexed_session_terminate_event = asyncio.Event() @@ -210,9 +211,405 @@ def mock_time(): from weakref import ref - with mock.patch("time.time", side_effect=mock_time): + async def mock_wait(event, timeout=None): + event.set() + + with mock.patch( + "google.cloud.spanner_v1._async.database_sessions_manager.time.monotonic", + side_effect=mock_time, + ): with mock.patch( - "google.cloud.spanner_v1._async.database_sessions_manager.CrossSync.sleep", - mock.AsyncMock(), - ): + "google.cloud.spanner_v1._async.database_sessions_manager.CrossSync.event_wait", + side_effect=mock_wait, + ) as mock_event_wait: await manager._maintain_multiplexed_session(ref(manager)) + mock_event_wait.assert_called_once() + + async def test_get_multiplexed_session_fast_path(self): + manager = DatabaseSessionsManager(self.database, self.pool) + mock_session = mock.Mock() + manager._multiplexed_session = mock_session + manager._init_lock = mock.Mock(wraps=manager._init_lock) + + session = await manager._get_multiplexed_session() + self.assertIs(session, mock_session) + manager._init_lock.acquire.assert_not_called() + self.assertIsNone(manager._multiplexed_session_lock) + + async def test_get_multiplexed_session_fast_path_lock_already_created(self): + manager = DatabaseSessionsManager(self.database, self.pool) + mock_session = mock.Mock() + manager._multiplexed_session = mock_session + mock_lock = mock.AsyncMock() + manager._multiplexed_session_lock = mock_lock + manager._init_lock = mock.Mock(wraps=manager._init_lock) + + session = await manager._get_multiplexed_session() + self.assertIs(session, mock_session) + manager._init_lock.acquire.assert_not_called() + mock_lock.acquire.assert_not_called() + mock_lock.__aenter__.assert_not_called() + + async def test_maintain_multiplexed_session_swaps_before_deleting_old_session( + self, + ): + from weakref import ref + + manager = DatabaseSessionsManager(self.database, self.pool) + manager._multiplexed_session_lock = asyncio.Lock() + manager._multiplexed_session_terminate_event = asyncio.Event() + + old_session = mock.AsyncMock() + new_session = mock.AsyncMock() + manager._multiplexed_session = old_session + + async def verify_swap_on_delete(): + self.assertIs(manager._multiplexed_session, new_session) + manager._multiplexed_session_terminate_event.set() + + old_session.delete.side_effect = verify_swap_on_delete + + refresh_interval = manager._MAINTENANCE_THREAD_REFRESH_INTERVAL.total_seconds() + call_count = 0 + + def mock_time(): + nonlocal call_count + call_count += 1 + if call_count == 1: + return 0 + return refresh_interval + 100 + + with mock.patch( + "google.cloud.spanner_v1._async.database_sessions_manager.time.monotonic", + side_effect=mock_time, + ): + with mock.patch.object( + manager, "_build_multiplexed_session", return_value=new_session + ) as mock_build: + await DatabaseSessionsManager._maintain_multiplexed_session( + ref(manager) + ) + mock_build.assert_called_once() + old_session.delete.assert_called_once() + self.assertIs(manager._multiplexed_session, new_session) + + async def test_maintain_multiplexed_session_handles_build_failure(self): + from weakref import ref + + manager = DatabaseSessionsManager(self.database, self.pool) + manager._multiplexed_session_lock = asyncio.Lock() + manager._multiplexed_session_terminate_event = asyncio.Event() + + current_session = mock.AsyncMock() + manager._multiplexed_session = current_session + + refresh_interval = manager._MAINTENANCE_THREAD_REFRESH_INTERVAL.total_seconds() + call_count = 0 + + def mock_time(): + nonlocal call_count + call_count += 1 + if call_count == 1: + return 0 + return refresh_interval + 100 + + async def mock_event_wait(event, timeout=None): + manager._multiplexed_session_terminate_event.set() + + with mock.patch( + "google.cloud.spanner_v1._async.database_sessions_manager.time.monotonic", + side_effect=mock_time, + ): + with mock.patch.object( + manager, + "_build_multiplexed_session", + side_effect=Exception("network error"), + ) as mock_build: + with mock.patch( + "google.cloud.spanner_v1._async.database_sessions_manager.CrossSync.event_wait", + side_effect=mock_event_wait, + ) as mock_event_wait_call: + await DatabaseSessionsManager._maintain_multiplexed_session( + ref(manager) + ) + mock_build.assert_called_once() + mock_event_wait_call.assert_called_once() + current_session.delete.assert_not_called() + self.assertIs(manager._multiplexed_session, current_session) + + async def test_maintain_multiplexed_session_handles_delete_failure(self): + from weakref import ref + + manager = DatabaseSessionsManager(self.database, self.pool) + manager._multiplexed_session_lock = asyncio.Lock() + manager._multiplexed_session_terminate_event = asyncio.Event() + + old_session = mock.AsyncMock() + old_session.delete.side_effect = Exception("delete failed") + new_session = mock.AsyncMock() + manager._multiplexed_session = old_session + + async def verify_swap_on_delete(): + self.assertIs(manager._multiplexed_session, new_session) + manager._multiplexed_session_terminate_event.set() + raise Exception("delete failed") + + old_session.delete.side_effect = verify_swap_on_delete + + refresh_interval = manager._MAINTENANCE_THREAD_REFRESH_INTERVAL.total_seconds() + call_count = 0 + + def mock_time(): + nonlocal call_count + call_count += 1 + if call_count == 1: + return 0 + return refresh_interval + 100 + + with mock.patch( + "google.cloud.spanner_v1._async.database_sessions_manager.time.monotonic", + side_effect=mock_time, + ): + with mock.patch.object( + manager, "_build_multiplexed_session", return_value=new_session + ): + await DatabaseSessionsManager._maintain_multiplexed_session( + ref(manager) + ) + old_session.delete.assert_called_once() + self.assertIs(manager._multiplexed_session, new_session) + + async def test_maintain_multiplexed_session_old_session_none(self): + from weakref import ref + + manager = DatabaseSessionsManager(self.database, self.pool) + manager._multiplexed_session_lock = asyncio.Lock() + manager._multiplexed_session_terminate_event = asyncio.Event() + + new_session = mock.AsyncMock() + manager._multiplexed_session = None + + refresh_interval = manager._MAINTENANCE_THREAD_REFRESH_INTERVAL.total_seconds() + call_count = 0 + + def mock_time(): + nonlocal call_count + call_count += 1 + if call_count == 1: + return 0 + return refresh_interval + 100 + + async def mock_event_wait(event, timeout=None): + manager._multiplexed_session_terminate_event.set() + + with mock.patch( + "google.cloud.spanner_v1._async.database_sessions_manager.time.monotonic", + side_effect=mock_time, + ): + with mock.patch.object( + manager, "_build_multiplexed_session", return_value=new_session + ): + with mock.patch( + "google.cloud.spanner_v1._async.database_sessions_manager.CrossSync.event_wait", + side_effect=mock_event_wait, + ): + + async def build_and_terminate(): + manager._multiplexed_session_terminate_event.set() + return new_session + + manager._build_multiplexed_session = build_and_terminate + await DatabaseSessionsManager._maintain_multiplexed_session( + ref(manager) + ) + self.assertIs(manager._multiplexed_session, new_session) + + async def test_get_multiplexed_session_initial_failure_allows_retry(self): + manager = DatabaseSessionsManager(self.database, self.pool) + manager._multiplexed_session = None + + with mock.patch.object( + manager, + "_build_multiplexed_session", + side_effect=Exception("create failed"), + ): + with self.assertRaises(Exception): + await manager._get_multiplexed_session() + self.assertIsNone(manager._multiplexed_session) + + # Subsequent call succeeds and creates session and maintenance thread + mock_session = mock.Mock() + with mock.patch.object( + manager, "_build_multiplexed_session", return_value=mock_session + ): + with mock.patch.object(manager, "_build_maintenance_thread"): + session = await manager._get_multiplexed_session() + self.assertIs(session, mock_session) + + async def test_get_multiplexed_session_concurrent_initialization_double_checked( + self, + ): + manager = DatabaseSessionsManager(self.database, self.pool) + manager._multiplexed_session = None + + concurrent_session = mock.Mock() + mock_lock = mock.AsyncMock() + + async def lock_enter(): + manager._multiplexed_session = concurrent_session + return mock_lock + + mock_lock.__aenter__.side_effect = lock_enter + + with mock.patch.object(manager, "_build_multiplexed_session") as mock_build: + with mock.patch( + "google.cloud.spanner_v1._async.database_sessions_manager.CrossSync.Lock", + return_value=mock_lock, + ): + session = await manager._get_multiplexed_session() + self.assertIs(session, concurrent_session) + mock_build.assert_not_called() + + async def test_get_multiplexed_session_maintenance_thread_failure_allows_retry( + self, + ): + manager = DatabaseSessionsManager(self.database, self.pool) + manager._multiplexed_session = None + + mock_session = mock.Mock() + with mock.patch.object( + manager, "_build_multiplexed_session", return_value=mock_session + ): + with mock.patch.object( + manager, + "_build_maintenance_thread", + side_effect=RuntimeError("thread spawn failed"), + ): + with self.assertRaises(RuntimeError): + await manager._get_multiplexed_session() + self.assertIsNone(manager._multiplexed_session) + + async def test_build_maintenance_thread(self): + manager = DatabaseSessionsManager(self.database, self.pool) + mock_session = mock.Mock() + mock_session.session_id = "test-session-123" + + task = manager._build_maintenance_thread(mock_session) + self.assertIsNotNone(task) + task.cancel() + + # Backward compatibility when session is omitted + manager._multiplexed_session = mock_session + task_default = manager._build_maintenance_thread() + self.assertIsNotNone(task_default) + task_default.cancel() + + async def test_build_multiplexed_session(self): + manager = DatabaseSessionsManager(self.database, self.pool) + with mock.patch( + "google.cloud.spanner_v1._async.database_sessions_manager.Session" + ) as mock_session_class: + mock_session_instance = mock.AsyncMock() + mock_session_class.return_value = mock_session_instance + session = await manager._build_multiplexed_session() + self.assertIs(session, mock_session_instance) + mock_session_instance.create.assert_called_once() + + async def test_put_session_multiplexed_and_regular(self): + manager = DatabaseSessionsManager(self.database, self.pool) + multiplexed_session = mock.Mock() + multiplexed_session.is_multiplexed = True + multiplexed_session.session_id = "multi-1" + await manager.put_session(multiplexed_session) + self.pool.put.assert_not_called() + + regular_session = mock.Mock() + regular_session.is_multiplexed = False + regular_session.session_id = "reg-1" + await manager.put_session(regular_session) + self.pool.put.assert_called_once_with(regular_session) + + def test_use_multiplexed_read_only(self): + from google.cloud.spanner_v1._async.database_sessions_manager import ( + TransactionType, + ) + + with mock.patch.dict( + "os.environ", {DatabaseSessionsManager._ENV_VAR_MULTIPLEXED: "false"} + ): + self.assertFalse( + DatabaseSessionsManager._use_multiplexed(TransactionType.READ_ONLY) + ) + with mock.patch.dict( + "os.environ", {DatabaseSessionsManager._ENV_VAR_MULTIPLEXED: "true"} + ): + self.assertTrue( + DatabaseSessionsManager._use_multiplexed(TransactionType.READ_ONLY) + ) + + def test_use_multiplexed_partitioned(self): + from google.cloud.spanner_v1._async.database_sessions_manager import ( + TransactionType, + ) + + with mock.patch.dict( + "os.environ", + {DatabaseSessionsManager._ENV_VAR_MULTIPLEXED_PARTITIONED: "false"}, + ): + self.assertFalse( + DatabaseSessionsManager._use_multiplexed(TransactionType.PARTITIONED) + ) + with mock.patch.dict( + "os.environ", + {DatabaseSessionsManager._ENV_VAR_MULTIPLEXED_PARTITIONED: "true"}, + ): + self.assertTrue( + DatabaseSessionsManager._use_multiplexed(TransactionType.PARTITIONED) + ) + + async def test_rotate_multiplexed_session_success(self): + manager = DatabaseSessionsManager(self.database, self.pool) + manager._multiplexed_session_lock = asyncio.Lock() + old_session = mock.AsyncMock() + new_session = mock.AsyncMock() + manager._multiplexed_session = old_session + + with mock.patch.object( + manager, "_build_multiplexed_session", return_value=new_session + ): + result = await manager._rotate_multiplexed_session() + self.assertTrue(result) + self.assertIs(manager._multiplexed_session, new_session) + old_session.delete.assert_called_once() + + async def test_rotate_multiplexed_session_build_failure(self): + manager = DatabaseSessionsManager(self.database, self.pool) + manager._multiplexed_session_lock = asyncio.Lock() + current_session = mock.AsyncMock() + manager._multiplexed_session = current_session + + with mock.patch.object( + manager, + "_build_multiplexed_session", + side_effect=Exception("network down"), + ): + result = await manager._rotate_multiplexed_session() + self.assertFalse(result) + self.assertIs(manager._multiplexed_session, current_session) + current_session.delete.assert_not_called() + + async def test_rotate_multiplexed_session_delete_failure(self): + manager = DatabaseSessionsManager(self.database, self.pool) + manager._multiplexed_session_lock = asyncio.Lock() + old_session = mock.AsyncMock() + old_session.delete.side_effect = Exception("delete failed") + new_session = mock.AsyncMock() + manager._multiplexed_session = old_session + + with mock.patch.object( + manager, "_build_multiplexed_session", return_value=new_session + ): + result = await manager._rotate_multiplexed_session() + self.assertTrue(result) + self.assertIs(manager._multiplexed_session, new_session) + old_session.delete.assert_called_once() diff --git a/packages/google-cloud-spanner/tests/unit/test_database_session_manager.py b/packages/google-cloud-spanner/tests/unit/test_database_session_manager.py index e5ff5f9b1882..297ac3c3817c 100644 --- a/packages/google-cloud-spanner/tests/unit/test_database_session_manager.py +++ b/packages/google-cloud-spanner/tests/unit/test_database_session_manager.py @@ -18,7 +18,7 @@ from unittest import TestCase from google.api_core.exceptions import BadRequest, FailedPrecondition -from mock import Mock, patch +from mock import MagicMock, Mock, patch from google.cloud.spanner_v1.database_sessions_manager import ( DatabaseSessionsManager, @@ -249,6 +249,313 @@ async def run_concurrent(): "test_concurrent_get_multiplexed_session_no_deadlock timed out (DEADLOCK)!" ) + def test_get_multiplexed_session_fast_path(self): + manager = self._manager + mock_session = Mock() + manager._multiplexed_session = mock_session + manager._init_lock = Mock(wraps=manager._init_lock) + + session = manager._get_multiplexed_session() + self.assertIs(session, mock_session) + manager._init_lock.acquire.assert_not_called() + self.assertIsNone(manager._multiplexed_session_lock) + + def test_get_multiplexed_session_fast_path_lock_already_created(self): + manager = self._manager + mock_session = Mock() + manager._multiplexed_session = mock_session + mock_lock = MagicMock() + manager._multiplexed_session_lock = mock_lock + manager._init_lock = Mock(wraps=manager._init_lock) + + session = manager._get_multiplexed_session() + self.assertIs(session, mock_session) + manager._init_lock.acquire.assert_not_called() + mock_lock.acquire.assert_not_called() + mock_lock.__enter__.assert_not_called() + + def test_maintain_multiplexed_session_swaps_before_deleting_old_session(self): + import threading + from weakref import ref + + manager = DatabaseSessionsManager(self._manager._database, self._manager._pool) + manager._multiplexed_session_lock = threading.Lock() + manager._multiplexed_session_terminate_event = Mock() + manager._multiplexed_session_terminate_event.is_set.side_effect = [False, True] + + old_session = Mock() + new_session = Mock() + manager._multiplexed_session = old_session + + def verify_swap_on_delete(): + self.assertIs(manager._multiplexed_session, new_session) + + old_session.delete.side_effect = verify_swap_on_delete + + call_count = 0 + + def mock_time(): + nonlocal call_count + call_count += 1 + if call_count == 1: + return 0 + return 1000000 + + with patch( + "google.cloud.spanner_v1.database_sessions_manager.time.monotonic", + side_effect=mock_time, + ): + with patch.object( + manager, "_build_multiplexed_session", return_value=new_session + ) as mock_build: + DatabaseSessionsManager._maintain_multiplexed_session(ref(manager)) + mock_build.assert_called_once() + old_session.delete.assert_called_once() + self.assertIs(manager._multiplexed_session, new_session) + + def test_close_branches(self): + manager = DatabaseSessionsManager(self._manager._database, self._manager._pool) + + # Branch where thread is None + manager.close() + + # Branch where thread is not None + mock_thread = Mock() + mock_session = Mock() + manager._multiplexed_session_thread = mock_thread + manager._multiplexed_session = mock_session + manager._multiplexed_session_terminate_event = Mock() + + manager.close() + manager._multiplexed_session_terminate_event.set.assert_called_once() + mock_thread.join.assert_called_once() + self.assertIsNone(manager._multiplexed_session) + mock_session.delete.assert_called_once() + + def test_maintain_multiplexed_session_handles_build_failure(self): + import threading + from weakref import ref + + manager = DatabaseSessionsManager(self._manager._database, self._manager._pool) + manager._multiplexed_session_lock = threading.Lock() + manager._multiplexed_session_terminate_event = Mock() + manager._multiplexed_session_terminate_event.is_set.side_effect = [ + False, + True, + ] + + current_session = Mock() + manager._multiplexed_session = current_session + + call_count = 0 + + def mock_time(): + nonlocal call_count + call_count += 1 + if call_count == 1: + return 0 + return 1000000 + + with patch( + "google.cloud.spanner_v1.database_sessions_manager.time.monotonic", + side_effect=mock_time, + ): + with patch.object( + manager, + "_build_multiplexed_session", + side_effect=Exception("network error"), + ) as mock_build: + with patch( + "google.cloud.spanner_v1.database_sessions_manager.CrossSync._Sync_Impl.event_wait" + ) as mock_event_wait: + DatabaseSessionsManager._maintain_multiplexed_session(ref(manager)) + mock_build.assert_called_once() + mock_event_wait.assert_called_once() + current_session.delete.assert_not_called() + self.assertIs(manager._multiplexed_session, current_session) + + def test_maintain_multiplexed_session_handles_delete_failure(self): + import threading + from weakref import ref + + manager = DatabaseSessionsManager(self._manager._database, self._manager._pool) + manager._multiplexed_session_lock = threading.Lock() + manager._multiplexed_session_terminate_event = Mock() + manager._multiplexed_session_terminate_event.is_set.side_effect = [False, True] + + old_session = Mock() + old_session.delete.side_effect = Exception("delete failed") + new_session = Mock() + manager._multiplexed_session = old_session + + call_count = 0 + + def mock_time(): + nonlocal call_count + call_count += 1 + if call_count == 1: + return 0 + return 1000000 + + with patch( + "google.cloud.spanner_v1.database_sessions_manager.time.monotonic", + side_effect=mock_time, + ): + with patch.object( + manager, "_build_multiplexed_session", return_value=new_session + ): + DatabaseSessionsManager._maintain_multiplexed_session(ref(manager)) + old_session.delete.assert_called_once() + self.assertIs(manager._multiplexed_session, new_session) + + def test_maintain_multiplexed_session_old_session_none(self): + import threading + from weakref import ref + + manager = DatabaseSessionsManager(self._manager._database, self._manager._pool) + manager._multiplexed_session_lock = threading.Lock() + manager._multiplexed_session_terminate_event = Mock() + manager._multiplexed_session_terminate_event.is_set.side_effect = [False, True] + + new_session = Mock() + manager._multiplexed_session = None + + call_count = 0 + + def mock_time(): + nonlocal call_count + call_count += 1 + if call_count == 1: + return 0 + return 1000000 + + with patch( + "google.cloud.spanner_v1.database_sessions_manager.time.monotonic", + side_effect=mock_time, + ): + with patch.object( + manager, "_build_multiplexed_session", return_value=new_session + ): + DatabaseSessionsManager._maintain_multiplexed_session(ref(manager)) + self.assertIs(manager._multiplexed_session, new_session) + + def test_get_multiplexed_session_initial_failure_allows_retry(self): + manager = DatabaseSessionsManager(self._manager._database, self._manager._pool) + manager._multiplexed_session = None + + with patch.object( + manager, + "_build_multiplexed_session", + side_effect=Exception("create failed"), + ): + with self.assertRaises(Exception): + manager._get_multiplexed_session() + self.assertIsNone(manager._multiplexed_session) + + # Subsequent call succeeds and creates session and maintenance thread + mock_session = Mock() + mock_thread = Mock() + mock_thread.is_alive.return_value = False + with patch.object( + manager, "_build_multiplexed_session", return_value=mock_session + ): + with patch.object( + manager, "_build_maintenance_thread", return_value=mock_thread + ): + session = manager._get_multiplexed_session() + self.assertIs(session, mock_session) + mock_thread.start.assert_called_once() + + def test_get_multiplexed_session_concurrent_initialization_double_checked( + self, + ): + manager = DatabaseSessionsManager(self._manager._database, self._manager._pool) + manager._multiplexed_session = None + + concurrent_session = Mock() + mock_lock = MagicMock() + + def lock_enter(): + manager._multiplexed_session = concurrent_session + return mock_lock + + mock_lock.__enter__.side_effect = lock_enter + + with patch.object(manager, "_build_multiplexed_session") as mock_build: + with patch( + "google.cloud.spanner_v1.database_sessions_manager.CrossSync._Sync_Impl.Lock", + return_value=mock_lock, + ): + session = manager._get_multiplexed_session() + self.assertIs(session, concurrent_session) + mock_build.assert_not_called() + + def test_get_multiplexed_session_maintenance_thread_failure_allows_retry( + self, + ): + manager = DatabaseSessionsManager(self._manager._database, self._manager._pool) + manager._multiplexed_session = None + + mock_session = Mock() + with patch.object( + manager, "_build_multiplexed_session", return_value=mock_session + ): + with patch.object( + manager, + "_build_maintenance_thread", + side_effect=RuntimeError("thread spawn failed"), + ): + with self.assertRaises(RuntimeError): + manager._get_multiplexed_session() + self.assertIsNone(manager._multiplexed_session) + + def test_build_maintenance_thread(self): + manager = DatabaseSessionsManager(self._manager._database, self._manager._pool) + mock_session = Mock() + mock_session.session_id = "sync-test-session-456" + + thread = manager._build_maintenance_thread(mock_session) + self.assertIsNotNone(thread) + self.assertEqual( + thread.name, "maintenance-multiplexed-session-sync-test-session-456" + ) + self.assertTrue(thread.daemon) + + # Backward compatibility when session is omitted + manager._multiplexed_session = mock_session + thread_default = manager._build_maintenance_thread() + self.assertIsNotNone(thread_default) + self.assertEqual( + thread_default.name, + "maintenance-multiplexed-session-sync-test-session-456", + ) + self.assertTrue(thread_default.daemon) + + def test_maintain_multiplexed_session_manager_gone(self): + from weakref import ref + + class Fake: + pass + + fake = Fake() + reference = ref(fake) + del fake + DatabaseSessionsManager._maintain_multiplexed_session(reference) + + def test_maintain_multiplexed_session_manager_collected_in_loop(self): + manager = DatabaseSessionsManager(self._manager._database, self._manager._pool) + call_count = 0 + + def mock_ref(): + nonlocal call_count + call_count += 1 + if call_count == 1: + return manager + return None + + DatabaseSessionsManager._maintain_multiplexed_session(mock_ref) + self.assertGreaterEqual(call_count, 2) + def test_exception_bad_request(self): manager = self._manager api = manager._database.spanner_api @@ -340,6 +647,59 @@ def test__getenv(self): DatabaseSessionsManager._use_multiplexed(TransactionType.READ_ONLY) ) + def test_rotate_multiplexed_session_success(self): + import threading + + manager = DatabaseSessionsManager(self._manager._database, self._manager._pool) + manager._multiplexed_session_lock = threading.Lock() + old_session = Mock() + new_session = Mock() + manager._multiplexed_session = old_session + + with patch.object( + manager, "_build_multiplexed_session", return_value=new_session + ): + result = manager._rotate_multiplexed_session() + self.assertTrue(result) + self.assertIs(manager._multiplexed_session, new_session) + old_session.delete.assert_called_once() + + def test_rotate_multiplexed_session_build_failure(self): + import threading + + manager = DatabaseSessionsManager(self._manager._database, self._manager._pool) + manager._multiplexed_session_lock = threading.Lock() + current_session = Mock() + manager._multiplexed_session = current_session + + with patch.object( + manager, + "_build_multiplexed_session", + side_effect=Exception("network down"), + ): + result = manager._rotate_multiplexed_session() + self.assertFalse(result) + self.assertIs(manager._multiplexed_session, current_session) + current_session.delete.assert_not_called() + + def test_rotate_multiplexed_session_delete_failure(self): + import threading + + manager = DatabaseSessionsManager(self._manager._database, self._manager._pool) + manager._multiplexed_session_lock = threading.Lock() + old_session = Mock() + old_session.delete.side_effect = Exception("delete failed") + new_session = Mock() + manager._multiplexed_session = old_session + + with patch.object( + manager, "_build_multiplexed_session", return_value=new_session + ): + result = manager._rotate_multiplexed_session() + self.assertTrue(result) + self.assertIs(manager._multiplexed_session, new_session) + old_session.delete.assert_called_once() + def _assert_true_with_timeout(self, condition: Callable) -> None: """Asserts that the given condition is met within a timeout period.