HBASE-27657 Connection and Request Attributes (#5332)

Signed-off-by: Bryan Beaudreault <bbeaudreault@apache.org>

Author: rmdmattingly

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

@@ -20,6 +20,7 @@
 import static org.apache.hadoop.hbase.util.FutureUtils.addListener;
 
 import java.io.IOException;
+import java.util.Collections;
 import java.util.concurrent.CompletableFuture;
 import org.apache.hadoop.hbase.ServerName;
 import org.apache.hadoop.hbase.ipc.HBaseRpcController;
@@ -44,7 +45,7 @@ public AsyncAdminRequestRetryingCaller(Timer retryTimer, AsyncConnectionImpl con
     long pauseNs, long pauseNsForServerOverloaded, int maxAttempts, long operationTimeoutNs,
     long rpcTimeoutNs, int startLogErrorsCnt, ServerName serverName, Callable<T> callable) {
     super(retryTimer, conn, priority, pauseNs, pauseNsForServerOverloaded, maxAttempts,
-      operationTimeoutNs, rpcTimeoutNs, startLogErrorsCnt);
+      operationTimeoutNs, rpcTimeoutNs, startLogErrorsCnt, Collections.emptyMap());
     this.serverName = serverName;
     this.callable = callable;
   }

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

@@ -114,6 +114,8 @@ class AsyncBatchRpcRetryingCaller<T> {
 
   private final HBaseServerExceptionPauseManager pauseManager;
 
+  private final Map<String, byte[]> requestAttributes;
+
   // we can not use HRegionLocation as the map key because the hashCode and equals method of
   // HRegionLocation only consider serverName.
   private static final class RegionRequest {
@@ -149,7 +151,8 @@ public int getPriority() {
 
   public AsyncBatchRpcRetryingCaller(Timer retryTimer, AsyncConnectionImpl conn,
     TableName tableName, List<? extends Row> actions, long pauseNs, long pauseNsForServerOverloaded,
-    int maxAttempts, long operationTimeoutNs, long rpcTimeoutNs, int startLogErrorsCnt) {
+    int maxAttempts, long operationTimeoutNs, long rpcTimeoutNs, int startLogErrorsCnt,
+    Map<String, byte[]> requestAttributes) {
     this.retryTimer = retryTimer;
     this.conn = conn;
     this.tableName = tableName;
@@ -180,6 +183,7 @@ public AsyncBatchRpcRetryingCaller(Timer retryTimer, AsyncConnectionImpl conn,
     }
     this.action2Errors = new IdentityHashMap<>();
     this.startNs = System.nanoTime();
+    this.requestAttributes = requestAttributes;
   }
 
   private static boolean hasIncrementOrAppend(Row action) {
@@ -392,6 +396,7 @@ private void sendToServer(ServerName serverName, ServerRequest serverReq, int tr
     HBaseRpcController controller = conn.rpcControllerFactory.newController();
     resetController(controller, Math.min(rpcTimeoutNs, remainingNs),
       calcPriority(serverReq.getPriority(), tableName));
+    controller.setRequestAttributes(requestAttributes);
     if (!cells.isEmpty()) {
       controller.setCellScanner(createCellScanner(cells));
     }

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

@@ -32,6 +32,7 @@
 import io.opentelemetry.api.trace.StatusCode;
 import io.opentelemetry.context.Scope;
 import java.io.IOException;
+import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -92,9 +93,12 @@ class AsyncClientScanner {
 
   private final Span span;
 
+  private final Map<String, byte[]> requestAttributes;
+
   public AsyncClientScanner(Scan scan, AdvancedScanResultConsumer consumer, TableName tableName,
     AsyncConnectionImpl conn, Timer retryTimer, long pauseNs, long pauseNsForServerOverloaded,
-    int maxAttempts, long scanTimeoutNs, long rpcTimeoutNs, int startLogErrorsCnt) {
+    int maxAttempts, long scanTimeoutNs, long rpcTimeoutNs, int startLogErrorsCnt,
+    Map<String, byte[]> requestAttributes) {
     if (scan.getStartRow() == null) {
       scan.withStartRow(EMPTY_START_ROW, scan.includeStartRow());
     }
@@ -113,6 +117,7 @@ public AsyncClientScanner(Scan scan, AdvancedScanResultConsumer consumer, TableN
     this.rpcTimeoutNs = rpcTimeoutNs;
     this.startLogErrorsCnt = startLogErrorsCnt;
     this.resultCache = createScanResultCache(scan);
+    this.requestAttributes = requestAttributes;
     if (scan.isScanMetricsEnabled()) {
       this.scanMetrics = new ScanMetrics();
       consumer.onScanMetricsCreated(scanMetrics);
@@ -191,15 +196,17 @@ private CompletableFuture<OpenScannerResponse> callOpenScanner(HBaseRpcControlle
   }
 
   private void startScan(OpenScannerResponse resp) {
-    addListener(conn.callerFactory.scanSingleRegion().id(resp.resp.getScannerId())
-      .location(resp.loc).remote(resp.isRegionServerRemote)
-      .scannerLeaseTimeoutPeriod(resp.resp.getTtl(), TimeUnit.MILLISECONDS).stub(resp.stub)
-      .setScan(scan).metrics(scanMetrics).consumer(consumer).resultCache(resultCache)
-      .rpcTimeout(rpcTimeoutNs, TimeUnit.NANOSECONDS)
-      .scanTimeout(scanTimeoutNs, TimeUnit.NANOSECONDS).pause(pauseNs, TimeUnit.NANOSECONDS)
-      .pauseForServerOverloaded(pauseNsForServerOverloaded, TimeUnit.NANOSECONDS)
-      .maxAttempts(maxAttempts).startLogErrorsCnt(startLogErrorsCnt)
-      .start(resp.controller, resp.resp), (hasMore, error) -> {
+    addListener(
+      conn.callerFactory.scanSingleRegion().id(resp.resp.getScannerId()).location(resp.loc)
+        .remote(resp.isRegionServerRemote)
+        .scannerLeaseTimeoutPeriod(resp.resp.getTtl(), TimeUnit.MILLISECONDS).stub(resp.stub)
+        .setScan(scan).metrics(scanMetrics).consumer(consumer).resultCache(resultCache)
+        .rpcTimeout(rpcTimeoutNs, TimeUnit.NANOSECONDS)
+        .scanTimeout(scanTimeoutNs, TimeUnit.NANOSECONDS).pause(pauseNs, TimeUnit.NANOSECONDS)
+        .pauseForServerOverloaded(pauseNsForServerOverloaded, TimeUnit.NANOSECONDS)
+        .maxAttempts(maxAttempts).startLogErrorsCnt(startLogErrorsCnt)
+        .setRequestAttributes(requestAttributes).start(resp.controller, resp.resp),
+      (hasMore, error) -> {
         try (Scope ignored = span.makeCurrent()) {
           if (error != null) {
             try {
@@ -231,8 +238,8 @@ private CompletableFuture<OpenScannerResponse> openScanner(int replicaId) {
         .priority(scan.getPriority()).rpcTimeout(rpcTimeoutNs, TimeUnit.NANOSECONDS)
         .operationTimeout(scanTimeoutNs, TimeUnit.NANOSECONDS).pause(pauseNs, TimeUnit.NANOSECONDS)
         .pauseForServerOverloaded(pauseNsForServerOverloaded, TimeUnit.NANOSECONDS)
-        .maxAttempts(maxAttempts).startLogErrorsCnt(startLogErrorsCnt).action(this::callOpenScanner)
-        .call();
+        .maxAttempts(maxAttempts).startLogErrorsCnt(startLogErrorsCnt)
+        .setRequestAttributes(requestAttributes).action(this::callOpenScanner).call();
     }
   }
 

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

@@ -30,6 +30,8 @@
 
 import io.opentelemetry.api.trace.Span;
 import java.io.IOException;
+import java.util.Collections;
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
@@ -122,6 +124,11 @@ public class AsyncConnectionImpl implements AsyncConnection {
 
   public AsyncConnectionImpl(Configuration conf, ConnectionRegistry registry, String clusterId,
     User user) {
+    this(conf, registry, clusterId, user, Collections.emptyMap());
+  }
+
+  public AsyncConnectionImpl(Configuration conf, ConnectionRegistry registry, String clusterId,
+    User user, Map<String, byte[]> connectionAttributes) {
     this.conf = conf;
     this.user = user;
     this.metricsScope = MetricsConnection.getScope(conf, clusterId, this);
@@ -137,7 +144,8 @@ public AsyncConnectionImpl(Configuration conf, ConnectionRegistry registry, Stri
     } else {
       this.metrics = Optional.empty();
     }
-    this.rpcClient = RpcClientFactory.createClient(conf, clusterId, metrics.orElse(null));
+    this.rpcClient =
+      RpcClientFactory.createClient(conf, clusterId, metrics.orElse(null), connectionAttributes);
     this.rpcControllerFactory = RpcControllerFactory.instantiate(conf);
     this.rpcTimeout =
       (int) Math.min(Integer.MAX_VALUE, TimeUnit.NANOSECONDS.toMillis(connConf.getRpcTimeoutNs()));

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

@@ -19,6 +19,7 @@
 
 import static org.apache.hadoop.hbase.util.FutureUtils.addListener;
 
+import java.util.Collections;
 import java.util.concurrent.CompletableFuture;
 import org.apache.hadoop.hbase.exceptions.ClientExceptionsUtil;
 import org.apache.hadoop.hbase.ipc.HBaseRpcController;
@@ -47,7 +48,7 @@ public AsyncMasterRequestRpcRetryingCaller(Timer retryTimer, AsyncConnectionImpl
     Callable<T> callable, int priority, long pauseNs, long pauseNsForServerOverloaded,
     int maxRetries, long operationTimeoutNs, long rpcTimeoutNs, int startLogErrorsCnt) {
     super(retryTimer, conn, priority, pauseNs, pauseNsForServerOverloaded, maxRetries,
-      operationTimeoutNs, rpcTimeoutNs, startLogErrorsCnt);
+      operationTimeoutNs, rpcTimeoutNs, startLogErrorsCnt, Collections.emptyMap());
     this.callable = callable;
   }
 

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

@@ -17,8 +17,10 @@
  */
 package org.apache.hadoop.hbase.client;
 
+import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.ExecutorService;
 import org.apache.hadoop.hbase.TableName;
 import org.apache.hadoop.hbase.client.coprocessor.Batch;
@@ -63,6 +65,7 @@ public static class Builder<T> {
     private int operationTimeout;
     private CancellableRegionServerCallable callable;
     private Object[] results;
+    private Map<String, byte[]> requestAttributes = Collections.emptyMap();
 
     private Builder() {
     }
@@ -124,9 +127,14 @@ Builder<T> setCallable(CancellableRegionServerCallable callable) {
       return this;
     }
 
+    Builder<T> setRequestAttributes(Map<String, byte[]> requestAttributes) {
+      this.requestAttributes = requestAttributes;
+      return this;
+    }
+
     public AsyncProcessTask<T> build() {
       return new AsyncProcessTask<>(pool, tableName, rows, submittedRows, callback, callable,
-        needResults, rpcTimeout, operationTimeout, results);
+        needResults, rpcTimeout, operationTimeout, results, requestAttributes);
     }
   }
 
@@ -140,16 +148,18 @@ public AsyncProcessTask<T> build() {
   private final int rpcTimeout;
   private final int operationTimeout;
   private final Object[] results;
+  private final Map<String, byte[]> requestAttributes;
 
   AsyncProcessTask(AsyncProcessTask<T> task) {
     this(task.getPool(), task.getTableName(), task.getRowAccess(), task.getSubmittedRows(),
       task.getCallback(), task.getCallable(), task.getNeedResults(), task.getRpcTimeout(),
-      task.getOperationTimeout(), task.getResults());
+      task.getOperationTimeout(), task.getResults(), task.getRequestAttributes());
   }
 
   AsyncProcessTask(ExecutorService pool, TableName tableName, RowAccess<? extends Row> rows,
     SubmittedRows size, Batch.Callback<T> callback, CancellableRegionServerCallable callable,
-    boolean needResults, int rpcTimeout, int operationTimeout, Object[] results) {
+    boolean needResults, int rpcTimeout, int operationTimeout, Object[] results,
+    Map<String, byte[]> requestAttributes) {
     this.pool = pool;
     this.tableName = tableName;
     this.rows = rows;
@@ -160,6 +170,7 @@ public AsyncProcessTask<T> build() {
     this.rpcTimeout = rpcTimeout;
     this.operationTimeout = operationTimeout;
     this.results = results;
+    this.requestAttributes = requestAttributes;
   }
 
   public int getOperationTimeout() {
@@ -190,6 +201,10 @@ CancellableRegionServerCallable getCallable() {
     return callable;
   }
 
+  public Map<String, byte[]> getRequestAttributes() {
+    return requestAttributes;
+  }
+
   Object[] getResults() {
     return results;
   }

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

@@ -284,6 +284,7 @@ public void run() {
   private final int operationTimeout;
   private final int rpcTimeout;
   private final AsyncProcess asyncProcess;
+  private final Map<String, byte[]> requestAttributes;
 
   /**
    * For {@link AsyncRequestFutureImpl#manageError(int, Row, Retry, Throwable, ServerName)}. Only
@@ -398,6 +399,7 @@ public AsyncRequestFutureImpl(AsyncProcessTask task, List<Action> actions, long
     if (task.getCallable() == null) {
       tracker = new RetryingTimeTracker().start();
     }
+    this.requestAttributes = task.getRequestAttributes();
   }
 
   protected Set<CancellableRegionServerCallable> getCallsInProgress() {
@@ -1316,7 +1318,8 @@ private ConnectionImplementation.ServerErrorTracker createServerErrorTracker() {
   private MultiServerCallable createCallable(final ServerName server, TableName tableName,
     final MultiAction multi) {
     return new MultiServerCallable(asyncProcess.connection, tableName, server, multi,
-      asyncProcess.rpcFactory.newController(), rpcTimeout, tracker, multi.getPriority());
+      asyncProcess.rpcFactory.newController(), rpcTimeout, tracker, multi.getPriority(),
+      requestAttributes);
   }
 
   private void updateResult(int index, Object result) {

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

@@ -22,6 +22,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.OptionalLong;
 import java.util.concurrent.CompletableFuture;
@@ -78,7 +79,7 @@ public abstract class AsyncRpcRetryingCaller<T> {
 
   public AsyncRpcRetryingCaller(Timer retryTimer, AsyncConnectionImpl conn, int priority,
     long pauseNs, long pauseNsForServerOverloaded, int maxAttempts, long operationTimeoutNs,
-    long rpcTimeoutNs, int startLogErrorsCnt) {
+    long rpcTimeoutNs, int startLogErrorsCnt, Map<String, byte[]> requestAttributes) {
     this.retryTimer = retryTimer;
     this.conn = conn;
     this.priority = priority;
@@ -89,6 +90,7 @@ public AsyncRpcRetryingCaller(Timer retryTimer, AsyncConnectionImpl conn, int pr
     this.future = new CompletableFuture<>();
     this.controller = conn.rpcControllerFactory.newController();
     this.controller.setPriority(priority);
+    this.controller.setRequestAttributes(requestAttributes);
     this.exceptions = new ArrayList<>();
     this.startNs = System.nanoTime();
     this.pauseManager =

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

@@ -23,7 +23,9 @@
 import static org.apache.hbase.thirdparty.com.google.common.base.Preconditions.checkArgument;
 import static org.apache.hbase.thirdparty.com.google.common.base.Preconditions.checkNotNull;
 
+import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import org.apache.hadoop.hbase.HRegionLocation;
@@ -83,6 +85,8 @@ public class SingleRequestCallerBuilder<T> extends BuilderBase {
 
     private int priority = PRIORITY_UNSET;
 
+    private Map<String, byte[]> requestAttributes = Collections.emptyMap();
+
     public SingleRequestCallerBuilder<T> table(TableName tableName) {
       this.tableName = tableName;
       return this;
@@ -144,6 +148,12 @@ public SingleRequestCallerBuilder<T> priority(int priority) {
       return this;
     }
 
+    public SingleRequestCallerBuilder<T>
+      setRequestAttributes(Map<String, byte[]> requestAttributes) {
+      this.requestAttributes = requestAttributes;
+      return this;
+    }
+
     private void preCheck() {
       checkArgument(replicaId >= 0, "invalid replica id %s", replicaId);
       checkNotNull(tableName, "tableName is null");
@@ -157,7 +167,7 @@ public AsyncSingleRequestRpcRetryingCaller<T> build() {
       preCheck();
       return new AsyncSingleRequestRpcRetryingCaller<>(retryTimer, conn, tableName, row, replicaId,
         locateType, callable, priority, pauseNs, pauseNsForServerOverloaded, maxAttempts,
-        operationTimeoutNs, rpcTimeoutNs, startLogErrorsCnt);
+        operationTimeoutNs, rpcTimeoutNs, startLogErrorsCnt, requestAttributes);
     }
 
     /**
@@ -201,6 +211,8 @@ public class ScanSingleRegionCallerBuilder extends BuilderBase {
 
     private int priority = PRIORITY_UNSET;
 
+    private Map<String, byte[]> requestAttributes = Collections.emptyMap();
+
     public ScanSingleRegionCallerBuilder id(long scannerId) {
       this.scannerId = scannerId;
       return this;
@@ -278,6 +290,12 @@ public ScanSingleRegionCallerBuilder startLogErrorsCnt(int startLogErrorsCnt) {
       return this;
     }
 
+    public ScanSingleRegionCallerBuilder
+      setRequestAttributes(Map<String, byte[]> requestAttributes) {
+      this.requestAttributes = requestAttributes;
+      return this;
+    }
+
     private void preCheck() {
       checkArgument(scannerId != null, "invalid scannerId %d", scannerId);
       checkNotNull(scan, "scan is null");
@@ -293,7 +311,7 @@ public AsyncScanSingleRegionRpcRetryingCaller build() {
       return new AsyncScanSingleRegionRpcRetryingCaller(retryTimer, conn, scan, scanMetrics,
         scannerId, resultCache, consumer, stub, loc, isRegionServerRemote, priority,
         scannerLeaseTimeoutPeriodNs, pauseNs, pauseNsForServerOverloaded, maxAttempts,
-        scanTimeoutNs, rpcTimeoutNs, startLogErrorsCnt);
+        scanTimeoutNs, rpcTimeoutNs, startLogErrorsCnt, requestAttributes);
     }
 
     /**
@@ -322,6 +340,8 @@ public class BatchCallerBuilder extends BuilderBase {
 
     private long rpcTimeoutNs = -1L;
 
+    private Map<String, byte[]> requestAttributes = Collections.emptyMap();
+
     public BatchCallerBuilder table(TableName tableName) {
       this.tableName = tableName;
       return this;
@@ -362,10 +382,15 @@ public BatchCallerBuilder startLogErrorsCnt(int startLogErrorsCnt) {
       return this;
     }
 
+    public BatchCallerBuilder setRequestAttributes(Map<String, byte[]> requestAttributes) {
+      this.requestAttributes = requestAttributes;
+      return this;
+    }
+
     public <T> AsyncBatchRpcRetryingCaller<T> build() {
       return new AsyncBatchRpcRetryingCaller<>(retryTimer, conn, tableName, actions, pauseNs,
         pauseNsForServerOverloaded, maxAttempts, operationTimeoutNs, rpcTimeoutNs,
-        startLogErrorsCnt);
+        startLogErrorsCnt, requestAttributes);
     }
 
     public <T> List<CompletableFuture<T>> call() {