Update StorageItemsTable.MAX_MUTATIONS logic to account for insert mu…

…tations
signalapp · Oct 24, 2023 · 7264b8e · 7264b8e
1 parent 56c2f8e
commit 7264b8e
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 5 deletions.
diff --git a/src/main/java/org/signal/storageservice/controllers/StorageController.java b/src/main/java/org/signal/storageservice/controllers/StorageController.java
@@ -85,7 +85,7 @@ public CompletableFuture<Response> write(@Auth User user, @HeaderParam(HttpHeade
     distributionSummary(INSERT_DISTRIBUTION_SUMMARY_NAME, userAgent).record(writeOperation.getInsertItemCount());
     distributionSummary(DELETE_DISTRIBUTION_SUMMARY_NAME, userAgent).record(writeOperation.getDeleteKeyCount());
 
-    if (writeOperation.getInsertItemCount() + writeOperation.getDeleteKeyCount() > StorageItemsTable.MAX_MUTATIONS) {
+    if (writeOperation.getInsertItemCount() * StorageItemsTable.MUTATIONS_PER_INSERT + writeOperation.getDeleteKeyCount() > StorageItemsTable.MAX_MUTATIONS) {
       return CompletableFuture.failedFuture(new WebApplicationException(Status.REQUEST_ENTITY_TOO_LARGE));
     }
 

diff --git a/src/main/java/org/signal/storageservice/storage/StorageItemsTable.java b/src/main/java/org/signal/storageservice/storage/StorageItemsTable.java
@@ -35,6 +35,7 @@ public class StorageItemsTable extends Table {
   public static final String COLUMN_KEY = "k";
 
   public static final int MAX_MUTATIONS = 100_000;
+  public static final int MUTATIONS_PER_INSERT = 2;
 
   private final MetricRegistry metricRegistry = SharedMetricRegistries.getOrCreate(StorageMetrics.NAME);
   private final Timer getTimer = metricRegistry.timer(name(StorageItemsTable.class, "get"));
@@ -52,6 +53,7 @@ public CompletableFuture<Void> set(User user, List<StorageItem> inserts, List<By
     for (StorageItem insert : inserts) {
       bulkMutation.add(getRowKeyFor(user, insert.getKey()),
           Mutation.create()
+              // each setCell() counts as mutation. If the below code changes, update MUTATIONS_PER_INSERT
               .setCell(FAMILY, ByteString.copyFromUtf8(COLUMN_DATA), 0, insert.getValue())
               .setCell(FAMILY, ByteString.copyFromUtf8(COLUMN_KEY), 0, insert.getKey()));
     }

diff --git a/src/test/java/org/signal/storageservice/controllers/StorageControllerTest.java b/src/test/java/org/signal/storageservice/controllers/StorageControllerTest.java
@@ -26,7 +26,6 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
-import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import javax.ws.rs.client.Entity;
 import javax.ws.rs.core.Response;
@@ -358,7 +357,7 @@ void testWriteOversizeList() {
         .setValue(ByteString.copyFromUtf8("A manifest"))
         .build();
 
-    final int insertCount = 1 + StorageItemsTable.MAX_MUTATIONS / 2;
+    final int insertCount = 1 + StorageItemsTable.MAX_MUTATIONS / 2 / StorageItemsTable.MUTATIONS_PER_INSERT;
     final int deleteCount = 1 + StorageItemsTable.MAX_MUTATIONS / 2;
 
    final WriteOperation.Builder builder = WriteOperation.newBuilder()

diff --git a/src/test/java/org/signal/storageservice/storage/StorageManagerTest.java b/src/test/java/org/signal/storageservice/storage/StorageManagerTest.java
@@ -49,7 +49,7 @@ class StorageManagerTest {
   private static final String MANIFESTS_TABLE_ID = "manifest-table";
 
   @RegisterExtension
-  private final BigtableEmulatorExtension bigtableEmulator = BigtableEmulatorExtension.create();
+  public final BigtableEmulatorExtension bigtableEmulator = BigtableEmulatorExtension.create();
 
   private BigtableDataClient client;
   private BigtableTableAdminClient tableAdminClient;
@@ -375,7 +375,9 @@ void testClearItemsLargeBatch() {
     for (int chunk = 0; chunk < 2; chunk++) {
       final BulkMutation bulkMutation = BulkMutation.create(CONTACTS_TABLE_ID);
 
-      for (int i = 0; i < StorageItemsTable.MAX_MUTATIONS; i++) {
+      // Each setCell() is a mutation
+      final int setCellCallsPerLoop = 2;
+      for (int i = 0; i < StorageItemsTable.MAX_MUTATIONS; i += setCellCallsPerLoop) {
         bulkMutation.add(String.format("%s#contact#somekey%d_%05d", userId, chunk, i),
             Mutation.create()
                 .setCell(StorageItemsTable.FAMILY, StorageItemsTable.COLUMN_DATA, "data" + String.format("%03d", i))