HBASE-22623 - Add RegionObserver coprocessor hook for preWALAppend

Author: gjacoby126

hbase-server/src/main/java/org/apache/hadoop/hbase/coprocessor/RegionObserver.java

@@ -1104,4 +1104,15 @@ default DeleteTracker postInstantiateDeleteTracker(
       throws IOException {
     return delTracker;
   }
+
+  /**
+   * Called just before the WAL Entry is appended to the WAL. Implementing this hook allows
+   * coprocessors to add extended attributes to the WALKey that then get persisted to the
+   * WAL, and are available to replication endpoints to use in processing WAL Entries.
+   * @param ctx the environment provided by the region server
+   * @param key the WALKey associated with a particular append to a WAL
+   */
+  default void preWALAppend(ObserverContext<RegionCoprocessorEnvironment> ctx, WALKey key)
+    throws IOException {
+  }
 }

hbase-server/src/main/java/org/apache/hadoop/hbase/regionserver/HRegion.java

@@ -7951,6 +7951,9 @@ private WriteEntry doWALAppend(WALEdit walEdit, Durability durability, List<UUID
     if (walEdit.isReplay()) {
       walKey.setOrigLogSeqNum(origLogSeqNum);
     }
+    if (this.coprocessorHost != null) {
+      this.coprocessorHost.preWALAppend(walKey);
+    }
     WriteEntry writeEntry = null;
     try {
       long txid = this.wal.append(this.getRegionInfo(), walKey, walEdit, true);

hbase-server/src/main/java/org/apache/hadoop/hbase/regionserver/RegionCoprocessorHost.java

@@ -1720,6 +1720,19 @@ public List<Pair<Cell, Cell>> call(RegionObserver observer) throws IOException {
         });
   }
 
