Refactor NestedSet to permit alternate NestedSetStore implementations.

PiperOrigin-RevId: 194787067

4 files changed, 180 insertions, 31 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 a1a5bb259d..725c8c8a73 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
@@ -23,6 +23,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization:constants",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec",
+        "//third_party:auto_value",
         "//third_party:guava",
         "//third_party:jsr305",
         "//third_party/protobuf:protobuf_java",
@@ -30,6 +31,19 @@ java_library(
 )
 
 java_library(
+    name = "testutils",
+    testonly = 1,
+    srcs = ["NestedSetCodecTestUtils.java"],
+    deps = [
+        ":nestedset",
+        "//src/main/java/com/google/devtools/build/lib/skyframe/serialization",
+        "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils",
+        "//third_party:guava",
+        "//third_party:truth",
+    ],
+)
+
+java_library(
     name = "fingerprint_cache",
     srcs = [
         "DigestMap.java",
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
new file mode 100644
index 0000000000..89faad4a8b
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecTestUtils.java
@@ -0,0 +1,89 @@
+// Copyright 2018 The Bazel Authors. 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.devtools.build.lib.collect.nestedset;
+
+import static com.google.common.truth.Truth.assertThat;
+
+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.testutils.SerializationTester;
+import java.io.IOException;
+
+/** Utilities for testing NestedSet serialization. */
+public class NestedSetCodecTestUtils {
+
+  private static final NestedSet<String> SHARED_NESTED_SET =
+      NestedSetBuilder.<String>stableOrder().add("e").build();
+
+  /** Perform serialization/deserialization checks for several simple NestedSet examples. */
+  public static void checkCodec(ObjectCodecs objectCodecs) throws Exception {
+    new SerializationTester(
+            NestedSetBuilder.emptySet(Order.STABLE_ORDER),
+            NestedSetBuilder.emptySet(Order.NAIVE_LINK_ORDER),
+            NestedSetBuilder.create(Order.STABLE_ORDER, "a"),
+            NestedSetBuilder.create(Order.STABLE_ORDER, "a", "b", "c"),
+            NestedSetBuilder.<String>stableOrder()
+                .add("a")
+                .add("b")
+                .addTransitive(
+                    NestedSetBuilder.<String>stableOrder()
+                        .add("c")
+                        .addTransitive(SHARED_NESTED_SET)
+                        .build())
+                .addTransitive(
+                    NestedSetBuilder.<String>stableOrder()
+                        .add("d")
+                        .addTransitive(SHARED_NESTED_SET)
+                        .build())
+                .addTransitive(NestedSetBuilder.emptySet(Order.STABLE_ORDER))
+                .build())
+        .setObjectCodecs(objectCodecs)
+        .setVerificationFunction(NestedSetCodecTestUtils::verifyDeserialization)
+        .runTests();
+  }
+
+  public static ListenableFuture<Void> writeToStoreFuture(
+      NestedSetStore store, NestedSet<?> nestedSet, SerializationContext serializationContext)
+      throws IOException, SerializationException {
+    return store
+        .computeFingerprintAndStore((Object[]) nestedSet.rawChildren(), serializationContext)
+        .writeStatus();
+  }
+
+  private static void verifyDeserialization(
+      NestedSet<String> subject, NestedSet<String> deserialized) {
+    assertThat(subject.getOrder()).isEqualTo(deserialized.getOrder());
+    assertThat(subject.toSet()).isEqualTo(deserialized.toSet());
+    verifyStructure(subject.rawChildren(), deserialized.rawChildren());
+  }
+
+  private static void verifyStructure(Object lhs, Object rhs) {
+    if (lhs == NestedSet.EMPTY_CHILDREN) {
+      assertThat(rhs).isSameAs(NestedSet.EMPTY_CHILDREN);
+    } else if (lhs instanceof Object[]) {
+      assertThat(rhs).isInstanceOf(Object[].class);
+      Object[] lhsArray = (Object[]) lhs;
+      Object[] rhsArray = (Object[]) rhs;
+      int n = lhsArray.length;
+      assertThat(rhsArray).hasLength(n);
+      for (int i = 0; i < n; ++i) {
+        verifyStructure(lhsArray[i], rhsArray[i]);
+      }
+    } else {
+      assertThat(lhs).isEqualTo(rhs);
+    }
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecWithStore.java b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecWithStore.java
index 0923404ee7..285bd6f230 100644
--- a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecWithStore.java
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecWithStore.java
@@ -13,6 +13,7 @@
 // limitations under the License.
 package com.google.devtools.build.lib.collect.nestedset;
 
+import com.google.devtools.build.lib.collect.nestedset.NestedSetStore.FingerprintComputationResult;
 import com.google.devtools.build.lib.skyframe.serialization.DeserializationContext;
 import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec;
 import com.google.devtools.build.lib.skyframe.serialization.SerializationConstants;
@@ -70,9 +71,9 @@ public class NestedSetCodecWithStore<T> implements ObjectCodec<NestedSet<T>> {
       context.serialize(obj.rawChildren(), codedOut);
     } else {
       context.serialize(NestedSetSize.GROUP, codedOut);
-      ByteString fingerprint =
+      FingerprintComputationResult fingerprintComputationResult =
           nestedSetStore.computeFingerprintAndStore((Object[]) obj.rawChildren(), context);
-      codedOut.writeByteArrayNoTag(fingerprint.toByteArray());
+      codedOut.writeByteArrayNoTag(fingerprintComputationResult.fingerprint().toByteArray());
     }
   }
 
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 7377daa336..fd96cc4a59 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
@@ -13,9 +13,14 @@
 // limitations under the License.
 package com.google.devtools.build.lib.collect.nestedset;
 
-import com.google.common.base.Preconditions;
+import com.google.auto.value.AutoValue;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.MapMaker;
 import com.google.common.hash.Hashing;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.MoreExecutors;
 import com.google.devtools.build.lib.skyframe.serialization.DeserializationContext;
 import com.google.devtools.build.lib.skyframe.serialization.SerializationConstants;
 import com.google.devtools.build.lib.skyframe.serialization.SerializationContext;
@@ -54,16 +59,18 @@ import javax.annotation.Nullable;
  * recursively retrieving B using its fingerprint.
  */
 public class NestedSetStore {
-
   /** Stores fingerprint -> NestedSet associations. */
-  interface NestedSetStorageEndpoint {
-    /** Associates a fingerprint with the serialized representation of some NestedSet contents. */
-    void put(ByteString fingerprint, byte[] serializedBytes);
+  public interface NestedSetStorageEndpoint {
+    /**
+     * Associates a fingerprint with the serialized representation of some NestedSet contents.
+     * Returns a future that completes when the write completes.
+     */
+    ListenableFuture<Void> put(ByteString fingerprint, byte[] serializedBytes) throws IOException;
 
     /**
      * Retrieves the serialized bytes for the NestedSet contents associated with this fingerprint.
      */
-    byte[] get(ByteString fingerprint);
+    byte[] get(ByteString fingerprint) throws IOException;
   }
 
   /** An in-memory {@link NestedSetStorageEndpoint} */
@@ -72,8 +79,9 @@ public class NestedSetStore {
         new ConcurrentHashMap<>();
 
     @Override
-    public void put(ByteString fingerprint, byte[] serializedBytes) {
+    public ListenableFuture<Void> put(ByteString fingerprint, byte[] serializedBytes) {
       fingerprintToContents.put(fingerprint, serializedBytes);
+      return Futures.immediateFuture(null);
     }
 
     @Override
@@ -91,7 +99,7 @@ public class NestedSetStore {
             .makeMap();
 
     /** Object/Object[] contents to fingerprint. Maintained for fast fingerprinting. */
-    private final Map<Object[], ByteString> contentsToFingerprint =
+    private final Map<Object[], FingerprintComputationResult> contentsToFingerprint =
         new MapMaker()
             .concurrencyLevel(SerializationConstants.DESERIALIZATION_POOL_SIZE)
             .weakKeys()
@@ -111,20 +119,44 @@ public class NestedSetStore {
      * contents are not known.
      */
     @Nullable
-    public ByteString fingerprintForContents(Object[] contents) {
+    public FingerprintComputationResult fingerprintForContents(Object[] contents) {
       return contentsToFingerprint.get(contents);
     }
 
     /** Associates the provided fingerprint and NestedSet contents. */
-    public void put(ByteString fingerprint, Object[] contents) {
-      contentsToFingerprint.put(contents, fingerprint);
-      fingerprintToContents.put(fingerprint, contents);
+    public void put(FingerprintComputationResult fingerprintComputationResult, Object[] contents) {
+      contentsToFingerprint.put(contents, fingerprintComputationResult);
+      fingerprintToContents.put(fingerprintComputationResult.fingerprint(), contents);
+    }
+  }
+
+  /** The result of a fingerprint computation, including the status of its storage. */
+  @VisibleForTesting
+  @AutoValue
+  public abstract static class FingerprintComputationResult {
+    static FingerprintComputationResult create(
+        ByteString fingerprint, ListenableFuture<Void> writeStatus) {
+      return new AutoValue_NestedSetStore_FingerprintComputationResult(fingerprint, writeStatus);
     }
+
+    abstract ByteString fingerprint();
+
+    @VisibleForTesting
+    public abstract ListenableFuture<Void> writeStatus();
   }
 
   private final NestedSetCache nestedSetCache = new NestedSetCache();
-  private final NestedSetStorageEndpoint nestedSetStorageEndpoint =
-      new InMemoryNestedSetStorageEndpoint();
+  private final NestedSetStorageEndpoint nestedSetStorageEndpoint;
+
+  /** Creates a NestedSetStore with the provided {@link NestedSetStorageEndpoint} as a backend. */
+  public NestedSetStore(NestedSetStorageEndpoint nestedSetStorageEndpoint) {
+    this.nestedSetStorageEndpoint = nestedSetStorageEndpoint;
+  }
+
+  /** Creates a NestedSetStore with an in-memory storage backend. */
+  public static NestedSetStore inMemory() {
+    return new NestedSetStore(new InMemoryNestedSetStorageEndpoint());
+  }
 
   /**
    * Computes and returns the fingerprint for the given NestedSet contents using the given {@link
@@ -132,9 +164,11 @@ public class NestedSetStore {
    * store. Recursively does the same for all transitive members (i.e. Object[] members) of the
    * provided contents.
    */
-  ByteString computeFingerprintAndStore(
-      Object[] contents, SerializationContext serializationContext) throws SerializationException {
-    ByteString priorFingerprint = nestedSetCache.fingerprintForContents(contents);
+  @VisibleForTesting
+  public FingerprintComputationResult computeFingerprintAndStore(
+      Object[] contents, SerializationContext serializationContext)
+      throws SerializationException, IOException {
+    FingerprintComputationResult priorFingerprint = nestedSetCache.fingerprintForContents(contents);
     if (priorFingerprint != null) {
       return priorFingerprint;
     }
@@ -149,13 +183,16 @@ public class NestedSetStore {
     ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
     CodedOutputStream codedOutputStream = CodedOutputStream.newInstance(byteArrayOutputStream);
 
+    ImmutableList.Builder<ListenableFuture<Void>> futureBuilder = ImmutableList.builder();
     try {
       codedOutputStream.writeInt32NoTag(contents.length);
       for (Object child : contents) {
         if (child instanceof Object[]) {
-          ByteString fingerprint =
+          FingerprintComputationResult fingerprintComputationResult =
               computeFingerprintAndStore((Object[]) child, serializationContext);
-          newSerializationContext.serialize(fingerprint, codedOutputStream);
+          futureBuilder.add(fingerprintComputationResult.writeStatus());
+          newSerializationContext.serialize(
+              fingerprintComputationResult.fingerprint(), codedOutputStream);
         } else {
           newSerializationContext.serialize(child, codedOutputStream);
         }
@@ -168,27 +205,33 @@ public class NestedSetStore {
     byte[] serializedBytes = byteArrayOutputStream.toByteArray();
     ByteString fingerprint =
         ByteString.copyFrom(Hashing.md5().hashBytes(serializedBytes).asBytes());
+    futureBuilder.add(nestedSetStorageEndpoint.put(fingerprint, serializedBytes));
 
-    nestedSetCache.put(fingerprint, contents);
-    nestedSetStorageEndpoint.put(fingerprint, serializedBytes);
+    ListenableFuture<Void> writeFuture =
+        Futures.whenAllComplete(futureBuilder.build())
+            .call(() -> null, MoreExecutors.directExecutor());
+    FingerprintComputationResult fingerprintComputationResult =
+        FingerprintComputationResult.create(fingerprint, writeFuture);
 
-    return fingerprint;
+    nestedSetCache.put(fingerprintComputationResult, contents);
+
+    return fingerprintComputationResult;
   }
 
   /** Retrieves and deserializes the NestedSet contents associated with the given fingerprint. */
   public Object[] getContentsAndDeserialize(
       ByteString fingerprint, DeserializationContext deserializationContext)
-      throws IOException, SerializationException {
+      throws SerializationException, IOException {
     Object[] contents = nestedSetCache.contentsForFingerprint(fingerprint);
     if (contents != null) {
       return contents;
     }
 
-    byte[] retrieved =
-        Preconditions.checkNotNull(
-            nestedSetStorageEndpoint.get(fingerprint),
-            "Fingerprint %s not found in NestedSetStore",
-            fingerprint);
+    byte[] retrieved = nestedSetStorageEndpoint.get(fingerprint);
+    if (retrieved == null) {
+      throw new AssertionError("Fingerprint " + fingerprint + " not found in NestedSetStore");
+    }
+
     CodedInputStream codedIn = CodedInputStream.newInstance(retrieved);
     DeserializationContext newDeserializationContext =
         deserializationContext.getNewMemoizingContext();
@@ -204,7 +247,9 @@ public class NestedSetStore {
               : deserializedElement;
     }
 
-    nestedSetCache.put(fingerprint, dereferencedContents);
+    FingerprintComputationResult fingerprintComputationResult =
+        FingerprintComputationResult.create(fingerprint, Futures.immediateFuture(null));
+    nestedSetCache.put(fingerprintComputationResult, dereferencedContents);
     return dereferencedContents;
   }
 }