HubSpot Backport HBASE-28001: Add request attribute support to BufferedMutator (apache#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;
 
@@ -103,6 +104,16 @@ default AsyncBufferedMutatorBuilder setMaxRetries(int maxRetries) {
    */
   AsyncBufferedMutatorBuilder setMaxKeyValueSize(int maxKeyValueSize);
 
+  /**
+   * Set a rpc request attribute.
+   */
+  AsyncBufferedMutatorBuilder setRequestAttribute(String key, byte[] value);
+
+  /**
+   * Set multiple rpc request attributes.
+   */
+  AsyncBufferedMutatorBuilder setRequestAttributes(Map<String, byte[]> requestAttributes);
+
   /**
    * Create the {@link AsyncBufferedMutator} instance.
    */

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;
@@ -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

@@ -22,12 +22,12 @@
 import static org.apache.hadoop.hbase.util.FutureUtils.allOf;
 
 import com.google.protobuf.RpcChannel;
+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/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;
@@ -194,6 +196,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/BufferedMutatorImpl.java

@@ -26,6 +26,7 @@
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Timer;
 import java.util.TimerTask;
@@ -89,6 +90,7 @@ public class BufferedMutatorImpl implements BufferedMutator {
   private final ExecutorService pool;
   private final AtomicInteger rpcTimeout;
   private final AtomicInteger operationTimeout;
+  private final Map<String, byte[]> requestAttributes;
   private final boolean cleanupPoolOnClose;
   private volatile boolean closed = false;
   private final AsyncProcess ap;
@@ -135,6 +137,9 @@ public class BufferedMutatorImpl implements BufferedMutator {
     this.operationTimeout = new AtomicInteger(params.getOperationTimeout() != UNSET
       ? params.getOperationTimeout()
       : conn.getConnectionConfiguration().getOperationTimeout());
+
+    this.requestAttributes = params.getRequestAttributes();
+
     this.ap = ap;
   }
 
@@ -252,7 +257,8 @@ public synchronized void close() throws IOException {
 
   private AsyncProcessTask createTask(QueueRowAccess access) {
     return new AsyncProcessTask(AsyncProcessTask.newBuilder().setPool(pool).setTableName(tableName)
-      .setRowAccess(access).setSubmittedRows(AsyncProcessTask.SubmittedRows.AT_LEAST_ONE).build()) {
+      .setRowAccess(access).setSubmittedRows(AsyncProcessTask.SubmittedRows.AT_LEAST_ONE)
+      .setRequestAttributes(requestAttributes).build()) {
       @Override
       public int getRpcTimeout() {
         return rpcTimeout.get();
@@ -391,6 +397,11 @@ public void setOperationTimeout(int operationTimeout) {
     this.operationTimeout.set(operationTimeout);
   }
 
+  @Override
+  public Map<String, byte[]> getRequestAttributes() {
+    return requestAttributes;
+  }
+
   long getCurrentWriteBufferSize() {
     return currentWriteBufferSize.get();
   }

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-examples/src/main/java/org/apache/hadoop/hbase/client/example/BufferedMutatorExample.java

@@ -19,15 +19,19 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.conf.Configured;
+import org.apache.hadoop.hbase.AuthUtil;
 import org.apache.hadoop.hbase.TableName;
 import org.apache.hadoop.hbase.client.BufferedMutator;
 import org.apache.hadoop.hbase.client.BufferedMutatorParams;
@@ -67,12 +71,18 @@ public void onException(RetriesExhaustedWithDetailsException e, BufferedMutator
         }
       }
     };
-    BufferedMutatorParams params = new BufferedMutatorParams(TABLE).listener(listener);
+    BufferedMutatorParams params = new BufferedMutatorParams(TABLE).listener(listener)
+      .setRequestAttribute("requestInfo", Bytes.toBytes("bar"));
 
     //
     // step 1: create a single Connection and a BufferedMutator, shared by all worker threads.
     //
-    try (final Connection conn = ConnectionFactory.createConnection(getConf());
+    Map<String, byte[]> connectionAttributes = new HashMap<>();
+    connectionAttributes.put("clientId", Bytes.toBytes("foo"));
+    Configuration conf = getConf();
+    try (
+      final Connection conn = ConnectionFactory.createConnection(conf, null,
+        AuthUtil.loginClient(conf), connectionAttributes);
       final BufferedMutator mutator = conn.getBufferedMutator(params)) {
 
       /** worker pool that operates on BufferedTable instances */

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

@@ -0,0 +1,123 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.hadoop.hbase.client;
+
+import static org.junit.Assert.assertEquals;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Random;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.hbase.AuthUtil;
+import org.apache.hadoop.hbase.Cell;
+import org.apache.hadoop.hbase.CellComparator;
+import org.apache.hadoop.hbase.HBaseClassTestRule;
+import org.apache.hadoop.hbase.HBaseTestingUtility;
+import org.apache.hadoop.hbase.HConstants;
+import org.apache.hadoop.hbase.MiniHBaseCluster;
+import org.apache.hadoop.hbase.TableName;
+import org.apache.hadoop.hbase.coprocessor.ObserverContext;
+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.RpcCall;
+import org.apache.hadoop.hbase.ipc.RpcServer;
+import org.apache.hadoop.hbase.testclassification.ClientTests;
+import org.apache.hadoop.hbase.testclassification.MediumTests;
+import org.apache.hadoop.hbase.util.Bytes;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.ClassRule;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+
+@Category({ ClientTests.class, MediumTests.class })
+public class TestConnectionAttributes {
+
+  @ClassRule
+  public static final HBaseClassTestRule CLASS_RULE =
+    HBaseClassTestRule.forClass(TestConnectionAttributes.class);
+
+  private static final Map<String, byte[]> CONNECTION_ATTRIBUTES = new HashMap<>();
+  static {
+    CONNECTION_ATTRIBUTES.put("clientId", Bytes.toBytes("foo"));
+  }
+  private static final byte[] FAMILY = Bytes.toBytes("0");
+  private static final TableName TABLE_NAME = TableName.valueOf("testConnectionAttributes");
+
+  private static final HBaseTestingUtility TEST_UTIL = new HBaseTestingUtility();
+  private static MiniHBaseCluster cluster;
+
+  @BeforeClass
+  public static void setUp() throws Exception {
+    cluster = TEST_UTIL.startMiniCluster(1);
+    Table table = TEST_UTIL.createTable(TABLE_NAME, new byte[][] { FAMILY }, 1,
+      HConstants.DEFAULT_BLOCKSIZE, TestConnectionAttributes.AttributesCoprocessor.class.getName());
+    table.close();
+  }
+
+  @AfterClass
+  public static void afterClass() throws Exception {
+    cluster.close();
+    TEST_UTIL.shutdownMiniCluster();
+  }
+
+  @Test
+  public void testConnectionHeaderOverwrittenAttributesRemain() throws IOException {
+    Configuration conf = TEST_UTIL.getConfiguration();
+    try (Connection conn = ConnectionFactory.createConnection(conf, null,
+      AuthUtil.loginClient(conf), CONNECTION_ATTRIBUTES); Table table = conn.getTable(TABLE_NAME)) {
+
+      // submit a 300 byte rowkey here to encourage netty's allocator to overwrite the connection
+      // header
+      byte[] bytes = new byte[300];
+      new Random().nextBytes(bytes);
+      Result result = table.get(new Get(bytes));
+
+      assertEquals(CONNECTION_ATTRIBUTES.size(), result.size());
+      for (Map.Entry<String, byte[]> attr : CONNECTION_ATTRIBUTES.entrySet()) {
+        byte[] val = result.getValue(FAMILY, Bytes.toBytes(attr.getKey()));
+        assertEquals(Bytes.toStringBinary(attr.getValue()), Bytes.toStringBinary(val));
+      }
+    }
+  }
+
+  public static class AttributesCoprocessor implements RegionObserver, RegionCoprocessor {
+
+    @Override
+    public Optional<RegionObserver> getRegionObserver() {
+      return Optional.of(this);
+    }
+
+    @Override
+    public void preGetOp(ObserverContext<RegionCoprocessorEnvironment> c, Get get,
+      List<Cell> result) throws IOException {
+      RpcCall rpcCall = RpcServer.getCurrentCall().get();
+      for (Map.Entry<String, byte[]> attr : rpcCall.getConnectionAttributes().entrySet()) {
+        result.add(c.getEnvironment().getCellBuilder().clear().setRow(get.getRow())
+          .setFamily(FAMILY).setQualifier(Bytes.toBytes(attr.getKey())).setValue(attr.getValue())
+          .setType(Cell.Type.Put).setTimestamp(1).build());
+      }
+      result.sort(CellComparator.getInstance());
+      c.bypass();
+    }
+  }
+}