Serialization of NestedSet<NestedSet> with NestedSetCodecWithStore waits on

backend writes for inner NestedSet serialization.

PiperOrigin-RevId: 195294676

4 files changed, 43 insertions, 2 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/collect/nestedset/BUILD b/src/main/java/com/google/devtools/build/lib/collect/nestedset/BUILD
index 725c8c8a73..e9a487f281 100644
--- a/src/main/java/com/google/devtools/build/lib/collect/nestedset/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/BUILD
@@ -37,6 +37,7 @@ java_library(
     deps = [
         ":nestedset",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization",
+        "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils",
         "//third_party:guava",
         "//third_party:truth",
diff --git a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecTestUtils.java b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecTestUtils.java
index 901b771c8b..ea455f87e1 100644
--- a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecTestUtils.java
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecTestUtils.java
@@ -15,10 +15,13 @@ package com.google.devtools.build.lib.collect.nestedset;
 
 import static com.google.common.truth.Truth.assertThat;
 
+import com.google.common.base.Objects;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.devtools.build.lib.skyframe.serialization.ObjectCodecs;
 import com.google.devtools.build.lib.skyframe.serialization.SerializationContext;
 import com.google.devtools.build.lib.skyframe.serialization.SerializationException;
+import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
+import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec.VisibleForSerialization;
 import com.google.devtools.build.lib.skyframe.serialization.testutils.SerializationTester;
 import java.io.IOException;
 
@@ -28,6 +31,33 @@ public class NestedSetCodecTestUtils {
   private static final NestedSet<String> SHARED_NESTED_SET =
       NestedSetBuilder.<String>stableOrder().add("e").build();
 
+  @AutoCodec
+  static class HasNestedSet {
+    private final NestedSet<String> nestedSetField;
+
+    @VisibleForSerialization
+    HasNestedSet(NestedSet<String> nestedSetField) {
+      this.nestedSetField = nestedSetField;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) {
+        return true;
+      }
+      if (o == null || getClass() != o.getClass()) {
+        return false;
+      }
+      HasNestedSet that = (HasNestedSet) o;
+      return Objects.equal(nestedSetField.rawChildren(), that.nestedSetField.rawChildren());
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hashCode(nestedSetField);
+    }
+  }
+
   /** Perform serialization/deserialization checks for several simple NestedSet examples. */
   public static void checkCodec(ObjectCodecs objectCodecs, boolean allowFutureBlocking)
       throws Exception {
@@ -50,7 +80,10 @@ public class NestedSetCodecTestUtils {
                         .addTransitive(SHARED_NESTED_SET)
                         .build())
                 .addTransitive(NestedSetBuilder.emptySet(Order.STABLE_ORDER))
-                .build())
+                .build(),
+            NestedSetBuilder.create(
+                Order.STABLE_ORDER,
+                new HasNestedSet(NestedSetBuilder.create(Order.STABLE_ORDER, "a"))))
         .setObjectCodecs(objectCodecs)
         .makeMemoizingAndAllowFutureBlocking(allowFutureBlocking)
         .setVerificationFunction(NestedSetCodecTestUtils::verifyDeserialization)
diff --git a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetStore.java b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetStore.java
index fd96cc4a59..d4f03de4d7 100644
--- a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetStore.java
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetStore.java
@@ -207,6 +207,13 @@ public class NestedSetStore {
         ByteString.copyFrom(Hashing.md5().hashBytes(serializedBytes).asBytes());
     futureBuilder.add(nestedSetStorageEndpoint.put(fingerprint, serializedBytes));
 
+    // If this is a NestedSet<NestedSet>, serialization of the contents will itself have writes.
+    ListenableFuture<Void> innerWriteFutures =
+        newSerializationContext.createFutureToBlockWritingOn();
+    if (innerWriteFutures != null) {
+      futureBuilder.add(innerWriteFutures);
+    }
+
     ListenableFuture<Void> writeFuture =
         Futures.whenAllComplete(futureBuilder.build())
             .call(() -> null, MoreExecutors.directExecutor());
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContext.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContext.java
index 2c19c6b8dc..468a362e24 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContext.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContext.java
@@ -123,7 +123,7 @@ public class SerializationContext {
    * getMemoizingContext, this method is not idempotent - the returned context will always be fresh.
    */
   public SerializationContext getNewMemoizingContext() {
-    return getNewMemoizingContext(/*allowFuturesToBlockWritingOn=*/ false);
+    return getNewMemoizingContext(allowFuturesToBlockWritingOn);
   }
 
   private SerializationContext getNewMemoizingContext(boolean allowFuturesToBlockWritingOn) {