Add Collection#upsertAndReturnOlderDocument for atomic pre-write snap…

document-store/src/main/java/org/hypertrace/core/documentstore/Collection.java

@@ -11,9 +11,16 @@
 import org.hypertrace.core.documentstore.model.exception.DuplicateDocumentException;
 import org.hypertrace.core.documentstore.model.options.QueryOptions;
 import org.hypertrace.core.documentstore.model.options.UpdateOptions;
+import org.hypertrace.core.documentstore.model.options.UpsertOptions;
 import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate;
 
-/** Interface spec for common operations on a collection of documents */
+/**
+ * Interface spec for common operations on a collection of documents.
+ *
+ * <p>Methods are not atomic, and give no isolation from concurrent writers, unless their javadoc
+ * explicitly says so. In particular, reading a document before or after a write with separate calls
+ * can observe other writers' changes (or stale replicas, depending on read preference).
+ */
 public interface Collection {
 
   /**
@@ -57,6 +64,26 @@ public interface Collection {
   @Deprecated(forRemoval = true)
   Document upsertAndReturn(Key key, Document document) throws IOException;
 
+  /**
+   * Atomically perform the same merge-upsert as {@link #upsert(Key, Document)} and return the
+   * document images selected by {@link UpsertOptions#getReturnDocuments()}.
+   *
+   * <p>Whichever images are requested come from the same atomic write, even under concurrent
+   * writers on the same key: the before-image is exactly the document this write replaced and the
+   * after-image is exactly the document this write produced.
+   *
+   * @param key Unique key of the document in the collection.
+   * @param document Document to be upserted.
+   * @param options Selects which images (BEFORE, AFTER or BOTH) to return.
+   * @return {@link DocumentChange#getBefore()} is empty if not requested or if this call created
+   *     the document; {@link DocumentChange#getAfter()} is empty if not requested.
+   * @throws IOException If the operation could not be performed
+   */
+  default DocumentChange upsertAndReturn(Key key, Document document, UpsertOptions options)
+      throws IOException {
+    throw new UnsupportedOperationException("upsertAndReturn with UpsertOptions is not supported");
+  }
+
   /**
    * Update a sub document
    *
@@ -307,6 +334,27 @@ default CloseableIterator<Document> bulkCreateOrReplaceReturnOlderDocuments(
    */
   Document createOrReplaceAndReturn(final Key key, final Document document) throws IOException;
 
+  /**
+   * Atomically perform the same full-replace write as {@link #createOrReplace(Key, Document)} and
+   * return the document images selected by {@link UpsertOptions#getReturnDocuments()}.
+   *
+   * <p>Whichever images are requested come from the same atomic write, even under concurrent
+   * writers on the same key: the before-image is exactly the document this write replaced and the
+   * after-image is exactly the document this write produced.
+   *
+   * @param key Unique key of the document in the collection.
+   * @param document The document to be created/replaced
+   * @param options Selects which images (BEFORE, AFTER or BOTH) to return.
+   * @return {@link DocumentChange#getBefore()} is empty if not requested or if this call created
+   *     the document; {@link DocumentChange#getAfter()} is empty if not requested.
+   * @throws IOException If the operation could not be performed
+   */
+  default DocumentChange createOrReplaceAndReturn(
+      final Key key, final Document document, final UpsertOptions options) throws IOException {
+    throw new UnsupportedOperationException(
+        "createOrReplaceAndReturn with UpsertOptions is not supported");
+  }
+
   /**
    * Updates existing documents if the corresponding Filter condition evaluates to true
    *

document-store/src/main/java/org/hypertrace/core/documentstore/DocumentChange.java

@@ -0,0 +1,61 @@
+package org.hypertrace.core.documentstore;
+
+import java.util.Objects;
+import java.util.Optional;
+import javax.annotation.Nullable;
+
+/**
+ * The document images returned by a single atomic write. Which images are present depends on the
+ * requested {@link org.hypertrace.core.documentstore.model.options.UpsertOptions.ReturnDocuments}.
+ */
+public final class DocumentChange {
+
+  @Nullable private final Document before;
+  @Nullable private final Document after;
+
+  private DocumentChange(@Nullable final Document before, @Nullable final Document after) {
+    this.before = before;
+    this.after = after;
+  }
+
+  /**
+   * @param before the document before the write; null if not requested or if the write created the
+   *     document
+   * @param after the document after the write; null if not requested
+   */
+  public static DocumentChange of(@Nullable final Document before, @Nullable final Document after) {
+    return new DocumentChange(before, after);
+  }
+
+  /** Empty when the before-image was not requested, or when the write created the document. */
+  public Optional<Document> getBefore() {
+    return Optional.ofNullable(before);
+  }
+
+  /** Empty when the after-image was not requested. */
+  public Optional<Document> getAfter() {
+    return Optional.ofNullable(after);
+  }
+
+  @Override
+  public boolean equals(final Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    DocumentChange that = (DocumentChange) o;
+    return Objects.equals(before, that.before) && Objects.equals(after, that.after);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(before, after);
+  }
+
+  @Override
+  public String toString() {
+    return "DocumentChange{before=" + before + ", after=" + after + '}';
+  }
+}

document-store/src/main/java/org/hypertrace/core/documentstore/model/options/UpsertOptions.java

@@ -0,0 +1,20 @@
+package org.hypertrace.core.documentstore.model.options;
+
+import lombok.Builder;
+import lombok.Builder.Default;
+import lombok.Value;
+
+@Value
+@Builder
+public class UpsertOptions {
+  public static UpsertOptions DEFAULT_UPSERT_OPTIONS =
+      UpsertOptions.builder().returnDocuments(ReturnDocuments.AFTER).build();
+
+  @Default ReturnDocuments returnDocuments = ReturnDocuments.AFTER;
+
+  public enum ReturnDocuments {
+    BEFORE,
+    AFTER,
+    BOTH,
+  }
+}

document-store/src/main/java/org/hypertrace/core/documentstore/mongo/MongoCollection.java

@@ -42,6 +42,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
+import javax.annotation.Nullable;
 import net.jodah.failsafe.Failsafe;
 import net.jodah.failsafe.RetryPolicy;
 import org.bson.BsonInt64;
@@ -56,12 +57,15 @@
 import org.hypertrace.core.documentstore.Collection;
 import org.hypertrace.core.documentstore.CreateResult;
 import org.hypertrace.core.documentstore.Document;
+import org.hypertrace.core.documentstore.DocumentChange;
 import org.hypertrace.core.documentstore.Filter;
 import org.hypertrace.core.documentstore.Key;
 import org.hypertrace.core.documentstore.Query;
 import org.hypertrace.core.documentstore.model.config.ConnectionConfig;
 import org.hypertrace.core.documentstore.model.exception.DuplicateDocumentException;
 import org.hypertrace.core.documentstore.model.options.QueryOptions;
+import org.hypertrace.core.documentstore.model.options.UpsertOptions;
+import org.hypertrace.core.documentstore.model.options.UpsertOptions.ReturnDocuments;
 import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate;
 import org.hypertrace.core.documentstore.mongo.query.MongoQueryExecutor;
 import org.hypertrace.core.documentstore.mongo.update.MongoUpdateExecutor;
@@ -300,6 +304,148 @@ public Document upsertAndReturn(Key key, Document document) throws IOException {
     return dbObjectToDocument(upsertResult);
   }
 
+  /**
+   * Performs one atomic server-side write (a single {@code findOneAndUpdate} with upsert).
+   *
+   * <ul>
+   *   <li>BEFORE / AFTER: the server returns the requested image from that write; the update is the
+   *       same as {@link #upsert(Key, Document)}.
+   *   <li>BOTH: the server returns the before-image from that write; the after-image is computed by
+   *       applying that same update to the before-image (no second read). This is exact because the
+   *       update only sets top-level fields to literal values. For that reason {@code
+   *       _lastUpdateTime} is set from the client clock instead of {@code $currentDate} in this
+   *       mode only.
+   * </ul>
+   */
+  @Override
+  public DocumentChange upsertAndReturn(Key key, Document document, UpsertOptions options)
+      throws IOException {
+    final ReturnDocuments returnDocuments = options.getReturnDocuments();
+    final BasicDBObject update =
+        returnDocuments == ReturnDocuments.BOTH
+            ? prepareUpsertWithExplicitLastUpdateTime(key, document, System.currentTimeMillis())
+            : prepareUpsert(key, document);
+    final ReturnDocument serverImage =
+        returnDocuments == ReturnDocuments.AFTER ? ReturnDocument.AFTER : ReturnDocument.BEFORE;
+    final BasicDBObject serverResult =
+        Failsafe.with(upsertRetryPolicy)
+            .get(
+                () ->
+                    collection.findOneAndUpdate(
+                        this.selectionCriteriaForKey(key),
+                        update,
+                        new FindOneAndUpdateOptions().upsert(true).returnDocument(serverImage)));
+
+    switch (returnDocuments) {
+      case BEFORE:
+        return DocumentChange.of(toDocumentOrNull(serverResult), null);
+      case AFTER:
+        if (serverResult == null) {
+          throw new IOException("Could not upsert the document with key: " + key);
+        }
+        return DocumentChange.of(null, dbObjectToDocument(serverResult));
+      case BOTH:
+        final BasicDBObject serverBefore = serverResult;
+        final BasicDBObject derivedAfter = applyUpsertToBeforeImage(key, serverBefore, update);
+        return DocumentChange.of(toDocumentOrNull(serverBefore), dbObjectToDocument(derivedAfter));
+      default:
+        throw new UnsupportedOperationException("Unhandled return documents: " + returnDocuments);
+    }
+  }
+
+  /**
+   * Performs one atomic server-side write (a single {@code findOneAndUpdate} with upsert).
+   *
+   * <ul>
+   *   <li>BEFORE / AFTER: the server returns the requested image from that write.
+   *   <li>BOTH: the server returns the before-image from that write; the after-image is computed by
+   *       applying the same replace to the before-image (no second read).
+   * </ul>
+   *
+   * <p>Unlike {@link #createOrReplace(Key, Document)}, every top-level value is wrapped in {@code
+   * $literal}, so values are stored exactly as given (e.g. a string starting with {@code $} is not
+   * read as a field path). That is what makes the derived after-image exact.
+   */
+  @Override
+  public DocumentChange createOrReplaceAndReturn(
+      final Key key, final Document document, final UpsertOptions options) throws IOException {
+    final ReturnDocuments returnDocuments = options.getReturnDocuments();
+    final long now = System.currentTimeMillis();
+    final BasicDBObject replacement = prepareDocumentForLiteralReplace(key, document, now);
+    final ReturnDocument serverImage =
+        returnDocuments == ReturnDocuments.AFTER ? ReturnDocument.AFTER : ReturnDocument.BEFORE;
+
+    try {
+      final BasicDBObject serverResult =
+          collection.findOneAndUpdate(
+              this.selectionCriteriaForKey(key),
+              prepareLiteralCreateOrReplace(replacement, now),
+              new FindOneAndUpdateOptions().upsert(true).returnDocument(serverImage));
+      LOGGER.debug("Create or replace result: {}", serverResult);
+
+      switch (returnDocuments) {
+        case BEFORE:
+          return DocumentChange.of(toDocumentOrNull(serverResult), null);
+        case AFTER:
+          return DocumentChange.of(null, dbObjectToDocument(requireNonNull(serverResult)));
+        case BOTH:
+          final BasicDBObject serverBefore = serverResult;
+          final BasicDBObject derivedAfter =
+              applyReplaceToBeforeImage(serverBefore, replacement, now);
+          return DocumentChange.of(
+              toDocumentOrNull(serverBefore), dbObjectToDocument(derivedAfter));
+        default:
+          throw new UnsupportedOperationException("Unhandled return documents: " + returnDocuments);
+      }
+    } catch (final Exception e) {
+      LOGGER.error("Exception creating/replacing document. key: {} content:{}", key, document, e);
+      throw e;
+    }
+  }
+
+  private BasicDBObject prepareDocumentForLiteralReplace(Key key, Document document, long now)
+      throws JsonProcessingException {
+    final BasicDBObject replacement = prepareDocument(key, document, now);
+    replacement.put(LAST_UPDATE_TIMESTAMP_ISO_8601, new Date(now));
+    return replacement;
+  }
+
+  /** Same stages as {@link #prepareForCreateOrReplace}, with every value wrapped in $literal. */
+  private List<BasicDBObject> prepareLiteralCreateOrReplace(BasicDBObject replacement, long now) {
+    final BasicDBObject project =
+        new BasicDBObject(
+            PROJECT_CLAUSE,
+            new BasicDBObject(
+                CREATED_TIME,
+                new BasicDBObject(
+                    IF_NULL_CLAUSE,
+                    new BsonValue[] {new BsonString(PREFIX + CREATED_TIME), new BsonInt64(now)})));
+    final BasicDBObject literalSet = new BasicDBObject();
+    replacement.forEach(
+        (field, value) -> literalSet.put(field, new BasicDBObject(PREFIX + "literal", value)));
+    return List.of(project, new BasicDBObject(SET_CLAUSE, literalSet));
+  }
+
+  /**
+   * Applies {@link #prepareLiteralCreateOrReplace} to the before-image the server returned for that
+   * same write: {@code $project} keeps only {@code _id} and {@code createdTime} (defaulting to
+   * {@code now} when missing or null), then {@code $set} adds the replacement fields.
+   */
+  private static BasicDBObject applyReplaceToBeforeImage(
+      @Nullable BasicDBObject before, BasicDBObject replacement, long now) {
+    final BasicDBObject after =
+        new BasicDBObject(ID_KEY, before == null ? replacement.get(ID_KEY) : before.get(ID_KEY));
+    final Object createdTime = before == null ? null : before.get(CREATED_TIME);
+    after.put(CREATED_TIME, createdTime == null ? now : createdTime);
+    after.putAll(replacement.toMap());
+    return after;
+  }
+
+  @Nullable
+  private static Document toDocumentOrNull(@Nullable final BasicDBObject dbObject) {
+    return dbObject == null ? null : dbObjectToDocument(dbObject);
+  }
+
   @SuppressWarnings("MismatchedQueryAndUpdateOfCollection")
   private List<BasicDBObject> prepareForCreateOrReplace(final Key key, final Document document)
       throws JsonProcessingException {
@@ -328,6 +474,37 @@ private BasicDBObject prepareUpsert(Key key, Document document) throws JsonProce
         .append("$setOnInsert", new BasicDBObject(CREATED_TIME, now));
   }
 
+  /**
+   * Same as {@link #prepareUpsert} except {@code _lastUpdateTime} is a literal client-clock value,
+   * so the stored document is fully determined by the before-image and this update.
+   */
+  private BasicDBObject prepareUpsertWithExplicitLastUpdateTime(
+      Key key, Document document, long now) throws JsonProcessingException {
+    BasicDBObject setObject = prepareDocument(key, document, now);
+    setObject.put(LAST_UPDATE_TIME, new Date(now));
+    return new BasicDBObject("$set", setObject)
+        .append("$setOnInsert", new BasicDBObject(CREATED_TIME, now));
+  }
+
+  /**
+   * Applies an update from {@link #prepareUpsertWithExplicitLastUpdateTime} to the before-image the
+   * server returned for that same write, yielding the stored document. Exact because the update
+   * only sets top-level fields to literal values ({@code $setOnInsert} applies only on insert).
+   */
+  @SuppressWarnings("unchecked")
+  private static BasicDBObject applyUpsertToBeforeImage(
+      Key key, @Nullable BasicDBObject before, BasicDBObject update) {
+    final BasicDBObject after =
+        before == null
+            ? new BasicDBObject(ID_KEY, key.toString())
+            : new BasicDBObject(before.toMap());
+    after.putAll((Map<String, Object>) update.get("$set"));
+    if (before == null) {
+      after.putAll((Map<String, Object>) update.get("$setOnInsert"));
+    }
+    return after;
+  }
+
   private BasicDBObject prepareInsert(Key key, Document document) throws JsonProcessingException {
     long now = System.currentTimeMillis();
     BasicDBObject insertDbObject = prepareDocument(key, document, now);