feat(cwl): "clear screen", "stop session" actions aws#5958

Users may want to clear their screen during a tailing session. The
liveTail document is read only.
Users may want to stop their session without closing the document.

Provide a codeLens to clear the screen (make document empty)
Provide a codeLens to stop the session without closing the document.

Moves Interval Timer for updating the StatusBar to the LiveTailSession
object so `stopLiveTailSession()` can interrupt it, and be guaranteed to
be cleaned up.

The TextDocument's URI and Session URI seem to not be equal. Looking up
in the LiveTailSessionRegistry with a document URI causes a session to
not be found. Converting with `toString` allows these URIs to match and
the registry to work as intended.

Improves Exception handling on-stream. Previously, stopping the session
in an expected fashion (codeLens, Closing editors) would cause an
exception to bubble up and appear to the User. New change recognizes
when the Abort Controller triggers, signifying an error has not occured,
and logs the event as opposed to surfacing an error. Exposing only
`isAborted` so that a caller has to use `stopLiveTailSession` to trigger
the abortController and guarantee that other clean up actions have taken
place.

Author: keeganirby

packages/core/src/awsService/cloudWatchLogs/activation.ts

@@ -19,15 +19,14 @@ import { searchLogGroup } from './commands/searchLogGroup'
 import { changeLogSearchParams } from './changeLogSearch'
 import { CloudWatchLogsNode } from './explorer/cloudWatchLogsNode'
 import { loadAndOpenInitialLogStreamFile, LogStreamCodeLensProvider } from './document/logStreamsCodeLensProvider'
+import { tailLogGroup } from './commands/tailLogGroup'
+import { LiveTailDocumentProvider } from './document/liveTailDocumentProvider'
+import { LiveTailSessionRegistry } from './registry/liveTailSessionRegistry'
 import { DeployedResourceNode } from '../appBuilder/explorer/nodes/deployedNode'
 import { isTreeNode } from '../../shared/treeview/resourceTreeDataProvider'
 import { getLogger } from '../../shared/logger/logger'
 import { ToolkitError } from '../../shared'
-import {  clearDocument, closeSession, tailLogGroup } from './commands/tailLogGroup'
-import { LiveTailDocumentProvider } from './document/liveTailDocumentProvider'
-import { LiveTailSessionRegistry } from './registry/liveTailSessionRegistry'
 import { LiveTailCodeLensProvider } from './document/liveTailCodeLensProvider'
