fixup! fixup! fix(server-utils): Ensure orchestrion instrumentation lazy registers on Node (#22518)

Author: isaacs

packages/server-utils/src/orchestrion/instrumentation.ts

@@ -58,7 +58,9 @@ export function invokeOrchestrionInstrumentation<Callback extends Instrumentatio
     return;
   }
 
-  const run = (): void => {
+  // cleanup function is passed in so that we only remove the listener
+  // once the callback is actually instrumented and called.
+  const run = (cleanup?: () => void): void => {
     // Mark as instrumented only if the channel binding actually lands.
     // Avoids a very narrow case where the binding never arrives within the
     // retry window. So, don't mark until we know we're actually calling.
@@ -67,6 +69,7 @@ export function invokeOrchestrionInstrumentation<Callback extends Instrumentatio
         return;
       }
       markInstrumented(callback);
+      cleanup?.();
       callback(...args);
     });
   };
@@ -88,8 +91,7 @@ export function invokeOrchestrionInstrumentation<Callback extends Instrumentatio
       return;
     }
     if (moduleNames.includes(moduleName)) {
-      run();
-      cleanup();
+      run(cleanup);
     }
   });
 }

packages/server-utils/test/orchestrion/instrumentation.test.ts

@@ -100,6 +100,29 @@ describe('invokeOrchestrionInstrumentation', () => {
     expect(callback).toHaveBeenCalledTimes(1);
   });
 
+  it('keeps the listener when the binding is not ready, so a later injection recovers', async () => {
+    // No async-context strategy, so the first injection finds no binding.
+    setAsyncContextStrategy(undefined);
+    const client = makeClient();
+    const callback = vi.fn();
+
+    invokeOrchestrionInstrumentation(client as never, ['pg', 'pg-pool'], callback, []);
+    expect(client.listenerCount()).toBe(1);
+
+    // First matching injection fires with no binding: the single retry bails,
+    // the callback stays unmarked, and the listener must remain for a retry.
+    client.inject('pg');
+    await new Promise(resolve => setTimeout(resolve, 10));
+    expect(callback).not.toHaveBeenCalled();
+    expect(client.listenerCount()).toBe(1);
+
+    // Binding now available; a later matching injection subscribes.
+    installBinding();
+    client.inject('pg-pool');
+    expect(callback).toHaveBeenCalledTimes(1);
+    expect(client.listenerCount()).toBe(0);
+  });
+
   it('matches on any of several module names', () => {
     const client = makeClient();
     const callback = vi.fn();