Python: route imports through a PythonImportService so Java ChangeType stops corrupting files

rewrite-core/src/main/java/org/openrewrite/rpc/RewriteRpc.java

@@ -42,6 +42,7 @@
 import java.time.Duration;
 import java.util.*;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
@@ -76,12 +77,16 @@ public class RewriteRpc {
      * Keeps track of the local and remote state of objects that are used in
      * visits and other operations for which incremental state sharing is useful
      * between two processes.
+     * <p>
+     * Concurrent because {@link GetObject.Handler} reads these from its own thread pool while
+     * multiple in-flight {@link #visit} calls write them, and a {@code HashMap} resize racing a
+     * {@code get} would drop a present key, which the peer surfaces as "Tree not found".
      */
     @VisibleForTesting
-    final Map<String, Object> remoteObjects = new HashMap<>();
+    final Map<String, Object> remoteObjects = new ConcurrentHashMap<>();
 
     @VisibleForTesting
-    final Map<String, Object> localObjects = new HashMap<>();
+    final Map<String, Object> localObjects = new ConcurrentHashMap<>();
 
     /* A reverse map of the objects back to their IDs */
     final Map<Object, String> localObjectIds = new IdentityHashMap<>();
@@ -372,6 +377,15 @@ public void evict(String id, int localRefsCheckpoint, int remoteRefsCheckpoint)
     }
 
     public <P> @Nullable Tree visit(Tree tree, String visitorName, P p, @Nullable Cursor cursor) {
+        return visit(tree, visitorName, null, p, cursor);
+    }
+
+    /**
+     * Run a remote visitor that takes constructor arguments (e.g. the peer's own {@code AddImport}),
+     * or {@code null} {@code visitorOptions} when it takes none.
+     */
+    public <P> @Nullable Tree visit(Tree tree, String visitorName, @Nullable Map<String, Object> visitorOptions,
+                                    P p, @Nullable Cursor cursor) {
         ensureDataTableStoreSent();
         // Set the local state of this tree, so that when the remote asks for it, we know what to send.
         localObjects.put(tree.getId().toString(), tree);
@@ -381,7 +395,7 @@ public void evict(String id, int localRefsCheckpoint, int remoteRefsCheckpoint)
 
         String sourceFileType = DynamicDispatchRpcCodec.canonicalSourceFileType(
                 (tree instanceof SourceFile ? tree : requireNonNull(cursor).firstEnclosingOrThrow(SourceFile.class)).getClass());
-        Supplier<VisitResponse> doSend = () -> send("Visit", new Visit(visitorName, sourceFileType, null,
+        Supplier<VisitResponse> doSend = () -> send("Visit", new Visit(visitorName, sourceFileType, visitorOptions,
                 tree.getId().toString(), pId, cursorIds), VisitResponse.class);
         VisitResponse response = p instanceof ExecutionContext
                 ? RewriteRpcExecutionContextView.view((ExecutionContext) p).withInFlightSlot(doSend)

rewrite-core/src/test/java/org/openrewrite/rpc/RewriteRpcTest.java

@@ -239,6 +239,18 @@ void evictDropsTreeFromBothPeers() {
         assertThat(server.remoteObjects).doesNotContainKey(id);
     }
 
+    /**
+     * The id-keyed object maps must stay concurrent, since {@code GetObject.Handler} reads them from
+     * its own thread pool while in-flight visits write them and a {@code HashMap} resize race dropped
+     * a present key as "Tree not found" (customer-requests#2858) -- an invariant rather than a
+     * behavioral test because the resize window is too small to hit deterministically.
+     */
+    @Test
+    void sharedObjectMapsAreThreadSafe() {
+        assertThat(client.localObjects).isInstanceOf(java.util.concurrent.ConcurrentMap.class);
+        assertThat(client.remoteObjects).isInstanceOf(java.util.concurrent.ConcurrentMap.class);
+    }
+
     @DocumentExample
     @Test
     void sendReceiveIdempotence() {

rewrite-java/src/main/java/org/openrewrite/java/ChangeType.java

@@ -241,6 +241,12 @@ private void addImport(JavaType.FullyQualified owningClass) {
                 j = ((TypedTree) tree).withType(updateType(((TypedTree) tree).getType()));
             } else if (tree instanceof JavaSourceFile) {
                 JavaSourceFile sf = (JavaSourceFile) tree;
+                // Languages that keep imports among the statements (Python) expose none through
+                // getImports(), so hand removal to their import service; on a genuinely import-less
+                // Java file there is nothing to remove and this is a no-op.
+                boolean removeThroughImportService = targetType instanceof JavaType.FullyQualified &&
+                        sf.getImports().isEmpty();
+
                 if (targetType instanceof JavaType.FullyQualified) {
                     for (J.Import anImport : sf.getImports()) {
                         if (anImport.isStatic()) {
@@ -280,6 +286,12 @@ private void addImport(JavaType.FullyQualified owningClass) {
                     }
                 }
 
+                if (removeThroughImportService) {
+                    // Queued after the addition so the import service sees the name already bound by
+                    // the new import, which is what makes the old one safe to retire.
+                    maybeRemoveImport(originalType.getFullyQualifiedName());
+                }
+
                 j = sf.withImports(ListUtils.map(sf.getImports(), i -> {
                     Cursor cursor = getCursor();
                     setCursor(new Cursor(cursor, i));

rewrite-java/src/main/java/org/openrewrite/java/JavaVisitor.java

@@ -199,7 +199,7 @@ public void maybeRemoveImport(JavaType.@Nullable FullyQualified clazz) {
     }
 
     public void maybeRemoveImport(String fullyQualifiedName) {
-        RemoveImport<P> op = new RemoveImport<>(fullyQualifiedName);
+        JavaVisitor<P> op = service(ImportService.class).removeImportVisitor(fullyQualifiedName);
         if (!getAfterVisit().contains(op)) {
             doAfterVisit(op);
         }

rewrite-java/src/main/java/org/openrewrite/java/service/ImportService.java

@@ -20,6 +20,7 @@
 import org.openrewrite.Incubating;
 import org.openrewrite.java.AddImport;
 import org.openrewrite.java.JavaVisitor;
+import org.openrewrite.java.RemoveImport;
 import org.openrewrite.java.ShortenFullyQualifiedTypeReferences;
 import org.openrewrite.java.tree.J;
 
@@ -34,6 +35,10 @@ public <P> JavaVisitor<P> addImportVisitor(@Nullable String packageName,
         return new AddImport<>(packageName, typeName, member,  alias, onlyIfReferenced);
     }
 
+    public <P> JavaVisitor<P> removeImportVisitor(String fullyQualifiedName) {
+        return new RemoveImport<>(fullyQualifiedName);
+    }
+
     public <J2 extends J> JavaVisitor<ExecutionContext> shortenAllFullyQualifiedTypeReferences() {
         return new ShortenFullyQualifiedTypeReferences().getVisitor();
     }

rewrite-python/rewrite/src/rewrite/java/extensions.py

@@ -2,7 +2,7 @@
 
 from rewrite import Cursor
 from rewrite.tree import Tree
-from rewrite.utils import list_map
+from rewrite.utils import list_map, replace_if_changed
 from .support_types import J, JRightPadded, JLeftPadded, JContainer, Space, P, T, J2
 
 if TYPE_CHECKING:
@@ -68,4 +68,4 @@ def visit_space(v: 'JavaVisitor', space: Optional[Space], p: P) -> Space:
 
 def with_name(method: 'MethodInvocation', name: 'Identifier') -> 'MethodInvocation':
     # FIXME add type attribution logic
-    return method if name is method.name else replace(method, _name=name)
+    return replace_if_changed(method, _name=name)