-import { tailLogGroup } from './commands/tailLogGroup'
 
 export async function activate(context: vscode.ExtensionContext, configuration: Settings): Promise<void> {
     const registry = LogDataRegistry.instance
@@ -50,6 +49,16 @@ export async function activate(context: vscode.ExtensionContext, configuration:
         vscode.workspace.registerTextDocumentContentProvider(CLOUDWATCH_LOGS_SCHEME, documentProvider)
     )
 
+    context.subscriptions.push(
+        vscode.languages.registerCodeLensProvider(
+            {
+                language: 'log',
+                scheme: cloudwatchLogsLiveTailScheme,
+            },
+            new LiveTailCodeLensProvider()
+        )
+    )
+
     context.subscriptions.push(
         vscode.workspace.registerTextDocumentContentProvider(cloudwatchLogsLiveTailScheme, liveTailDocumentProvider)
     )
@@ -113,7 +122,16 @@ export async function activate(context: vscode.ExtensionContext, configuration:
                     ? { regionName: node.regionCode, groupName: node.logGroup.logGroupName! }
                     : undefined
             await tailLogGroup(liveTailRegistry, logGroupInfo)
-        })
+        }),
+
+        Commands.register('aws.cwl.stopTailingLogGroup', async (document: vscode.TextDocument) => {
+            closeSession(document.uri, liveTailRegistry)
+        }),
+
+        Commands.register('aws.cwl.clearDocument', async (document: vscode.TextDocument) => {
+            await clearDocument(document)
+        }),
+
         Commands.register('aws.appBuilder.searchLogs', async (node: DeployedResourceNode) => {
             try {
                 const logGroupInfo = isTreeNode(node)

packages/core/src/awsService/cloudWatchLogs/commands/tailLogGroup.ts

@@ -13,7 +13,8 @@ import {
     LiveTailSessionUpdate,
     StartLiveTailResponseStream,
 } from '@aws-sdk/client-cloudwatch-logs'
-import { globals, ToolkitError } from '../../../shared'
+import { getLogger, ToolkitError } from '../../../shared'
+import { uriToKey } from '../cloudWatchLogsUtils'
 
 export async function tailLogGroup(
     registry: LiveTailSessionRegistry,
@@ -32,32 +33,29 @@ export async function tailLogGroup(
         region: wizardResponse.regionLogGroupSubmenuResponse.region,
     }
     const session = new LiveTailSession(liveTailSessionConfig)
-    if (registry.has(session.uri)) {
+    if (registry.has(uriToKey(session.uri))) {
         await prepareDocument(session)
         return
     }
-    registry.set(session.uri, session)
+    registry.set(uriToKey(session.uri), session)
 
     const document = await prepareDocument(session)
-    const timer = globals.clock.setInterval(() => {
-        session.updateStatusBarItemText()
-    }, 500)
+
     hideShowStatusBarItemsOnActiveEditor(session, document)
-    registerTabChangeCallback(session, registry, document, timer)
+    registerTabChangeCallback(session, registry, document)
 
     const stream = await session.startLiveTailSession()
 
-    await handleSessionStream(stream, document, session, timer)
+    await handleSessionStream(stream, document, session)
 }
 
-export function closeSession(sessionUri: vscode.Uri, registry: LiveTailSessionRegistry, timer: NodeJS.Timer) {
-    globals.clock.clearInterval(timer)
-    const session = registry.get(sessionUri)
+export function closeSession(sessionUri: vscode.Uri, registry: LiveTailSessionRegistry) {
+    const session = registry.get(uriToKey(sessionUri))
     if (session === undefined) {
         throw new ToolkitError(`No LiveTail session found for URI: ${sessionUri.toString()}`)
     }
     session.stopLiveTailSession()
-    registry.delete(sessionUri)
+    registry.delete(uriToKey(sessionUri))
 }
 
 export async function clearDocument(textDocument: vscode.TextDocument) {
@@ -80,8 +78,7 @@ async function prepareDocument(session: LiveTailSession): Promise<vscode.TextDoc
 async function handleSessionStream(
     stream: AsyncIterable<StartLiveTailResponseStream>,
     document: vscode.TextDocument,
-    session: LiveTailSession,
-    timer: NodeJS.Timer
+    session: LiveTailSession
 ) {
     try {
         for await (const event of stream) {
@@ -100,8 +97,21 @@ async function handleSessionStream(
                 session.isSampled = isSampled(event.sessionUpdate)
             }
         }
-    } finally {
-        globals.clock.clearInterval(timer)
+    } catch (e) {
+        if (session.isAborted) {
+            //Expected case. User action cancelled stream (CodeLens, Close Editor, etc.).
+            //AbortSignal interrupts the LiveTail stream, causing error to be thrown here.
+            //Can assume that stopLiveTailSession() has already been called - AbortSignal is only
+            //exposed through that method.
+            getLogger().info(`Session stopped: ${uriToKey(session.uri)}`)
+        } else {
+            //Unexpected exception.
+            session.stopLiveTailSession()
+            throw ToolkitError.chain(
+                e,
+                `Unexpected on-stream exception while tailing session: ${session.uri.toString()}`
+            )
+        }
     }
 }
 
@@ -196,13 +206,12 @@ function hideShowStatusBarItemsOnActiveEditor(session: LiveTailSession, document
 function registerTabChangeCallback(
     session: LiveTailSession,
     registry: LiveTailSessionRegistry,
-    document: vscode.TextDocument,
-    timer: NodeJS.Timer
+    document: vscode.TextDocument
 ) {
     vscode.window.tabGroups.onDidChangeTabs((tabEvent) => {
         const isOpen = isLiveTailSessionOpenInAnyTab(session)
         if (!isOpen) {
-            closeSession(session.uri, registry, timer)
+            closeSession(session.uri, registry)
             void clearDocument(document)
         }
     })

packages/core/src/awsService/cloudWatchLogs/document/liveTailCodeLensProvider.ts

@@ -0,0 +1,52 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as vscode from 'vscode'
+import { cloudwatchLogsLiveTailScheme } from '../../../shared/constants'
+
+export class LiveTailCodeLensProvider implements vscode.CodeLensProvider {
+    onDidChangeCodeLenses?: vscode.Event<void> | undefined
+
+    provideCodeLenses(
+        document: vscode.TextDocument,
+        token: vscode.CancellationToken
+    ): vscode.ProviderResult<vscode.CodeLens[]> {
+        const uri = document.uri
+        if (uri.scheme !== cloudwatchLogsLiveTailScheme) {
+            return []
+        }
+        const codeLenses: vscode.CodeLens[] = []
+        codeLenses.push(this.buildClearDocumentCodeLens(document))
+        codeLenses.push(this.buildStopTailingCodeLens(document))
+        return codeLenses
+    }
+
+    private buildClearDocumentCodeLens(document: vscode.TextDocument): vscode.CodeLens {
+        const range = this.getBottomOfDocumentRange(document)
+        const command: vscode.Command = {
+            title: 'Clear document',
+            command: 'aws.cwl.clearDocument',
+            arguments: [document],
+        }
+        return new vscode.CodeLens(range, command)
+    }
+
+    private buildStopTailingCodeLens(document: vscode.TextDocument): vscode.CodeLens {
+        const range = this.getBottomOfDocumentRange(document)
+        const command: vscode.Command = {
+            title: 'Stop tailing',
+            command: 'aws.cwl.stopTailingLogGroup',
+            arguments: [document],
+        }
+        return new vscode.CodeLens(range, command)
+    }
+
+    private getBottomOfDocumentRange(document: vscode.TextDocument): vscode.Range {
+        return new vscode.Range(
+            new vscode.Position(document.lineCount - 1, 0),
+            new vscode.Position(document.lineCount - 1, 0)
+        )
+    }
+}

packages/core/src/awsService/cloudWatchLogs/registry/liveTailSession.ts

@@ -10,7 +10,7 @@ import {
 } from '@aws-sdk/client-cloudwatch-logs'
 import { LogStreamFilterResponse } from '../wizard/liveTailLogStreamSubmenu'
 import { CloudWatchLogsSettings } from '../cloudWatchLogsUtils'
-import { convertToTimeString, Settings, ToolkitError } from '../../../shared'
+import { convertToTimeString, globals, Settings, ToolkitError } from '../../../shared'
 import { createLiveTailURIFromArgs } from './liveTailSessionRegistry'
 import { getUserAgent } from '../../../shared/telemetry/util'
 
@@ -39,6 +39,9 @@ export class LiveTailSession {
     private _eventRate: number
     private _isSampled: boolean
 
+    //While session is running, used to update the StatusBar each half second.
+    private statusBarUpdateTimer: NodeJS.Timer | undefined
+
     static settings = new CloudWatchLogsSettings(Settings.instance)
 
     public constructor(configuration: LiveTailSessionConfiguration) {
@@ -89,6 +92,10 @@ export class LiveTailSession {
             }
             this.startTime = Date.now()
             this.endTime = undefined
+            this.statusBarUpdateTimer = globals.clock.setInterval(() => {
+                this.updateStatusBarItemText()
+            }, 500)
+
             return commandOutput.responseStream
         } catch (e) {
             throw new ToolkitError('Encountered error while trying to start LiveTail session.')
@@ -98,6 +105,7 @@ export class LiveTailSession {
     public stopLiveTailSession() {
         this.endTime = Date.now()
         this.statusBarItem.dispose()
+        globals.clock.clearInterval(this.statusBarUpdateTimer)
         this.liveTailClient.abortController.abort()
         this.liveTailClient.cwlClient.destroy()
     }
@@ -145,4 +153,8 @@ export class LiveTailSession {
         const sampledString = this._isSampled ? 'Yes' : 'No'
         this.statusBarItem.text = `Tailing: ${timeString}, ${this._eventRate} events/sec, Sampled: ${sampledString}`
     }
+
+    public get isAborted() {
+        return this.liveTailClient.abortController.signal.aborted
+    }
 }

packages/core/src/awsService/cloudWatchLogs/registry/liveTailSessionRegistry.ts

@@ -6,7 +6,7 @@ import * as vscode from 'vscode'
 import { cloudwatchLogsLiveTailScheme } from '../../../shared/constants'
 import { LiveTailSession, LiveTailSessionConfiguration } from './liveTailSession'
 
-export class LiveTailSessionRegistry extends Map<vscode.Uri, LiveTailSession> {
+export class LiveTailSessionRegistry extends Map<string, LiveTailSession> {
     static #instance: LiveTailSessionRegistry
 
     public static get instance() {

packages/core/src/test/awsService/cloudWatchLogs/commands/tailLogGroup.test.ts

@@ -18,7 +18,7 @@ import {
     TailLogGroupWizardResponse,
 } from '../../../../awsService/cloudWatchLogs/wizard/tailLogGroupWizard'
 import { getTestWindow } from '../../../shared/vscode/window'
-import { CloudWatchLogsSettings } from '../../../../awsService/cloudWatchLogs/cloudWatchLogsUtils'
+import { CloudWatchLogsSettings, uriToKey } from '../../../../awsService/cloudWatchLogs/cloudWatchLogsUtils'
 import { installFakeClock } from '../../../testUtil'
 
 describe('TailLogGroup', function () {
@@ -125,15 +125,14 @@ describe('TailLogGroup', function () {
             .callsFake(async function () {
                 return
             })
-        // const fakeClock = installFakeClock()
-        const timer = setInterval(() => {}, 1000)
+
         const session = new LiveTailSession({
             logGroupName: testLogGroup,
             region: testRegion,
         })
-        registry.set(session.uri, session)
+        registry.set(uriToKey(session.uri), session)
 
-        closeSession(session.uri, registry, timer)
+        closeSession(session.uri, registry)
         assert.strictEqual(0, registry.size)
         assert.strictEqual(true, stopLiveTailSessionSpy.calledOnce)
         assert.strictEqual(0, clock.countTimers())