fix(bigtable): count mutations wrapped from protos towards the row limits

Mutation.MAX_MUTATIONS and MAX_BYTE_SIZE are enforced in addMutation, against
counters only addMutation maintains. fromProtoUnsafe(List),
fromProtoUnsafe(Iterable) and fromProto(List) add to the mutation list
directly and leave both counters at zero, so mutations wrapped from existing
protos count towards neither limit.

The mutation count is backstopped by RowMutationEntry.toProto() and
BulkMutation.add, which re-check the real list size, so it surfaces late and as
a different exception type. The byte size is not backstopped anywhere, so a
Mutation seeded from protos can exceed 200 MB with nothing client-side
objecting.

Count wrapped protos in the three factories, so the counters describe the whole
row. Deliberately not a checkState in the factories themselves: that would make
wrapping an already-over-limit proto throw where it currently does not.

Author: laughingman7743

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/models/Mutation.java

@@ -87,7 +87,7 @@ public static Mutation createUnsafe() {
   @BetaApi
   public static Mutation fromProtoUnsafe(List<com.google.bigtable.v2.Mutation> protos) {
     Mutation mutation = new Mutation(true);
-    mutation.mutations.addAll(protos);
+    mutation.addAllFromProto(protos);
     return mutation;
   }
 
@@ -100,7 +100,7 @@ public static Mutation fromProtoUnsafe(List<com.google.bigtable.v2.Mutation> pro
   @BetaApi
   public static Mutation fromProtoUnsafe(Iterable<com.google.bigtable.v2.Mutation> protos) {
     Mutation mutation = new Mutation(true);
-    mutation.mutations.addAll(protos);
+    mutation.addAllFromProto(protos);
     return mutation;
   }
 
@@ -116,7 +116,7 @@ public static Mutation fromProtoUnsafe(Iterable<com.google.bigtable.v2.Mutation>
    */
   static Mutation fromProto(List<com.google.bigtable.v2.Mutation> protos) {
     Mutation mutation = new Mutation(false);
-    mutation.mutations.addAll(protos);
+    mutation.addAllFromProto(protos);
     return mutation;
   }
 
@@ -346,10 +346,22 @@ private void addMutation(com.google.bigtable.v2.Mutation mutation) {
         byteSize + mutation.getSerializedSize() <= MAX_BYTE_SIZE,
         "Byte size of mutations is too large");
 
+    countTowardsLimits(mutation);
+
+    mutations.add(mutation);
+  }
+
+  private void countTowardsLimits(com.google.bigtable.v2.Mutation mutation) {
     numMutations++;
     byteSize += mutation.getSerializedSize();
+  }
 
-    mutations.add(mutation);
+  private void addAllFromProto(Iterable<com.google.bigtable.v2.Mutation> protos) {
+    // One traversal: protos may be a non-repeatable Iterable.
+    for (com.google.bigtable.v2.Mutation proto : protos) {
+      mutations.add(proto);
+      countTowardsLimits(proto);
+    }
   }
 
   private static ByteString wrapByteString(String str) {

java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/models/MutationTest.java

@@ -24,6 +24,7 @@
 import com.google.bigtable.v2.Mutation.MergeToCell;
 import com.google.bigtable.v2.Mutation.TimestampOrigin;
 import com.google.cloud.bigtable.data.v2.models.Range.TimestampRange;
+import com.google.common.collect.ImmutableList;
 import com.google.common.primitives.Longs;
 import com.google.protobuf.ByteString;
 import java.io.ByteArrayInputStream;
@@ -32,6 +33,7 @@
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
 import java.time.Instant;
+import java.util.Iterator;
 import java.util.List;
 import org.junit.Before;
 import org.junit.Test;
@@ -270,6 +272,105 @@ public void tooLargeRequest() {
     assertThat(actualError).isInstanceOf(IllegalStateException.class);
   }
 
+  @Test
+  public void tooManyMutationsCountsWrappedProtosTest() {
+    Mutation mutation =
+        Mutation.fromProtoUnsafe(
+            ImmutableList.of(
+                com.google.bigtable.v2.Mutation.newBuilder()
+                    .setDeleteFromRow(com.google.bigtable.v2.Mutation.DeleteFromRow.newBuilder())
+                    .build()));
+
+    for (int i = 0; i < Mutation.MAX_MUTATIONS - 1; i++) {
+      mutation.setCell("f", "", "");
+    }
+
+    Exception actualError = null;
+    try {
+      mutation.setCell("f", "", "");
+    } catch (Exception e) {
+      actualError = e;
+    }
+
+    assertThat(actualError).isInstanceOf(IllegalStateException.class);
+    assertThat(mutation.getMutations()).hasSize(Mutation.MAX_MUTATIONS);
+  }
+
+  @Test
+  public void tooLargeRequestCountsWrappedProtosTest() {
+    Mutation mutation =
+        Mutation.fromProtoUnsafe(
+            ImmutableList.of(
+                com.google.bigtable.v2.Mutation.newBuilder()
+                    .setSetCell(
+                        com.google.bigtable.v2.Mutation.SetCell.newBuilder()
+                            .setFamilyName("f")
+                            .setValue(ByteString.copyFrom(new byte[Mutation.MAX_BYTE_SIZE / 2])))
+                    .build(),
+                com.google.bigtable.v2.Mutation.newBuilder()
+                    .setSetCell(
+                        com.google.bigtable.v2.Mutation.SetCell.newBuilder()
+                            .setFamilyName("f")
+                            .setValue(ByteString.copyFrom(new byte[Mutation.MAX_BYTE_SIZE / 2])))
+                    .build()));
+
+    Exception actualError = null;
+    try {
+      mutation.setCell("f", "", "");
+    } catch (Exception e) {
+      actualError = e;
+    }
+
+    assertThat(actualError).isInstanceOf(IllegalStateException.class);
+  }
+
+  @Test
+  public void fromProtoUnsafeTraversesTheIterableOnceTest() {
+    Iterable<com.google.bigtable.v2.Mutation> oneShot =
+        oneShot(
+            com.google.bigtable.v2.Mutation.newBuilder()
+                .setDeleteFromRow(com.google.bigtable.v2.Mutation.DeleteFromRow.newBuilder())
+                .build(),
+            com.google.bigtable.v2.Mutation.newBuilder()
+                .setDeleteFromRow(com.google.bigtable.v2.Mutation.DeleteFromRow.newBuilder())
+                .build());
+
+    Mutation mutation = Mutation.fromProtoUnsafe(oneShot);
+
+    assertThat(mutation.getMutations()).hasSize(2);
+
+    for (int i = 0; i < Mutation.MAX_MUTATIONS - 2; i++) {
+      mutation.setCell("f", "", "");
+    }
+
+    Exception actualError = null;
+    try {
+      mutation.setCell("f", "", "");
+    } catch (Exception e) {
+      actualError = e;
+    }
+
+    assertThat(actualError).isInstanceOf(IllegalStateException.class);
+  }
+
+  /** An Iterable that refuses a second traversal. */
+  private static Iterable<com.google.bigtable.v2.Mutation> oneShot(
+      com.google.bigtable.v2.Mutation... protos) {
+    List<com.google.bigtable.v2.Mutation> source = ImmutableList.copyOf(protos);
+    return new Iterable<com.google.bigtable.v2.Mutation>() {
+      private boolean consumed;
+
+      @Override
+      public Iterator<com.google.bigtable.v2.Mutation> iterator() {
+        if (consumed) {
+          throw new IllegalStateException("Iterable traversed more than once");
+        }
+        consumed = true;
+        return source.iterator();
+      }
+    };
+  }
+
   @Test
   public void testWithLongValue() {
     Mutation mutation =