Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

import asyncio
import threading
import time
from datetime import timedelta
from enum import Enum
from os import getenv
Expand Down Expand Up @@ -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()
Expand All @@ -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
Expand All @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -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()

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note: This is strictly speaking a bug. Multiplexed sessions cannot be deleted. I will fix this in a separate follow-up pull request. I kept it for now to keep the change as small as possible.

Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -126,16 +127,22 @@ 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()
if self._multiplexed_session_terminate_event is None:
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:
Expand All @@ -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.
Expand All @@ -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:
Expand All @@ -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()
Loading
Loading