2828using AdbcDrivers . Databricks . Telemetry . TagDefinitions ;
2929using AdbcDrivers . HiveServer2 ;
3030using AdbcDrivers . HiveServer2 . Hive2 ;
31+ using Apache . Arrow . Adbc . Tracing ;
3132using Apache . Hive . Service . Rpc . Thrift ;
3233
3334namespace AdbcDrivers . Databricks . Reader
@@ -42,6 +43,7 @@ internal class DatabricksOperationStatusPoller : IOperationStatusPoller
4243 private readonly int _heartbeatIntervalSeconds ;
4344 private readonly int _requestTimeoutSeconds ;
4445 private readonly IResponse _response ;
46+ private readonly IActivityTracer ? _activityTracer ;
4547 // internal cancellation token source - won't affect the external token
4648 private CancellationTokenSource ? _internalCts ;
4749 private Task ? _operationStatusPollingTask ;
@@ -58,12 +60,14 @@ public DatabricksOperationStatusPoller(
5860 IHiveServer2Statement statement ,
5961 IResponse response ,
6062 int heartbeatIntervalSeconds = DatabricksConstants . DefaultOperationStatusPollingIntervalSeconds ,
61- int requestTimeoutSeconds = DatabricksConstants . DefaultOperationStatusRequestTimeoutSeconds )
63+ int requestTimeoutSeconds = DatabricksConstants . DefaultOperationStatusRequestTimeoutSeconds ,
64+ IActivityTracer ? activityTracer = null )
6265 {
6366 _statement = statement ?? throw new ArgumentNullException ( nameof ( statement ) ) ;
6467 _response = response ;
6568 _heartbeatIntervalSeconds = heartbeatIntervalSeconds ;
6669 _requestTimeoutSeconds = requestTimeoutSeconds ;
70+ _activityTracer = activityTracer ;
6771 }
6872
6973 public bool IsStarted => _operationStatusPollingTask != null ;
@@ -82,10 +86,22 @@ public void Start(CancellationToken externalToken = default)
8286 _internalCts = new CancellationTokenSource ( ) ;
8387 // create a linked token to the external token so that the external token can cancel the operation status polling task if needed
8488 var linkedToken = CancellationTokenSource . CreateLinkedTokenSource ( _internalCts . Token , externalToken ) . Token ;
85- _operationStatusPollingTask = Task . Run ( ( ) => PollOperationStatus ( linkedToken ) ) ;
89+ _operationStatusPollingTask = Task . Run ( async ( ) =>
90+ {
91+ if ( _activityTracer != null )
92+ {
93+ await _activityTracer . Trace . TraceActivityAsync (
94+ activity => PollOperationStatus ( linkedToken , activity ) ,
95+ activityName : "PollOperationStatus" ) . ConfigureAwait ( false ) ;
96+ }
97+ else
98+ {
99+ await PollOperationStatus ( linkedToken , null ) . ConfigureAwait ( false ) ;
100+ }
101+ } ) ;
86102 }
87103
88- private async Task PollOperationStatus ( CancellationToken cancellationToken )
104+ private async Task PollOperationStatus ( CancellationToken cancellationToken , Activity ? activity )
89105 {
90106 int consecutiveFailures = 0 ;
91107
@@ -108,6 +124,13 @@ private async Task PollOperationStatus(CancellationToken cancellationToken)
108124 // Track poll count for telemetry
109125 _pollCount ++ ;
110126
127+ activity ? . AddEvent ( new ActivityEvent ( "operation_status_poller.poll_success" ,
128+ tags : new ActivityTagsCollection
129+ {
130+ { "poll_count" , _pollCount } ,
131+ { "operation_state" , response . OperationState . ToString ( ) }
132+ } ) ) ;
133+
111134 // end the heartbeat if the command has terminated
112135 if ( response . OperationState == TOperationState . CANCELED_STATE ||
113136 response . OperationState == TOperationState . ERROR_STATE ||
@@ -130,7 +153,7 @@ private async Task PollOperationStatus(CancellationToken cancellationToken)
130153 // Log the error but continue polling. Transient errors (e.g. ObjectDisposedException
131154 // from TLS connection recycling) should not kill the heartbeat poller, as that would
132155 // cause the server-side command inactivity timeout to expire and terminate the query.
133- Activity . Current ? . AddEvent ( new ActivityEvent ( "operation_status_poller.poll_error" ,
156+ activity ? . AddEvent ( new ActivityEvent ( "operation_status_poller.poll_error" ,
134157 tags : new ActivityTagsCollection
135158 {
136159 { "error.type" , ex . GetType ( ) . Name } ,
@@ -141,7 +164,7 @@ private async Task PollOperationStatus(CancellationToken cancellationToken)
141164
142165 if ( consecutiveFailures >= MaxConsecutiveFailures )
143166 {
144- Activity . Current ? . AddEvent ( new ActivityEvent ( "operation_status_poller.max_failures_reached" ,
167+ activity ? . AddEvent ( new ActivityEvent ( "operation_status_poller.max_failures_reached" ,
145168 tags : new ActivityTagsCollection
146169 {
147170 { "consecutive_failures" , consecutiveFailures } ,
@@ -151,13 +174,20 @@ private async Task PollOperationStatus(CancellationToken cancellationToken)
151174 }
152175 }
153176
154- // Wait before next poll. On cancellation this throws OperationCanceledException
155- // which propagates up to the caller (Dispose catches it).
156- await Task . Delay ( TimeSpan . FromSeconds ( _heartbeatIntervalSeconds ) , cancellationToken ) . ConfigureAwait ( false ) ;
177+ // Wait before next poll.
178+ try
179+ {
180+ await Task . Delay ( TimeSpan . FromSeconds ( _heartbeatIntervalSeconds ) , cancellationToken ) . ConfigureAwait ( false ) ;
181+ }
182+ catch ( OperationCanceledException ) when ( cancellationToken . IsCancellationRequested )
183+ {
184+ // Normal shutdown — don't let this propagate as an error through TraceActivityAsync
185+ break ;
186+ }
157187 }
158188
159189 // Add telemetry tags to current activity when polling completes
160- Activity . Current ? . SetTag ( StatementExecutionEvent . PollCount , _pollCount ) ;
190+ activity ? . SetTag ( StatementExecutionEvent . PollCount , _pollCount ) ;
161191 }
162192
163193 public void Stop ( )
0 commit comments