[branch-2] HBASE-27657: Connection and Request Attributes

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() {