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 8274efa6c41..d6581df893c 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 @@ -84,6 +84,10 @@ protected List handle(GetObject request) throws Exception { RpcSendQueue sendQueue = new RpcSendQueue(batchSize.get(), batch::put, localRefs, request.getSourceFileType(), traceGetObject.get()); TREE_TRAVERSAL_POOL.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); @@ -96,6 +100,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 47b2062a1e0..42de8d1b81b 100644 --- a/rewrite-core/src/test/java/org/openrewrite/rpc/RewriteRpcTest.java +++ b/rewrite-core/src/test/java/org/openrewrite/rpc/RewriteRpcTest.java @@ -181,9 +181,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 31386195cd4..3b50e68e09d 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 cb9ce154dfa..797e14642e1 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);