Add clear screen and stop session code lens

Author: Keegan Irby

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

@@ -19,13 +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 { clearDocument, closeSession, 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 { LiveTailCodeLensProvider } from './document/liveTailCodeLensProvider'
 
 export async function activate(context: vscode.ExtensionContext, configuration: Settings): Promise<void> {
     const registry = LogDataRegistry.instance
@@ -48,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)
     )
@@ -112,6 +123,14 @@ export async function activate(context: vscode.ExtensionContext, configuration:
             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,7 @@ import {
     LiveTailSessionUpdate,
     StartLiveTailResponseStream,
 } from '@aws-sdk/client-cloudwatch-logs'
-import { globals, ToolkitError } from '../../../shared'
+import { getLogger, ToolkitError } from '../../../shared'
 
 export async function tailLogGroup(
     registry: LiveTailSessionRegistry,
@@ -32,32 +32,29 @@ export async function tailLogGroup(
         region: wizardResponse.regionLogGroupSubmenuResponse.region,
     }
     const session = new LiveTailSession(liveTailSessionConfig)
-    if (registry.has(session.uri)) {
+    if (registry.has(session.uri.toString())) {
         await prepareDocument(session)
         return
     }
-    registry.set(session.uri, session)
+    registry.set(session.uri.toString(), 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(sessionUri.toString())
     if (session === undefined) {
         throw new ToolkitError(`No LiveTail session found for URI: ${sessionUri.toString()}`)
     }
     session.stopLiveTailSession()
-    registry.delete(sessionUri)
+    registry.delete(sessionUri.toString())
 }
 
 export async function clearDocument(textDocument: vscode.TextDocument) {
@@ -80,8 +77,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 +96,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 ${session.uri.toString()} stopped.`)
+        } else {
+            //Unexpected exception.
+            session.stopLiveTailSession()
+            throw ToolkitError.chain(
+                e,
+                `Unexpected on-stream execption while tailing session: ${session.uri.toString()}`
+            )
+        }
     }
 }
 
@@ -196,13 +205,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,51 @@
+/*!
+ * 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 = new vscode.Range(
+            new vscode.Position(document.lineCount - 1, 0),
+            new vscode.Position(document.lineCount - 1, 0)
+        )
+        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 = new vscode.Range(
+            new vscode.Position(document.lineCount - 1, 0),
+            new vscode.Position(document.lineCount - 1, 0)
+        )
+        const command: vscode.Command = {
+            title: 'Stop tailing',
+            command: 'aws.cwl.stopTailingLogGroup',
+            arguments: [document],
+        }
+        return new vscode.CodeLens(range, command)
+    }
+}

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

@@ -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(session.uri.toString(), 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())