feat(bigtable): expose the row key, mutation count and serialized size on RowMutationEntry

RowMutationEntry exposed no accessor for its serialized size, its row key or
its mutation count, so a caller that needs any of them had to call the
@internalapi toProto() and build a whole MutateRowsRequest.Entry to read one
number. Bounding in-flight memory by bytes is the motivating case: a single row
mutation may be megabytes, so bounding by entry count bounds nothing useful.

Mutation.byteSize could not simply be exposed - it sums each child mutation's
own serialized size and omits the row-key field and the per-mutation tag and
length prefixes. Mutation now accumulates the framed size alongside it, and
RowMutationEntry.getSerializedSize() adds the row-key field to it. Reads are
O(1) and allocate nothing.

The counters are maintained on the fromProto factories too, which bypass
addMutation, and are transient and recomputed in readObject so an instance
restored from an older serialized form is correct rather than reporting zero.
numMutations and byteSize are deliberately left alone.

Author: laughingman7743

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

@@ -17,6 +17,7 @@
 
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
+import com.google.bigtable.v2.MutateRowsRequest;
 import com.google.bigtable.v2.Mutation.AddToCell;
 import com.google.bigtable.v2.Mutation.DeleteFromColumn;
 import com.google.bigtable.v2.Mutation.DeleteFromFamily;
@@ -28,6 +29,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.primitives.Longs;
 import com.google.protobuf.ByteString;
