branch-2 connection and request attributes impl

Author: Ray Mattingly

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

@@ -153,12 +153,13 @@ public void waitUntilDone() throws InterruptedIOException {
   private final int periodToLog;
 
   AsyncProcess(ClusterConnection hc, Configuration conf, RpcRetryingCallerFactory rpcCaller,
-    RpcControllerFactory rpcFactory) {
-    this(hc, conf, rpcCaller, rpcFactory, hc.getConnectionConfiguration().getRetriesNumber());
+    RpcControllerFactory rpcFactory, Map<String, byte[]> requestAttributes) {
+    this(hc, conf, rpcCaller, rpcFactory, hc.getConnectionConfiguration().getRetriesNumber(),
+      requestAttributes);
   }
 
   AsyncProcess(ClusterConnection hc, Configuration conf, RpcRetryingCallerFactory rpcCaller,
-    RpcControllerFactory rpcFactory, int retriesNumber) {
+    RpcControllerFactory rpcFactory, int retriesNumber, Map<String, byte[]> requestAttributes) {
     if (hc == null) {
       throw new IllegalArgumentException("ClusterConnection cannot be null.");
     }
@@ -189,6 +190,7 @@ public void waitUntilDone() throws InterruptedIOException {
 
     this.rpcCallerFactory = rpcCaller;
     this.rpcFactory = rpcFactory;
+    this.rpcFactory.setRequestAttributes(requestAttributes);
     this.logBatchErrorDetails = conf.getBoolean(LOG_DETAILS_FOR_BATCH_ERROR, false);
 
     this.requestController = RequestControllerFactory.create(conf);
@@ -218,6 +220,10 @@ public <CResult> AsyncRequestFuture submit(AsyncProcessTask<CResult> task)
     }
   }
 
+  public void setRequestAttributes(Map<String, byte[]> requestAttributes) {
+    this.rpcFactory.setRequestAttributes(requestAttributes);
+  }
+
   /**
    * Extract from the rows list what we can submit. The rows we can not submit are kept in the list.
    * Does not send requests to replicas (not currently used for anything other than streaming puts

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

@@ -142,7 +142,9 @@ public class BufferedMutatorImpl implements BufferedMutator {
     RpcControllerFactory rpcFactory, BufferedMutatorParams params) {
     this(conn, params,
       // puts need to track errors globally due to how the APIs currently work.
-      new AsyncProcess(conn, conn.getConfiguration(), rpcCallerFactory, rpcFactory));
+      // todo rmattingly support buffered mutator request attributes
+      new AsyncProcess(conn, conn.getConfiguration(), rpcCallerFactory, rpcFactory,
+        Collections.emptyMap()));
   }
 
   private void checkClose() {

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

@@ -230,10 +230,10 @@ public static Connection createConnection(Configuration conf, ExecutorService po
       try {
         // Default HCM#HCI is not accessible; make it so before invoking.
         Constructor<?> constructor = clazz.getDeclaredConstructor(Configuration.class,
-          ExecutorService.class, User.class, Map.class);
+          ExecutorService.class, User.class, ConnectionRegistry.class, Map.class);
         constructor.setAccessible(true);
         return user.runAs((PrivilegedExceptionAction<Connection>) () -> (Connection) constructor
-          .newInstance(conf, pool, user, connectionAttributes));
+          .newInstance(conf, pool, user, null, connectionAttributes));
       } catch (Exception e) {
         throw new IOException(e);
       }

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

@@ -39,6 +39,7 @@
 import java.util.Collections;
 import java.util.Date;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
@@ -292,6 +293,14 @@ public class ConnectionImplementation implements ClusterConnection, Closeable {
    */
   ConnectionImplementation(Configuration conf, ExecutorService pool, User user,
     ConnectionRegistry registry) throws IOException {
+    this(conf, pool, user, registry, Collections.emptyMap());
+  }
+
+  /**
+   * Constructor, for creating cluster connection with provided ConnectionRegistry.
+   */
+  ConnectionImplementation(Configuration conf, ExecutorService pool, User user,
+    ConnectionRegistry registry, Map<String, byte[]> connectionAttributes) throws IOException {
     this.conf = conf;
     this.user = user;
     if (user != null && user.isLoginFromKeytab()) {
@@ -348,11 +357,13 @@ public class ConnectionImplementation implements ClusterConnection, Closeable {
       }
       this.metaCache = new MetaCache(this.metrics);
 
-      this.rpcClient = RpcClientFactory.createClient(this.conf, this.clusterId, this.metrics);
+      this.rpcClient = RpcClientFactory.createClient(this.conf, this.clusterId, this.metrics,
+        connectionAttributes);
       this.rpcControllerFactory = RpcControllerFactory.instantiate(conf);
       this.rpcCallerFactory = RpcRetryingCallerFactory.instantiate(conf, connectionConfig,
         interceptor, this.stats, this.metrics);
-      this.asyncProcess = new AsyncProcess(this, conf, rpcCallerFactory, rpcControllerFactory);
+      this.asyncProcess = new AsyncProcess(this, conf, rpcCallerFactory, rpcControllerFactory,
+        Collections.emptyMap());
 
       // Do we publish the status?
       if (shouldListen) {
@@ -476,7 +487,7 @@ public TableBuilder getTableBuilder(TableName tableName, ExecutorService pool) {
       @Override
       public Table build() {
         return new HTable(ConnectionImplementation.this, this, rpcCallerFactory,
-          rpcControllerFactory, pool);
+          rpcControllerFactory, pool, requestAttributes);
       }
     };
   }

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

@@ -59,6 +59,7 @@
 import org.apache.hadoop.hbase.filter.Filter;
 import org.apache.hadoop.hbase.io.TimeRange;
 import org.apache.hadoop.hbase.ipc.CoprocessorRpcChannel;
+import org.apache.hadoop.hbase.ipc.HBaseRpcControllerImpl;
 import org.apache.hadoop.hbase.ipc.RpcControllerFactory;
 import org.apache.hadoop.hbase.trace.HBaseSemanticAttributes;
 import org.apache.hadoop.hbase.trace.TraceUtil;
@@ -130,6 +131,8 @@ public class HTable implements Table {
   private final RpcRetryingCallerFactory rpcCallerFactory;
   private final RpcControllerFactory rpcControllerFactory;
 
+  private final Map<String, byte[]> requestAttributes;
+
   // Marked Private @since 1.0
   @InterfaceAudience.Private
   public static ThreadPoolExecutor getDefaultExecutor(Configuration conf) {
@@ -165,7 +168,8 @@ public static ThreadPoolExecutor getDefaultExecutor(Configuration conf) {
   @InterfaceAudience.Private
   protected HTable(final ConnectionImplementation connection, final TableBuilderBase builder,
     final RpcRetryingCallerFactory rpcCallerFactory,
-    final RpcControllerFactory rpcControllerFactory, final ExecutorService pool) {
+    final RpcControllerFactory rpcControllerFactory, final ExecutorService pool,
+    final Map<String, byte[]> requestAttributes) {
     this.connection = Preconditions.checkNotNull(connection, "connection is null");
     this.configuration = connection.getConfiguration();
     this.connConfiguration = connection.getConnectionConfiguration();
@@ -197,9 +201,11 @@ protected HTable(final ConnectionImplementation connection, final TableBuilderBa
     this.scanTimeout = builder.scanTimeout;
     this.scannerCaching = connConfiguration.getScannerCaching();
     this.scannerMaxResultSize = connConfiguration.getScannerMaxResultSize();
+    this.requestAttributes = requestAttributes;
 
     // puts need to track errors globally due to how the APIs currently work.
     multiAp = this.connection.getAsyncProcess();
+    multiAp.setRequestAttributes(requestAttributes);
     this.locator = new HRegionLocator(tableName, connection);
   }
 
@@ -1308,7 +1314,7 @@ public <R extends Message> void batchCoprocessorService(
       AsyncProcess asyncProcess = new AsyncProcess(connection, configuration,
         RpcRetryingCallerFactory.instantiate(configuration, connConfiguration,
           connection.getStatisticsTracker(), connection.getConnectionMetrics()),
-        RpcControllerFactory.instantiate(configuration));
+        RpcControllerFactory.instantiate(configuration), requestAttributes);
 
       Batch.Callback<ClientProtos.CoprocessorServiceResult> resultsCallback =
         (byte[] region, byte[] row, ClientProtos.CoprocessorServiceResult serviceResult) -> {
@@ -1433,6 +1439,11 @@ public boolean thenMutate(RowMutations mutation) throws IOException {
     }
   }
 
+  @Override
+  public Map<String, byte[]> getRequestAttributes() {
+    return requestAttributes;
+  }
+
   private class CheckAndMutateWithFilterBuilderImpl implements CheckAndMutateWithFilterBuilder {
 
     private final byte[] row;
@@ -1478,4 +1489,9 @@ public boolean thenMutate(RowMutations mutation) throws IOException {
           .isSuccess(), supplier);
     }
   }
+
+  private void maybePopulateRequestAttributes(ClientServiceCallable<?> callable) {
+    HBaseRpcControllerImpl.configureRequestAttributes(callable.getRpcController(),
+      requestAttributes);
+  }
 }

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

@@ -434,7 +434,8 @@ public FlushWorker(Configuration conf, ClusterConnection conn, HRegionLocation a
         HConstants.DEFAULT_HBASE_CLIENT_OPERATION_TIMEOUT);
       // Specify 0 retries in AsyncProcess because we need to reassign puts to different queues
       // if regions are moved.
-      this.ap = new AsyncProcess(conn, conf, rpcCallerFactory, rpcControllerFactory, 0);
+      this.ap = new AsyncProcess(conn, conf, rpcCallerFactory, rpcControllerFactory, 0,
+        Collections.emptyMap());
       this.executor = executor;
       this.maxRetryInQueue = conf.getInt(TABLE_MULTIPLEXER_MAX_RETRIES_IN_QUEUE, 10000);
       this.pool = pool;

hbase-client/src/main/java/org/apache/hadoop/hbase/ipc/HBaseRpcControllerImpl.java

@@ -31,6 +31,7 @@
 import org.apache.yetus.audience.InterfaceAudience;
 
 import org.apache.hbase.thirdparty.com.google.protobuf.RpcCallback;
+import org.apache.hbase.thirdparty.com.google.protobuf.RpcController;
 
 /**
  * Get instances via {@link RpcControllerFactory} on client-side.
@@ -281,4 +282,15 @@ public String toString() {
       + exception + ", regionInfo=" + regionInfo + ", priority=" + priority + ", cellScanner="
       + cellScanner + '}';
   }
+
+  public static void configureRequestAttributes(RpcController rpcController,
+    Map<String, byte[]> requestAttributes) {
+    if (
+      !requestAttributes.isEmpty() && rpcController != null
+        && rpcController instanceof HBaseRpcController
+    ) {
+      HBaseRpcController controller = (HBaseRpcController) rpcController;
+      controller.setRequestAttributes(requestAttributes);
+    }
+  }
 }

hbase-client/src/main/java/org/apache/hadoop/hbase/ipc/RpcControllerFactory.java

@@ -17,7 +17,11 @@
  */
 package org.apache.hadoop.hbase.ipc;
 
+import static org.apache.hadoop.hbase.ipc.HBaseRpcControllerImpl.configureRequestAttributes;
+
+import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.CellScannable;
 import org.apache.hadoop.hbase.CellScanner;
@@ -45,30 +49,42 @@ public class RpcControllerFactory {
   public static final String CUSTOM_CONTROLLER_CONF_KEY = "hbase.rpc.controllerfactory.class";
   protected final Configuration conf;
 
+  private Map<String, byte[]> requestAttributes = Collections.emptyMap();
+
   public RpcControllerFactory(Configuration conf) {
     this.conf = conf;
   }
 
   public HBaseRpcController newController() {
     // TODO: Set HConstants default rpc timeout here rather than nothing?
-    return new HBaseRpcControllerImpl();
+    HBaseRpcController controller = new HBaseRpcControllerImpl();
+    configureRequestAttributes(controller, requestAttributes);
+    return controller;
   }
 
   public HBaseRpcController newController(CellScanner cellScanner) {
-    return new HBaseRpcControllerImpl(null, cellScanner);
+    HBaseRpcController controller = new HBaseRpcControllerImpl(null, cellScanner);
+    configureRequestAttributes(controller, requestAttributes);
+    return controller;
   }
 
   public HBaseRpcController newController(RegionInfo regionInfo, CellScanner cellScanner) {
-    return new HBaseRpcControllerImpl(regionInfo, cellScanner);
+    HBaseRpcController controller = new HBaseRpcControllerImpl(regionInfo, cellScanner);
+    configureRequestAttributes(controller, requestAttributes);
+    return controller;
   }
 
   public HBaseRpcController newController(final List<CellScannable> cellIterables) {
-    return new HBaseRpcControllerImpl(null, cellIterables);
+    HBaseRpcController controller = new HBaseRpcControllerImpl(null, cellIterables);
+    configureRequestAttributes(controller, requestAttributes);
+    return controller;
   }
 
   public HBaseRpcController newController(RegionInfo regionInfo,
     final List<CellScannable> cellIterables) {
-    return new HBaseRpcControllerImpl(regionInfo, cellIterables);
+    HBaseRpcController controller = new HBaseRpcControllerImpl(regionInfo, cellIterables);
+    configureRequestAttributes(controller, requestAttributes);
+    return controller;
   }
 
   public static RpcControllerFactory instantiate(Configuration configuration) {
@@ -91,4 +107,9 @@ public static RpcControllerFactory instantiate(Configuration configuration) {
       return new RpcControllerFactory(configuration);
     }
   }
+
+  public RpcControllerFactory setRequestAttributes(Map<String, byte[]> requestAttributes) {
+    this.requestAttributes = requestAttributes;
+    return this;
+  }
 }

hbase-client/src/test/java/org/apache/hadoop/hbase/client/TestAsyncProcess.java

@@ -185,14 +185,14 @@ public TableName getTableName() {
 
     public MyAsyncProcess(ClusterConnection hc, Configuration conf) {
       super(hc, conf, new RpcRetryingCallerFactory(conf, hc.getConnectionConfiguration()),
-        new RpcControllerFactory(conf));
+        new RpcControllerFactory(conf), Collections.emptyMap());
       service = Executors.newFixedThreadPool(5);
       this.conf = conf;
     }
 
     public MyAsyncProcess(ClusterConnection hc, Configuration conf, AtomicInteger nbThreads) {
       super(hc, conf, new RpcRetryingCallerFactory(conf, hc.getConnectionConfiguration()),
-        new RpcControllerFactory(conf));
+        new RpcControllerFactory(conf), Collections.emptyMap());
       service = new ThreadPoolExecutor(1, 20, 60, TimeUnit.SECONDS, new SynchronousQueue<>(),
         new CountingThreadFactory(nbThreads));
     }
@@ -1705,7 +1705,7 @@ public Future submit(Runnable runnable) {
   static class AsyncProcessForThrowableCheck extends AsyncProcess {
     public AsyncProcessForThrowableCheck(ClusterConnection hc, Configuration conf) {
       super(hc, conf, new RpcRetryingCallerFactory(conf, hc.getConnectionConfiguration()),
-        new RpcControllerFactory(conf));
+        new RpcControllerFactory(conf), Collections.emptyMap());
     }
   }
 

hbase-client/src/test/java/org/apache/hadoop/hbase/client/TestAsyncProcessWithRegionException.java

@@ -199,7 +199,7 @@ private static class MyAsyncProcess extends AsyncProcess {
 
     MyAsyncProcess(ClusterConnection hc, Configuration conf) {
       super(hc, conf, new RpcRetryingCallerFactory(conf, hc.getConnectionConfiguration()),
-        new RpcControllerFactory(conf));
+        new RpcControllerFactory(conf), Collections.emptyMap());
     }
 
     public AsyncRequestFuture submit(TableName tableName, List<? extends Row> rows)