Skip to content

Commit

Permalink
Add ack result, ack end event to subscribe span and end it
Browse files Browse the repository at this point in the history
  • Loading branch information
mukund-ananthu committed Sep 18, 2024
1 parent 8827117 commit 1703dd2
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 0 deletions.
11 changes: 11 additions & 0 deletions google/cloud/pubsub_v1/open_telemetry/subscribe_opentelemetry.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,3 +77,14 @@ def add_subscribe_span_event(self, event: str) -> None:
"timestamp": str(datetime.now()),
},
)

def end_subscribe_span(self) -> None:
assert self._subscribe_span is not None
self._subscribe_span.end()

def set_subscribe_span_result(self, result: str) -> None:
assert self._subscribe_span is not None
self._subscribe_span.set_attribute(
key="messaging.gcp_pubsub.result",
value=result,
)
13 changes: 13 additions & 0 deletions google/cloud/pubsub_v1/subscriber/_protocol/dispatcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -243,6 +243,12 @@ def ack(self, items: Sequence[requests.AckRequest]) -> None:
ack_reqs_dict=ack_reqs_dict,
)

for completed_ack in requests_completed:
if completed_ack.opentelemetry_data:
completed_ack.opentelemetry_data.set_subscribe_span_result("acked")
completed_ack.opentelemetry_data.add_subscribe_span_event("ack end")
completed_ack.opentelemetry_data.end_subscribe_span()

# Remove the completed messages from lease management.
self.drop(requests_completed)

Expand Down Expand Up @@ -286,6 +292,13 @@ def _retry_acks(self, requests_to_retry):
ack_ids=[req.ack_id for req in requests_to_retry],
ack_reqs_dict=ack_reqs_dict,
)

for completed_ack in requests_completed:
if completed_ack.opentelemetry_data:
completed_ack.opentelemetry_data.set_subscribe_span_result("acked")
completed_ack.opentelemetry_data.add_subscribe_span_event("ack end")
completed_ack.opentelemetry_data.end_subscribe_span()

assert (
len(requests_to_retry) <= _ACK_IDS_BATCH_SIZE
), "Too many requests to be retried."
Expand Down

0 comments on commit 1703dd2

Please sign in to comment.