Fix polling mechanism for DocDB explorer

Muzaffer Aydin · Muzaffer Aydin · commit bbe560aa6a0a · 2024-11-29T13:10:43.000-08:00
diff --git a/packages/core/src/docdb/activation.ts b/packages/core/src/docdb/activation.ts
@@ -23,56 +23,78 @@ import { renameInstance } from './commands/renameInstance'
 import { addTag, listTags, removeTag } from './commands/tagCommands'
 import { Uri } from 'vscode'
 import { openUrl } from '../shared/utilities/vsCodeUtils'
+import { getLogger } from '../shared/logger'
 
 /**
- * Activates DocumentDB components.
+ * A utility function to automatically invoke trackChanges after a command.
  */
 
+function withTrackChanges<T extends DBResourceNode>(
+    command: (node: T) => Promise<void>,
+    commandName: string = 'UnnamedCommand'
+): (node: T) => Promise<void> {
+    return async (node: T) => {
+        const arn = node.arn || 'UnknownARN'
+        const startTime = new Date().toISOString()
+
+        getLogger().info(
+            `[${startTime}] Executing command "${commandName}" for resource with ARN: ${arn}. Tracking changes will be invoked post-execution.`
+        )
+
+        await command(node)
+
+        const endTime = new Date().toISOString()
+        getLogger().info(
+            `[${endTime}] Successfully executed command "${commandName}" for resource with ARN: ${arn}. Invoking trackChanges now.`
+        )
+
+        await node.trackChangesWithWaitProcessingStatus()
+    }
+}
+
+/**
+ * Activates DocumentDB components.
+ */
 export async function activate(ctx: ExtContext): Promise<void> {
     ctx.extensionContext.subscriptions.push(
         Commands.register('aws.docdb.createCluster', async (node?: DocumentDBNode) => {
             await createCluster(node)
         }),
 
-        Commands.register('aws.docdb.deleteCluster', async (node: DBClusterNode) => {
-            await deleteCluster(node)
-        }),
+        Commands.register('aws.docdb.deleteCluster', withTrackChanges<DBClusterNode>(deleteCluster, 'deleteCluster')),
 
-        Commands.register('aws.docdb.renameCluster', async (node: DBClusterNode) => {
-            await renameCluster(node)
-        }),
+        Commands.register('aws.docdb.renameCluster', withTrackChanges<DBClusterNode>(renameCluster, 'renameCluster')),
 
-        Commands.register('aws.docdb.startCluster', async (node?: DBClusterNode) => {
-            await startCluster(node)
-        }),
+        Commands.register('aws.docdb.startCluster', withTrackChanges<DBClusterNode>(startCluster, 'startCluster')),
 
-        Commands.register('aws.docdb.stopCluster', async (node?: DBClusterNode) => {
-            await stopCluster(node)
-        }),
+        Commands.register('aws.docdb.stopCluster', withTrackChanges<DBClusterNode>(stopCluster, 'stopCluster')),
 
-        Commands.register('aws.docdb.addRegion', async (node: DBClusterNode) => {
-            await addRegion(node)
-        }),
+        Commands.register('aws.docdb.addRegion', withTrackChanges<DBClusterNode>(addRegion, 'addRegion')),
 
-        Commands.register('aws.docdb.createInstance', async (node: DBClusterNode) => {
-            await createInstance(node)
-        }),
+        Commands.register(
+            'aws.docdb.createInstance',
+            withTrackChanges<DBClusterNode>(createInstance, 'createInstance')
+        ),
 
-        Commands.register('aws.docdb.deleteInstance', async (node: DBInstanceNode) => {
-            await deleteInstance(node)
-        }),
+        Commands.register(
+            'aws.docdb.deleteInstance',
+            withTrackChanges<DBInstanceNode>(deleteInstance, 'deleteInstance')
+        ),
 
-        Commands.register('aws.docdb.modifyInstance', async (node: DBInstanceNode) => {
-            await modifyInstance(node)
-        }),
+        Commands.register(
+            'aws.docdb.modifyInstance',
+            withTrackChanges<DBInstanceNode>(modifyInstance, 'modifyInstance')
+        ),
 
-        Commands.register('aws.docdb.rebootInstance', async (node: DBInstanceNode) => {
-            await rebootInstance(node)
-        }),
+        Commands.register(
+            'aws.docdb.rebootInstance',
+            withTrackChanges<DBInstanceNode>(rebootInstance, 'rebootInstance')
+        ),
 
-        Commands.register('aws.docdb.renameInstance', async (node: DBInstanceNode) => {
-            await renameInstance(node)
-        }),
+        Commands.register(
+            'aws.docdb.renameInstance',
+            withTrackChanges<DBInstanceNode>(renameInstance, 'renameInstance')
+        ),
 
         Commands.register('aws.docdb.listTags', async (node: DBResourceNode) => {
             await listTags(node)
@@ -90,7 +112,7 @@ export async function activate(ctx: ExtContext): Promise<void> {
             await node?.openInBrowser()
         }),
 
-        Commands.register('aws.docdb.viewDocs', async (node?: DBResourceNode) => {
+        Commands.register('aws.docdb.viewDocs', async () => {
             await openUrl(
                 Uri.parse('https://docs.aws.amazon.com/documentdb/latest/developerguide/get-started-guide.html')
             )
diff --git a/packages/core/src/docdb/explorer/dbClusterNode.ts b/packages/core/src/docdb/explorer/dbClusterNode.ts
@@ -50,16 +50,17 @@ export class DBClusterNode extends DBResourceNode {
         )
         this.description = this.getDescription()
         this.tooltip = `${this.name}${os.EOL}Engine: ${this.cluster.EngineVersion}${os.EOL}Status: ${this.cluster.Status}`
-        this.trackChanges()
+        if (this.isStatusRequiringPolling()) {
+            getLogger().info(`${this.arn} requires polling.`)
+            this.trackChanges()
+        } else {
+            getLogger().info(`${this.arn} does NOT require polling.`)
+        }
     }
 
     public override async getChildren(): Promise<AWSTreeNodeBase[]> {
+        getLogger().info(`DBClusterNode.getChildren() called`)
         return telemetry.docdb_listInstances.run(async () => {
-            this.childNodes.forEach((node) => {
-                getLogger().info(`(getChildren) Removing Polling from node: ${node.arn}`)
-                node.pollingSet.delete(node.arn)
-                node.pollingSet.clearTimer()
-            })
             return await makeChildrenNodes({
                 getChildNodes: async () => {
                     this.instances = (await this.client.listInstances([this.arn])).map((i) => {
@@ -143,7 +144,7 @@ export class DBClusterNode extends DBResourceNode {
 
     override async getStatus() {
         const [cluster] = await this.client.listClusters(this.arn)
-        getLogger().info(`Updating Status: new status ${cluster?.Status} for cluster ${this.arn}`)
+        getLogger().info(`Get Status: status ${cluster?.Status} for cluster ${this.arn}`)
         this.cluster.Status = cluster?.Status
         return cluster.Status
     }
@@ -162,7 +163,6 @@ export class DBClusterNode extends DBResourceNode {
     }
 
     override refreshTree(): void {
-        this.clearTimer()
         this.refresh()
         this.parent.refresh()
     }
diff --git a/packages/core/src/docdb/explorer/dbElasticClusterNode.ts b/packages/core/src/docdb/explorer/dbElasticClusterNode.ts
@@ -12,6 +12,7 @@ import { DocDBContext } from './docdbContext'
 import { copyToClipboard } from '../../shared/utilities/messages'
 import { localize } from '../../shared/utilities/vsCodeUtils'
 import { getAwsConsoleUrl } from '../../shared/awsConsole'
+import { getLogger } from '../../shared/logger'
 
 /**
  * An AWS Explorer node representing DocumentDB elastic clusters.
@@ -32,6 +33,12 @@ export class DBElasticClusterNode extends DBResourceNode {
         )
         this.description = this.getDescription()
         this.tooltip = `${this.name}\nStatus: ${this.status}`
+        if (this.isStatusRequiringPolling()) {
+            getLogger().info(`${this.arn} requires polling.`)
+            this.trackChanges()
+        } else {
+            getLogger().info(`${this.arn} does NOT require polling.`)
+        }
     }
 
     private getContext() {
@@ -70,7 +77,9 @@ export class DBElasticClusterNode extends DBResourceNode {
 
     override async getStatus() {
         const cluster = await this.client.getElasticCluster(this.arn)
-        return cluster?.status
+        getLogger().info(`Get Status: status ${cluster?.status} for elastic-cluster ${this.arn}`)
+        this.cluster.status = cluster?.status
+        return cluster?.status?.toLowerCase()
     }
 
     override get isAvailable() {
diff --git a/packages/core/src/docdb/explorer/dbGlobalClusterNode.ts b/packages/core/src/docdb/explorer/dbGlobalClusterNode.ts
@@ -17,6 +17,7 @@ import { DBResourceNode } from './dbResourceNode'
 import { DocDBContext } from './docdbContext'
 import { copyToClipboard } from '../../shared/utilities/messages'
 import { getAwsConsoleUrl } from '../../shared/awsConsole'
+import { getLogger } from '../../shared/logger'
 
 function getRegionFromArn(arn: string) {
     const match = arn.match(/:rds:([^:]+):.*:cluster:/)
@@ -46,6 +47,12 @@ export class DBGlobalClusterNode extends DBResourceNode {
         this.iconPath = new vscode.ThemeIcon('globe') //TODO: determine icon for global cluster
         this.description = 'Global cluster'
         this.tooltip = `${this.name}\nEngine: ${this.cluster.EngineVersion}\nStatus: ${this.cluster.Status} (read-only)`
+        if (this.isStatusRequiringPolling()) {
+            getLogger().info(`${this.arn} requires polling.`)
+            this.trackChanges()
+        } else {
+            getLogger().info(`${this.arn} does NOT require polling.`)
+        }
     }
 
     public override async getChildren(): Promise<AWSTreeNodeBase[]> {
diff --git a/packages/core/src/docdb/explorer/dbInstanceNode.ts b/packages/core/src/docdb/explorer/dbInstanceNode.ts
@@ -32,7 +32,25 @@ export class DBInstanceNode extends DBResourceNode {
         this.contextValue = this.getContext()
         this.iconPath = this.isAvailable || this.isStopped ? undefined : new vscode.ThemeIcon('loading~spin')
         this.tooltip = `${this.name}\nClass: ${this.instance.DBInstanceClass}\nStatus: ${this.status}`
-        this.trackChanges()
+        getLogger().info(`Parent of ${instance.DBInstanceArn} is ${parent.arn}`)
+        if (this.isStatusRequiringPolling()) {
+            getLogger().info(`${instance.DBInstanceArn} requires polling.`)
+            this.trackChanges()
+        } else {
+            getLogger().info(`${instance.DBInstanceArn} does NOT require polling.`)
+        }
+    }
+
+    public override isStatusRequiringPolling(): boolean {
+        const instanceRequiresPolling = super.isStatusRequiringPolling()
+        const parentRequiresPolling = this.parent.isStatusRequiringPolling()
+        const requiresPolling = instanceRequiresPolling || parentRequiresPolling
+
+        getLogger().info(
+            `isStatusRequiringPolling (DBInstanceNode): Instance ${this.arn} requires polling: ${instanceRequiresPolling}, Parent ${this.parent.arn} requires polling: ${parentRequiresPolling}, Combined result: ${requiresPolling}`
+        )
+
+        return requiresPolling
     }
 
     private makeDescription(): string {
@@ -85,7 +103,6 @@ export class DBInstanceNode extends DBResourceNode {
     }
 
     override refreshTree(): void {
-        this.clearTimer()
         this.refresh()
         this.parent.refresh()
     }
diff --git a/packages/core/src/docdb/explorer/dbResourceNode.ts b/packages/core/src/docdb/explorer/dbResourceNode.ts
@@ -17,7 +17,16 @@ export abstract class DBResourceNode extends AWSTreeNodeBase implements AWSResou
     public override readonly regionCode: string
     public abstract readonly arn: string
     public abstract readonly name: string
-    public readonly pollingSet: PollingSet<string> = new PollingSet(10000, this.updateNodeStatus.bind(this))
+    public readonly pollingSet: PollingSet<string> = new PollingSet(30000, this.updateNodeStatus.bind(this))
+    private static readonly globalPollingArns: Set<string> = new Set<string>()
+    public processingStatuses = new Set<string>([
+        'creating',
+        'modifying',
+        'rebooting',
+        'starting',
+        'stopping',
+        'renaming',
+    ])
 
     protected constructor(
         public readonly client: DocumentDBClient,
@@ -29,6 +38,15 @@ export abstract class DBResourceNode extends AWSTreeNodeBase implements AWSResou
         getLogger().info(`NEW DBResourceNode`)
     }
 
+    public isStatusRequiringPolling(): boolean {
+        const currentStatus = this.status?.toLowerCase()
+        const isProcessingStatus = currentStatus !== undefined && this.processingStatuses.has(currentStatus)
+        getLogger().info(
+            `isStatusRequiringPolling (DBResourceNode):: Checking if status "${currentStatus}" for ARN: ${this.arn} requires polling: ${isProcessingStatus}`
+        )
+        return isProcessingStatus
+    }
+
     public [inspect.custom](): string {
         return 'DBResourceNode'
     }
@@ -49,30 +67,75 @@ export abstract class DBResourceNode extends AWSTreeNodeBase implements AWSResou
         return this.status === 'stopped'
     }
 
-    public async waitUntilStatusChanged(): Promise<boolean> {
-        const currentStatus = this.status
+    public get isPolling(): boolean {
+        const isPolling = DBResourceNode.globalPollingArns.has(this.arn)
+        getLogger().info(`isPolling: ARN ${this.arn} is ${isPolling ? '' : 'not '}being polled.`)
+        return isPolling
+    }
 
+    public set isPolling(value: boolean) {
+        if (value) {
+            if (!this.isPolling) {
+                DBResourceNode.globalPollingArns.add(this.arn)
+                getLogger().info(`Polling started for ARN: ${this.arn}`)
+            } else {
+                getLogger().info(`Polling already active for ARN: ${this.arn}`)
+            }
+        } else {
+            if (this.isPolling) {
+                DBResourceNode.globalPollingArns.delete(this.arn)
+                getLogger().info(`Polling stopped for ARN: ${this.arn}`)
+            } else {
+                getLogger().info(`Polling was not active for ARN: ${this.arn}`)
+            }
+        }
+    }
+
+    public async waitUntilStatusChanged(
+        checkProcessingStatuses: boolean = false,
+        timeout: number = 1200000,
+        interval: number = 5000
+    ): Promise<boolean> {
         await waitUntil(
             async () => {
                 const status = await this.getStatus()
-                getLogger().info('docdb: waitUntilStatusChanged (status): %O', status)
-                return status !== currentStatus
+                if (checkProcessingStatuses) {
+                    const isProcessingStatus = status !== undefined && this.processingStatuses.has(status.toLowerCase())
+                    getLogger().info('docdb: waitUntilStatusChangedToProcessingStatus: %O', isProcessingStatus)
+                    return isProcessingStatus
+                } else {
+                    const hasStatusChanged = status !== this.status
+                    getLogger().info('docdb: waitUntilStatusChanged (status): %O', hasStatusChanged)
+                    return hasStatusChanged
+                }
             },
-            { timeout: 1200000, interval: 5000, truthy: true }
+            { timeout, interval, truthy: true }
         )
-
+        this.refreshTree()
         return false
     }
 
-    public trackChanges() {
+    public async trackChangesWithWaitProcessingStatus() {
         getLogger().info(
-            `Preparing to track changes for ARN: ${this.arn}; pollingSet: ${this.pollingSet.has(this.arn)}; condition: ${this.pollingSet.has(this.arn) === false}`
+            `Preparing to track changes with waiting a processing status for ARN: ${this.arn}; condition: ${this.isPolling};`
         )
-        if (this.pollingSet.has(this.arn) === false) {
+        if (!this.isPolling) {
+            this.isPolling = true
+            await this.waitUntilStatusChanged(true, 60000, 1000)
+            getLogger().info(`Tracking changes for a processing status wait is over`)
             this.pollingSet.start(this.arn)
-            getLogger().info(
-                `Tracking changes for ARN: ${this.arn}; pollingSet: ${this.pollingSet.has(this.arn)}; condition: ${this.pollingSet.has(this.arn) === false}`
-            )
+            getLogger().info(`Tracking changes for ARN: ${this.arn}; condition: ${this.isPolling};`)
+        } else {
+            getLogger().info(`ARN: ${this.arn} already being tracked`)
+        }
+    }
+
+    public trackChanges() {
+        getLogger().info(`Preparing to track immdiately for ARN: ${this.arn}; condition: ${this.isPolling};`)
+        if (!this.isPolling) {
+            this.isPolling = true
+            this.pollingSet.start(this.arn)
+            getLogger().info(`Tracking changes for ARN: ${this.arn}; condition: ${this.isPolling};`)
         } else {
             getLogger().info(`ARN: ${this.arn} already being tracked`)
         }
@@ -98,8 +161,14 @@ export abstract class DBResourceNode extends AWSTreeNodeBase implements AWSResou
         )
         if (currentStatus !== newStatus) {
             getLogger().info(`docdb: ${this.arn} updateNodeStatus - refreshing UI`)
+            this.refreshTree()
+        }
+        if (!this.isStatusRequiringPolling()) {
+            getLogger().info(`docdb: ${this.arn} updateNodeStatus - refreshing UI`)
+            getLogger().info(`pollingSet delete ${this.arn} updateNodeStatus`)
             this.pollingSet.delete(this.arn)
             this.pollingSet.clearTimer()
+            this.isPolling = false
             this.refreshTree()
         }
     }
diff --git a/packages/core/src/docdb/explorer/docdbNode.ts b/packages/core/src/docdb/explorer/docdbNode.ts
