googleapis · harshachinta · Dec 6, 2024 · Nov 18, 2024 · harshachinta · Nov 29, 2024
@@ -463,6 +463,7 @@ def _retry(
     retry_count=5,
     delay=2,
     allowed_exceptions=None,
+    beforeNextRetry=None,
 ):
     """
     Retry a function with a specified number of retries, delay between retries, and list of allowed exceptions.
@@ -479,6 +480,9 @@ def _retry(
     """
     retries = 0
     while retries <= retry_count:
+        if retries > 0 and beforeNextRetry:
+            beforeNextRetry(retries, delay)
+
         try:
             return func()
         except Exception as exc:

@@ -81,10 +81,11 @@ def trace_call(name, session, extra_attributes=None, observability_options=None)
     tracer = get_tracer(tracer_provider)
 
     # Set base attributes that we know for every trace created
+    db = session._database
     attributes = {
         "db.type": "spanner",
         "db.url": SpannerClient.DEFAULT_ENDPOINT,
-        "db.instance": session._database.name,
+        "db.instance": "" if not db else db.name,
         "net.host.name": SpannerClient.DEFAULT_ENDPOINT,
         OTEL_SCOPE_NAME: TRACER_NAME,
         OTEL_SCOPE_VERSION: TRACER_VERSION,
@@ -106,7 +107,10 @@ def trace_call(name, session, extra_attributes=None, observability_options=None)
             yield span
         except Exception as error:
             span.set_status(Status(StatusCode.ERROR, str(error)))
-            span.record_exception(error)
+            # OpenTelemetry-Python imposes invoking span.record_exception on __exit__
+            # on any exception. We should file a bug later on with them to only
+            # invoke .record_exception if not already invoked, hence we should not
+            # invoke .record_exception on our own else we shall have 2 exceptions.
             raise
         else:
             if (not span._status) or span._status.status_code == StatusCode.UNSET:
@@ -116,3 +120,14 @@ def trace_call(name, session, extra_attributes=None, observability_options=None)
                 # it wasn't previously set otherwise.
                 # https://github.com/googleapis/python-spanner/issues/1246
                 span.set_status(Status(StatusCode.OK))
+
+
+def get_current_span():
+    if not HAS_OPENTELEMETRY_INSTALLED:
+        return None
+    return trace.get_current_span()
+
+
+def add_span_event(span, event_name, event_attributes=None):
+    if span:
+        span.add_event(event_name, event_attributes)
@@ -67,6 +67,10 @@
     SpannerGrpcTransport,
 )
 from google.cloud.spanner_v1.table import Table
+from google.cloud.spanner_v1._opentelemetry_tracing import (
+    add_span_event,
+    get_current_span,
+)
 
 
 SPANNER_DATA_SCOPE = "https://www.googleapis.com/auth/spanner.data"
@@ -1164,7 +1168,9 @@ def __init__(
 
     def __enter__(self):
         """Begin ``with`` block."""
+        current_span = get_current_span()
         session = self._session = self._database._pool.get()
+        add_span_event(current_span, "Using session", {"id": session.session_id})
         batch = self._batch = Batch(session)
         if self._request_options.transaction_tag:
             batch.transaction_tag = self._request_options.transaction_tag
@@ -1187,6 +1193,12 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                     extra={"commit_stats": self._batch.commit_stats},
                 )
             self._database._pool.put(self._session)
+            current_span = get_current_span()
+            add_span_event(
+                current_span,
+                "Returned session to pool",
+                {"id": self._session.session_id},
+            )
 
 
 class MutationGroupsCheckout(object):

@@ -16,6 +16,7 @@
 
 import datetime
 import queue
+import time
 
 from google.cloud.exceptions import NotFound
 from google.cloud.spanner_v1 import BatchCreateSessionsRequest
@@ -24,6 +25,10 @@
     _metadata_with_prefix,
     _metadata_with_leader_aware_routing,
 )
+from google.cloud.spanner_v1._opentelemetry_tracing import (
+    add_span_event,
+    get_current_span,
+)
 from warnings import warn
 
 _NOW = datetime.datetime.utcnow  # unit tests may replace
@@ -196,20 +201,50 @@ def bind(self, database):
                          when needed.
         """
         self._database = database
+        requested_session_count = self.size - self._sessions.qsize()
+        span = get_current_span()
+        span_event_attributes = {"kind": type(self).__name__}
+
+        if requested_session_count <= 0:
+            add_span_event(
+                span,
+                f"Invalid session pool size({requested_session_count}) <= 0",
+                span_event_attributes,
+            )
+            return
+
         api = database.spanner_api
         metadata = _metadata_with_prefix(database.name)
         if database._route_to_leader_enabled:
             metadata.append(
                 _metadata_with_leader_aware_routing(database._route_to_leader_enabled)
             )
         self._database_role = self._database_role or self._database.database_role
+        if requested_session_count > 0:
+            add_span_event(
+                span,
+                f"Requesting {requested_session_count} sessions",
+                span_event_attributes,
+            )
+
+        if self._sessions.full():
+            add_span_event(span, "Session pool is already full", span_event_attributes)
+            return
+
         request = BatchCreateSessionsRequest(
             database=database.name,
-            session_count=self.size - self._sessions.qsize(),
+            session_count=requested_session_count,
             session_template=Session(creator_role=self.database_role),
         )
 
+        returned_session_count = 0
         while not self._sessions.full():
+            request.session_count = requested_session_count - self._sessions.qsize()
+            add_span_event(
+                span,
+                f"Creating {request.session_count} sessions",
+                span_event_attributes,
+            )
             resp = api.batch_create_sessions(
                 request=request,
                 metadata=metadata,
@@ -218,6 +253,13 @@ def bind(self, database):
                 session = self._new_session()
                 session._session_id = session_pb.name.split("/")[-1]
                 self._sessions.put(session)
+                returned_session_count += 1
+
+        add_span_event(
+            span,
+            f"Requested for {requested_session_count} sessions, returned {returned_session_count}",
+            span_event_attributes,
+        )
 
     def get(self, timeout=None):
         """Check a session out from the pool.
