From 4c4a883a08df821c380b6f0d56a749a568792d35 Mon Sep 17 00:00:00 2001
From: Craigory Coppola <craigorycoppola@gmail.com>
Date: Mon, 24 Jun 2024 11:32:11 -0400
Subject: [PATCH] fix(core): isolated plugins should provide cleanup function

---
 .../project-graph/plugins/isolation/index.ts  | 22 ++++++++++++++-----
 .../plugins/isolation/plugin-pool.ts          | 22 ++++++++++++-------
 .../nx/src/project-graph/project-graph.ts     |  8 ++++++-
 3 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/packages/nx/src/project-graph/plugins/isolation/index.ts b/packages/nx/src/project-graph/plugins/isolation/index.ts
index 5e422cb64b763..19a5ba5abac9f 100644
--- a/packages/nx/src/project-graph/plugins/isolation/index.ts
+++ b/packages/nx/src/project-graph/plugins/isolation/index.ts
@@ -6,20 +6,30 @@ import { loadRemoteNxPlugin } from './plugin-pool';
 /**
  * Used to ensure 1 plugin : 1 worker
  */
-const remotePluginCache = new Map<string, Promise<LoadedNxPlugin>>();
+const remotePluginCache = new Map<
+  string,
+  readonly [Promise<LoadedNxPlugin>, () => void]
+>();
 
 export function loadNxPluginInIsolation(
   plugin: PluginConfiguration,
   root = workspaceRoot
-): [Promise<LoadedNxPlugin>, () => void] {
+): readonly [Promise<LoadedNxPlugin>, () => void] {
   const cacheKey = JSON.stringify(plugin);
 
   if (remotePluginCache.has(cacheKey)) {
-    return [remotePluginCache.get(cacheKey), () => {}];
+    return remotePluginCache.get(cacheKey);
   }
 
-  const loadingPlugin = loadRemoteNxPlugin(plugin, root);
-  remotePluginCache.set(cacheKey, loadingPlugin);
+  const [loadingPlugin, cleanup] = loadRemoteNxPlugin(plugin, root);
   // We clean up plugin workers when Nx process completes.
-  return [loadingPlugin, () => {}];
+  const val = [
+    loadingPlugin,
+    () => {
+      cleanup();
+      remotePluginCache.delete(cacheKey);
+    },
+  ] as const;
+  remotePluginCache.set(cacheKey, val);
+  return val;
 }
diff --git a/packages/nx/src/project-graph/plugins/isolation/plugin-pool.ts b/packages/nx/src/project-graph/plugins/isolation/plugin-pool.ts
index a34c8f6dc6ae9..ecfc269de746e 100644
--- a/packages/nx/src/project-graph/plugins/isolation/plugin-pool.ts
+++ b/packages/nx/src/project-graph/plugins/isolation/plugin-pool.ts
@@ -22,7 +22,7 @@ interface PendingPromise {
 export function loadRemoteNxPlugin(
   plugin: PluginConfiguration,
   root: string
-): Promise<LoadedNxPlugin> {
+): [Promise<LoadedNxPlugin>, () => void] {
   // this should only really be true when running unit tests within
   // the Nx repo. We still need to start the worker in this case,
   // but its typescript.
@@ -66,13 +66,19 @@ export function loadRemoteNxPlugin(
 
   cleanupFunctions.add(cleanupFunction);
 
-  return new Promise<LoadedNxPlugin>((res, rej) => {
-    worker.on(
-      'message',
-      createWorkerHandler(worker, pendingPromises, res, rej)
-    );
-    worker.on('exit', exitHandler);
-  });
+  return [
+    new Promise<LoadedNxPlugin>((res, rej) => {
+      worker.on(
+        'message',
+        createWorkerHandler(worker, pendingPromises, res, rej)
+      );
+      worker.on('exit', exitHandler);
+    }),
+    () => {
+      cleanupFunction();
+      cleanupFunctions.delete(cleanupFunction);
+    },
+  ];
 }
 
 function shutdownPluginWorker(worker: ChildProcess) {
diff --git a/packages/nx/src/project-graph/project-graph.ts b/packages/nx/src/project-graph/project-graph.ts
index 05e2086d5a7ad..81a5d51b7e6d5 100644
--- a/packages/nx/src/project-graph/project-graph.ts
+++ b/packages/nx/src/project-graph/project-graph.ts
@@ -148,7 +148,13 @@ export async function buildProjectGraphAndSourceMapsWithoutDaemon() {
       throw e;
     }
   } finally {
-    cleanup();
+    // When plugins are isolated we don't clean them up during
+    // a single run of the CLI. They are cleaned up when the CLI
+    // process exits. Cleaning them here could cause issues if pending
+    // promises are not resolved.
+    if (process.env.NX_ISOLATE_PLUGINS !== 'true') {
+      cleanup();
+    }
   }
 
   const { projectGraph, projectFileMapCache } = projectGraphResult;