HBASE-28001: Add request attribute support to BufferedMutator

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

@@ -20,6 +20,7 @@
 import java.io.Closeable;
 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.conf.Configuration;
@@ -93,4 +94,11 @@ default CompletableFuture<Void> mutate(Mutation mutation) {
   default long getPeriodicalFlushTimeout(TimeUnit unit) {
     throw new UnsupportedOperationException("Not implemented");
   }
+
+  /**
+   * Returns the rpc request attributes.
+   */
+  default Map<String, byte[]> getRequestAttributes() {
+    throw new UnsupportedOperationException("Not implemented");
+  }
 }

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

@@ -19,6 +19,7 @@
 
 import static org.apache.hadoop.hbase.client.ConnectionUtils.retries2Attempts;
 
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import org.apache.yetus.audience.InterfaceAudience;
 
@@ -38,6 +39,16 @@ public interface AsyncBufferedMutatorBuilder {
    */
   AsyncBufferedMutatorBuilder setRpcTimeout(long timeout, TimeUnit unit);
 
+  /**
+   * Set a rpc request attribute.
+   */
+  AsyncBufferedMutatorBuilder setRequestAttribute(String key, byte[] value);
+
+  /**
+   * Set multiple rpc request attributes.
+   */
+  AsyncBufferedMutatorBuilder setRequestAttributes(Map<String, byte[]> requestAttributes);
+
   /**
    * Set the base pause time for retrying. We use an exponential policy to generate sleep time when
    * retrying.

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

@@ -17,6 +17,7 @@
  */
 package org.apache.hadoop.hbase.client;
 
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import org.apache.yetus.audience.InterfaceAudience;
 
@@ -78,6 +79,20 @@ public AsyncBufferedMutatorBuilder setStartLogErrorsCnt(int startLogErrorsCnt) {
     return this;
   }
 
+  @Override
+  public AsyncBufferedMutatorBuilder setRequestAttribute(String key, byte[] value) {
+    tableBuilder.setRequestAttribute(key, value);
+    return this;
+  }
+
+  @Override
+  public AsyncBufferedMutatorBuilder setRequestAttributes(Map<String, byte[]> requestAttributes) {
+    for (Map.Entry<String, byte[]> requestAttribute : requestAttributes.entrySet()) {
+      tableBuilder.setRequestAttribute(requestAttribute.getKey(), requestAttribute.getValue());
+    }
+    return this;
+  }
+
   @Override
   public AsyncBufferedMutatorBuilder setWriteBufferSize(long writeBufferSize) {
     Preconditions.checkArgument(writeBufferSize > 0, "writeBufferSize %d must be > 0",

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

@@ -24,6 +24,7 @@
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
@@ -130,7 +131,7 @@ Stream.<CompletableFuture<Void>> generate(CompletableFuture::new).limit(mutation
         periodicFlushTask = periodicalFlushTimer.newTimeout(timeout -> {
           synchronized (AsyncBufferedMutatorImpl.this) {
             // confirm that we are still valid, if there is already an internalFlush call before us,
-            // then we should not execute any more. And in internalFlush we will set periodicFlush
+            // then we should not execute anymore. And in internalFlush we will set periodicFlush
             // to null, and since we may schedule a new one, so here we check whether the references
             // are equal.
             if (timeout == periodicFlushTask) {
@@ -170,4 +171,9 @@ public long getWriteBufferSize() {
   public long getPeriodicalFlushTimeout(TimeUnit unit) {
     return unit.convert(periodicFlushTimeoutNs, TimeUnit.NANOSECONDS);
   }
+
+  @Override
+  public Map<String, byte[]> getRequestAttributes() {
+    return table.getRequestAttributes();
+  }
 }

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

@@ -129,10 +129,10 @@ public AsyncTableBuilderBase<C> setStartLogErrorsCnt(int startLogErrorsCnt) {
 
   @Override
   public AsyncTableBuilder<C> setRequestAttribute(String key, byte[] value) {
-    if (this.requestAttributes.isEmpty()) {
-      this.requestAttributes = new HashMap<>();
+    if (requestAttributes.isEmpty()) {
+      requestAttributes = new HashMap<>();
     }
-    this.requestAttributes.put(key, value);
+    requestAttributes.put(key, value);
     return this;
   }
 }

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

@@ -20,6 +20,7 @@
 import java.io.Closeable;
 import java.io.IOException;
 import java.util.List;
+import java.util.Map;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.TableName;
 import org.apache.yetus.audience.InterfaceAudience;
@@ -204,6 +205,14 @@ default void setOperationTimeout(int timeout) {
       "The BufferedMutator::setOperationTimeout has not been implemented");
   }
 
+  /**
+   * Returns the rpc request attributes.
+   */
+  default Map<String, byte[]> getRequestAttributes() {
+    throw new UnsupportedOperationException(
+      "The BufferedMutator::setOperationTimeout has not been implemented");
+  }
+
   /**
    * Listens for asynchronous exceptions on a {@link BufferedMutator}.
    */

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

@@ -23,6 +23,7 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
@@ -186,4 +187,9 @@ public void setRpcTimeout(int timeout) {
   public void setOperationTimeout(int timeout) {
     // no effect
   }
+
+  @Override
+  public Map<String, byte[]> getRequestAttributes() {
+    return mutator.getRequestAttributes();
+  }
 }

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

@@ -17,6 +17,9 @@
  */
 package org.apache.hadoop.hbase.client;
 
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.concurrent.ExecutorService;
 import org.apache.hadoop.hbase.TableName;
 import org.apache.yetus.audience.InterfaceAudience;
@@ -38,6 +41,7 @@ public class BufferedMutatorParams implements Cloneable {
   private String implementationClassName = null;
   private int rpcTimeout = UNSET;
   private int operationTimeout = UNSET;
+  protected Map<String, byte[]> requestAttributes = Collections.emptyMap();
   private BufferedMutator.ExceptionListener listener = new BufferedMutator.ExceptionListener() {
     @Override
     public void onException(RetriesExhaustedWithDetailsException exception,
@@ -85,6 +89,18 @@ public int getOperationTimeout() {
     return operationTimeout;
   }
 
+  public BufferedMutatorParams setRequestAttribute(String key, byte[] value) {
+    if (requestAttributes.isEmpty()) {
+      requestAttributes = new HashMap<>();
+    }
+    requestAttributes.put(key, value);
+    return this;
+  }
+
+  public Map<String, byte[]> getRequestAttributes() {
+    return requestAttributes;
+  }
+
   /**
    * Override the write buffer size specified by the provided {@link Connection}'s
    * {@link org.apache.hadoop.conf.Configuration} instance, via the configuration key

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

@@ -107,6 +107,10 @@ public BufferedMutator getBufferedMutator(BufferedMutatorParams params) throws I
     if (params.getMaxKeyValueSize() != BufferedMutatorParams.UNSET) {
       builder.setMaxKeyValueSize(params.getMaxKeyValueSize());
     }
+    if (!params.getRequestAttributes().isEmpty()) {
+
+      builder.setRequestAttributes(params.getRequestAttributes());
+    }
     return new BufferedMutatorOverAsyncBufferedMutator(builder.build(), params.getListener());
   }
 

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

@@ -214,8 +214,8 @@ private <T> SingleRequestCallerBuilder<T> newCaller(byte[] row, int priority, lo
       .operationTimeout(operationTimeoutNs, TimeUnit.NANOSECONDS)
       .pause(pauseNs, TimeUnit.NANOSECONDS)
       .pauseForServerOverloaded(pauseNsForServerOverloaded, TimeUnit.NANOSECONDS)
-      .maxAttempts(maxAttempts).setRequestAttributes(requestAttributes)
-      .startLogErrorsCnt(startLogErrorsCnt).setRequestAttributes(requestAttributes);
+      .maxAttempts(maxAttempts).startLogErrorsCnt(startLogErrorsCnt)
+      .setRequestAttributes(requestAttributes);
   }
 
   private <T, R extends OperationWithAttributes & Row> SingleRequestCallerBuilder<T>

hbase-server/src/test/java/org/apache/hadoop/hbase/client/TestRequestAndConnectionAttributes.java

@@ -36,6 +36,8 @@
 import org.apache.hadoop.hbase.AuthUtil;
 import org.apache.hadoop.hbase.Cell;
 import org.apache.hadoop.hbase.CellComparator;
+import org.apache.hadoop.hbase.ExtendedCellScannable;
+import org.apache.hadoop.hbase.ExtendedCellScanner;
 import org.apache.hadoop.hbase.HBaseClassTestRule;
 import org.apache.hadoop.hbase.HBaseTestingUtil;
 import org.apache.hadoop.hbase.HConstants;
@@ -44,7 +46,10 @@
 import org.apache.hadoop.hbase.coprocessor.RegionCoprocessor;
 import org.apache.hadoop.hbase.coprocessor.RegionCoprocessorEnvironment;
 import org.apache.hadoop.hbase.coprocessor.RegionObserver;
+import org.apache.hadoop.hbase.ipc.DelegatingHBaseRpcController;
+import org.apache.hadoop.hbase.ipc.HBaseRpcController;
 import org.apache.hadoop.hbase.ipc.RpcCall;
+import org.apache.hadoop.hbase.ipc.RpcControllerFactory;
 import org.apache.hadoop.hbase.ipc.RpcServer;
 import org.apache.hadoop.hbase.regionserver.InternalScanner;
 import org.apache.hadoop.hbase.testclassification.ClientTests;
@@ -230,6 +235,47 @@ public void testRequestAttributesMultiPut() throws IOException {
     assertTrue(REQUEST_ATTRIBUTES_VALIDATED.get());
   }
 
+  @Test
+  public void testRequestAttributesBufferedMutate() throws IOException, InterruptedException {
+    assertFalse(REQUEST_ATTRIBUTES_VALIDATED.get());
+    addRandomRequestAttributes();
+
+    Configuration conf = TEST_UTIL.getConfiguration();
+    try (
+      Connection conn = ConnectionFactory.createConnection(conf, null, AuthUtil.loginClient(conf),
+        CONNECTION_ATTRIBUTES);
+      BufferedMutator bufferedMutator = conn.getBufferedMutator(
+        configureRequestAttributes(new BufferedMutatorParams(REQUEST_ATTRIBUTES_TEST_TABLE)));) {
+      Put put = new Put(Bytes.toBytes("a"));
+      put.addColumn(REQUEST_ATTRIBUTES_TEST_TABLE_CF, Bytes.toBytes("c"), Bytes.toBytes("v"));
+      bufferedMutator.mutate(put);
+      bufferedMutator.flush();
+    }
+
+    assertTrue(REQUEST_ATTRIBUTES_VALIDATED.get());
+  }
+
+  @Test
+  public void testRequestAttributesFromRpcController() throws IOException, InterruptedException {
+    assertFalse(REQUEST_ATTRIBUTES_VALIDATED.get());
+    addRandomRequestAttributes();
+
+    Configuration conf = TEST_UTIL.getConfiguration();
+    conf.setClass(RpcControllerFactory.CUSTOM_CONTROLLER_CONF_KEY,
+      RequestMetadataControllerFactory.class, RpcControllerFactory.class);
+    try (
+      Connection conn = ConnectionFactory.createConnection(conf, null, AuthUtil.loginClient(conf),
+        CONNECTION_ATTRIBUTES);
+      BufferedMutator bufferedMutator = conn.getBufferedMutator(REQUEST_ATTRIBUTES_TEST_TABLE);) {
+      Put put = new Put(Bytes.toBytes("a"));
+      put.addColumn(REQUEST_ATTRIBUTES_TEST_TABLE_CF, Bytes.toBytes("c"), Bytes.toBytes("v"));
+      bufferedMutator.mutate(put);
+      bufferedMutator.flush();
+    }
+    conf.unset(RpcControllerFactory.CUSTOM_CONTROLLER_CONF_KEY);
+    assertTrue(REQUEST_ATTRIBUTES_VALIDATED.get());
+  }
+
   @Test
   public void testNoRequestAttributes() throws IOException {
     assertFalse(REQUEST_ATTRIBUTES_VALIDATED.get());
@@ -262,6 +308,64 @@ private static TableBuilder configureRequestAttributes(TableBuilder tableBuilder
     return tableBuilder;
   }
 
+  private static BufferedMutatorParams configureRequestAttributes(BufferedMutatorParams params) {
+    REQUEST_ATTRIBUTES.forEach(params::setRequestAttribute);
+    return params;
+  }
+
+  public static class RequestMetadataControllerFactory extends RpcControllerFactory {
+
+    public RequestMetadataControllerFactory(Configuration conf) {
+      super(conf);
+    }
+
+    @Override
+    public HBaseRpcController newController() {
+      return new RequestMetadataController(super.newController(), REQUEST_ATTRIBUTES);
+    }
+
+    @Override
+    public HBaseRpcController newController(ExtendedCellScanner cellScanner) {
+      return new RequestMetadataController(super.newController(null, cellScanner),
+        REQUEST_ATTRIBUTES);
+    }
+
+    @Override
+    public HBaseRpcController newController(RegionInfo regionInfo,
+      ExtendedCellScanner cellScanner) {
+      return new RequestMetadataController(super.newController(regionInfo, cellScanner),
+        REQUEST_ATTRIBUTES);
+    }
+
+    @Override
+    public HBaseRpcController newController(final List<ExtendedCellScannable> cellIterables) {
+      return new RequestMetadataController(super.newController(null, cellIterables),
+        REQUEST_ATTRIBUTES);
+    }
+
+    @Override
+    public HBaseRpcController newController(RegionInfo regionInfo,
+      final List<ExtendedCellScannable> cellIterables) {
+      return new RequestMetadataController(super.newController(regionInfo, cellIterables),
+        REQUEST_ATTRIBUTES);
+    }
+
+    public static class RequestMetadataController extends DelegatingHBaseRpcController {
+      private final Map<String, byte[]> requestAttributes;
+
+      RequestMetadataController(HBaseRpcController delegate,
+        Map<String, byte[]> requestAttributes) {
+        super(delegate);
+        this.requestAttributes = requestAttributes;
+      }
+
+      @Override
+      public Map<String, byte[]> getRequestAttributes() {
+        return requestAttributes;
+      }
+    }
+  }
+
   public static class AttributesCoprocessor implements RegionObserver, RegionCoprocessor {
 
     @Override