Repository navigation
Conversation
…shots Callers that need the previous version of a document in order to emit a change event currently have to issue a separate read before the write. That read is not atomic with respect to the write, so a concurrent upsert on the same key can land in between and the "previous" snapshot ends up already holding the new value. Add a merge-upsert that hands back the pre-write document as part of the write itself. Mongo uses findOneAndUpdate with ReturnDocument.BEFORE; both Postgres collections do the read and the upsert inside a single transaction, with the read taking a row lock. An empty Optional means the document did not exist and this call created it. Co-authored-by: Cursor <cursoragent@cursor.com>
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #328 +/- ##
============================================
- Coverage 81.23% 81.12% -0.11%
- Complexity 1698 1736 +38
============================================
Files 243 245 +2
Lines 7780 7964 +184
Branches 771 798 +27
============================================
+ Hits 6320 6461 +141
- Misses 962 993 +31
- Partials 498 510 +12
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
| * @return the pre-write document, or {@link Optional#empty()} if the document did not exist (i.e. | ||
| * this call created it) | ||
| */ | ||
| Optional<Document> upsertAndReturnOlderDocument(Key key, Document document) throws IOException; |
There was a problem hiding this comment.
This'll break existing clients. Add a default impl throwing UnsupportedException
There was a problem hiding this comment.
Added a default upsertAndReturnOlderDocument implementation that throws UnsupportedOperationException, matching the other optional Collection APIs, so existing clients are not forced to implement the new method.
…oid a Postgres deadlock. Add a default UnsupportedOperationException implementation so existing Collection clients still compile, and reuse the locked before-image when preparing the JSON upsert so the follow-up createdTime read cannot block on FOR UPDATE. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
| } | ||
|
|
||
| @Override | ||
| public Optional<Document> upsertAndReturnOlderDocument(Key key, Document document) |
There was a problem hiding this comment.
Why do we need this method when we have the new one?
There was a problem hiding this comment.
We don't. Removed upsertAndReturnOlderDocument from Collection and all implementations in 46ae7db; it was never released. upsertAndReturn(key, doc, UpsertOptions) with BEFORE covers that use case.
| } | ||
|
|
||
| @Override | ||
| public Optional<Document> upsertAndReturnOlderDocument(Key key, Document document) |
There was a problem hiding this comment.
Same, why do we need this with the new method?
There was a problem hiding this comment.
Removed here too in 46ae7db, so PostgresCollection is back to what's on main. The nested-JSON Postgres collection doesn't implement the new UpsertOptions overloads (default UnsupportedOperationException); only Mongo and flat Postgres do, which is what entity-service uses.
| } | ||
|
|
||
| @Override | ||
| public DocumentChange upsertAndReturnChange(Key key, Document document) throws IOException { |
There was a problem hiding this comment.
Is this atomic? Because the method contract says atomic but it doesn't look like it.
There was a problem hiding this comment.
Fair point. The write was one atomic findOneAndUpdate, but the after-image was computed client-side and could drift from what was stored: _lastUpdateTime came from $currentDate on the server while the computed image used the client clock.
In 46ae7db:
BEFOREandAFTERare each a singlefindOneAndUpdateand return exactly the document the server hands back.BOTHis still a singlefindOneAndUpdate(ReturnDocument.BEFORE); the after-image is built by applying that same update to the server-returned before-image, with no second read. That's now exact because the update only sets top-level fields to literal values;_lastUpdateTimeis set explicitly in$setinstead of via$currentDate.- An integration test reads the document back after a
BOTHupsert (insert and update) and asserts it equals the returned after-image field for field. A 20-thread test asserts concurrent writes on one key form a single gap-free before → after chain. - The javadoc now states exactly this.
Returning both images from the server would need a Mongo transaction (replica set only), which I didn't want to make a requirement. createOrReplaceAndReturn with BOTH on Mongo throws UnsupportedOperationException, because the replace pipeline evaluates values as expressions, so its result can't be derived exactly.
| * The before and after images always reflect the same atomic write, even under concurrent writers | ||
| * on the same key. | ||
| */ | ||
| default DocumentChange createOrReplaceAndReturnChange(Key key, Document document) |
There was a problem hiding this comment.
I feel a better API contract would be to add UpsertOptions - And let the user control if he wants BEFORE, AFTER or BOTH.
There was a problem hiding this comment.
Agreed, done in 46ae7db. Added UpsertOptions (@Value @Builder, like UpdateOptions) with ReturnDocuments { BEFORE, AFTER, BOTH }, defaulting to AFTER, and replaced the *AndReturnChange methods with:
DocumentChange upsertAndReturn(Key key, Document document, UpsertOptions options)
DocumentChange createOrReplaceAndReturn(Key key, Document document, UpsertOptions options)Both are default methods throwing UnsupportedOperationException, so existing implementations aren't affected. DocumentChange.getBefore() is empty when not requested or when the write created the document; getAfter() is empty when not requested. On flat Postgres all three modes run in one transaction (SELECT ... FOR UPDATE on the real PK column, write, SELECT), retrying if two first inserts race.
…rtAndReturnOlderDocument Co-authored-by: Cursor <cursoragent@cursor.com>
|
|
||
| switch (returnDocuments) { | ||
| case BEFORE: | ||
| return DocumentChange.of(toDocumentOrNull(result), null); |
There was a problem hiding this comment.
What does result contain? Doesn't look like it returns the existing doc - Can you confirm?
There was a problem hiding this comment.
It's the server's response to that single findOneAndUpdate. For BEFORE and BOTH we request ReturnDocument.BEFORE, so it's the existing document, or null on insert. For AFTER we request ReturnDocument.AFTER, so it's the stored document. b96ad4a makes this explicit: it's named serverResult, and in the BOTH branch it's assigned to serverBefore before the after-image is derived as derivedAfter.
| if (result == null) { | ||
| throw new IOException("Could not upsert the document with key: " + key); | ||
| } | ||
| return DocumentChange.of(null, dbObjectToDocument(result)); |
There was a problem hiding this comment.
We're using the same result for both BEFORE and AFTER?
There was a problem hiding this comment.
One variable, but it never holds both images in one call. Which image the server returns depends on the requested mode (serverImage): BEFORE for BEFORE/BOTH, AFTER for AFTER. b96ad4a renames it to serverResult and adds explicit serverBefore / derivedAfter variables in the BOTH branch so this is clearer.
| } | ||
|
|
||
| /** Asserts the update uses a literal _lastUpdateTime and returns its $set clause. */ | ||
| private BasicDBObject assertExplicitLastUpdateTime() { |
There was a problem hiding this comment.
It's recommended to move private methods towards the end of the file.
There was a problem hiding this comment.
Done in b96ad4a: the private helpers in UpsertAndReturnWithOptionsTest are now at the end of the class.
| Key key, | ||
| Document document, | ||
| UpsertOptions options, | ||
| BiFunction<List<String>, String, String> sqlBuilder, |
There was a problem hiding this comment.
Is this param really needed? Can't it refer to buildUpsertSql directly?
There was a problem hiding this comment.
It's needed because writeAndReturnChange is shared by two operations with different SQL. upsertAndReturn passes buildUpsertSql, which merges (COALESCE(EXCLUDED.col, table.col)). createOrReplaceAndReturn passes buildCreateOrReplaceSql, which full-replaces (EXCLUDED.col). Everything else is the same: the before-image read, locking, parameter binding, RETURNING, and the schema-refresh retry. Passing the builder avoids duplicating that.
| final Query keyQuery = buildPrimaryKeyQuery(tableName, key); | ||
|
|
||
| try { | ||
| for (int attempt = 0; attempt < MAX_FIRST_INSERT_RACE_RETRIES; attempt++) { |
There was a problem hiding this comment.
Shall we leave the retries to the client rather?
There was a problem hiding this comment.
b96ad4a removes the retry loop. The race it covered is now prevented rather than retried: if SELECT ... FOR UPDATE finds no row (so it locks nothing), we take pg_advisory_xact_lock on the table and key and re-read. Concurrent first writers through these methods are serialised, so the second one sees the first one's row as its before-image. The 20-thread concurrent first-insert integration test still forms a single before → after chain.
One case remains: a writer that doesn't take the lock (e.g. plain upsert()) inserts the same key between our read and our write. Then we roll back and throw IOException without writing anything, and retrying is left to the client, as you suggested.
suresh-prakash
left a comment
There was a problem hiding this comment.
Summary
Reviewed 4 findings. Posted 4 inline comments.
| Severity | Count |
|---|---|
| High | 1 |
| Medium | 3 |
| Low | 0 |
Reviewed using @claude traceable:traceable-java-code-review
💡 Want this skill on your own PRs? Install the Traceable plugin from the Harness internal marketplace:
/plugin marketplace add https://git0.harness.io/l7B_kbSEQD2wjrM7PShm5w/PROD/Harness_Commons/claude-code-marketplace.git /plugin install traceable@harness-internal-claude-pluginsThen run
/skill traceable:traceable-java-code-review <pr-url>.
| try (Connection connection = client.getTransactionalConnection()) { | ||
| try { | ||
| final Document before = | ||
| readBefore ? selectFirstDocument(connection, keyQuery).orElse(null) : null; |
There was a problem hiding this comment.
[high] Postgres before/after images include _implicit_id
writeAndReturnChange reads both images through selectFirstDocument, and that query is buildSelectQueryForUpdate(), which appends id AS _implicit_id. A normal find() does not return that column. MongoPostgresWriteConsistencyTest strips _implicit_id before asserting the snapshot equals the stored document, so a change event built from getBefore() or getAfter() will not match a later read. Read the row with the same projection as find(), or drop _implicit_id before wrapping the Document.
Rule: snapshot-field
Reviewed using @claude traceable:traceable-java-code-review
There was a problem hiding this comment.
Fixed in b96ad4a. The after-image now comes from the write itself (INSERT ... ON CONFLICT ... RETURNING *), which has the same columns as find(). _implicit_id is dropped from the SELECT ... FOR UPDATE before-image. The integration test no longer strips _implicit_id: it asserts both images equal what find() returns.
| * 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( |
There was a problem hiding this comment.
[medium] New Mongo upsert stores client _lastUpdateTime instead of $currentDate
upsertAndReturn(Key, Document, UpsertOptions) is documented as the same merge-upsert as upsert(). prepareUpsert sets _lastUpdateTime with $currentDate (Mongo server clock). prepareUpsertWithExplicitLastUpdateTime stores new Date(now) from the client so the BOTH after-image can be derived locally. Callers that switch to this method change the stored field, and mixed with rows still written by upsert(), _lastUpdateTime can move backwards when the client clock and the server clock disagree.
Rule: upsert-semantics
Reviewed using @claude traceable:traceable-java-code-review
There was a problem hiding this comment.
Addressed in b96ad4a. BEFORE and AFTER now send exactly the same update as upsert() (prepareUpsert, with $currentDate). Only BOTH sets _lastUpdateTime from the client clock, because that's what lets the after-image be derived exactly from the server-returned before-image without a second read. The javadoc states this, and unit tests assert $currentDate for BEFORE/AFTER and the literal value for BOTH.
| final Query keyQuery = buildPrimaryKeyQuery(tableName, key); | ||
|
|
||
| try { | ||
| for (int attempt = 0; attempt < MAX_FIRST_INSERT_RACE_RETRIES; attempt++) { |
There was a problem hiding this comment.
[medium] New Postgres write holds FOR UPDATE across extra reads
writeAndReturnChange runs SELECT … FOR UPDATE, the upsert, and a second SELECT … FOR UPDATE when an after-image is requested, all in one transaction. BOTH is three round trips. Today's upsertAndReturn(Key, Document) is a single statement. The row lock is held until commit, so an existing upsert of the same key waits for those extra reads. A first-insert race rolls the transaction back and retries it up to three times.
Rule: lock-latency
Reviewed using @claude traceable:traceable-java-code-review
There was a problem hiding this comment.
Reworked in b96ad4a. The after-image comes from RETURNING * on the write, so there is no second read and no retry loop. Statements per call:
- AFTER: 1
- BEFORE / BOTH on an existing row: 2 (
SELECT ... FOR UPDATE, write) - First insert: 4 (select, advisory lock, re-select, write), once per key
I benchmarked the previous commit (46ae7db) against b96ad4a on flat Postgres 13.1 (Testcontainers, local, 3000 ops per row). p50 shown; two runs agreed:
| case | before | after |
|---|---|---|
| upsert AFTER, existing row | 1.00 ms | 0.57 ms |
| createOrReplace AFTER, existing row | 1.01 ms | 0.72 ms |
| upsert BOTH, existing row | 1.17 ms | 0.78 ms |
| createOrReplace BOTH, existing row | 1.35 ms | 1.04 ms |
| upsert BOTH, first insert | 1.11 ms | 1.10 ms |
| upsert BEFORE, first insert | 0.90 ms | 1.27 ms |
| createOrReplace BOTH, first insert | 1.29 ms | 1.61 ms |
| upsert BOTH, 8 threads on one key (p50 / p99) | 8.9 / 17.8 ms | 6.3 / 13.1 ms |
| createOrReplace BOTH, 8 threads on one key (p50 / p99) | 9.0 / 21.0 ms | 6.9 / 15.5 ms |
| upsert BOTH, 8 threads racing first inserts (p50 / p99) | 1.74 / 3.7 ms | 1.90 / 4.0 ms |
The row lock is now held for one statement less, which shows up as lower latency for contended writes on the same key. The cost is about +0.3 ms on a key's first insert in BEFORE/BOTH mode. These methods stay opt-in; upsert() and the existing upsertAndReturn(Key, Document) are unchanged and remain a single statement.
| public DocumentChange createOrReplaceAndReturn( | ||
| final Key key, final Document document, final UpsertOptions options) throws IOException { | ||
| final ReturnDocuments returnDocuments = options.getReturnDocuments(); | ||
| if (returnDocuments == ReturnDocuments.BOTH) { |
There was a problem hiding this comment.
[medium] Mongo createOrReplaceAndReturn rejects BOTH
Collection.createOrReplaceAndReturn documents BEFORE, AFTER, and BOTH as the images returned from one atomic write. This method throws UnsupportedOperationException for ReturnDocuments.BOTH before it writes. Flat Postgres accepts BOTH, so the same call is not portable across collection implementations.
Rule: api-contract
Reviewed using @claude traceable:traceable-java-code-review
There was a problem hiding this comment.
Supported now (b96ad4a). createOrReplaceAndReturn is still a single findOneAndUpdate. The difference from createOrReplace is that every top-level value in the pipeline's $set is wrapped in $literal, so values are stored exactly as given; a string like "$price" is no longer read as a field path. That makes the after-image derivable from the server-returned before-image: keep _id, keep createdTime (or now if missing or null, matching $ifNull), then the replacement fields. The new integration test replaces with "$price" and an empty object and asserts the returned after-image equals the stored document. The concurrent-chain test now also runs for Mongo createOrReplace BOTH.
- Postgres: replace first-insert retry loop with pg_advisory_xact_lock on the key, take the after-image from INSERT ... RETURNING, strip _implicit_id from images - Mongo: BEFORE/AFTER use $currentDate like upsert(); explicit _lastUpdateTime only for BOTH - Mongo: support BOTH in createOrReplaceAndReturn by deriving the after-image from a fully $literal-wrapped replace - Document non-atomicity on Collection Co-authored-by: Cursor <cursoragent@cursor.com>
…shots
Callers that need the previous version of a document in order to emit a change event currently have to issue a separate read before the write. That read is not atomic with respect to the write, so a concurrent upsert on the same key can land in between and the "previous" snapshot ends up already holding the new value.
Add a merge-upsert that hands back the pre-write document as part of the write itself. Mongo uses findOneAndUpdate with ReturnDocument.BEFORE; both Postgres collections do the read and the upsert inside a single transaction, with the read taking a row lock. An empty Optional means the document did not exist and this call created it.
Description
Please include a summary of the change, motivation and context.
Testing
Please describe the tests that you ran to verify your changes. Please summarize what did you test and what needs to be tested e.g. deployed and tested helm chart locally.
Checklist:
Documentation
Make sure that you have documented corresponding changes in this repository or hypertrace docs repo if required.