fix(cwl): Dispose event listeners when command exits (aws#6095)

keeganirby · karanA-aws · commit 7a0d456fa52a · 2025-01-17T14:42:54.000-08:00
## Problem
TailLogGroup command is not being fully cleaned up properly. The event
listeners for:
* showing the StatusBar item
*  closing a session when exiting its editors  

are not being disposed of. As users stop sessions, these listeners
remain active whilst doing no meaningful work. This is effectively a
memory leak.

Additionally during testing, I noticed that we are adding a session to
the LiveTailRegistry before the session is actually started. This could
cause an issue where a session stream isn't successfully created, but
the session is still registered.

## Solution
* Create an array of disposables and dispose of them as we exit the
TailLogGroup command
* Only register a session after startLiveTail returns


Because now we dispose of the "tab close" listener when tailLogGroup
exits, changes are needed in the unit tests to keep the mock response
stream "open" as the test executes. This keeps tailLogGroup from
returning/disposing while we perform test assertions and close the
editors.
diff --git a/packages/core/src/awsService/cloudWatchLogs/commands/tailLogGroup.ts b/packages/core/src/awsService/cloudWatchLogs/commands/tailLogGroup.ts
@@ -52,25 +52,28 @@ export async function tailLogGroup(
             })
             return
         }
-
-        registry.set(uriToKey(session.uri), session)
-
         const document = await prepareDocument(session)
 
-        hideShowStatusBarItemsOnActiveEditor(session, document)
-        registerTabChangeCallback(session, registry, document, codeLensProvider)
+        const disposables: vscode.Disposable[] = []
+        disposables.push(hideShowStatusBarItemsOnActiveEditor(session, document))
+        disposables.push(closeSessionWhenAllEditorsClosed(session, registry, document, codeLensProvider))
 
-        const stream = await session.startLiveTailSession()
-        getLogger().info(`LiveTail session started: ${uriToKey(session.uri)}`)
-
-        span.record({
-            source: source,
-            result: 'Succeeded',
-            sessionAlreadyStarted: false,
-            hasLogEventFilterPattern: Boolean(wizardResponse.filterPattern),
-            logStreamFilterType: wizardResponse.logStreamFilter.type,
-        })
-        await handleSessionStream(stream, document, session)
+        try {
+            const stream = await session.startLiveTailSession()
+            registry.set(uriToKey(session.uri), session)
+            codeLensProvider.refresh()
+            getLogger().info(`LiveTail session started: ${uriToKey(session.uri)}`)
+            span.record({
+                source: source,
+                result: 'Succeeded',
+                sessionAlreadyStarted: false,
+                hasLogEventFilterPattern: Boolean(wizardResponse.filterPattern),
+                logStreamFilterType: wizardResponse.logStreamFilter.type,
+            })
+            await handleSessionStream(stream, document, session)
+        } finally {
+            disposables.forEach((disposable) => disposable.dispose())
+        }
     })
 }
 
@@ -224,8 +227,11 @@ function eventRate(event: LiveTailSessionUpdate): number {
     return event.sessionResults === undefined ? 0 : event.sessionResults.length
 }
 
