From 53aed251d3e45c76eea46aaefd721120f36eb7e5 Mon Sep 17 00:00:00 2001 From: Jonathan Schneider Date: Wed, 26 Feb 2025 10:39:16 -0500 Subject: [PATCH] RewriteRpc no longer uses identity hash codes --- .../java/org/openrewrite/rpc/RewriteRpc.java | 44 +++++++++++++------ .../org/openrewrite/rpc/request/Generate.java | 17 ++++++- .../org/openrewrite/rpc/request/Visit.java | 2 +- 3 files changed, 48 insertions(+), 15 deletions(-) diff --git a/rewrite-core/src/main/java/org/openrewrite/rpc/RewriteRpc.java b/rewrite-core/src/main/java/org/openrewrite/rpc/RewriteRpc.java index a95ea84f4fd..355bb856085 100644 --- a/rewrite-core/src/main/java/org/openrewrite/rpc/RewriteRpc.java +++ b/rewrite-core/src/main/java/org/openrewrite/rpc/RewriteRpc.java @@ -18,6 +18,8 @@ import io.moderne.jsonrpc.JsonRpc; import io.moderne.jsonrpc.JsonRpcMethod; import io.moderne.jsonrpc.JsonRpcRequest; +import io.moderne.jsonrpc.internal.SnowflakeId; +import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.VisibleForTesting; import org.jspecify.annotations.Nullable; import org.openrewrite.*; @@ -50,6 +52,8 @@ public class RewriteRpc { @VisibleForTesting final Map localObjects = new HashMap<>(); + /* A reverse map of the objects back to their IDs */ + final Map localObjectIds = new IdentityHashMap<>(); private final Map remoteRefs = new IdentityHashMap<>(); @@ -117,17 +121,8 @@ public

VisitResponse scan(SourceFile sourceFile, String visitorName, P p, // Set the local state of this tree, so that when the remote // asks for it, we know what to send. localObjects.put(sourceFile.getId().toString(), sourceFile); - String pId = Integer.toString(System.identityHashCode(p)); - - Object p2 = p; - while (p2 instanceof DelegatingExecutionContext) { - p2 = ((DelegatingExecutionContext) p2).getDelegate(); - } - if (p2 instanceof ExecutionContext) { - ((ExecutionContext) p2).putMessage("org.openrewrite.rpc.id", pId); - } - localObjects.put(pId, p2); + String pId = maybeUnwrapExecutionContext(p); List cursorIds = getCursorIds(cursor); return send("Visit", new Visit(visitorName, null, sourceFile.getId().toString(), pId, cursorIds), @@ -135,8 +130,9 @@ public

VisitResponse scan(SourceFile sourceFile, String visitorName, P p, } public Collection generate(String remoteRecipeId, ExecutionContext ctx) { - List generated = send("Generate", new Generate(remoteRecipeId, - Integer.toString(System.identityHashCode(ctx))), GenerateResponse.class); + String ctxId = maybeUnwrapExecutionContext(ctx); + List generated = send("Generate", new Generate(remoteRecipeId, ctxId), + GenerateResponse.class); if (!generated.isEmpty()) { return generated.stream() .map(this::getObject) @@ -145,6 +141,27 @@ public Collection generate(String remoteRecipeId, Executio return emptyList(); } + /** + * If the p object is a DelegatingExecutionContext, unwrap it to get the underlying + * ExecutionContext + * + * @param p A visitor parameter, which may or may not be an ExecutionContext + * @param

The type of p + * @return The ID of p as represented in the local object cache. + */ + private

String maybeUnwrapExecutionContext(P p) { + Object p2 = p; + while (p2 instanceof DelegatingExecutionContext) { + p2 = ((DelegatingExecutionContext) p2).getDelegate(); + } + String pId = localObjectIds.computeIfAbsent(p2, p3 -> SnowflakeId.generateId()); + if (p2 instanceof ExecutionContext) { + ((ExecutionContext) p2).putMessage("org.openrewrite.rpc.id", pId); + } + localObjects.put(pId, p2); + return pId; + } + public List getRecipes() { return send("GetRecipes", null, GetRecipesResponse.class); } @@ -171,7 +188,7 @@ List getCursorIds(@Nullable Cursor cursor) { cursorIds = cursor.getPathAsStream().map(c -> { String id = c instanceof Tree ? ((Tree) c).getId().toString() - : Integer.toString(System.identityHashCode(c)); + : localObjectIds.computeIfAbsent(c, c2 -> SnowflakeId.generateId()); localObjects.put(id, c); return id; }).collect(Collectors.toList()); @@ -195,6 +212,7 @@ public T getObject(String id) { } // We are now in sync with the remote state of the object. remoteObjects.put(id, remoteObject); + localObjects.put(id, remoteObject); //noinspection unchecked return (T) remoteObject; diff --git a/rewrite-core/src/main/java/org/openrewrite/rpc/request/Generate.java b/rewrite-core/src/main/java/org/openrewrite/rpc/request/Generate.java index 55b21c200cd..97f3cbb2962 100644 --- a/rewrite-core/src/main/java/org/openrewrite/rpc/request/Generate.java +++ b/rewrite-core/src/main/java/org/openrewrite/rpc/request/Generate.java @@ -19,6 +19,11 @@ import lombok.RequiredArgsConstructor; import lombok.Value; import org.openrewrite.*; +import org.openrewrite.scheduling.RecipeRunCycle; +import org.openrewrite.scheduling.WatchableExecutionContext; +import org.openrewrite.table.RecipeRunStats; +import org.openrewrite.table.SourcesFileErrors; +import org.openrewrite.table.SourcesFileResults; import java.util.Collection; import java.util.Map; @@ -26,6 +31,7 @@ import java.util.stream.Collectors; import static java.util.Collections.emptyList; +import static org.openrewrite.ExecutionContext.CURRENT_RECIPE; @Value public class Generate implements RpcRequest { @@ -45,8 +51,17 @@ public static class Handler extends JsonRpcMethod { @Override protected Object handle(Generate request) throws Exception { - ExecutionContext ctx = (ExecutionContext) getObject.apply(request.getP()); Recipe recipe = preparedRecipes.get(request.getId()); + + ExecutionContext ctx = (ExecutionContext) getObject.apply(request.getP()); + if (ctx.getMessage(CURRENT_RECIPE) == null) { + WatchableExecutionContext wctx = new WatchableExecutionContext((ExecutionContext) ctx); + wctx.putCycle(new RecipeRunCycle<>(recipe, 0, new Cursor(null, Cursor.ROOT_VALUE), wctx, + new RecipeRunStats(Recipe.noop()), new SourcesFileResults(Recipe.noop()), + new SourcesFileErrors(Recipe.noop()), LargeSourceSet::edit)); + ctx.putCurrentRecipe(recipe); + } + if (recipe instanceof ScanningRecipe) { //noinspection unchecked ScanningRecipe scanningRecipe = (ScanningRecipe) recipe; diff --git a/rewrite-core/src/main/java/org/openrewrite/rpc/request/Visit.java b/rewrite-core/src/main/java/org/openrewrite/rpc/request/Visit.java index 9c62f525ae3..d2c38f9cbc3 100644 --- a/rewrite-core/src/main/java/org/openrewrite/rpc/request/Visit.java +++ b/rewrite-core/src/main/java/org/openrewrite/rpc/request/Visit.java @@ -127,12 +127,12 @@ private Object getVisitorP(Visit request) { String visitorName = request.getVisitor(); if (visitorName.startsWith("scan:") || visitorName.startsWith("edit:")) { - WatchableExecutionContext ctx = new WatchableExecutionContext((ExecutionContext) p); Recipe recipe = preparedRecipes.get(visitorName.substring( "edit:".length() /* 'scan:' has same length*/)); // This is really probably particular to the Java implementation, // because we are carrying forward the legacy of cycles that are likely to be // removed from OpenRewrite in the future. + WatchableExecutionContext ctx = new WatchableExecutionContext((ExecutionContext) p); ctx.putCycle(new RecipeRunCycle<>(recipe, 0, new Cursor(null, Cursor.ROOT_VALUE), ctx, new RecipeRunStats(Recipe.noop()), new SourcesFileResults(Recipe.noop()), new SourcesFileErrors(Recipe.noop()), LargeSourceSet::edit));