Migrate waitForConnectionToBeUsable() to react to pushed connection event state changes

.versions/ide-sidecar.txt

@@ -1 +1 @@
-v0.134.1
+v0.139.0

src/authn/ccloudProvider.test.ts

@@ -59,7 +59,7 @@ describe("ConfluentCloudAuthProvider", () => {
 
     // assume the connection is immediately usable for most tests
     sandbox
-      .stub(connections, "waitForConnectionToBeUsable")
+      .stub(connections, "waitForConnectionToBeStable")
       .resolves(TEST_AUTHENTICATED_CCLOUD_CONNECTION);
 
     // don't handle the progress notification, openExternal, etc in this test suite

src/authn/ccloudProvider.ts

@@ -13,7 +13,7 @@ import {
   createCCloudConnection,
   deleteCCloudConnection,
   getCCloudConnection,
-  waitForConnectionToBeUsable,
+  waitForConnectionToBeStable,
 } from "../sidecar/connections";
 import { getStorageManager } from "../storage";
 import { SecretStorageKeys } from "../storage/constants";
@@ -152,7 +152,7 @@ export class ConfluentCloudAuthProvider implements vscode.AuthenticationProvider
       throw e;
     }
 
-    const authenticatedConnection = await waitForConnectionToBeUsable(CCLOUD_CONNECTION_ID);
+    const authenticatedConnection = await waitForConnectionToBeStable(CCLOUD_CONNECTION_ID);
     if (!authenticatedConnection) {
       throw new Error("CCloud connection failed to become usable after authentication.");
     }

src/directConnectManager.test.ts

@@ -71,7 +71,7 @@ describe("DirectConnectionManager behavior", () => {
       .stub(connections, "tryToUpdateConnection")
       .resolves({} as any);
     // assume the connection is immediately usable for most tests
-    sandbox.stub(connections, "waitForConnectionToBeUsable").resolves(TEST_DIRECT_CONNECTION);
+    sandbox.stub(connections, "waitForConnectionToBeStable").resolves(TEST_DIRECT_CONNECTION);
   });
 
   afterEach(() => {

src/directConnectManager.ts

@@ -18,7 +18,7 @@ import {
   tryToCreateConnection,
   tryToDeleteConnection,
   tryToUpdateConnection,
-  waitForConnectionToBeUsable,
+  waitForConnectionToBeStable,
 } from "./sidecar/connections";
 import { SecretStorageKeys } from "./storage/constants";
 import { DirectResourceLoader } from "./storage/directResourceLoader";
@@ -47,6 +47,8 @@ export class DirectConnectionManager {
    * cleaned up on extension deactivation. */
   disposables: Disposable[] = [];
 
+  // Map of connection ID to {current from-websocket-event state, event emitter when said state changes}
+
   // singleton instance to prevent multiple listeners and single source of connection management
   private static instance: DirectConnectionManager | null = null;
   private constructor() {
@@ -218,6 +220,7 @@ export class DirectConnectionManager {
   ): Promise<{ connection: Connection | null; errorMessage: string | null }> {
     let connection: Connection | null = null;
     let errorMessage: string | null = null;
+
     try {
       connection = update
         ? await tryToUpdateConnection(spec)
@@ -230,7 +233,7 @@ export class DirectConnectionManager {
             title: `Waiting for "${connection.spec.name}" to be usable...`,
           },
           async () => {
-            await waitForConnectionToBeUsable(connectionId);
+            await waitForConnectionToBeStable(connectionId);
           },
         );
       }
@@ -298,7 +301,7 @@ export class DirectConnectionManager {
       // wait for all new connections to be created before checking their status
       await Promise.all(newConnectionPromises);
       // kick off background checks to ensure the new connections are usable
-      connectionIdsToCheck.map((id) => waitForConnectionToBeUsable(id));
+      connectionIdsToCheck.map((id) => waitForConnectionToBeStable(id));
       logger.debug(
         `created and checked ${connectionIdsToCheck.length} new connection(s), firing event`,
       );

src/extension.ts

@@ -74,7 +74,7 @@ import { createConfigChangeListener } from "./preferences/listener";
 import { updatePreferences } from "./preferences/updates";
 import { registerProjectGenerationCommand } from "./scaffold";
 import { getSidecarManager, sidecarOutputChannel } from "./sidecar";
-import { getCCloudAuthSession } from "./sidecar/connections";
+import { ConnectionStateWatcher, getCCloudAuthSession } from "./sidecar/connections";
 import { WebsocketManager } from "./sidecar/websocketManager";
 import { getStorageManager, StorageManager } from "./storage";
 import { SecretStorageKeys } from "./storage/constants";
@@ -224,6 +224,10 @@ async function _activateExtension(
   // reset the Docker credentials secret so `src/docker/configs.ts` can pull it fresh
   getStorageManager().deleteSecret(SecretStorageKeys.DOCKER_CREDS_SECRET_KEY);
 
+  // Watch for sidecar pushing connection state changes over websocket.
+  // (side effect of causing the watcher to be created)
+  ConnectionStateWatcher.getInstance();
+
   const directConnectionManager = DirectConnectionManager.getInstance();
   context.subscriptions.push(...directConnectionManager.disposables);
 

src/models/resource.ts

@@ -1,8 +1,19 @@
 import { ConnectionType } from "../clients/sidecar";
-import { IconNames } from "../constants";
+import { CCLOUD_CONNECTION_ID, IconNames, LOCAL_CONNECTION_ID } from "../constants";
 
 /** A uniquely-branded string-type for a connection ID. */
 export type ConnectionId = string & { readonly brand: unique symbol };
+
+/** Given a ConnectionId, return the corresponding ConnectionType */
+export function connectionTypeFromId(id: ConnectionId): ConnectionType {
+  // CCloud and local are defined by constants. Direct are then arbitrary uuid strings.
+  return id === CCLOUD_CONNECTION_ID
+    ? ConnectionType.Ccloud
+    : id === LOCAL_CONNECTION_ID
+      ? ConnectionType.Local
+      : ConnectionType.Direct;
+}
+
 // TODO: use other branded resource ID types here
 
 export interface IResourceBase {

src/sidecar/connections.test.ts

@@ -22,21 +22,23 @@ import {
 } from "../emitters";
 import { ConnectionId } from "../models/resource";
 import { getResourceManager } from "../storage/resourceManager";
+import { ConnectionEventAction, Message, MessageType, newMessageHeaders } from "../ws/messageTypes";
 import {
   clearCurrentCCloudResources,
+  ConnectionStateWatcher,
   getLocalConnection,
   hasCCloudAuthSession,
   tryToCreateConnection,
   tryToDeleteConnection,
   tryToUpdateConnection,
-  waitForConnectionToBeUsable,
+  waitForConnectionToBeStable,
 } from "./connections";
 
+import * as connectionStatusUtils from "./connectionStatusUtils";
+
 describe("sidecar/connections.ts", () => {
   let sandbox: sinon.SinonSandbox;
   let stubConnectionsResourceApi: sinon.SinonStubbedInstance<ConnectionsResourceApi>;
-  let clock: sinon.SinonFakeTimers;
-  let connectionUsableFireStub: sinon.SinonStub;
 
   beforeEach(() => {
     sandbox = sinon.createSandbox();
@@ -47,8 +49,6 @@ describe("sidecar/connections.ts", () => {
     stubSidecarHandle.getConnectionsResourceApi.returns(stubConnectionsResourceApi);
     // stub the getSidecar function to return the stub sidecar handle
     sandbox.stub(sidecar, "getSidecar").resolves(stubSidecarHandle);
-    // stub the event emitter
-    connectionUsableFireStub = sandbox.stub(connectionUsable, "fire");
   });
 
   afterEach(() => {
@@ -133,6 +133,40 @@ describe("sidecar/connections.ts", () => {
     setContextValue(ContextValues.ccloudConnectionAvailable, true);
     assert.strictEqual(hasCCloudAuthSession(), true);
   });
+});
+
+describe("sidecar/connections.ts waitForConnectionToBeUsable() tests", () => {
+  const connectionStateWatcher = ConnectionStateWatcher.getInstance();
+
+  let sandbox: sinon.SinonSandbox;
+  let clock: sinon.SinonFakeTimers;
+  let connectionUsableFireStub: sinon.SinonStub;
+
+  beforeEach(() => {
+    sandbox = sinon.createSandbox();
+
+    // stub the event emitter
+    connectionUsableFireStub = sandbox.stub(connectionUsable, "fire");
+  });
+
+  afterEach(() => {
+    sandbox.restore();
+  });
+
+  function announceConnectionState(connection: Connection): void {
+    // inject a event to updated the connection state
+    // as if sent from sidecar from websocket.
+    const websocketMessage: Message<MessageType.CONNECTION_EVENT> = {
+      headers: newMessageHeaders(MessageType.CONNECTION_EVENT),
+      body: {
+        action: ConnectionEventAction.UPDATED,
+        connection: connection,
+      },
+    };
+
+    // As if had been just sent from sidecar.
+    connectionStateWatcher.handleConnectionUpdateEvent(websocketMessage);
+  }
 
   // dynamically set up tests for `waitForConnectionToBeUsable()` using different connections and states
   type ConnectionStateMatches = [
@@ -182,9 +216,10 @@ describe("sidecar/connections.ts", () => {
           ...testAuthStatus,
         },
       };
-      stubConnectionsResourceApi.gatewayV1ConnectionsIdGet.resolves(testConnection);
 
-      const connection = await waitForConnectionToBeUsable(testConnectionId);
+      announceConnectionState(testConnection);
+
+      const connection = await waitForConnectionToBeStable(testConnectionId);
 
       assert.deepStrictEqual(connection, testConnection);
     });
@@ -202,14 +237,13 @@ describe("sidecar/connections.ts", () => {
           ...testAuthStatus,
         },
       };
-      stubConnectionsResourceApi.gatewayV1ConnectionsIdGet.resolves(testConnection);
+      announceConnectionState(testConnection);
 
       // set a short timeout, even though we're using fake timers
       const shortTimeoutMs = 10;
-      const connectionPromise: Promise<Connection | null> = waitForConnectionToBeUsable(
+      const connectionPromise: Promise<Connection | null> = waitForConnectionToBeStable(
         testConnectionId,
         shortTimeoutMs,
-        shortTimeoutMs / 2,
       );
       // "wait" for the timeout to occur
       await clock.tickAsync(300);
@@ -220,7 +254,7 @@ describe("sidecar/connections.ts", () => {
       assert.ok(connectionUsableFireStub.calledOnce);
     });
 
-    it(`${baseConnection.spec.type}: waitForConnectionToBeUsable() should continue polling if the connection is not found initially`, async () => {
+    it(`${baseConnection.spec.type}: waitForConnectionToBeUsable() should wait for websocket event if the connection is not found initially`, async () => {
       const testConnection = {
         ...baseConnection,
         status: {
@@ -230,46 +264,43 @@ describe("sidecar/connections.ts", () => {
           ...testAuthStatus,
         },
       };
-      stubConnectionsResourceApi.gatewayV1ConnectionsIdGet
-        .onFirstCall()
-        .rejects(new ResponseError(new Response(null, { status: 404 })));
-      stubConnectionsResourceApi.gatewayV1ConnectionsIdGet.onSecondCall().resolves(testConnection);
 
-      const connection = await waitForConnectionToBeUsable(testConnectionId);
+      // wrap a spy around isConnectionStable so we can check when it's called
+      const isConnectionStableSpy = sandbox.spy(connectionStatusUtils, "isConnectionStable");
 
-      assert.deepStrictEqual(connection, testConnection);
-    });
+      // clear out prior connection state so that top of ConnectionStateWatcher.waitForConnectionUpdate will be a cache
+      // miss and it has to wait for an update.
+      connectionStateWatcher.purgeCachedConnectionState(testConnectionId);
 
-    it(`${baseConnection.spec.type}: waitForConnectionToBeUsable() should wait for a connection to be usable`, async () => {
-      const pendingConnection = {
-        ...baseConnection,
-        status: {
-          kafka_cluster: { state: pendingState },
-          schema_registry: { state: pendingState },
-          ccloud: { state: pendingState },
-          ...testAuthStatus,
-        },
-      };
-      stubConnectionsResourceApi.gatewayV1ConnectionsIdGet
-        .onFirstCall()
-        .resolves(pendingConnection);
-
-      const usableConnection = {
-        ...pendingConnection,
-        status: {
-          kafka_cluster: { state: usableKafkaClusterState },
-          schema_registry: { state: usableSchemaRegistryState },
-          ccloud: { state: usableCcloudState },
-          ...testAuthStatus,
-        },
-      };
-      stubConnectionsResourceApi.gatewayV1ConnectionsIdGet
-        .onSecondCall()
-        .resolves(usableConnection);
-
-      const connection = await waitForConnectionToBeUsable(testConnectionId, 10, 5);
+      clock = sandbox.useFakeTimers(Date.now());
 
-      assert.deepStrictEqual(connection, usableConnection);
+      async function scriptEventFlow() {
+        // let time pass some ...
+        await clock.tickAsync(100);
+
+        // isConnectionStableSpy should not have been called yet.
+        assert.ok(isConnectionStableSpy.notCalled);
+
+        // simulate a websocket event that updates the connection state to this new stable state.
+        announceConnectionState(testConnection);
+
+        // And now isConnectionStable should have been called.
+        assert.ok(isConnectionStableSpy.calledOnce);
+        // and it should have returned true
+        assert.ok(isConnectionStableSpy.returned(true));
+      }
+
+      // await both the script and waitForConnectionToBeStable
+      const results = await Promise.all([
+        scriptEventFlow(),
+        waitForConnectionToBeStable(testConnectionId),
+      ]);
+
+      // waitForConnectionToBeStable should have returned the connection
+      assert.deepStrictEqual(results[1], testConnection);
+      // and that isDirectConnectionStable was called (after the websocket event)
+      // log number of calls.
+      assert.ok(isConnectionStableSpy.calledOnce);
     });
   }
 });