-function hideShowStatusBarItemsOnActiveEditor(session: LiveTailSession, document: vscode.TextDocument) {
-    vscode.window.onDidChangeActiveTextEditor((editor) => {
+function hideShowStatusBarItemsOnActiveEditor(
+    session: LiveTailSession,
+    document: vscode.TextDocument
+): vscode.Disposable {
+    return vscode.window.onDidChangeActiveTextEditor((editor) => {
         session.showStatusBarItem(editor?.document === document)
     })
 }
@@ -242,13 +248,13 @@ function hideShowStatusBarItemsOnActiveEditor(session: LiveTailSession, document
  * from view, will not be returned. Meaning a Tab that is created (shown in top bar), but not open, will not be returned. Even if
  * the tab isn't visible, we want to continue writing to the doc, and keep the session alive.
  */
-function registerTabChangeCallback(
+function closeSessionWhenAllEditorsClosed(
     session: LiveTailSession,
     registry: LiveTailSessionRegistry,
     document: vscode.TextDocument,
     codeLensProvider: LiveTailCodeLensProvider
-) {
-    vscode.window.tabGroups.onDidChangeTabs((tabEvent) => {
+): vscode.Disposable {
+    return vscode.window.tabGroups.onDidChangeTabs((tabEvent) => {
         const isOpen = isLiveTailSessionOpenInAnyTab(session)
         if (!isOpen) {
             closeSession(session.uri, registry, 'ClosedEditors', codeLensProvider)
diff --git a/packages/core/src/test/awsService/cloudWatchLogs/commands/tailLogGroup.test.ts b/packages/core/src/test/awsService/cloudWatchLogs/commands/tailLogGroup.test.ts
@@ -4,23 +4,20 @@
  */
 
 import * as sinon from 'sinon'
-import * as FakeTimers from '@sinonjs/fake-timers'
 import * as vscode from 'vscode'
 
 import assert from 'assert'
 import { clearDocument, closeSession, tailLogGroup } from '../../../../awsService/cloudWatchLogs/commands/tailLogGroup'
-import { LiveTailSessionLogEvent, StartLiveTailResponseStream } from '@aws-sdk/client-cloudwatch-logs'
+import { StartLiveTailResponseStream } from '@aws-sdk/client-cloudwatch-logs'
 import { LiveTailSessionRegistry } from '../../../../awsService/cloudWatchLogs/registry/liveTailSessionRegistry'
 import { LiveTailSession } from '../../../../awsService/cloudWatchLogs/registry/liveTailSession'
-import { asyncGenerator } from '../../../../shared/utilities/collectionUtils'
 import {
     TailLogGroupWizard,
     TailLogGroupWizardResponse,
 } from '../../../../awsService/cloudWatchLogs/wizard/tailLogGroupWizard'
 import { getTestWindow } from '../../../shared/vscode/window'
 import { CloudWatchLogsSettings, uriToKey } from '../../../../awsService/cloudWatchLogs/cloudWatchLogsUtils'
-import { installFakeClock } from '../../../testUtil'
-import { DefaultAwsContext, ToolkitError } from '../../../../shared'
+import { DefaultAwsContext, ToolkitError, waitUntil } from '../../../../shared'
 import { LiveTailCodeLensProvider } from '../../../../awsService/cloudWatchLogs/document/liveTailCodeLensProvider'
 
 describe('TailLogGroup', function () {
@@ -39,23 +36,12 @@ describe('TailLogGroup', function () {
     let cloudwatchSettingsSpy: sinon.SinonSpy
     let wizardSpy: sinon.SinonSpy
 
-    let clock: FakeTimers.InstalledClock
-
-    before(function () {
-        clock = installFakeClock()
-    })
-
     beforeEach(function () {
-        clock.reset()
         sandbox = sinon.createSandbox()
         registry = new LiveTailSessionRegistry()
         codeLensProvider = new LiveTailCodeLensProvider(registry)
     })
 
-    after(function () {
-        clock.uninstall()
-    })
-
     afterEach(function () {
         sandbox.restore()
     })
@@ -64,47 +50,46 @@ describe('TailLogGroup', function () {
         sandbox.stub(DefaultAwsContext.prototype, 'getCredentialAccountId').returns(testAwsAccountId)
         sandbox.stub(DefaultAwsContext.prototype, 'getCredentials').returns(Promise.resolve(testAwsCredentials))
 
-        wizardSpy = sandbox.stub(TailLogGroupWizard.prototype, 'run').callsFake(async function () {
-            return getTestWizardResponse()
-        })
-        const testMessage2 = `${testMessage}-2`
-        const testMessage3 = `${testMessage}-3`
+        const startTimestamp = 1732276800000 // 11-22-2024 12:00:00PM GMT
+        const updateFrames: StartLiveTailResponseStream[] = [
+            getSessionUpdateFrame(false, `${testMessage}-1`, startTimestamp + 1000),
+            getSessionUpdateFrame(false, `${testMessage}-2`, startTimestamp + 2000),
+            getSessionUpdateFrame(false, `${testMessage}-3`, startTimestamp + 3000),
+        ]
+        //Returns the configured update frames and then indefinitely blocks.
+        //This keeps the stream 'open', simulating an open network stream waiting for new events.
+        //If the stream were to close, the event listeners in the TailLogGroup command would dispose,
+        //breaking the 'closes tab closes session' assertions this test makes.
+        async function* generator(): AsyncIterable<StartLiveTailResponseStream> {
+            for (const frame of updateFrames) {
+                yield frame
+            }
+            await new Promise(() => {})
+        }
+
         startLiveTailSessionSpy = sandbox
             .stub(LiveTailSession.prototype, 'startLiveTailSession')
-            .callsFake(async function () {
-                return getTestResponseStream([
-                    {
-                        message: testMessage,
-                        timestamp: 876830400000,
-                    },
-                    {
-                        message: testMessage2,
-                        timestamp: 876830402000,
-                    },
-                    {
-                        message: testMessage3,
-                        timestamp: 876830403000,
-                    },
-                ])
-            })
+            .returns(Promise.resolve(generator()))
         stopLiveTailSessionSpy = sandbox
             .stub(LiveTailSession.prototype, 'stopLiveTailSession')
             .callsFake(async function () {
                 return
             })
-
+        wizardSpy = sandbox.stub(TailLogGroupWizard.prototype, 'run').callsFake(async function () {
+            return getTestWizardResponse()
+        })
         //Set maxLines to 1.
         cloudwatchSettingsSpy = sandbox.stub(CloudWatchLogsSettings.prototype, 'get').callsFake(() => {
             return 1
         })
-        await tailLogGroup(registry, testSource, codeLensProvider, {
+
+        //The mock stream doesn't 'close', causing tailLogGroup to not return. If we `await`, it will never resolve.
+        //Run it in the background and use waitUntil to poll its state.
+        void tailLogGroup(registry, testSource, codeLensProvider, {
             groupName: testLogGroup,
             regionName: testRegion,
         })
-        assert.strictEqual(wizardSpy.calledOnce, true)
-        assert.strictEqual(cloudwatchSettingsSpy.calledOnce, true)
-        assert.strictEqual(startLiveTailSessionSpy.calledOnce, true)
-        assert.strictEqual(registry.size, 1)
+        await waitUntil(async () => registry.size !== 0, { interval: 100, timeout: 1000 })
 
         //registry is asserted to have only one entry, so this is assumed to be the session that was
         //started in this test.
@@ -113,13 +98,24 @@ describe('TailLogGroup', function () {
         if (sessionUri === undefined) {
             throw Error
         }
-        const document = getTestWindow().activeTextEditor?.document
+
+        assert.strictEqual(wizardSpy.calledOnce, true)
+        assert.strictEqual(cloudwatchSettingsSpy.calledOnce, true)
+        assert.strictEqual(startLiveTailSessionSpy.calledOnce, true)
+        assert.strictEqual(registry.size, 1)
+
+        //Validate writing to the document.
+        //MaxLines is set to 1, and "testMessage3" is the last event in the stream, its contents should be the only thing in the doc.
+        const window = getTestWindow()
+        const document = window.activeTextEditor?.document
         assert.strictEqual(sessionUri.toString(), document?.uri.toString())
-        //Test responseStream has 3 events, maxLines is set to 1. Only 3rd event should be in doc.
-        assert.strictEqual(document?.getText().trim(), `12:00:03\t${testMessage3}`)
+        const doesDocumentContainExpectedContent = await waitUntil(
+            async () => document?.getText().trim() === `12:00:03\t${testMessage}-3`,
+            { interval: 100, timeout: 1000 }
+        )
+        assert.strictEqual(doesDocumentContainExpectedContent, true)
 
         //Test that closing all tabs the session's document is open in will cause the session to close
-        const window = getTestWindow()
         let tabs: vscode.Tab[] = []
         window.tabGroups.all.forEach((tabGroup) => {
             tabs = tabs.concat(getLiveTailSessionTabsFromTabGroup(tabGroup, sessionUri!))
@@ -159,7 +155,6 @@ describe('TailLogGroup', function () {
         closeSession(session.uri, registry, testSource, codeLensProvider)
         assert.strictEqual(0, registry.size)
         assert.strictEqual(true, stopLiveTailSessionSpy.calledOnce)
-        assert.strictEqual(0, clock.countTimers())
     })
 
     it('clearDocument clears all text from document', async function () {
@@ -200,26 +195,23 @@ describe('TailLogGroup', function () {
         }
     }
 
-    //Creates a test response stream. Each log event provided will be its own "frame" of the input stream.
-    function getTestResponseStream(logEvents: LiveTailSessionLogEvent[]): AsyncIterable<StartLiveTailResponseStream> {
-        const sessionStartFrame: StartLiveTailResponseStream = {
-            sessionStart: {
-                logGroupIdentifiers: [testLogGroup],
+    function getSessionUpdateFrame(
+        isSampled: boolean,
+        message: string,
+        timestamp: number
+    ): StartLiveTailResponseStream {
+        return {
+            sessionUpdate: {
+                sessionMetadata: {
+                    sampled: isSampled,
+                },
+                sessionResults: [
+                    {
+                        message: message,
+                        timestamp: timestamp,
+                    },
+                ],
             },
-            sessionUpdate: undefined,
         }
-
-        const updateFrames: StartLiveTailResponseStream[] = logEvents.map((event) => {
-            return {
-                sessionUpdate: {
-                    sessionMetadata: {
-                        sampled: false,
-                    },
-                    sessionResults: [event],
-                },
-            }
-        })
-
-        return asyncGenerator([sessionStartFrame, ...updateFrames])
     }
 })
