Skip to content

perf(spanner): prune lock and begin event allocations on single-use snapshots #18420

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
olavloite wants to merge 1 commit into
base: main
Choose a base branch
from
Open

perf(spanner): prune lock and begin event allocations on single-use snapshots #18420

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
51 changes: 38 additions & 13 deletions packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -213,31 +213,43 @@ class _SnapshotBase(_SessionWrapper):

Allows reuse of API request methods with different transaction selector.

.. note::
Single-use snapshots (``multi_use=False``) are designed for a single read
or query operation and are thread-confined; they are not safe for concurrent
invocation across multiple threads. Multi-use snapshots and transactions
synchronize concurrent operations using internal locks.

:type session: :class:`~google.cloud.spanner_v1.session.Session`
:param session: the session used to perform transaction operations.
"""

_read_only: bool = True
_multi_use: bool = False

def __init__(self, session, client_context=None):
def __init__(self, session, client_context=None, multi_use: Optional[bool] = None):
super().__init__(session)
self._client_context = _validate_client_context(client_context)
self._execute_sql_request_count: int = 0
self._read_request_count: int = 0
self._begin_request_sent: bool = False

if multi_use is not None:
self._multi_use = multi_use

# Identifier for the transaction.
self._transaction_id: Optional[bytes] = None
self._precommit_token: Optional[MultiplexedSessionPrecommitToken] = None

# Operations within a transaction can be performed concurrently, so we
# need to use a lock when updating the transaction.
self._lock: CrossSync.Lock = CrossSync.Lock()

# Event to coordinate concurrent requests beginning the transaction.
# This is used to prevent the "Transaction has not begun" race condition.
self._transaction_begin_event: CrossSync.Event = CrossSync.Event()
# Single-use snapshots never execute concurrently or begin transactions,
# so we skip allocating Lock and Event for them.
if self._multi_use:
self._lock: Optional[CrossSync.Lock] = CrossSync.Lock()
self._transaction_begin_event: Optional[CrossSync.Event] = CrossSync.Event()
else:
self._lock = None
self._transaction_begin_event = None

@property
def _resource_info(self):
Expand All @@ -258,15 +270,23 @@ async def _wait_for_transaction_begin(self) -> None:
id is available, must wait for that first request to complete instead
of assuming that the transaction has not begun.

For single-use snapshots, this method checks and enforces sequential
reuse prevention without synchronization, as single-use snapshots are
thread-confined.

:raises ValueError: if the transaction has already been used to execute
a request, but is not a multi-use transaction, or if the concurrent
request that began the transaction did not complete in time.
"""

if not self._multi_use:
if self._begin_request_sent or self._read_request_count > 0:
raise ValueError("Cannot re-use single-use snapshot.")
self._begin_request_sent = True
return