+import com.google.protobuf.CodedOutputStream;
 import java.io.IOException;
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
@@ -60,6 +62,23 @@ public final class Mutation implements MutationApi<Mutation>, Serializable {
   private int numMutations;
   private long byteSize;
 
+  /**
+   * Number of accumulated mutations, maintained on every path that adds one — unlike {@link
+   * #numMutations}, which counts only what went through {@link #addMutation} and stays at zero for
+   * the {@code fromProto} factories.
+   */
+  private transient int mutationCount;
+
+  /**
+   * Size these mutations occupy inside a request proto that carries them as a repeated field: each
+   * mutation's tag, length prefix and body. Backs {@link RowMutationEntry#getSerializedSize()} so
+   * callers can size an entry without building one.
+   *
+   * <p>Transient and recomputed by {@code readObject}, so an instance restored from a form written
+   * by an older version is correct too.
+   */
+  private transient long mutationsSerializedSize;
+
   /** Creates new instance of Mutation object. */
   public static Mutation create() {
     return new Mutation(false);
@@ -86,6 +105,7 @@ public static Mutation createUnsafe() {
   public static Mutation fromProtoUnsafe(List<com.google.bigtable.v2.Mutation> protos) {
     Mutation mutation = new Mutation(true);
     mutation.mutations.addAll(protos);
+    mutation.countAllTowardsSerializedSize(protos);
     return mutation;
   }
 
@@ -99,6 +119,7 @@ public static Mutation fromProtoUnsafe(List<com.google.bigtable.v2.Mutation> pro
   public static Mutation fromProtoUnsafe(Iterable<com.google.bigtable.v2.Mutation> protos) {
     Mutation mutation = new Mutation(true);
     mutation.mutations.addAll(protos);
+    mutation.countAllTowardsSerializedSize(protos);
     return mutation;
   }
 
@@ -115,6 +136,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.countAllTowardsSerializedSize(protos);
     return mutation;
   }
 
@@ -129,6 +151,7 @@ private void readObject(ObjectInputStream input) throws IOException, ClassNotFou
     ImmutableList<com.google.bigtable.v2.Mutation> deserialized =
         (ImmutableList<com.google.bigtable.v2.Mutation>) input.readObject();
     this.mutations = ImmutableList.<com.google.bigtable.v2.Mutation>builder().addAll(deserialized);
+    countAllTowardsSerializedSize(deserialized);
   }
 
   private void writeObject(ObjectOutputStream output) throws IOException {
@@ -335,10 +358,38 @@ private void addMutation(com.google.bigtable.v2.Mutation mutation) {
 
     numMutations++;
     byteSize += mutation.getSerializedSize();
+    countTowardsSerializedSize(mutation);
 
     mutations.add(mutation);
   }
 
+  /** Accumulates one mutation into {@link #mutationCount} and {@link #mutationsSerializedSize}. */
+  private void countTowardsSerializedSize(com.google.bigtable.v2.Mutation mutation) {
+    mutationCount++;
+    // The field number only decides the tag's own width, and every request proto that carries
+    // these mutations numbers the field below 16, so this is the same one byte in all of them.
+    mutationsSerializedSize +=
+        CodedOutputStream.computeMessageSize(
+            MutateRowsRequest.Entry.MUTATIONS_FIELD_NUMBER, mutation);
+  }
+
+  /** Accumulates mutations added without going through {@link #addMutation}. */
+  private void countAllTowardsSerializedSize(Iterable<com.google.bigtable.v2.Mutation> protos) {
+    for (com.google.bigtable.v2.Mutation proto : protos) {
+      countTowardsSerializedSize(proto);
+    }
+  }
+
+  /** Returns the number of accumulated mutations, on every construction path. */
+  int getMutationCount() {
+    return mutationCount;
+  }
+
+  /** Returns {@link #mutationsSerializedSize}. */
+  long getMutationsSerializedSize() {
+    return mutationsSerializedSize;
+  }
+
   private static ByteString wrapByteString(String str) {
     if (str == null) {
       return null;

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

@@ -20,6 +20,7 @@
 import com.google.bigtable.v2.MutateRowsRequest;
 import com.google.common.base.Preconditions;
 import com.google.protobuf.ByteString;
+import com.google.protobuf.CodedOutputStream;
 import java.io.Serializable;
 import javax.annotation.Nonnull;
 
@@ -200,6 +201,35 @@ public RowMutationEntry mergeToCell(
     return this;
   }
 
+  /** Returns the row key these mutations apply to. */
+  public ByteString getRowKey() {
+    return key;
+  }
+
+  /** Returns the number of mutations accumulated in this entry. */
+  public int getMutationCount() {
+    return mutation.getMutationCount();
+  }
+
+  /**
+   * Returns the serialized size, in bytes, of the request entry this becomes on the wire — the same
+   * number as {@code toProto().getSerializedSize()}, without building the proto.
+   *
+   * <p>Intended for callers that bound their own in-flight memory or report a bytes-written metric.
+   * A single row mutation may legitimately be megabytes, so bounding by entry count bounds nothing
+   * useful; this is the number that bounds memory. It is accumulated as mutations are added rather
+   * than computed on each call.
+   */
+  public long getSerializedSize() {
+    long size = mutation.getMutationsSerializedSize();
+    if (!key.isEmpty()) {
+      // Implicit presence: an empty row key is not written at all, so it contributes no tag and
+      // no length prefix either.
+      size += CodedOutputStream.computeBytesSize(MutateRowsRequest.Entry.ROW_KEY_FIELD_NUMBER, key);
+    }
+    return size;
+  }
+
   @InternalApi
   public MutateRowsRequest.Entry toProto() {
     Preconditions.checkArgument(

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

@@ -165,4 +165,120 @@ public void testWithLongValue() {
                 .setValue(ByteString.copyFrom(Longs.toByteArray(100_000L)))
                 .build());
   }
+
+  @Test
+  public void getRowKeyTest() {
+    ByteString rowKey = ByteString.copyFromUtf8("row-key");
+
+    assertThat(RowMutationEntry.create(rowKey).getRowKey()).isEqualTo(rowKey);
+    assertThat(RowMutationEntry.create("row-key").getRowKey()).isEqualTo(rowKey);
+    assertThat(RowMutationEntry.create(ByteString.EMPTY).getRowKey()).isEqualTo(ByteString.EMPTY);
+  }
+
+  @Test
+  public void getMutationCountTest() {
+    assertThat(RowMutationEntry.create("row-key").getMutationCount()).isEqualTo(0);
+    assertThat(
+            RowMutationEntry.create("row-key")
+                .setCell("fake-family", "q1", 10_000L, "v1")
+                .setCell("fake-family", "q2", 10_000L, "v2")
+                .deleteFamily("fake-family")
+                .deleteRow()
+                .getMutationCount())
+        .isEqualTo(4);
+  }
+
+  @Test
+  public void getSerializedSizeMatchesProtoTest() {
+    // Every shape that changes the framing: no mutations, an empty row key, a value long enough
+    // to need a two-byte length prefix, and each mutation kind.
+    assertSizeMatchesProto(RowMutationEntry.create("row-key"));
+    assertSizeMatchesProto(RowMutationEntry.create(ByteString.EMPTY));
+    assertSizeMatchesProto(
+        RowMutationEntry.create("row-key").setCell("fake-family", "q", 10_000L, "v"));
+    assertSizeMatchesProto(
+        RowMutationEntry.create("row-key").setCell("fake-family", "q", 10_000L, repeat("v", 200)));
+    assertSizeMatchesProto(
+        RowMutationEntry.create(ByteString.copyFromUtf8(repeat("k", 5_000)))
+            .setCell("fake-family", "q", 10_000L, repeat("v", 70_000)));
+    assertSizeMatchesProto(
+        RowMutationEntry.create("row-key")
+            .setCell("fake-family", "q1", 10_000L, "v1")
+            .setCell("fake-family", "q2", 10_000L, 100_000L)
+            .deleteCells("fake-family", "q1")
+            .deleteFamily("fake-family")
+            .deleteRow());
+  }
+
+  @Test
+  public void getSerializedSizeMatchesProtoForManyMutationsTest() {
+    // Crosses the point where the repeated field's own count changes nothing but the accumulated
+    // per-mutation length prefixes do.
+    RowMutationEntry underTest = RowMutationEntry.create("row-key");
+    for (int i = 0; i < 300; i++) {
+      underTest.setCell("fake-family", "q" + i, 10_000L, "value-" + i);
+    }
+
+    assertThat(underTest.getMutationCount()).isEqualTo(300);
+    assertSizeMatchesProto(underTest);
+  }
+
+  @Test
+  public void getSerializedSizeFromMutationUnsafeTest() {
+    // createFromMutationUnsafe bypasses addMutation, so the accumulated size has to be maintained
+    // by the fromProto factories too.
+    com.google.cloud.bigtable.data.v2.models.Mutation mutation =
+        com.google.cloud.bigtable.data.v2.models.Mutation.fromProtoUnsafe(
+            ImmutableList.of(
+                Mutation.newBuilder()
+                    .setSetCell(
+                        Mutation.SetCell.newBuilder()
+                            .setFamilyName("fake-family")
+                            .setColumnQualifier(ByteString.copyFromUtf8("q"))
+                            .setTimestampMicros(10_000L)
+                            .setValue(ByteString.copyFromUtf8(repeat("v", 500))))
+                    .build(),
+                Mutation.newBuilder()
+                    .setDeleteFromRow(Mutation.DeleteFromRow.getDefaultInstance())
+                    .build()));
+    RowMutationEntry underTest =
+        RowMutationEntry.createFromMutationUnsafe(ByteString.copyFromUtf8("row-key"), mutation);
+
+    assertThat(underTest.getMutationCount()).isEqualTo(2);
+    assertSizeMatchesProto(underTest);
+  }
+
+  @Test
+  public void getSerializedSizeSurvivesSerializationTest()
+      throws IOException, ClassNotFoundException {
+    RowMutationEntry underTest =
+        RowMutationEntry.create("row-key")
+            .setCell("fake-family", "q", 10_000L, repeat("v", 500))
+            .deleteFamily("fake-family");
+
+    ByteArrayOutputStream bos = new ByteArrayOutputStream();
+    ObjectOutputStream oos = new ObjectOutputStream(bos);
+    oos.writeObject(underTest);
+    oos.close();
+    ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(bos.toByteArray()));
+    RowMutationEntry actual = (RowMutationEntry) ois.readObject();
+
+    assertThat(actual.getRowKey()).isEqualTo(underTest.getRowKey());
+    assertThat(actual.getMutationCount()).isEqualTo(underTest.getMutationCount());
+    assertThat(actual.getSerializedSize()).isEqualTo(underTest.getSerializedSize());
+    assertSizeMatchesProto(actual);
+  }
+
+  private static void assertSizeMatchesProto(RowMutationEntry entry) {
+    assertThat(entry.getSerializedSize()).isEqualTo(entry.toProto().getSerializedSize());
+    assertThat(entry.getMutationCount()).isEqualTo(entry.toProto().getMutationsCount());
+  }
+
+  private static String repeat(String unit, int times) {
+    StringBuilder sb = new StringBuilder(unit.length() * times);
+    for (int i = 0; i < times; i++) {
+      sb.append(unit);
+    }
+    return sb.toString();
+  }
 }