HBASE-26545 Implement tracing of scan

* on `AsyncTable`, both `scan` and `scanAll` methods should result in `SCAN` table operations.
* the span of the `SCAN` table operation should have children representing all the RPC calls
  involved in servicing the scan.
* when a user provides custom implementation of `AdvancedScanResultConsumer`, any spans emitted
  from the callback methods should also be tied to the span that represents the `SCAN` table
  operation. This is easily done because these callbacks are executed on the RPC thread.
* when a user provides a custom implementation of `ScanResultConsumer`, any spans emitted from the
  callback methods should be also be tied to the span that represents the `SCAN` table
  operation. This accomplished by carefully passing the span instance around after it is created.

Signed-off-by: Andrew Purtell <[email protected]>
Signed-off-by: Duo Zhang <[email protected]>

Author: ndimiduk

hbase-client/src/main/java/org/apache/hadoop/hbase/client/AsyncClientScanner.java

@@ -1,4 +1,4 @@
-/**
+/*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
@@ -27,19 +27,21 @@
 import static org.apache.hadoop.hbase.client.ConnectionUtils.isRemote;
 import static org.apache.hadoop.hbase.client.ConnectionUtils.timelineConsistentRead;
 import static org.apache.hadoop.hbase.util.FutureUtils.addListener;
-
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.StatusCode;
+import io.opentelemetry.context.Scope;
 import java.io.IOException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.hadoop.hbase.HRegionLocation;
 import org.apache.hadoop.hbase.TableName;
 import org.apache.hadoop.hbase.client.metrics.ScanMetrics;
+import org.apache.hadoop.hbase.client.trace.TableOperationSpanBuilder;
 import org.apache.hadoop.hbase.ipc.HBaseRpcController;
+import org.apache.hadoop.hbase.trace.TraceUtil;
 import org.apache.yetus.audience.InterfaceAudience;
-
 import org.apache.hbase.thirdparty.io.netty.util.Timer;
-
 import org.apache.hadoop.hbase.shaded.protobuf.RequestConverter;
 import org.apache.hadoop.hbase.shaded.protobuf.generated.ClientProtos.ClientService;
 import org.apache.hadoop.hbase.shaded.protobuf.generated.ClientProtos.ClientService.Interface;
@@ -85,6 +87,8 @@ class AsyncClientScanner {
 
   private final ScanResultCache resultCache;
 
+  private final Span span;
+
   public AsyncClientScanner(Scan scan, AdvancedScanResultConsumer consumer, TableName tableName,
       AsyncConnectionImpl conn, Timer retryTimer, long pauseNs, long pauseForCQTBENs,
       int maxAttempts, long scanTimeoutNs, long rpcTimeoutNs, int startLogErrorsCnt) {
@@ -112,6 +116,21 @@ public AsyncClientScanner(Scan scan, AdvancedScanResultConsumer consumer, TableN
     } else {
       this.scanMetrics = null;
     }
+
+    /*
+     * Assumes that the `start()` method is called immediately after construction. If this is no
+     * longer the case, for tracing correctness, we should move the start of the span into the
+     * `start()` method. The cost of doing so would be making access to the `span` safe for
+     * concurrent threads.
+     */
+    span = new TableOperationSpanBuilder(conn)
+      .setTableName(tableName)
+      .setOperation(scan)
+      .build();
+    if (consumer instanceof AsyncTableResultScanner) {
+      AsyncTableResultScanner scanner = (AsyncTableResultScanner) consumer;
+      scanner.setSpan(span);
+    }
   }
 
   private static final class OpenScannerResponse {
@@ -140,26 +159,35 @@ public OpenScannerResponse(HRegionLocation loc, boolean isRegionServerRemote, In
 
   private CompletableFuture<OpenScannerResponse> callOpenScanner(HBaseRpcController controller,
       HRegionLocation loc, ClientService.Interface stub) {
-    boolean isRegionServerRemote = isRemote(loc.getHostname());
-    incRPCCallsMetrics(scanMetrics, isRegionServerRemote);
-    if (openScannerTries.getAndIncrement() > 1) {
-      incRPCRetriesMetrics(scanMetrics, isRegionServerRemote);
+    try (Scope ignored = span.makeCurrent()) {
+      boolean isRegionServerRemote = isRemote(loc.getHostname());
+      incRPCCallsMetrics(scanMetrics, isRegionServerRemote);
+      if (openScannerTries.getAndIncrement() > 1) {
+        incRPCRetriesMetrics(scanMetrics, isRegionServerRemote);
+      }
+      CompletableFuture<OpenScannerResponse> future = new CompletableFuture<>();
+      try {
+        ScanRequest request = RequestConverter.buildScanRequest(
+          loc.getRegion().getRegionName(), scan, scan.getCaching(), false);
+        stub.scan(controller, request, resp -> {
+          try (Scope ignored1 = span.makeCurrent()) {
+            if (controller.failed()) {
+              final IOException e = controller.getFailed();
+              future.completeExceptionally(e);
+              TraceUtil.setError(span, e);
+              span.end();
+              return;
+            }
+            future.complete(new OpenScannerResponse(
+              loc, isRegionServerRemote, stub, controller, resp));
+          }
+        });
+      } catch (IOException e) {
+        // span is closed by listener attached to the Future in `openScanner()`
+        future.completeExceptionally(e);
+      }
+      return future;
     }
-    CompletableFuture<OpenScannerResponse> future = new CompletableFuture<>();
-    try {
-      ScanRequest request = RequestConverter.buildScanRequest(loc.getRegion().getRegionName(), scan,
-        scan.getCaching(), false);
-      stub.scan(controller, request, resp -> {
-        if (controller.failed()) {
-          future.completeExceptionally(controller.getFailed());
-          return;
-        }
-        future.complete(new OpenScannerResponse(loc, isRegionServerRemote, stub, controller, resp));
-      });
-    } catch (IOException e) {
-      future.completeExceptionally(e);
-    }
-    return future;
   }
 
   private void startScan(OpenScannerResponse resp) {
@@ -173,26 +201,40 @@ private void startScan(OpenScannerResponse resp) {
         .pauseForCQTBE(pauseForCQTBENs, TimeUnit.NANOSECONDS).maxAttempts(maxAttempts)
         .startLogErrorsCnt(startLogErrorsCnt).start(resp.controller, resp.resp),
       (hasMore, error) -> {
-        if (error != null) {
-          consumer.onError(error);
-          return;
-        }
-        if (hasMore) {
-          openScanner();
-        } else {
-          consumer.onComplete();
+        try (Scope ignored = span.makeCurrent()) {
+          if (error != null) {
+            try {
+              consumer.onError(error);
+              return;
+            } finally {
+              TraceUtil.setError(span, error);
+              span.end();
+            }
+          }
+          if (hasMore) {
+            openScanner();
+          } else {
+            try {
+              consumer.onComplete();
+            } finally {
+              span.setStatus(StatusCode.OK);
+              span.end();
+            }
+          }
         }
       });
   }
 
   private CompletableFuture<OpenScannerResponse> openScanner(int replicaId) {
-    return conn.callerFactory.<OpenScannerResponse> single().table(tableName)
-      .row(scan.getStartRow()).replicaId(replicaId).locateType(getLocateType(scan))
-      .priority(scan.getPriority())
-      .rpcTimeout(rpcTimeoutNs, TimeUnit.NANOSECONDS)
-      .operationTimeout(scanTimeoutNs, TimeUnit.NANOSECONDS).pause(pauseNs, TimeUnit.NANOSECONDS)
-      .pauseForCQTBE(pauseForCQTBENs, TimeUnit.NANOSECONDS).maxAttempts(maxAttempts)
-      .startLogErrorsCnt(startLogErrorsCnt).action(this::callOpenScanner).call();
+    try (Scope ignored = span.makeCurrent()) {
+      return conn.callerFactory.<OpenScannerResponse> single().table(tableName)
+        .row(scan.getStartRow()).replicaId(replicaId).locateType(getLocateType(scan))
+        .priority(scan.getPriority())
+        .rpcTimeout(rpcTimeoutNs, TimeUnit.NANOSECONDS)
+        .operationTimeout(scanTimeoutNs, TimeUnit.NANOSECONDS).pause(pauseNs, TimeUnit.NANOSECONDS)
+        .pauseForCQTBE(pauseForCQTBENs, TimeUnit.NANOSECONDS).maxAttempts(maxAttempts)
+        .startLogErrorsCnt(startLogErrorsCnt).action(this::callOpenScanner).call();
+    }
   }
 
   private long getPrimaryTimeoutNs() {
@@ -206,15 +248,24 @@ private void openScanner() {
     addListener(timelineConsistentRead(conn.getLocator(), tableName, scan, scan.getStartRow(),
       getLocateType(scan), this::openScanner, rpcTimeoutNs, getPrimaryTimeoutNs(), retryTimer,
       conn.getConnectionMetrics()), (resp, error) -> {
-        if (error != null) {
-          consumer.onError(error);
-          return;
+        try (Scope ignored = span.makeCurrent()) {
+          if (error != null) {
+            try {
+              consumer.onError(error);
+              return;
+            } finally {
+              TraceUtil.setError(span, error);
+              span.end();
+            }
+          }
+          startScan(resp);
         }
-        startScan(resp);
       });
   }
 
   public void start() {
-    openScanner();
+    try (Scope ignored = span.makeCurrent()) {
+      openScanner();
+    }
   }
 }

hbase-client/src/main/java/org/apache/hadoop/hbase/client/AsyncScanSingleRegionRpcRetryingCaller.java

@@ -27,7 +27,8 @@
 import static org.apache.hadoop.hbase.client.ConnectionUtils.translateException;
 import static org.apache.hadoop.hbase.client.ConnectionUtils.updateResultsMetrics;
 import static org.apache.hadoop.hbase.client.ConnectionUtils.updateServerSideMetrics;
-
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
@@ -50,11 +51,9 @@
 import org.apache.yetus.audience.InterfaceAudience;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-
 import org.apache.hbase.thirdparty.com.google.common.base.Preconditions;
 import org.apache.hbase.thirdparty.io.netty.util.Timeout;
 import org.apache.hbase.thirdparty.io.netty.util.Timer;
-
 import org.apache.hadoop.hbase.shaded.protobuf.ProtobufUtil;
 import org.apache.hadoop.hbase.shaded.protobuf.RequestConverter;
 import org.apache.hadoop.hbase.shaded.protobuf.ResponseConverter;
@@ -573,7 +572,12 @@ private void call() {
     resetController(controller, callTimeoutNs, priority);
     ScanRequest req = RequestConverter.buildScanRequest(scannerId, scan.getCaching(), false,
       nextCallSeq, scan.isScanMetricsEnabled(), false, scan.getLimit());
-    stub.scan(controller, req, resp -> onComplete(controller, resp));
+    final Context context = Context.current();
+    stub.scan(controller, req, resp -> {
+      try (Scope ignored = context.makeCurrent()) {
+        onComplete(controller, resp);
+      }
+    });
   }
 
   private void next() {

hbase-client/src/main/java/org/apache/hadoop/hbase/client/AsyncTableImpl.java

@@ -18,7 +18,9 @@
 package org.apache.hadoop.hbase.client;
 
 import static java.util.stream.Collectors.toList;
+import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
 import java.io.IOException;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
@@ -231,22 +233,29 @@ public ResultScanner getScanner(Scan scan) {
   }
 
   private void scan0(Scan scan, ScanResultConsumer consumer) {
-    try (ResultScanner scanner = getScanner(scan)) {
-      consumer.onScanMetricsCreated(scanner.getScanMetrics());
-      for (Result result; (result = scanner.next()) != null;) {
-        if (!consumer.onNext(result)) {
-          break;
+    Span span = null;
+    try (AsyncTableResultScanner scanner = rawTable.getScanner(scan)) {
+      span = scanner.getSpan();
+      try (Scope ignored = span.makeCurrent()) {
+        consumer.onScanMetricsCreated(scanner.getScanMetrics());
+        for (Result result; (result = scanner.next()) != null; ) {
+          if (!consumer.onNext(result)) {
+            break;
+          }
         }
+        consumer.onComplete();
       }
-      consumer.onComplete();
     } catch (IOException e) {
-      consumer.onError(e);
+      try (Scope ignored = span.makeCurrent()) {
+        consumer.onError(e);
+      }
     }
   }
 
   @Override
   public void scan(Scan scan, ScanResultConsumer consumer) {
-    pool.execute(() -> scan0(scan, consumer));
+    final Context context = Context.current();
+    pool.execute(context.wrap(() -> scan0(scan, consumer)));
   }
 
   @Override

hbase-client/src/main/java/org/apache/hadoop/hbase/client/AsyncTableResultScanner.java

@@ -18,7 +18,7 @@
 package org.apache.hadoop.hbase.client;
 
 import static org.apache.hadoop.hbase.client.ConnectionUtils.calcEstimatedSize;
-
+import io.opentelemetry.api.trace.Span;
 import java.io.IOException;
 import java.io.InterruptedIOException;
 import java.util.ArrayDeque;
@@ -58,6 +58,9 @@ class AsyncTableResultScanner implements ResultScanner, AdvancedScanResultConsum
 
   private ScanResumer resumer;
 
+  // Used to pass the span instance to the `AsyncTableImpl` from its underlying `rawAsyncTable`.
+  private Span span = null;
+
   public AsyncTableResultScanner(TableName tableName, Scan scan, long maxCacheSize) {
     this.tableName = tableName;
     this.maxCacheSize = maxCacheSize;
@@ -79,6 +82,14 @@ private void stopPrefetch(ScanController controller) {
     resumer = controller.suspend();
   }
 
+  Span getSpan() {
+    return span;
+  }
+
+  void setSpan(final Span span) {
+    this.span = span;
+  }
+
   @Override
   public synchronized void onNext(Result[] results, ScanController controller) {
     assert results.length > 0;

hbase-client/src/main/java/org/apache/hadoop/hbase/client/RawAsyncTableImpl.java

@@ -640,30 +640,26 @@ public AsyncTableResultScanner getScanner(Scan scan) {
 
   @Override
   public CompletableFuture<List<Result>> scanAll(Scan scan) {
-    final Supplier<Span> supplier = newTableOperationSpanBuilder()
-      .setOperation(scan);
-    return tracedFuture(() -> {
-      CompletableFuture<List<Result>> future = new CompletableFuture<>();
-      List<Result> scanResults = new ArrayList<>();
-      scan(scan, new AdvancedScanResultConsumer() {
+    CompletableFuture<List<Result>> future = new CompletableFuture<>();
+    List<Result> scanResults = new ArrayList<>();
+    scan(scan, new AdvancedScanResultConsumer() {
 
-        @Override
-        public void onNext(Result[] results, ScanController controller) {
-          scanResults.addAll(Arrays.asList(results));
-        }
+      @Override
+      public void onNext(Result[] results, ScanController controller) {
+        scanResults.addAll(Arrays.asList(results));
+      }
 
-        @Override
-        public void onError(Throwable error) {
-          future.completeExceptionally(error);
-        }
+      @Override
+      public void onError(Throwable error) {
+        future.completeExceptionally(error);
+      }
 
-        @Override
-        public void onComplete() {
-          future.complete(scanResults);
-        }
-      });
-      return future;
-    }, supplier);
+      @Override
+      public void onComplete() {
+        future.complete(scanResults);
+      }
+    });
+    return future;
   }
 
   @Override