async with self._lock:
if self._begin_request_sent or self._read_request_count > 0:
if not self._multi_use:
raise ValueError("Cannot re-use single-use snapshot.")
wait_needed = self._transaction_id is None
else:
wait_needed = False
Expand Down Expand Up @@ -663,7 +683,7 @@ async def _get_streamed_result_set(
trace_name = f"CloudSpanner.{type(self).__name__}.{trace_method_name}"

is_inline_begin = False
if self._transaction_id is None:
if self._multi_use and self._transaction_id is None:
is_inline_begin = True
await self._lock.acquire()

Expand Down Expand Up @@ -983,7 +1003,8 @@ def _update_for_transaction_pb(self, transaction_pb: Transaction) -> None:
if self._transaction_id is None and transaction_pb.id:
self._transaction_id = transaction_pb.id
# Release any request waiting for this transaction to begin.
self._transaction_begin_event.set()
if self._transaction_begin_event is not None:
self._transaction_begin_event.set()

if transaction_pb._pb.HasField("precommit_token"):
self._update_for_precommit_token_pb_unsafe(transaction_pb.precommit_token)
Expand All @@ -993,7 +1014,10 @@ async def _update_for_precommit_token_pb(
self, precommit_token_pb: MultiplexedSessionPrecommitToken
) -> None:
"""Updates the snapshot for the given multiplexed session precommit token."""
async with self._lock:
if self._lock is not None:
async with self._lock:
self._update_for_precommit_token_pb_unsafe(precommit_token_pb)
else:
self._update_for_precommit_token_pb_unsafe(precommit_token_pb)

def _update_for_precommit_token_pb_unsafe(
Expand Down Expand Up @@ -1029,7 +1053,9 @@ def __init__(
transaction_id=None,
client_context=None,
):
super(Snapshot, self).__init__(session, client_context=client_context)
super(Snapshot, self).__init__(
session, client_context=client_context, multi_use=multi_use
)
opts = [read_timestamp, min_read_timestamp, max_staleness, exact_staleness]
flagged = [opt for opt in opts if opt is not None]
if len(flagged) > 1:
Expand All @@ -1047,7 +1073,6 @@ def __init__(
self._min_read_timestamp = min_read_timestamp
self._max_staleness = max_staleness
self._exact_staleness = exact_staleness
self._multi_use = multi_use
self._transaction_id = transaction_id

def _build_transaction_options_pb(self) -> TransactionOptions:
Expand Down
4 changes: 3 additions & 1 deletion packages/google-cloud-spanner/google/cloud/spanner_v1/_async/transaction.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,9 @@ class Transaction(_SnapshotBase, _BatchBase):
_read_only: bool = False

def __init__(self, session, client_context=None):
super(Transaction, self).__init__(session, client_context=client_context)
super(Transaction, self).__init__(
session, client_context=client_context, multi_use=True
)
self.rolled_back: bool = False

# If this transaction is used to retry a previous aborted transaction with a
Expand Down
50 changes: 38 additions & 12 deletions packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -187,25 +187,39 @@ class _SnapshotBase(_SessionWrapper):

Allows reuse of API request methods with different transaction selector.

.. note::
Single-use snapshots (``multi_use=False``) are designed for a single read
or query operation and are thread-confined; they are not safe for concurrent
invocation across multiple threads. Multi-use snapshots and transactions
synchronize concurrent operations using internal locks.

:type session: :class:`~google.cloud.spanner_v1.session.Session`
:param session: the session used to perform transaction operations.
"""

_read_only: bool = True
_multi_use: bool = False

def __init__(self, session, client_context=None):
def __init__(self, session, client_context=None, multi_use: Optional[bool] = None):
super().__init__(session)
self._client_context = _validate_client_context(client_context)
self._execute_sql_request_count: int = 0
self._read_request_count: int = 0
self._begin_request_sent: bool = False
if multi_use is not None:
self._multi_use = multi_use
self._transaction_id: Optional[bytes] = None
self._precommit_token: Optional[MultiplexedSessionPrecommitToken] = None
self._lock: CrossSync._Sync_Impl.Lock = CrossSync._Sync_Impl.Lock()
self._transaction_begin_event: CrossSync._Sync_Impl.Event = (
CrossSync._Sync_Impl.Event()
)
if self._multi_use:
self._lock: Optional[CrossSync._Sync_Impl.Lock] = (
CrossSync._Sync_Impl.Lock()
)
self._transaction_begin_event: Optional[CrossSync._Sync_Impl.Event] = (
CrossSync._Sync_Impl.Event()
)
else:
self._lock = None
self._transaction_begin_event = None

@property
def _resource_info(self):
Expand All @@ -225,13 +239,20 @@ def _wait_for_transaction_begin(self) -> None:
id is available, must wait for that first request to complete instead
of assuming that the transaction has not begun.

For single-use snapshots, this method checks and enforces sequential
reuse prevention without synchronization, as single-use snapshots are
thread-confined.

:raises ValueError: if the transaction has already been used to execute
a request, but is not a multi-use transaction, or if the concurrent
request that began the transaction did not complete in time."""
if not self._multi_use:
if self._begin_request_sent or self._read_request_count > 0:
raise ValueError("Cannot re-use single-use snapshot.")
self._begin_request_sent = True
return
with self._lock:
if self._begin_request_sent or self._read_request_count > 0:
if not self._multi_use:
raise ValueError("Cannot re-use single-use snapshot.")
wait_needed = self._transaction_id is None
else:
wait_needed = False
Expand Down Expand Up @@ -596,7 +617,7 @@ def _get_streamed_result_set(
trace_method_name = "execute_sql" if is_execute_sql_request else "read"
trace_name = f"CloudSpanner.{type(self).__name__}.{trace_method_name}"
is_inline_begin = False
if self._transaction_id is None:
if self._multi_use and self._transaction_id is None:
is_inline_begin = True
self._lock.acquire()
try:
Expand Down Expand Up @@ -882,15 +903,19 @@ def _update_for_transaction_pb(self, transaction_pb: Transaction) -> None:
"""Updates the snapshot for the given transaction."""
if self._transaction_id is None and transaction_pb.id:
self._transaction_id = transaction_pb.id
self._transaction_begin_event.set()
if self._transaction_begin_event is not None:
self._transaction_begin_event.set()
if transaction_pb._pb.HasField("precommit_token"):
self._update_for_precommit_token_pb_unsafe(transaction_pb.precommit_token)

def _update_for_precommit_token_pb(
self, precommit_token_pb: MultiplexedSessionPrecommitToken
) -> None:
"""Updates the snapshot for the given multiplexed session precommit token."""
with self._lock:
if self._lock is not None:
with self._lock:
self._update_for_precommit_token_pb_unsafe(precommit_token_pb)
else:
self._update_for_precommit_token_pb_unsafe(precommit_token_pb)

def _update_for_precommit_token_pb_unsafe(
Expand Down Expand Up @@ -918,7 +943,9 @@ def __init__(
transaction_id=None,
client_context=None,
):
super(Snapshot, self).__init__(session, client_context=client_context)
super(Snapshot, self).__init__(
session, client_context=client_context, multi_use=multi_use
)
opts = [read_timestamp, min_read_timestamp, max_staleness, exact_staleness]
flagged = [opt for opt in opts if opt is not None]
if len(flagged) > 1:
Expand All @@ -934,7 +961,6 @@ def __init__(
self._min_read_timestamp = min_read_timestamp
self._max_staleness = max_staleness
self._exact_staleness = exact_staleness
self._multi_use = multi_use
self._transaction_id = transaction_id

def _build_transaction_options_pb(self) -> TransactionOptions:
Expand Down
4 changes: 3 additions & 1 deletion packages/google-cloud-spanner/google/cloud/spanner_v1/transaction.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,9 @@ class Transaction(_SnapshotBase, _BatchBase):
_read_only: bool = False

def __init__(self, session, client_context=None):
super(Transaction, self).__init__(session, client_context=client_context)
super(Transaction, self).__init__(
session, client_context=client_context, multi_use=True
)
self.rolled_back: bool = False
self._multiplexed_session_previous_transaction_id: Optional[bytes] = None

Expand Down
Loading
Loading