chore: add Span Links for asynchronous pagination

Author: keshavdandeva

java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java

@@ -59,6 +59,8 @@
 import com.google.common.util.concurrent.Uninterruptibles;
 import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanBuilder;
+import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.api.trace.StatusCode;
 import io.opentelemetry.api.trace.Tracer;
 import io.opentelemetry.context.Context;
@@ -1057,72 +1059,78 @@ Thread runNextPageTaskAsync(
     // calls
     populateFirstPage(result, rpcResponseQueue);
 
-    Context asyncContext = (this.otelContext != null) ? this.otelContext : Context.current();
+    SpanContext parentSpanContext = null;
+    if (this.otelContext != null) {
+      parentSpanContext = Span.fromContext(this.otelContext).getSpanContext();
+    }
+    final SpanContext finalParentSpanContext = parentSpanContext;
 
     // This thread makes the RPC calls and paginates
     Runnable nextPageTask =
-        asyncContext.wrap(
-            () -> {
-              Tracer tracer = getSafeTracer();
-              String currentPageToken = firstPageToken;
-              TableResult currentResults = result;
-              TableId destinationTable = null;
-              if (firstPageToken != null) {
-                destinationTable = getDestinationTable(jobId);
-              }
+        () -> {
+          Tracer tracer = getSafeTracer();
+          String currentPageToken = firstPageToken;
+          TableResult currentResults = result;
+          TableId destinationTable = null;
+          if (firstPageToken != null) {
+            destinationTable = getDestinationTable(jobId);
+          }
 
-              try {
-                while (currentPageToken != null) {
-                  // do not process further pages and shutdown
-                  if (Thread.currentThread().isInterrupted() || queryTaskExecutor.isShutdown()) {
-                    LOG.warning(
-                        "%s Interrupted @ runNextPageTaskAsync", Thread.currentThread().getName());
-                    break;
-                  }
+          try {
+            while (currentPageToken != null) {
+              // do not process further pages and shutdown
+              if (Thread.currentThread().isInterrupted() || queryTaskExecutor.isShutdown()) {
+                LOG.warning(
+                    "%s Interrupted @ runNextPageTaskAsync", Thread.currentThread().getName());
+                break;
+              }
 
-                  Span paginationSpan =
-                      tracer.spanBuilder("BigQueryStatement.pagination").startSpan();
-                  try (Scope scope = paginationSpan.makeCurrent()) {
-                    paginationSpan.setAttribute("db.pagination.page_token", currentPageToken);
-
-                    long startTime = System.nanoTime();
-                    currentResults =
-                        this.bigQuery.listTableData(
-                            destinationTable,
-                            TableDataListOption.pageSize(querySettings.getMaxResultPerPage()),
-                            TableDataListOption.pageToken(currentPageToken));
-
-                    long duration = (System.nanoTime() - startTime) / 1000000;
-                    paginationSpan.setAttribute("db.pagination.duration_ms", duration);
-                    paginationSpan.setAttribute(
-                        "db.pagination.rows_fetched", querySettings.getMaxResultPerPage());
-
-                    currentPageToken = currentResults.getNextPageToken();
-                    // this will be parsed asynchronously without blocking the current thread
-                    Uninterruptibles.putUninterruptibly(
-                        rpcResponseQueue, Tuple.of(currentResults, true));
-                    LOG.fine(
-                        "Fetched %d results from the server in %d ms.",
-                        querySettings.getMaxResultPerPage(), (int) duration);
-                  } catch (Exception e) {
-                    paginationSpan.recordException(e);
-                    paginationSpan.setStatus(StatusCode.ERROR, e.getMessage());
-                    throw e;
-                  } finally {
-                    paginationSpan.end();
-                  }
-                }
-              } catch (Exception ex) {
+              SpanBuilder spanBuilder = tracer.spanBuilder("BigQueryStatement.pagination");
+              if (finalParentSpanContext != null && finalParentSpanContext.isValid()) {
+                spanBuilder.addLink(finalParentSpanContext);
+              }
+              Span paginationSpan = spanBuilder.startSpan();
+              try (Scope scope = paginationSpan.makeCurrent()) {
+                paginationSpan.setAttribute("db.pagination.page_token", currentPageToken);
+
+                long startTime = System.nanoTime();
+                currentResults =
+                    this.bigQuery.listTableData(
+                        destinationTable,
+                        TableDataListOption.pageSize(querySettings.getMaxResultPerPage()),
+                        TableDataListOption.pageToken(currentPageToken));
+
+                long duration = (System.nanoTime() - startTime) / 1000000;
+                paginationSpan.setAttribute("db.pagination.duration_ms", duration);
+                paginationSpan.setAttribute(
+                    "db.pagination.rows_fetched", querySettings.getMaxResultPerPage());
+
+                currentPageToken = currentResults.getNextPageToken();
+                // this will be parsed asynchronously without blocking the current thread
                 Uninterruptibles.putUninterruptibly(
-                    bigQueryFieldValueListWrapperBlockingQueue,
-                    BigQueryFieldValueListWrapper.ofError(new BigQueryJdbcRuntimeException(ex)));
+                    rpcResponseQueue, Tuple.of(currentResults, true));
+                LOG.fine(
+                    "Fetched %d results from the server in %d ms.",
+                    querySettings.getMaxResultPerPage(), (int) duration);
+              } catch (Exception e) {
+                paginationSpan.recordException(e);
+                paginationSpan.setStatus(StatusCode.ERROR, e.getMessage());
+                throw e;
               } finally {
-                // this will stop the parseDataTask as well when the pagination completes
-                Uninterruptibles.putUninterruptibly(rpcResponseQueue, Tuple.of(null, false));
+                paginationSpan.end();
               }
-              // We cannot do queryTaskExecutor.shutdownNow() here as populate buffer method may not
-              // have finished processing the records and even that will be interrupted
-            });
+            }
+          } catch (Exception ex) {
+            Uninterruptibles.putUninterruptibly(
+                bigQueryFieldValueListWrapperBlockingQueue,
+                BigQueryFieldValueListWrapper.ofError(new BigQueryJdbcRuntimeException(ex)));
+          } finally {
+            // this will stop the parseDataTask as well when the pagination completes
+            Uninterruptibles.putUninterruptibly(rpcResponseQueue, Tuple.of(null, false));
+          }
+          // We cannot do queryTaskExecutor.shutdownNow() here as populate buffer method may not
+          // have finished processing the records and even that will be interrupted
+        };
 
     Thread nextPageWorker = JDBC_THREAD_FACTORY.newThread(nextPageTask);
     nextPageWorker.start();