From d76ce30cb7f610df85dcb2a55bcf91ec4ee84760 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kevin=20Carpenter=E2=84=A2=EF=B8=8F?= Date: Tue, 10 Mar 2026 19:37:28 -0300 Subject: [PATCH] Roll back localRefs on failed GetObject exchange When a GetObject exchange fails mid-serialization, the sender's localRefs retains refs assigned during the failed exchange. On subsequent exchanges, the sender sends pure references for objects the remote never received, causing "Received a reference to an object that was not previously sent". Fix by snapshotting the ref count before the exchange and rolling back any refs added during a failed exchange, on both Java and TypeScript sender sides. Fixes https://github.com/moderneinc/customer-requests/issues/1977 --- .../openrewrite/rpc/request/GetObject.java | 11 +++++++++++ .../org/openrewrite/rpc/RewriteRpcTest.java | 5 +++++ rewrite-javascript/rewrite/src/reference.ts | 17 ++++++++++++++++- .../rewrite/src/rpc/request/get-object.ts | 19 ++++++++++++++----- 4 files changed, 46 insertions(+), 6 deletions(-) diff --git a/rewrite-core/src/main/java/org/openrewrite/rpc/request/GetObject.java b/rewrite-core/src/main/java/org/openrewrite/rpc/request/GetObject.java index cefcf84176..6af1e24451 100644 --- a/rewrite-core/src/main/java/org/openrewrite/rpc/request/GetObject.java +++ b/rewrite-core/src/main/java/org/openrewrite/rpc/request/GetObject.java @@ -78,6 +78,10 @@ protected List handle(GetObject request) throws Exception { RpcSendQueue sendQueue = new RpcSendQueue(batchSize.get(), batch::put, localRefs, request.getSourceFileType(), traceGetObject.get()); forkJoin.submit(() -> { + // Snapshot the current ref count so we can roll back on failure. + // Ref IDs are assigned sequentially as localRefs.size() + 1, + // so any ref > savedRefCount was added during this exchange. + int savedRefCount = localRefs.size(); try { sendQueue.send(after, before, null); @@ -90,6 +94,13 @@ protected List handle(GetObject request) throws Exception { // forces a full object sync (ADD) instead of a delta (CHANGE) // against the stale, partially-sent baseline. remoteObjects.remove(id); + + // Roll back localRefs to remove refs assigned during this failed + // exchange. Without this, subsequent exchanges would send pure + // references for objects the remote never received, causing + // "Received a reference to an object that was not previously sent". + localRefs.values().removeIf(ref -> ref > savedRefCount); + PrintStream logFile = log.get(); //noinspection ConstantValue if (logFile != null) { diff --git a/rewrite-core/src/test/java/org/openrewrite/rpc/RewriteRpcTest.java b/rewrite-core/src/test/java/org/openrewrite/rpc/RewriteRpcTest.java index a07432e3e9..6336ad6806 100644 --- a/rewrite-core/src/test/java/org/openrewrite/rpc/RewriteRpcTest.java +++ b/rewrite-core/src/test/java/org/openrewrite/rpc/RewriteRpcTest.java @@ -177,9 +177,14 @@ void sendFailureCleansUpRemoteObjects() { } // Step 4: verify the sender cleaned up its stale remoteObjects entry + // and rolled back any refs assigned during the failed exchange assertThat(server.remoteObjects) .describedAs("Sender should remove stale remoteObjects entry after send failure") .doesNotContainKey(id); + int refsAfterFailure = server.localRefs.size(); + assertThat(refsAfterFailure) + .describedAs("Sender should roll back localRefs assigned during failed exchange") + .isEqualTo(0); // Step 5: put back a valid tree and retry — should succeed via full ADD PlainText fixed = original.withText("Fixed"); diff --git a/rewrite-javascript/rewrite/src/reference.ts b/rewrite-javascript/rewrite/src/reference.ts index 31386195cd..3b50e68e09 100644 --- a/rewrite-javascript/rewrite/src/reference.ts +++ b/rewrite-javascript/rewrite/src/reference.ts @@ -84,4 +84,19 @@ export class ReferenceMap { this.refs.set(ref, refId); this.refsById.set(refId, ref); } -} \ No newline at end of file + + snapshot(): number { + return this.refCount; + } + + rollbackTo(savedRefCount: number): void { + for (let i = savedRefCount; i < this.refCount; i++) { + const obj = this.refsById.get(i); + if (obj) { + this.refs.delete(obj); + this.refsById.delete(i); + } + } + this.refCount = savedRefCount; + } +} diff --git a/rewrite-javascript/rewrite/src/rpc/request/get-object.ts b/rewrite-javascript/rewrite/src/rpc/request/get-object.ts index cb9ce154df..797e14642e 100644 --- a/rewrite-javascript/rewrite/src/rpc/request/get-object.ts +++ b/rewrite-javascript/rewrite/src/rpc/request/get-object.ts @@ -63,11 +63,20 @@ export class GetObject { const after = obj; const before = remoteObjects.get(objId); - allData = await new RpcSendQueue(localRefs, request.sourceFileType, trace()) - .generate(after, before); - pendingData.set(objId, allData); - - remoteObjects.set(objId, after); + // Snapshot ref count so we can roll back on failure. + // Ref IDs are assigned sequentially, so any ref >= savedRefCount + // was added during this exchange. + const savedRefCount = localRefs.snapshot(); + try { + allData = await new RpcSendQueue(localRefs, request.sourceFileType, trace()) + .generate(after, before); + pendingData.set(objId, allData); + remoteObjects.set(objId, after); + } catch (e) { + remoteObjects.delete(objId); + localRefs.rollbackTo(savedRefCount); + throw e; + } } const batch = allData.splice(0, batchSize);