Repository: GoogleCloudPlatform/cloud-bigtable-client
Pull request #1671 (2018-02-14 15:26:27)

Adding batch operations in BigtableAsyncTable.

spollapally: LGTM

Patch:
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/BulkMutation.java b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/BulkMutation.java
index 7a83930e7..a477d24d4 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/BulkMutation.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/main/java/com/google/cloud/bigtable/grpc/async/BulkMutation.java
@@ -404,7 +404,7 @@ public class BulkMutation {
 
     boolean didSend = false;
     if (currentBatch != null && currentBatch.wouldBeFull(entry)) {
-      send();
+      sendUnsent();
       if (scheduledFlush != null) {
         scheduledFlush.cancel(true);
         scheduledFlush = null;
@@ -428,7 +428,7 @@ public class BulkMutation {
           public void run() {
             synchronized (BulkMutation.this) {
               scheduledFlush = null;
-              send();
+              sendUnsent();
             }
           }
         }, autoflushMs, TimeUnit.MILLISECONDS);
@@ -442,12 +442,11 @@ public class BulkMutation {
    * Send any outstanding {@link MutateRowRequest}s and wait until all requests are complete.
    */
   public void flush() throws InterruptedException {
-    send();
+    sendUnsent();
     operationAccountant.awaitCompletion();
   }
 
-  @VisibleForTesting
-  synchronized void send() {
+  public synchronized void sendUnsent() {
     if (currentBatch != null) {
       try {
         currentBatch.run();
diff --git a/bigtable-hbase-2.x-parent/bigtable-hbase-2.x/src/main/java/com/google/cloud/bigtable/hbase2_x/BigtableAsyncTable.java b/bigtable-hbase-2.x-parent/bigtable-hbase-2.x/src/main/java/com/google/cloud/bigtable/hbase2_x/BigtableAsyncTable.java
index 2b9824a53..98caf2ad9 100644
--- a/bigtable-hbase-2.x-parent/bigtable-hbase-2.x/src/main/java/com/google/cloud/bigtable/hbase2_x/BigtableAsyncTable.java
+++ b/bigtable-hbase-2.x-parent/bigtable-hbase-2.x/src/main/java/com/google/cloud/bigtable/hbase2_x/BigtableAsyncTable.java
@@ -15,10 +15,14 @@ d * Copyright 2017 Google Inc. All Rights Reserved.
  */
 package com.google.cloud.bigtable.hbase2_x;
 
+import static com.google.cloud.bigtable.hbase2_x.FutureUtils.toCompletableFuture;
+import static java.util.stream.Collectors.toList;
+
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Stream;
 
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.TableName;
@@ -40,7 +44,7 @@ import org.apache.hadoop.hbase.filter.KeyOnlyFilter;
 
 import com.google.bigtable.v2.MutateRowRequest;
 import com.google.bigtable.v2.MutateRowResponse;
-import com.google.cloud.bigtable.config.Logger;
+import com.google.bigtable.v2.ReadRowsRequest;
 import com.google.cloud.bigtable.grpc.BigtableDataClient;
 import com.google.cloud.bigtable.grpc.BigtableSession;
 import com.google.cloud.bigtable.grpc.scanner.FlatRow;
@@ -55,7 +59,6 @@ import com.google.common.util.concurrent.ListenableFuture;
  * @author spollapally
  */
 public class BigtableAsyncTable implements AsyncTable {
-  private final Logger LOG = new Logger(getClass());
 
   private final BigtableAsyncConnection asyncConnection;
   private final BigtableDataClient client;
@@ -86,9 +89,16 @@ public class BigtableAsyncTable implements AsyncTable {
     throw new UnsupportedOperationException("append"); // TODO
   }
 
+  @SuppressWarnings("unchecked")
   @Override
   public <T> List<CompletableFuture<T>> batch(List<? extends Row> actions) {
-    throw new UnsupportedOperationException("batch"); // TODO
+    List<? extends Row> updatedActions = actions.stream()
+        .map(row -> row instanceof Get ? fromHB2Get((Get) row) : row)
+        .collect(toList());
+    // TODO: The CompletableFutures need to return Void for Put/Delete.
+    return asyncRequests(updatedActions).stream()
+        .map(lfuture -> (CompletableFuture<T>) toCompletableFuture(lfuture))
+        .collect(toList());
   }
 
   @Override
@@ -101,23 +111,41 @@ public class BigtableAsyncTable implements AsyncTable {
     // figure out how to time this with Opencensus
     MutateRowRequest request = hbaseAdapter.adapt(delete);
     ListenableFuture<MutateRowResponse> future = client.mutateRowAsync(request);
-    return FutureUtils.toCompletableFuture(future, executorService)
+    return toCompletableFuture(future, executorService)
         .thenApply(r -> null);
   }
 
   @Override
   public List<CompletableFuture<Void>> delete(List<Delete> deletes) {
-    throw new UnsupportedOperationException("delete list"); // TODO
+    Stream<CompletableFuture<Void>> stream = asyncRequests(deletes)
+          .stream()
+          .map(lfuture -> toCompletableFuture(lfuture)
+                .thenApply(r -> null));
+    return stream.collect(toList());
+  }
+
+  private List<ListenableFuture<?>> asyncRequests(List<? extends Row> actions) {
+    Object[] results = new Object[actions.size()];
+    return getBatchExecutor().issueAsyncRowRequests(actions, results, null);
   }
 
   @Override
   public CompletableFuture<Result> get(Get get) {
-    ListenableFuture<List<FlatRow>> future = client.readFlatRowsAsync(hbaseAdapter.adapt(get));
-    return FutureUtils.toCompletableFuture(future, (list -> toResult("get", list)),
+    ReadRowsRequest request = hbaseAdapter.adapt(fromHB2Get(get));
+    ListenableFuture<List<FlatRow>> future = client.readFlatRowsAsync(request);
+    return toCompletableFuture(future, (list -> toResult("get", list)),
       executorService);
   }
+  
+  private static Get fromHB2Get(Get get) {
+    if (get.isCheckExistenceOnly()) {
+      return addKeyOnlyFilter(get);
+    } else {
+      return get;
+    }
+  }
 
-  private Get addKeyOnlyFilter(Get get) {
+  private static Get addKeyOnlyFilter(Get get) {
     Get existsGet = new Get(get);
     if (get.getFilter() == null) {
       existsGet.setFilter(new KeyOnlyFilter());
@@ -147,9 +175,25 @@ public class BigtableAsyncTable implements AsyncTable {
     }
   }
 
+  @SuppressWarnings("unchecked")
   @Override
-  public List<CompletableFuture<Result>> get(List<Get> arg0) {
-    throw new UnsupportedOperationException("get"); // TODO
+  public List<CompletableFuture<Result>> get(List<Get> gets) {
+    List<Get> hb1Gets = gets.stream()
+        .map(BigtableAsyncTable::fromHB2Get)
+        .collect(toList());
+    return asyncRequests(hb1Gets).stream()
+        .map(lfuture ->  (CompletableFuture<Result>) toCompletableFuture(lfuture))
+        .collect(toList());
+  }
+
+  @Override
+  public List<CompletableFuture<Boolean>> exists(List<Get> gets) {
+    List<Get> existGets = gets.stream()
+          .map(BigtableAsyncTable::addKeyOnlyFilter)
+          .collect(toList());
+    return get(existGets).stream()
+          .map(cf -> cf.thenApply(r -> !r.isEmpty()))
+          .collect(toList());
   }
 
   @Override
@@ -196,7 +240,7 @@ public class BigtableAsyncTable implements AsyncTable {
   public CompletableFuture<Void> mutateRow(RowMutations rowMutations) {
     MutateRowRequest request = hbaseAdapter.adapt(rowMutations);
     ListenableFuture<MutateRowResponse> future = client.mutateRowAsync(request);
-    return FutureUtils.toCompletableFuture(future, executorService)
+    return toCompletableFuture(future, executorService)
         .thenApply(r -> null);
   }
 
@@ -204,13 +248,17 @@ public class BigtableAsyncTable implements AsyncTable {
   public CompletableFuture<Void> put(Put put) {
     // figure out how to time this with Opencensus
     MutateRowRequest request = hbaseAdapter.adapt(put);
-    ListenableFuture<MutateRowResponse> future = client.mutateRowAsync(request);
-    return FutureUtils.toCompletableFuture(future, executorService).thenApply(r -> null);
+    ListenableFuture<?> future = client.mutateRowAsync(request);
+    return toCompletableFuture(future, executorService).thenApply(r -> null);
   }
 
   @Override
-  public List<CompletableFuture<Void>> put(List<Put> arg0) {
-    throw new UnsupportedOperationException("put"); // TODO
+  public List<CompletableFuture<Void>> put(List<Put> puts) {
+    Stream<CompletableFuture<Void>> stream = asyncRequests(puts)
+        .stream()
+        .map(lfuture -> toCompletableFuture(lfuture)
+              .thenApply(r -> null));
+    return stream.collect(toList());
   }
 
   @Override
@@ -227,5 +275,4 @@ public class BigtableAsyncTable implements AsyncTable {
   public void scan(Scan scan, ScanResultConsumer consumer) {
     throw new UnsupportedOperationException("scan"); // TODO
   }
-
 }
diff --git a/bigtable-hbase-2.x-parent/bigtable-hbase-2.x/src/main/java/com/google/cloud/bigtable/hbase2_x/FutureUtils.java b/bigtable-hbase-2.x-parent/bigtable-hbase-2.x/src/main/java/com/google/cloud/bigtable/hbase2_x/FutureUtils.java
index 45dee4aa4..f6c32ac0e 100644
--- a/bigtable-hbase-2.x-parent/bigtable-hbase-2.x/src/main/java/com/google/cloud/bigtable/hbase2_x/FutureUtils.java
+++ b/bigtable-hbase-2.x-parent/bigtable-hbase-2.x/src/main/java/com/google/cloud/bigtable/hbase2_x/FutureUtils.java
@@ -37,8 +37,7 @@ public class FutureUtils {
   private static final ExecutorService DIRECT_EXECUTOR = MoreExecutors.newDirectExecutorService();
   static Logger logger = new Logger(FutureUtils.class);
 
-  public static <T> CompletableFuture<T>
-      toCompletableFuture(final ListenableFuture<T> listenableFuture) {
+  public static <T> CompletableFuture<T> toCompletableFuture(ListenableFuture<T> listenableFuture) {
     return toCompletableFuture(listenableFuture, DIRECT_EXECUTOR);
   }
 
diff --git a/bigtable-hbase-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BatchExecutor.java b/bigtable-hbase-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BatchExecutor.java
index 9c6c7fef6..b9c791d47 100644
--- a/bigtable-hbase-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BatchExecutor.java
+++ b/bigtable-hbase-parent/bigtable-hbase/src/main/java/com/google/cloud/bigtable/hbase/BatchExecutor.java
@@ -145,9 +145,9 @@ public class BatchExecutor {
       this.bulkMutation = session.createBulkMutation(tableName);
     }
 
-    protected void flush() throws InterruptedException {
+    protected void flush() {
       // If there is a bulk mutation in progress, then send it.
-      bulkMutation.flush();
+      bulkMutation.sendUnsent();
       bulkRead.flush();
     }
   }
@@ -240,8 +240,8 @@ public class BatchExecutor {
     batchCallback(actions, results, null);
   }
 
-  private <R> List<ListenableFuture<?>> issueAsyncRowRequests(List<? extends Row> actions,
-      Object[] results, Batch.Callback<R> callback) throws InterruptedException {
+  public <R> List<ListenableFuture<?>> issueAsyncRowRequests(List<? extends Row> actions,
+      Object[] results, Batch.Callback<R> callback) {
     BulkOperation bulkOperation = new BulkOperation(session, requestAdapter.getBigtableTableName());
     try {
       List<ListenableFuture<?>> resultFutures = new ArrayList<>(actions.size());


Test patch:
diff --git a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestBulkMutation.java b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestBulkMutation.java
index 570e437ca..fa17a699a 100644
--- a/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestBulkMutation.java
+++ b/bigtable-client-core-parent/bigtable-client-core/src/test/java/com/google/cloud/bigtable/grpc/async/TestBulkMutation.java
@@ -124,7 +124,7 @@ public class TestBulkMutation {
   public void testAdd() {
     MutateRowsRequest.Entry entry = createRequestEntry();
     underTest.add(entry);
-    underTest.send();
+    underTest.sendUnsent();
     MutateRowsRequest expected = MutateRowsRequest.newBuilder()
         .setTableName(TABLE_NAME.toString())
         .addEntries(entry)
@@ -256,7 +256,7 @@ public class TestBulkMutation {
         for (int i = 0; i < batchCount * MAX_ROW_COUNT; i++) {
           futures.add(bulkMutation.add(createRequestEntry()));
         }
-        bulkMutation.send();
+        bulkMutation.sendUnsent();
       }
     };
     ExecutorService pool = Executors.newFixedThreadPool(100);
@@ -389,6 +389,6 @@ public class TestBulkMutation {
         .getStatusBuilder()
             .setCode(code.getCode().value());
     future.set(Arrays.asList(responseBuilder.build()));
-    underTest.send();
+    underTest.sendUnsent();
   }
 }
diff --git a/bigtable-hbase-2.x-parent/bigtable-hbase-2.x-integration-tests/src/test/java/com/google/cloud/bigtable/hbase/IntegrationTests.java b/bigtable-hbase-2.x-parent/bigtable-hbase-2.x-integration-tests/src/test/java/com/google/cloud/bigtable/hbase/IntegrationTests.java
index 286351839..e820cd594 100644
--- a/bigtable-hbase-2.x-parent/bigtable-hbase-2.x-integration-tests/src/test/java/com/google/cloud/bigtable/hbase/IntegrationTests.java
+++ b/bigtable-hbase-2.x-parent/bigtable-hbase-2.x-integration-tests/src/test/java/com/google/cloud/bigtable/hbase/IntegrationTests.java
@@ -16,6 +16,7 @@
 package com.google.cloud.bigtable.hbase;
 
 import com.google.cloud.bigtable.hbase.async.TestAsyncAdmin;
+import com.google.cloud.bigtable.hbase.async.TestAsyncBatch;
 import com.google.cloud.bigtable.hbase.async.TestAsyncConnection;
 import com.google.cloud.bigtable.hbase.async.TestBasicAsyncOps;
 import com.google.cloud.bigtable.hbase.test_env.SharedTestEnvRule;
@@ -57,7 +58,8 @@ import org.junit.runners.Suite;
     TestTruncateTable.class,
     TestAsyncAdmin.class,
     TestAsyncConnection.class,
-    TestBasicAsyncOps.class
+    TestBasicAsyncOps.class,
+    TestAsyncBatch.class
 })
 public class IntegrationTests {
 
diff --git a/bigtable-hbase-2.x-parent/bigtable-hbase-2.x-integration-tests/src/test/java/com/google/cloud/bigtable/hbase/async/TestAsyncBatch.java b/bigtable-hbase-2.x-parent/bigtable-hbase-2.x-integration-tests/src/test/java/com/google/cloud/bigtable/hbase/async/TestAsyncBatch.java
new file mode 100644
index 000000000..2d8f7aa43
--- /dev/null
+++ b/bigtable-hbase-2.x-parent/bigtable-hbase-2.x-integration-tests/src/test/java/com/google/cloud/bigtable/hbase/async/TestAsyncBatch.java
@@ -0,0 +1,145 @@
+/*
+ * Copyright 2018 Google Inc. All Rights Reserved.
+ *
+ * Licensed 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 com.google.cloud.bigtable.hbase.async;
+
+import static com.google.cloud.bigtable.hbase.test_env.SharedTestEnvRule.COLUMN_FAMILY;
+import static java.util.stream.Collectors.toList;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+
+import org.apache.hadoop.hbase.Cell;
+import org.apache.hadoop.hbase.CellUtil;
+import org.apache.hadoop.hbase.client.AsyncTable;
+import org.apache.hadoop.hbase.client.Delete;
+import org.apache.hadoop.hbase.client.Get;
+import org.apache.hadoop.hbase.client.Put;
+import org.apache.hadoop.hbase.client.Result;
+import org.junit.AfterClass;
+import org.junit.Assert;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+import com.google.common.util.concurrent.ThreadFactoryBuilder;
+
+/**
+ * Test to make sure that basic {@link AsyncTable} operations work
+ * @author sduskis
+ */
+@RunWith(JUnit4.class)
+public class TestAsyncBatch extends AbstractAsyncTest {
+
+  private static ExecutorService executor;
+
+  @BeforeClass
+  public static void setup() {
+    executor = Executors.newCachedThreadPool(
+      new ThreadFactoryBuilder().setDaemon(true).setNameFormat("table-deleter").build());
+  }
+
+  @AfterClass
+  public static void shutdown() {
+    if (executor != null) {
+      executor.shutdownNow();
+    }
+    executor = null;
+  }
+
+  @Test
+  public void testBasicAsyncOps() throws Exception {
+    System.out.println("TestBasicAsyncOps");
+    byte[][] rowKeys = dataHelper.randomData("TestBasicAsyncOps-", 5);
+    byte[][] testQualifiers = dataHelper.randomData("testQualifier-", 5);
+    byte[][] testValues = dataHelper.randomData("testValue-", 5);
+    testPutsGetsDeletes(true, rowKeys, testQualifiers, testValues);
+  }
+
+  static <T> CompletableFuture<List<T>> allOf(List<CompletableFuture<T>> futures) {
+    return CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
+        .thenApply(v -> futures.stream().map(f -> f.getNow(null)).collect(toList()));
+  }
+
+  private void testPutsGetsDeletes(boolean doGet, byte[][] rowKeys, byte[][] testQualifiers,
+      byte[][] testValues) throws IOException, InterruptedException, ExecutionException {
+    AsyncTable table = getAsyncConnection().getTable(sharedTestEnv.getDefaultTableName(), executor);
+
+    // setup
+    List<Get> gets = new ArrayList<>();
+    List<Put> puts = new ArrayList<>();
+    List<Delete> deletes = new ArrayList<>();
+    for (int i = 0; i < rowKeys.length; i++) {
+      puts.add(new Put(rowKeys[i])
+        .addColumn(COLUMN_FAMILY, testQualifiers[i], testValues[i]));
+
+      gets.add(new Get(rowKeys[i])
+        .addColumn(COLUMN_FAMILY, testQualifiers[i]));
+
+      deletes.add(new Delete(rowKeys[i]));
+    }
+
+    Stopwatch stopwatch = new Stopwatch();
+    // Put
+    allOf(table.put(puts)).get();
+    stopwatch.print("Put took %d ms");
+
+    // Get
+
+    // Do the get on some tests, but not others.  The rationale for that is to do performance
+    // testing on large values.
+    if (doGet) {
+      stopwatch.print("Get took %d ms");
+      List<Result> results = allOf(table.get(gets)).get();
+      Assert.assertEquals(rowKeys.length, results.size()); 
+      for (int i = 0; i < rowKeys.length; i++) {
+        Result result = results.get(i);
+        Assert.assertTrue(result.containsColumn(COLUMN_FAMILY, testQualifiers[i]));
+        List<Cell> cells = result.getColumnCells(COLUMN_FAMILY, testQualifiers[i]);
+        Assert.assertEquals(1, cells.size());
+        Assert.assertTrue(CellUtil.matchingValue(cells.get(0), testValues[i]));
+      }
+      stopwatch.print("Verifying took %d ms");
+    }
+    // Delete
+    allOf(table.delete(deletes)).get();
+    stopwatch.print("Delete took %d ms");
+
+    // Confirm deleted
+    List<Boolean> exists = table.existsAll(gets).get();
+    Assert.assertNotNull(exists);
+    Assert.assertEquals(rowKeys.length, exists.size());
+    for (Boolean b : exists) {
+      Assert.assertFalse(b);
+    }
+    stopwatch.print("Exists took %d ms");
+  }
+
+  private class Stopwatch {
+    long lastCheckin = System.currentTimeMillis();
+
+    private void print(String string) {
+      long now = System.currentTimeMillis();
+      logger.info(string, now - lastCheckin);
+      lastCheckin = now;
+    }
+  }
+}