| FazBrowse GitHub Viewer | Trending | | Home |
| Tools: [Download Repo ZIP] [Original HTTPS Page] |
1 parent cb8a2b7 commit eb9dd5a
11 files changed
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -78,8 +78,17 @@ def trace_call(name, session, extra_attributes=None): | |||
| 78 | 78 | try: | |
| 79 | 79 | yield span | |
| 80 | 80 | except Exception as error: | |
| 81 | - span.set_status(Status(StatusCode.ERROR, str(error))) | ||
| 82 | - span.record_exception(error) | ||
| 81 | + set_span_error_and_record_exception(span, error) | ||
| 83 | 82 | raise | |
| 84 | 83 | else: | |
| 85 | 84 | span.set_status(Status(StatusCode.OK)) | |
| 85 | + | ||
| 86 | + | ||
| 87 | + def set_span_error_and_record_exception(span, exc): | ||
| 88 | + if exc and span: | ||
| 89 | + span.set_status(Status(StatusCode.ERROR, str(exc))) | ||
| 90 | + span.record_exception(exc) | ||
| 91 | + | ||
| 92 | + | ||
| 93 | + def get_current_span(): | ||
| 94 | + return trace.get_current_span() | ||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -16,6 +16,7 @@ | |||
| 16 | 16 | ||
| 17 | 17 | import datetime | |
| 18 | 18 | import queue | |
| 19 | + import time | ||
| 19 | 20 | ||
| 20 | 21 | from google.cloud.exceptions import NotFound | |
| 21 | 22 | from google.cloud.spanner_v1 import BatchCreateSessionsRequest | |
@@ -24,6 +25,9 @@ | |||
| 24 | 25 | _metadata_with_prefix, | |
| 25 | 26 | _metadata_with_leader_aware_routing, | |
| 26 | 27 | ) | |
| 28 | + from google.cloud.spanner_v1._opentelemetry_tracing import ( | ||
| 29 | + get_current_span, | ||
| 30 | + ) | ||
| 27 | 31 | from warnings import warn | |
| 28 | 32 | ||
| 29 | 33 | _NOW = datetime.datetime.utcnow # unit tests may replace | |
@@ -199,13 +203,32 @@ def bind(self, database): | |||
| 199 | 203 | _metadata_with_leader_aware_routing(database._route_to_leader_enabled) | |
| 200 | 204 | ) | |
| 201 | 205 | self._database_role = self._database_role or self._database.database_role | |
| 206 | + requested_session_count = self.size - self._sessions.qsize() | ||
| 202 | 207 | request = BatchCreateSessionsRequest( | |
| 203 | 208 | database=database.name, | |
| 204 | - session_count=self.size - self._sessions.qsize(), | ||
| 209 | + session_count=requested_session_count, | ||
| 205 | 210 | session_template=Session(creator_role=self.database_role), | |
| 206 | 211 | ) | |
| 207 | 212 | ||
| 213 | + current_span = get_current_span() | ||
| 214 | + if requested_session_count > 0: | ||
| 215 | + current_span.add_event( | ||
| 216 | + f"Requesting {requested_session_count} sessions", | ||
| 217 | + {"kind": "fixed_size_pool"}, | ||
| 218 | + ) | ||
| 219 | + | ||
| 220 | + if self._sessions.full(): | ||
| 221 | + current_span.add_event( | ||
| 222 | + "Session pool is already full", {"kind": "fixed_size_pool"} | ||
| 223 | + ) | ||
| 224 | + return | ||
| 225 | + | ||
| 226 | + returned_session_count = 0 | ||
| 208 | 227 | while not self._sessions.full(): | |
| 228 | + current_span.add_event( | ||
| 229 | + f"Creating {request.session_count} sessions", | ||
| 230 | + {"kind": "fixed_size_pool"}, | ||
| 231 | + ) | ||
| 209 | 232 | resp = api.batch_create_sessions( | |
| 210 | 233 | request=request, | |
| 211 | 234 | metadata=metadata, | |
@@ -214,6 +237,12 @@ def bind(self, database): | |||
| 214 | 237 | session = self._new_session() | |
| 215 | 238 | session._session_id = session_pb.name.split("/")[-1] | |
| 216 | 239 | self._sessions.put(session) | |
| 240 | + returned_session_count += 1 | ||
| 241 | + | ||
| 242 | + current_span.add_event( | ||
| 243 | + f"Requested for {requested_session_count}, returned {returned_session_count}", | ||
| 244 | + {"kind": "fixed_size_pool"}, | ||
| 245 | + ) | ||
| 217 | 246 | ||
| 218 | 247 | def get(self, timeout=None): | |
| 219 | 248 | """Check a session out from the pool. | |
@@ -229,12 +258,23 @@ def get(self, timeout=None): | |||
| 229 | 258 | if timeout is None: | |
| 230 | 259 | timeout = self.default_timeout | |
| 231 | 260 | ||
| 261 | + start_time = time.time() | ||
| 262 | + current_span = get_current_span() | ||
| 263 | + current_span.add_event("Acquiring session", {"kind": type(self).__name__}) | ||
| 232 | 264 | session = self._sessions.get(block=True, timeout=timeout) | |
| 233 | 265 | ||
| 234 | 266 | if not session.exists(): | |
| 235 | 267 | session = self._database.session() | |
| 236 | 268 | session.create() | |
| 237 | 269 | ||
| 270 | + current_span.add_event( | ||
| 271 | + "Acquired session", | ||
| 272 | + { | ||
| 273 | + "time.elapsed": time.time() - start_time, | ||
| 274 | + "session.id": session.session_id, | ||
| 275 | + "kind": type(self).__name__, | ||
| 276 | + }, | ||
| 277 | + ) | ||
| 238 | 278 | return session | |
| 239 | 279 | ||
| 240 | 280 | def put(self, session): | |
@@ -307,6 +347,10 @@ def get(self): | |||
| 307 | 347 | :returns: an existing session from the pool, or a newly-created | |
| 308 | 348 | session. | |
| 309 | 349 | """ | |
| 350 | + start_time = time.time() | ||
| 351 | + current_span = get_current_span() | ||
| 352 | + current_span.add_event("Acquiring session", {"kind": type(self).__name__}) | ||
| 353 | + | ||
| 310 | 354 | try: | |
| 311 | 355 | session = self._sessions.get_nowait() | |
| 312 | 356 | except queue.Empty: | |
@@ -316,6 +360,15 @@ def get(self): | |||
| 316 | 360 | if not session.exists(): | |
| 317 | 361 | session = self._new_session() | |
| 318 | 362 | session.create() | |
| 363 | + else: | ||
| 364 | + current_span.add_event( | ||
| 365 | + "Cache hit: has usable session", | ||
| 366 | + { | ||
| 367 | + "id": session.session_id, | ||
| 368 | + "kind": type(self).__name__, | ||
| 369 | + }, | ||
| 370 | + ) | ||
| 371 | + | ||
| 319 | 372 | return session | |
| 320 | 373 | ||
| 321 | 374 | def put(self, session): | |
@@ -422,6 +475,18 @@ def bind(self, database): | |||
| 422 | 475 | session_template=Session(creator_role=self.database_role), | |
| 423 | 476 | ) | |
| 424 | 477 | ||
| 478 | + requested_session_count = request.session_count | ||
| 479 | + current_span = get_current_span() | ||
| 480 | + current_span.add_event(f"Requesting {requested_session_count} sessions") | ||
| 481 | + | ||
| 482 | + if created_session_count >= self.size: | ||
| 483 | + current_span.add_event( | ||
| 484 | + "Created no new sessions as sessionPool is full", | ||
| 485 | + {"kind": type(self).__name__}, | ||
| 486 | + ) | ||
| 487 | + return | ||
| 488 | + | ||
| 489 | + returned_session_count = 0 | ||
| 425 | 490 | while created_session_count < self.size: | |
| 426 | 491 | resp = api.batch_create_sessions( | |
| 427 | 492 | request=request, | |
@@ -431,8 +496,17 @@ def bind(self, database): | |||
| 431 | 496 | session = self._new_session() | |
| 432 | 497 | session._session_id = session_pb.name.split("/")[-1] | |
| 433 | 498 | self.put(session) | |
| 499 | + returned_session_count += 1 | ||
| 500 | + | ||
| 434 | 501 | created_session_count += len(resp.session) | |
| 435 | 502 | ||
| 503 | + current_span.add_event( | ||
| 504 | + "Requested for {requested_session_count} sessions, return {returned_session_count}", | ||
| 505 | + { | ||
| 506 | + "kind": "pinging_pool", | ||
| 507 | + }, | ||
| 508 | + ) | ||
| 509 | + | ||
| 436 | 510 | def get(self, timeout=None): | |
| 437 | 511 | """Check a session out from the pool. | |
| 438 | 512 | ||
@@ -447,6 +521,12 @@ def get(self, timeout=None): | |||
| 447 | 521 | if timeout is None: | |
| 448 | 522 | timeout = self.default_timeout | |
| 449 | 523 | ||
| 524 | + start_time = time.time() | ||
| 525 | + current_span = get_current_span() | ||
| 526 | + current_span.add_event( | ||
| 527 | + "Waiting for a session to become available", {"kind": "pinging_pool"} | ||
| 528 | + ) | ||
| 529 | + | ||
| 450 | 530 | ping_after, session = self._sessions.get(block=True, timeout=timeout) | |
| 451 | 531 | ||
| 452 | 532 | if _NOW() > ping_after: | |
@@ -457,6 +537,14 @@ def get(self, timeout=None): | |||
| 457 | 537 | session = self._new_session() | |
| 458 | 538 | session.create() | |
| 459 | 539 | ||
| 540 | + current_span.add_event( | ||
| 541 | + "Acquired session", | ||
| 542 | + { | ||
| 543 | + "time.elapsed": time.time() - start_time, | ||
| 544 | + "session.id": session.session_id, | ||
| 545 | + "kind": "pinging_pool", | ||
| 546 | + }, | ||
| 547 | + ) | ||
| 460 | 548 | return session | |
| 461 | 549 | ||
| 462 | 550 | def put(self, session): | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -30,7 +30,11 @@ | |||
| 30 | 30 | _metadata_with_prefix, | |
| 31 | 31 | _metadata_with_leader_aware_routing, | |
| 32 | 32 | ) | |
| 33 | - from google.cloud.spanner_v1._opentelemetry_tracing import trace_call | ||
| 33 | + from google.cloud.spanner_v1._opentelemetry_tracing import ( | ||
| 34 | + get_current_span, | ||
| 35 | + set_span_error_and_record_exception, | ||
| 36 | + trace_call, | ||
| 37 | + ) | ||
| 34 | 38 | from google.cloud.spanner_v1.batch import Batch | |
| 35 | 39 | from google.cloud.spanner_v1.snapshot import Snapshot | |
| 36 | 40 | from google.cloud.spanner_v1.transaction import Transaction | |
@@ -113,6 +117,10 @@ def name(self): | |||
| 113 | 117 | :raises ValueError: if session is not yet created | |
| 114 | 118 | """ | |
| 115 | 119 | if self._session_id is None: | |
| 120 | + err = "No session available" | ||
| 121 | + current_span = get_current_span() | ||
| 122 | + current_span.add_event(err) | ||
| 123 | + set_span_error_and_record_exception(current_span, err) | ||
| 116 | 124 | raise ValueError("No session ID set by back-end") | |
| 117 | 125 | return self._database.name + "/sessions/" + self._session_id | |
| 118 | 126 | ||
@@ -124,8 +132,14 @@ def create(self): | |||
| 124 | 132 | ||
| 125 | 133 | :raises ValueError: if :attr:`session_id` is already set. | |
| 126 | 134 | """ | |
| 135 | + current_span = get_current_span() | ||
| 136 | + current_span.add_event("Creating Session") | ||
| 137 | + | ||
| 127 | 138 | if self._session_id is not None: | |
| 128 | - raise ValueError("Session ID already set by back-end") | ||
| 139 | + err = "Session ID already set by back-end" | ||
| 140 | + current_span.add_event(err) | ||
| 141 | + set_span_error_and_record_exception(current_span, err) | ||
| 142 | + raise ValueError(err) | ||
| 129 | 143 | api = self._database.spanner_api | |
| 130 | 144 | metadata = _metadata_with_prefix(self._database.name) | |
| 131 | 145 | if self._database._route_to_leader_enabled: | |
@@ -148,6 +162,7 @@ def create(self): | |||
| 148 | 162 | metadata=metadata, | |
| 149 | 163 | ) | |
| 150 | 164 | self._session_id = session_pb.name.split("/")[-1] | |
| 165 | + current_span.add_event("Using Session", {"id": self._session_id}) | ||
| 151 | 166 | ||
| 152 | 167 | def exists(self): | |
| 153 | 168 | """Test for the existence of this session. | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -92,3 +92,13 @@ def assertSpanAttributes( | |||
| 92 | 92 | self.assertEqual(span.name, name) | |
| 93 | 93 | self.assertEqual(span.status.status_code, status) | |
| 94 | 94 | self.assertEqual(dict(span.attributes), attributes) | |
| 95 | + | ||
| 96 | + def assertSpanEvents(self, name, wantEventNames=[], span=None): | ||
| 97 | + if HAS_OPENTELEMETRY_INSTALLED: | ||
| 98 | + if not span: | ||
| 99 | + span_list = self.ot_exporter.get_finished_spans() | ||
| 100 | + self.assertEqual(len(span_list) > 0, True) | ||
| 101 | + span = span_list[0] | ||
| 102 | + | ||
| 103 | + self.assertEqual(span.name, name) | ||
| 104 | + self.assertEqual(len(span.events), len(wantEventNames)) | ||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -611,6 +611,10 @@ def __init__(self, database=None, name=TestBatch.SESSION_NAME): | |||
| 611 | 611 | self._database = database | |
| 612 | 612 | self.name = name | |
| 613 | 613 | ||
| 614 | + @property | ||
| 615 | + def session_id(self): | ||
| 616 | + return self.name | ||
| 617 | + | ||
| 614 | 618 | ||
| 615 | 619 | class _Database(object): | |
| 616 | 620 | name = "testing" | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -3188,6 +3188,10 @@ def run_in_transaction(self, func, *args, **kw): | |||
| 3188 | 3188 | self._retried = (func, args, kw) | |
| 3189 | 3189 | return self._committed | |
| 3190 | 3190 | ||
| 3191 | + @property | ||
| 3192 | + def session_id(self): | ||
| 3193 | + return self.name | ||
| 3194 | + | ||
| 3191 | 3195 | ||
| 3192 | 3196 | class _MockIterator(object): | |
| 3193 | 3197 | def __init__(self, *values, **kw): | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -14,6 +14,7 @@ | |||
| 14 | 14 | ||
| 15 | 15 | ||
| 16 | 16 | from functools import total_ordering | |
| 17 | + import time | ||
| 17 | 18 | import unittest | |
| 18 | 19 | ||
| 19 | 20 | import mock | |
@@ -923,6 +924,8 @@ def __init__(self, database, exists=True, transaction=None): | |||
| 923 | 924 | self.create = mock.Mock() | |
| 924 | 925 | self._deleted = False | |
| 925 | 926 | self._transaction = transaction | |
| 927 | + # Generate a faux id. | ||
| 928 | + self._session_id = f"time.time()" | ||
| 926 | 929 | ||
| 927 | 930 | def __lt__(self, other): | |
| 928 | 931 | return id(self) < id(other) | |
@@ -949,6 +952,10 @@ def transaction(self): | |||
| 949 | 952 | txn = self._transaction = _make_transaction(self) | |
| 950 | 953 | return txn | |
| 951 | 954 | ||
| 955 | + @property | ||
| 956 | + def session_id(self): | ||
| 957 | + return self._session_id | ||
| 958 | + | ||
| 952 | 959 | ||
| 953 | 960 | class _Database(object): | |
| 954 | 961 | def __init__(self, name): | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -15,6 +15,7 @@ | |||
| 15 | 15 | ||
| 16 | 16 | import google.api_core.gapic_v1.method | |
| 17 | 17 | from google.cloud.spanner_v1 import RequestOptions | |
| 18 | + from google.cloud.spanner_v1._opentelemetry_tracing import trace_call | ||
| 18 | 19 | import mock | |
| 19 | 20 | from tests._helpers import ( | |
| 20 | 21 | OpenTelemetryBase, | |
@@ -174,6 +175,43 @@ def test_create_w_database_role(self): | |||
| 174 | 175 | "CloudSpanner.CreateSession", attributes=TestSession.BASE_ATTRIBUTES | |
| 175 | 176 | ) | |
| 176 | 177 | ||
| 178 | + def test_create_session_span_annotations(self): | ||
| 179 | + from google.cloud.spanner_v1 import CreateSessionRequest | ||
| 180 | + from google.cloud.spanner_v1 import Session as SessionRequestProto | ||
| 181 | + | ||
| 182 | + session_pb = self._make_session_pb( | ||
| 183 | + self.SESSION_NAME, database_role=self.DATABASE_ROLE | ||
| 184 | + ) | ||
| 185 | + | ||
| 186 | + gax_api = self._make_spanner_api() | ||
| 187 | + gax_api.create_session.return_value = session_pb | ||
| 188 | + database = self._make_database(database_role=self.DATABASE_ROLE) | ||
| 189 | + database.spanner_api = gax_api | ||
| 190 | + session = self._make_one(database, database_role=self.DATABASE_ROLE) | ||
| 191 | + | ||
| 192 | + with trace_call("TestSessionSpan", session): | ||
| 193 | + session.create() | ||
| 194 | + | ||
| 195 | + self.assertEqual(session.session_id, self.SESSION_ID) | ||
| 196 | + self.assertEqual(session.database_role, self.DATABASE_ROLE) | ||
| 197 | + session_template = SessionRequestProto(creator_role=self.DATABASE_ROLE) | ||
| 198 | + | ||
| 199 | + request = CreateSessionRequest( | ||
| 200 | + database=database.name, | ||
| 201 | + session=session_template, | ||
| 202 | + ) | ||
| 203 | + | ||
| 204 | + gax_api.create_session.assert_called_once_with( | ||
| 205 | + request=request, | ||
| 206 | + metadata=[ | ||
| 207 | + ("google-cloud-resource-prefix", database.name), | ||
| 208 | + ("x-goog-spanner-route-to-leader", "true"), | ||
| 209 | + ], | ||
| 210 | + ) | ||
| 211 | + | ||
| 212 | + wantEventNames = ["Acquiring session", "Creating Session", "Using Session"] | ||
| 213 | + self.assertSpanEvents("CloudSpanner.CreateSession", wantEventNames) | ||
| 214 | + | ||
| 177 | 215 | def test_create_wo_database_role(self): | |
| 178 | 216 | from google.cloud.spanner_v1 import CreateSessionRequest | |
| 179 | 217 | ||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -1822,6 +1822,10 @@ def __init__(self, database=None, name=TestSnapshot.SESSION_NAME): | |||
| 1822 | 1822 | self._database = database | |
| 1823 | 1823 | self.name = name | |
| 1824 | 1824 | ||
| 1825 | + @property | ||
| 1826 | + def session_id(self): | ||
| 1827 | + return self.name | ||
| 1828 | + | ||
| 1825 | 1829 | ||
| 1826 | 1830 | class _MockIterator(object): | |
| 1827 | 1831 | def __init__(self, *values, **kw): | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -1082,6 +1082,10 @@ def __init__(self, database=None, name=TestTransaction.SESSION_NAME): | |||
| 1082 | 1082 | self._database = database | |
| 1083 | 1083 | self.name = name | |
| 1084 | 1084 | ||
| 1085 | + @property | ||
| 1086 | + def session_id(self): | ||
| 1087 | + return self.name | ||
| 1088 | + | ||
| 1085 | 1089 | ||
| 1086 | 1090 | class _MockIterator(object): | |
| 1087 | 1091 | def __init__(self, *values, **kw): | |
| Back | FazBrowse Home | New Git URL |
0 commit comments