firestore: update ObjectValue.buildProto to return the memoized proto (#7370)

Author: MarkDuckworth

firebase-firestore/CHANGELOG.md

@@ -6,6 +6,8 @@
 - [changed] Improve query performance in large result sets by replacing the deprecated AsyncTask
   thread pool with a self-managed thread pool.
   [#7376](//github.com/firebase/firebase-android-sdk/issues/7376)
+- [changed] Improve query performance via internal memoization of calculated document data.
+  [#7370](//github.com/firebase/firebase-android-sdk/issues/7370)
 
 # 26.0.0
 

firebase-firestore/src/main/java/com/google/firebase/firestore/model/ObjectValue.java

@@ -16,6 +16,7 @@
 
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 
+import androidx.annotation.GuardedBy;
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import com.google.firebase.firestore.model.mutation.FieldMask;
@@ -28,17 +29,39 @@
 
 /** A structured object value stored in Firestore. */
 public final class ObjectValue implements Cloneable {
+
+  private final Object lock = new Object();
+
+  /**
+   * The immutable Value proto for this object with all overlays applied.
+   * <p>
+   * The value for this member is calculated _lazily_ and is null if the overlays have not yet been
+   * applied.
+   * <p>
+   * This member MAY be READ concurrently from multiple threads without acquiring any particular
+   * locks; however, UPDATING it MUST have the `lock` lock held.
+   * <p>
+   * Internal Invariant: Exactly one of `mergedValue` and `partialValue` must be null, with the
+   * other being non-null.
+   */
+  @Nullable private volatile Value mergedValue;
+
   /**
    * The immutable Value proto for this object. Local mutations are stored in `overlayMap` and only
    * applied when {@link #buildProto()} is invoked.
+   * <p>
+   * Internal Invariant: Exactly one of `mergedValue` and `partialValue` must be null, with the
+   * other being non-null.
    */
+  @GuardedBy("lock")
   private Value partialValue;
 
   /**
    * A nested map that contains the accumulated changes that haven't yet been applied to {@link
    * #partialValue}. Values can either be {@link Value} protos, {@code Map<String, Object>} values
    * (to represent additional nesting) or {@code null} (to represent field deletes).
    */
+  @GuardedBy("lock")
   private final Map<String, Object> overlayMap = new HashMap<>();
 
   public static ObjectValue fromMap(Map<String, Value> value) {
@@ -53,7 +76,7 @@ public ObjectValue(Value value) {
     hardAssert(
         !ServerTimestamps.isServerTimestamp(value),
         "ServerTimestamps should not be used as an ObjectValue");
-    this.partialValue = value;
+    this.mergedValue = value;
   }
 
   public ObjectValue() {
@@ -103,7 +126,7 @@ private FieldMask extractFieldMask(MapValue value) {
   }
 
   @Nullable
-  private Value extractNestedValue(Value value, FieldPath fieldPath) {
+  private static Value extractNestedValue(Value value, FieldPath fieldPath) {
     if (fieldPath.isEmpty()) {
       return value;
     } else {
@@ -121,17 +144,40 @@ private Value extractNestedValue(Value value, FieldPath fieldPath) {
    * Returns the Protobuf that backs this ObjectValue.
    *
    * <p>This method applies any outstanding modifications and memoizes the result. Further
-   * invocations are based on this memoized result.
+   * invocations are based on this memoized result until overlays are applied, at which point the
+   * memoized result is marked as "stale" and a new result is calculated and memoized upon the next
+   * invocation of this method.
    */
   private Value buildProto() {
-    synchronized (overlayMap) {
-      MapValue mergedResult = applyOverlay(FieldPath.EMPTY_PATH, overlayMap);
-      if (mergedResult != null) {
-        partialValue = Value.newBuilder().setMapValue(mergedResult).build();
-        overlayMap.clear();
+    // Use double-checked locking to avoid acquiring a lock in the cases where the memoized result
+    // has already been calculated (https://en.wikipedia.org/wiki/Double-checked_locking).
+    Value value = this.mergedValue;
+
+    if (value == null) {
+      synchronized (lock) {
+        value = mergedValue;
+        if (value == null) {
+          assert (partialValue != null);
+          if (overlayMap.isEmpty()) {
+            value = partialValue;
+          } else {
+            MapValue mergedResult = applyOverlay(partialValue, FieldPath.EMPTY_PATH, overlayMap);
+            if (mergedResult == null) {
+              value = partialValue;
+            } else {
+              value = Value.newBuilder().setMapValue(mergedResult).build();
+            }
+          }
+
+          assert (value != null);
+          mergedValue = value;
+          partialValue = null;
+          overlayMap.clear();
+        }
       }
     }
-    return partialValue;
+
+    return value;
   }
 
   /**
@@ -171,6 +217,17 @@ public void setAll(Map<FieldPath, Value> data) {
    * Adds {@code value} to the overlay map at {@code path}. Creates nested map entries if needed.
    */
   private void setOverlay(FieldPath path, @Nullable Value value) {
+    synchronized (lock) {
+      if (mergedValue != null) {
+        partialValue = mergedValue;
+        mergedValue = null;
+      }
+      setOverlay(overlayMap, path, value);
+    }
+  }
+
+  private static void setOverlay(
+      Map<String, Object> overlayMap, FieldPath path, @Nullable Value value) {
     Map<String, Object> currentLevel = overlayMap;
 
     for (int i = 0; i < path.length() - 1; ++i) {
@@ -208,8 +265,8 @@ private void setOverlay(FieldPath path, @Nullable Value value) {
    *     overlayMap}.
    * @return The merged data at `currentPath` or null if no modifications were applied.
    */
-  private @Nullable MapValue applyOverlay(
-      FieldPath currentPath, Map<String, Object> currentOverlays) {
+  private static @Nullable MapValue applyOverlay(
+      Value partialValue, FieldPath currentPath, Map<String, Object> currentOverlays) {
     boolean modified = false;
 
     @Nullable Value existingValue = extractNestedValue(partialValue, currentPath);
@@ -227,7 +284,8 @@ private void setOverlay(FieldPath path, @Nullable Value value) {
       if (value instanceof Map) {
         @Nullable
         MapValue nested =
-            applyOverlay(currentPath.append(pathSegment), (Map<String, Object>) value);
+            applyOverlay(
+                partialValue, currentPath.append(pathSegment), (Map<String, Object>) value);
         if (nested != null) {
           resultAtPath.putFields(pathSegment, Value.newBuilder().setMapValue(nested).build());
           modified = true;