@@ -233,12 +275,43 @@ def get(self, timeout=None):
         if timeout is None:
             timeout = self.default_timeout
 
-        session = self._sessions.get(block=True, timeout=timeout)
-        age = _NOW() - session.last_use_time
+        start_time = time.time()
+        current_span = get_current_span()
+        span_event_attributes = {"kind": type(self).__name__}
+        add_span_event(current_span, "Acquiring session", span_event_attributes)
 
-        if age >= self._max_age and not session.exists():
-            session = self._database.session()
-            session.create()
+        session = None
+        try:
+            add_span_event(
+                current_span,
+                "Waiting for a session to become available",
+                span_event_attributes,
+            )
+
+            session = self._sessions.get(block=True, timeout=timeout)
+            age = _NOW() - session.last_use_time
+
+            if age >= self._max_age and not session.exists():
+                if not session.exists():
+                    add_span_event(
+                        current_span,
+                        "Session is not valid, recreating it",
+                        span_event_attributes,
+                    )
+                session = self._database.session()
+                session.create()
+                # Replacing with the updated session.id.
+                span_event_attributes["session.id"] = session._session_id
+
+            span_event_attributes["session.id"] = session._session_id
+            span_event_attributes["time.elapsed"] = time.time() - start_time
+            add_span_event(current_span, "Acquired session", span_event_attributes)
+
+        except queue.Empty as e:
+            add_span_event(
+                current_span, "No sessions available in the pool", span_event_attributes
+            )
+            raise e
 
         return session
 
@@ -312,13 +385,32 @@ def get(self):
         :returns: an existing session from the pool, or a newly-created
                   session.
         """
+        current_span = get_current_span()
+        span_event_attributes = {"kind": type(self).__name__}
+        add_span_event(current_span, "Acquiring session", span_event_attributes)
+
         try:
+            add_span_event(
+                current_span,
+                "Waiting for a session to become available",
+                span_event_attributes,
+            )
             session = self._sessions.get_nowait()
         except queue.Empty:
+            add_span_event(
+                current_span,
+                "No sessions available in pool. Creating session",
+                span_event_attributes,
+            )
             session = self._new_session()
             session.create()
         else:
             if not session.exists():
+                add_span_event(
+                    current_span,
+                    "Session is not valid, recreating it",
+                    span_event_attributes,
+                )
                 session = self._new_session()
                 session.create()
         return session
@@ -427,6 +519,38 @@ def bind(self, database):
             session_template=Session(creator_role=self.database_role),
         )
 
+        span_event_attributes = {"kind": type(self).__name__}
+        current_span = get_current_span()
+        requested_session_count = request.session_count
+        if requested_session_count <= 0:
+            add_span_event(
+                current_span,
+                f"Invalid session pool size({requested_session_count}) <= 0",
+                span_event_attributes,
+            )
+            return
+
+        add_span_event(
+            current_span,
+            f"Requesting {requested_session_count} sessions",
+            span_event_attributes,
+        )
+
+        if created_session_count >= self.size:
+            add_span_event(
+                current_span,
+                "Created no new sessions as sessionPool is full",
+                span_event_attributes,
+            )
+            return
+
+        add_span_event(
+            current_span,
+            f"Creating {request.session_count} sessions",
+            span_event_attributes,
+        )
+
+        returned_session_count = 0
         while created_session_count < self.size:
             resp = api.batch_create_sessions(
                 request=request,
@@ -436,8 +560,16 @@ def bind(self, database):
                 session = self._new_session()
                 session._session_id = session_pb.name.split("/")[-1]
                 self.put(session)
+                returned_session_count += 1
+
             created_session_count += len(resp.session)
 
+        add_span_event(
+            current_span,
+            f"Requested for {requested_session_count} sessions, return {returned_session_count}",
+            span_event_attributes,
+        )
+
     def get(self, timeout=None):
         """Check a session out from the pool.
 
@@ -452,7 +584,26 @@ def get(self, timeout=None):
         if timeout is None:
             timeout = self.default_timeout
 
-        ping_after, session = self._sessions.get(block=True, timeout=timeout)
+        start_time = time.time()
+        span_event_attributes = {"kind": type(self).__name__}
+        current_span = get_current_span()
+        add_span_event(
+            current_span,
+            "Waiting for a session to become available",
+            span_event_attributes,
+        )
+
+        ping_after = None
+        session = None
+        try:
+            ping_after, session = self._sessions.get(block=True, timeout=timeout)
+        except queue.Empty as e:
+            add_span_event(
+                current_span,
+                "No sessions available in the pool within the specified timeout",
+                span_event_attributes,
+            )
+            raise e
 
         if _NOW() > ping_after:
             # Using session.exists() guarantees the returned session exists.
@@ -462,6 +613,14 @@ def get(self, timeout=None):
                 session = self._new_session()
                 session.create()
 
+        span_event_attributes.update(
+            {
+                "time.elapsed": time.time() - start_time,
+                "session.id": session._session_id,
+                "kind": "pinging_pool",
+            }
+        )
+        add_span_event(current_span, "Acquired session", span_event_attributes)
         return session
 
     def put(self, session):