+  public void preWALAppend(WALKey key) throws IOException {
+    if (this.coprocEnvironments.isEmpty()){
+      return;
+    }
+    execOperation(new RegionObserverOperationWithoutResult()
+     {
+      @Override
+      public void call(RegionObserver observer) throws IOException {
+        observer.preWALAppend(this, key);
+      }
+    });
+  }
+
   public Message preEndpointInvocation(final Service service, final String methodName,
       Message request) throws IOException {
     if (coprocEnvironments.isEmpty()) {

hbase-server/src/main/java/org/apache/hadoop/hbase/wal/WALKey.java

@@ -86,6 +86,13 @@ default long getNonce() {
    */
   long getOrigLogSeqNum();
 
+  /**
+   * Add a named String value to this WALKey to be persisted into the WAL
+   * @param attributeKey Name of the attribute
+   * @param attributeValue Value of the attribute
+   */
+  void addExtendedAttribute(String attributeKey, byte[] attributeValue);
+
     /**
      * Return a named String value injected into the WALKey during processing, such as by a
      * coprocessor

hbase-server/src/main/java/org/apache/hadoop/hbase/wal/WALKeyImpl.java

@@ -195,6 +195,37 @@ public WALKeyImpl(final byte[] encodedRegionName,
         mvcc, null, null);
   }
 
+  /**
+   * Copy constructor that takes in an existing WALKeyImpl plus some extended attributes.
+   * Intended for coprocessors to add annotations to a system-generated WALKey
+   * for persistence to the WAL.
+   * @param key Key to be copied into this new key
+   * @param extendedAttributes Extra attributes to copy into the new key
+   */
+  public WALKeyImpl(WALKeyImpl key,
+                    Map<String, byte[]> extendedAttributes){
+    init(key.getEncodedRegionName(), key.getTableName(), key.getSequenceId(),
+        key.getWriteTime(), key.getClusterIds(), key.getNonceGroup(), key.getNonce(),
+        key.getMvcc(), key.getReplicationScopes(), extendedAttributes);
+
+  }
+
+  /**
+   * Copy constructor that takes in an existing WALKey, the extra WALKeyImpl fields that the
+   * parent interface is missing, plus some extended attributes. Intended
+   * for coprocessors to add annotations to a system-generated WALKey for
+   * persistence to the WAL.
+   */
+  public WALKeyImpl(WALKey key,
+                    List<UUID> clusterIds,
+                    MultiVersionConcurrencyControl mvcc,
+                    final NavigableMap<byte[], Integer> replicationScopes,
+                    Map<String, byte[]> extendedAttributes){
+    init(key.getEncodedRegionName(), key.getTableName(), key.getSequenceId(),
+        key.getWriteTime(), clusterIds, key.getNonceGroup(), key.getNonce(),
+        mvcc, replicationScopes, extendedAttributes);
+
+  }
   /**
    * Create the log key for writing to somewhere.
    * We maintain the tablename mainly for debugging purposes.
@@ -464,6 +495,14 @@ public UUID getOriginatingClusterId(){
     return clusterIds.isEmpty()? HConstants.DEFAULT_CLUSTER_ID: clusterIds.get(0);
   }
 
+  @Override
+  public void addExtendedAttribute(String attributeKey, byte[] attributeValue){
+    if (extendedAttributes == null){
+      extendedAttributes = new HashMap<String, byte[]>();
+    }
+    extendedAttributes.put(attributeKey, attributeValue);
+  }
+
   @Override
   public byte[] getExtendedAttribute(String attributeKey){
     return extendedAttributes != null ? extendedAttributes.get(attributeKey) : null;

hbase-server/src/test/java/org/apache/hadoop/hbase/coprocessor/SimpleRegionObserver.java

@@ -25,6 +25,7 @@
 import static org.junit.Assert.assertTrue;
 
 import java.io.IOException;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -124,7 +125,11 @@ public class SimpleRegionObserver implements RegionCoprocessor, RegionObserver {
   final AtomicInteger ctPostStartRegionOperation = new AtomicInteger(0);
   final AtomicInteger ctPostCloseRegionOperation = new AtomicInteger(0);
   final AtomicBoolean throwOnPostFlush = new AtomicBoolean(false);
+  final AtomicInteger ctPreWALAppend = new AtomicInteger(0);
+
   static final String TABLE_SKIPPED = "SKIPPED_BY_PREWALRESTORE";
+  Map<String, byte[]> extendedAttributes = new HashMap<String,byte[]>();
+  static final byte[] WAL_EXTENDED_ATTRIBUTE_BYTES = Bytes.toBytes("foo");
 
   public void setThrowOnPostFlush(Boolean val){
     throwOnPostFlush.set(val);
@@ -631,6 +636,15 @@ public StoreFileReader postStoreFileReaderOpen(ObserverContext<RegionCoprocessor
     return reader;
   }
 
+  @Override
+  public void preWALAppend(ObserverContext<RegionCoprocessorEnvironment> ctx,
+                                 WALKey key) throws IOException {
+    ctPreWALAppend.incrementAndGet();
+
+    key.addExtendedAttribute(Integer.toString(ctPreWALAppend.get()),
+        Bytes.toBytes("foo"));
+  }
+
   public boolean hadPreGet() {
     return ctPreGet.get() > 0;
   }
@@ -864,6 +878,10 @@ public int getCtPostWALRestore() {
     return ctPostWALRestore.get();
   }
 
+  public int getCtPreWALAppend() {
+    return ctPreWALAppend.get();
+  }
+
   public boolean wasStoreFileReaderOpenCalled() {
     return ctPreStoreFileReaderOpen.get() > 0 && ctPostStoreFileReaderOpen.get() > 0;
   }

hbase-server/src/test/java/org/apache/hadoop/hbase/coprocessor/TestRegionObserverInterface.java

@@ -25,6 +25,7 @@
 import java.io.IOException;
 import java.lang.reflect.Method;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Optional;
 import org.apache.hadoop.conf.Configuration;
@@ -70,20 +71,26 @@
 import org.apache.hadoop.hbase.regionserver.StoreFile;
 import org.apache.hadoop.hbase.regionserver.compactions.CompactionLifeCycleTracker;
 import org.apache.hadoop.hbase.regionserver.compactions.CompactionRequest;
+import org.apache.hadoop.hbase.regionserver.wal.WALActionsListener;
 import org.apache.hadoop.hbase.testclassification.CoprocessorTests;
 import org.apache.hadoop.hbase.testclassification.MediumTests;
 import org.apache.hadoop.hbase.tool.BulkLoadHFiles;
 import org.apache.hadoop.hbase.util.Bytes;
 import org.apache.hadoop.hbase.util.EnvironmentEdgeManager;
 import org.apache.hadoop.hbase.util.JVMClusterUtil;
 import org.apache.hadoop.hbase.util.Threads;
+import org.apache.hadoop.hbase.wal.WALEdit;
+import org.apache.hadoop.hbase.wal.WALKey;
+import org.apache.hadoop.hbase.wal.WALKeyImpl;
 import org.junit.AfterClass;
+import org.junit.Assert;
 import org.junit.BeforeClass;
 import org.junit.ClassRule;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 import org.junit.rules.TestName;
+import org.mockito.Mockito;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -663,6 +670,65 @@ public void testPreWALRestoreSkip() throws Exception {
     table.close();
   }
 
+  //called from testPreWALAppendIsWrittenToWAL
+  private void testPreWALAppendHook(Table table, TableName tableName) throws IOException {
+    int expectedCalls = 0;
+    String [] methodArray = new String[1];
+    methodArray[0] = "getCtPreWALAppend";
+    Object[] resultArray = new Object[1];
+
+    Put p = new Put(ROW);
+    p.addColumn(A, A, A);
+    table.put(p);
+    resultArray[0] = ++expectedCalls;
+    verifyMethodResult(SimpleRegionObserver.class, methodArray, tableName, resultArray);
+
+    Append a = new Append(ROW);
+    a.addColumn(B, B, B);
+    table.append(a);
+    resultArray[0] = ++expectedCalls;
+    verifyMethodResult(SimpleRegionObserver.class, methodArray, tableName, resultArray);
+
+    Increment i = new Increment(ROW);
+    i.addColumn(C, C, 1);
+    table.increment(i);
+    resultArray[0] = ++expectedCalls;
+    verifyMethodResult(SimpleRegionObserver.class, methodArray, tableName, resultArray);
+
+    Delete d = new Delete(ROW);
+    table.delete(d);
+    resultArray[0] = ++expectedCalls;
+    verifyMethodResult(SimpleRegionObserver.class, methodArray, tableName, resultArray);
+  }
+
+  @Test
+  public void testPreWALAppend() throws Exception {
+    SimpleRegionObserver sro = new SimpleRegionObserver();
+    ObserverContext ctx = Mockito.mock(ObserverContext.class);
+    WALKey key = new WALKeyImpl(Bytes.toBytes("region"), TEST_TABLE,
+        EnvironmentEdgeManager.currentTime());
+    sro.preWALAppend(ctx, key);
+    Assert.assertEquals(1, key.getExtendedAttributes().size());
+    Assert.assertArrayEquals(SimpleRegionObserver.WAL_EXTENDED_ATTRIBUTE_BYTES,
+        key.getExtendedAttribute(Integer.toString(sro.getCtPreWALAppend())));
+  }
+
+  @Test
+  public void testPreWALAppendIsWrittenToWAL() throws Exception {
+    final TableName tableName = TableName.valueOf(TEST_TABLE.getNameAsString() +
+        "." + name.getMethodName());
+    Table table = util.createTable(tableName, new byte[][] { A, B, C });
+
+    PreWALAppendWALActionsListener listener = new PreWALAppendWALActionsListener();
+    List<HRegion> regions = util.getHBaseCluster().getRegions(tableName);
+    //should be only one region
+    HRegion region = regions.get(0);
+    region.getWAL().registerWALActionsListener(listener);
+    testPreWALAppendHook(table, tableName);
+    boolean[] expectedResults = {true, true, true, true};
+    Assert.assertArrayEquals(expectedResults, listener.getWalKeysCorrectArray());
+
+  }
   // check each region whether the coprocessor upcalls are called or not.
   private void verifyMethodResult(Class<?> coprocessor, String methodName[], TableName tableName,
       Object value[]) throws IOException {
@@ -711,4 +777,23 @@ private static void createHFile(Configuration conf, FileSystem fs, Path path, by
       writer.close();
     }
   }
+
+  private static class PreWALAppendWALActionsListener implements WALActionsListener {
+    boolean[] walKeysCorrect = {false, false, false, false};
+
+    @Override
+    public void postAppend(long entryLen, long elapsedTimeMillis,
+                           WALKey logKey, WALEdit logEdit) throws IOException {
+      for (int k = 0; k < 4; k++) {
+        if (!walKeysCorrect[k]) {
+          walKeysCorrect[k] = Arrays.equals(SimpleRegionObserver.WAL_EXTENDED_ATTRIBUTE_BYTES,
+              logKey.getExtendedAttribute(Integer.toString(k + 1)));
+        }
+      }
+    }
+
+    boolean[] getWalKeysCorrectArray() {
+      return walKeysCorrect;
+    }
+  }
 }

hbase-server/src/test/java/org/apache/hadoop/hbase/regionserver/TestHRegion.java

@@ -104,6 +104,7 @@
 import org.apache.hadoop.hbase.Waiter;
 import org.apache.hadoop.hbase.client.Append;
 import org.apache.hadoop.hbase.client.ColumnFamilyDescriptor;
+import org.apache.hadoop.hbase.client.ColumnFamilyDescriptorBuilder;
 import org.apache.hadoop.hbase.client.Delete;
 import org.apache.hadoop.hbase.client.Durability;
 import org.apache.hadoop.hbase.client.Get;
@@ -165,7 +166,6 @@
 import org.apache.hadoop.hbase.wal.WALProvider;
 import org.apache.hadoop.hbase.wal.WALProvider.Writer;
 import org.apache.hadoop.hbase.wal.WALSplitUtil;
-import org.apache.hadoop.hbase.wal.WALSplitter;
 import org.apache.hadoop.metrics2.MetricsExecutor;
 import org.junit.After;
 import org.junit.Assert;
@@ -401,6 +401,7 @@ public void testMemstoreSizeAccountingWithFailedPostBatchMutate() throws IOExcep
     String testName = "testMemstoreSizeAccountingWithFailedPostBatchMutate";
     FileSystem fs = FileSystem.get(CONF);
     Path rootDir = new Path(dir + testName);
+    ChunkCreator.initialize(MemStoreLABImpl.CHUNK_SIZE_DEFAULT, false, 0, 0, 0, null);
     FSHLog hLog = new FSHLog(fs, rootDir, testName, CONF);
     hLog.init();
     region = initHRegion(tableName, null, null, false, Durability.SYNC_WAL, hLog,
@@ -2427,7 +2428,16 @@ public Void answer(InvocationOnMock invocation) throws Throwable {
         return null;
       }
     }).when(mockedCPHost).preBatchMutate(Mockito.isA(MiniBatchOperationInProgress.class));
+    ColumnFamilyDescriptorBuilder builder = ColumnFamilyDescriptorBuilder.
+        newBuilder(COLUMN_FAMILY_BYTES);
+    ScanInfo info = new ScanInfo(CONF, builder.build(), Long.MAX_VALUE,
+        Long.MAX_VALUE, region.getCellComparator());
+    Mockito.when(mockedCPHost.preFlushScannerOpen(Mockito.any(HStore.class),
+        Mockito.any())).thenReturn(info);
+    Mockito.when(mockedCPHost.preFlush(Mockito.any(), Mockito.any(StoreScanner.class),
+        Mockito.any())).thenAnswer(i -> i.getArgument(1));
     region.setCoprocessorHost(mockedCPHost);
+
     region.put(originalPut);
     region.setCoprocessorHost(normalCPHost);
     final long finalSize = region.getDataInMemoryWithoutWAL();