HBASE-28001: Add request attribute support to BufferedMutator (#6076)

Co-authored-by: Evie Boland <[email protected]>
Signed-off-by: Nick Dimiduk <[email protected]>

Author: eab148

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() {
+    return Collections.emptyMap();
+  }
 }

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/AsyncTable.java

@@ -21,12 +21,12 @@
 import static org.apache.hadoop.hbase.client.ConnectionUtils.toCheckExistenceOnly;
 import static org.apache.hadoop.hbase.util.FutureUtils.allOf;
 
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
-import org.apache.commons.lang3.NotImplementedException;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.CompareOperator;
 import org.apache.hadoop.hbase.TableName;
@@ -117,7 +117,7 @@ public interface AsyncTable<C extends ScanResultConsumerBase> {
    * @return a map of request attributes supplied by the client
    */
   default Map<String, byte[]> getRequestAttributes() {
-    throw new NotImplementedException("Add an implementation!");
+    return Collections.emptyMap();
   }
 
   /**

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

@@ -19,7 +19,9 @@
 
 import java.io.Closeable;
 import java.io.IOException;
+import java.util.Collections;
 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 +206,13 @@ default void setOperationTimeout(int timeout) {
       "The BufferedMutator::setOperationTimeout has not been implemented");
   }
 
+  /**
+   * Returns the rpc request attributes.
+   */
+  default Map<String, byte[]> getRequestAttributes() {
+    return Collections.emptyMap();
+  }
+
   /**
    * 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());
   }