Fix trace explorer / welcome page refresh corner cases

vscode-trace-common/src/client/tsp-client-provider-impl.ts

@@ -10,13 +10,22 @@ export class TspClientProvider implements ITspClientProvider {
     private _traceManager: TraceManager;
     private _experimentManager: ExperimentManager;
     private _listeners: ((tspClient: TspClient) => void)[] = [];
+    private _initialized = false;
 
     constructor(
         private _url: string,
         private _signalHandler: VsCodeMessageManager | undefined
     ) {
         this.updateClients();
-        RestClient.addConnectionStatusListener(status => this._signalHandler?.notifyConnection(status));
+
+        RestClient.addConnectionStatusListener(status => {
+            // Ignore the first update that is sent when calling addConnectionStatusListener
+            if (!this._initialized) {
+                this._initialized = true;
+                return;
+            }
+            this._signalHandler?.notifyConnection(status);
+        });
         this._tspClient.checkHealth(); // When this is called in the remote use-case, it will block the port-forwarding service-worker.
     }
 

vscode-trace-extension/src/extension.ts

@@ -80,7 +80,7 @@ export function activate(context: vscode.ExtensionContext): ExternalAPI {
             await startTraceServerIfAvailable(file.fsPath);
             if (await isTraceServerUp()) {
                 fileOpenHandler(context, file);
-                vscode.commands.executeCommand('setContext', 'trace-explorer.noExperiments', false);
+                vscode.commands.executeCommand('trace-explorer.refreshContext');
             }
         })
     );
@@ -106,6 +106,9 @@ export function activate(context: vscode.ExtensionContext): ExternalAPI {
 
                 // Signal the change to all trace panels
                 TraceViewerPanel.updateTraceServerUrl(newTspClientURL);
+
+                // Refresh so that either trace explorer or welcome page is rendered
+                vscode.commands.executeCommand('trace-explorer.refreshContext');
             }
         })
     );
@@ -181,6 +184,7 @@ export function activate(context: vscode.ExtensionContext): ExternalAPI {
             await startTraceServerIfAvailable(traceUri.fsPath);
             if (await isTraceServerUp()) {
                 fileOpenHandler(context, traceUri);
+                serverStatusService.updateServerStatus(true);
                 vscode.commands.executeCommand('setContext', 'trace-explorer.noExperiments', false);
             }
         })
@@ -193,24 +197,28 @@ export function activate(context: vscode.ExtensionContext): ExternalAPI {
     );
 
     context.subscriptions.push(
-        vscode.commands.registerCommand('serverStatus.started', () => {
-            serverStatusService.checkAndUpdateServerStatus();
+        vscode.commands.registerCommand('serverStatus.started', async () => {
+            await serverStatusService.updateServerStatus(true);
             if (tracesProvider) {
+                // Trigger webview refresh
                 tracesProvider.postMessagetoWebview(VSCODE_MESSAGES.TRACE_SERVER_STARTED, undefined);
             }
+            // Refresh so that either trace explorer or welcome page is rendered
+            updateNoExperimentsContext();
         })
     );
 
     context.subscriptions.push(
-        vscode.commands.registerCommand('serverStatus.stopped', () => {
-            serverStatusService.checkAndUpdateServerStatus();
+        vscode.commands.registerCommand('serverStatus.stopped', async () => {
+            await serverStatusService.updateServerStatus(false);
         })
     );
 
     context.subscriptions.push(
         vscode.commands.registerCommand('trace-explorer.refreshContext', async () => {
+            // Refresh so that either trace explorer or welcome page is rendered
             const isUp = await isTraceServerUp();
-            vscode.commands.executeCommand('setContext', 'traceViewer.serverUp', isUp);
+            await serverStatusService.updateServerStatus(isUp);
             if (isUp) {
                 await updateNoExperimentsContext();
             }
@@ -219,9 +227,13 @@ export function activate(context: vscode.ExtensionContext): ExternalAPI {
 
     vscode.commands.executeCommand('setContext', 'traceViewer.markerSetsPresent', false);
     vscode.commands.executeCommand('setContext', 'traceViewer.markerCategoriesPresent', false);
-    vscode.commands.executeCommand('setContext', 'trace-explorer.noExperiments', true);
+
+    // Initialize noExperiments/serverUp in a way so that trace explorer webviews are initialized
+    vscode.commands.executeCommand('setContext', 'trace-explorer.noExperiments', false);
+    vscode.commands.executeCommand('setContext', 'traceViewer.serverUp', true);
+
+    // Refresh to trigger rendering trace explorer or welcome page
     vscode.commands.executeCommand('trace-explorer.refreshContext');
-    serverStatusService.checkAndUpdateServerStatus();
     return traceExtensionAPI;
 }
 

vscode-trace-extension/src/trace-explorer/abstract-trace-explorer-provider.ts

@@ -6,7 +6,7 @@ import { traceExtensionWebviewManager } from 'vscode-trace-extension/src/extensi
 import { VSCODE_MESSAGES } from 'vscode-trace-common/lib/messages/vscode-message-manager';
 
 export abstract class AbstractTraceExplorerProvider implements vscode.WebviewViewProvider {
-    protected _view: vscode.WebviewView;
+    protected _view: vscode.WebviewView | undefined;
     protected _disposables: vscode.Disposable[] = [];
 
     /**
@@ -62,6 +62,11 @@ export abstract class AbstractTraceExplorerProvider implements vscode.WebviewVie
         traceExtensionWebviewManager.fireWebviewCreated(webviewView);
 
         this.init(webviewView, _context, _token);
+        webviewView.onDidDispose(_event => this.dispose(), undefined, this._disposables);
+    }
+
+    protected dispose() {
+        this._view = undefined;
     }
 
     /* eslint-disable max-len */

vscode-trace-extension/src/trace-explorer/available-views/trace-explorer-available-views-webview-provider.ts

@@ -1,13 +1,13 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import JSONBigConfig from 'json-bigint';
-import * as vscode from 'vscode';
 import { signalManager, Signals } from 'traceviewer-base/lib/signals/signal-manager';
 import { Experiment } from 'tsp-typescript-client/lib/models/experiment';
 import { OutputDescriptor } from 'tsp-typescript-client/lib/models/output-descriptor';
+import * as vscode from 'vscode';
+import { VSCODE_MESSAGES } from 'vscode-trace-common/lib/messages/vscode-message-manager';
+import { convertSignalExperiment } from 'vscode-trace-common/lib/signals/vscode-signal-converter';
 import { TraceViewerPanel } from '../../trace-viewer-panel/trace-viewer-webview-panel';
 import { getTspClientUrl } from '../../utils/backend-tsp-client-provider';
-import { convertSignalExperiment } from 'vscode-trace-common/lib/signals/vscode-signal-converter';
-import { VSCODE_MESSAGES } from 'vscode-trace-common/lib/messages/vscode-message-manager';
 import { AbstractTraceExplorerProvider } from '../abstract-trace-explorer-provider';
 
 const JSONBig = JSONBigConfig({
@@ -39,13 +39,14 @@ export class TraceExplorerAvailableViewsProvider extends AbstractTraceExplorerPr
                 const data: any = message.data;
                 switch (command) {
                     case VSCODE_MESSAGES.CONNECTION_STATUS:
-                        if (data && data.status) {
-                            this._statusService.checkAndUpdateServerStatus();
+                        if (data?.status) {
+                            const status: boolean = JSON.parse(message.data.status);
+                            this._statusService.updateServerStatus(status);
                         }
                         return;
                     case VSCODE_MESSAGES.WEBVIEW_READY:
                         // Post the tspTypescriptClient
-                        this._view.webview.postMessage({
+                        this._view?.webview.postMessage({
                             command: VSCODE_MESSAGES.SET_TSP_CLIENT,
                             data: getTspClientUrl()
                         });
@@ -86,13 +87,11 @@ export class TraceExplorerAvailableViewsProvider extends AbstractTraceExplorerPr
         );
 
         signalManager().on(Signals.EXPERIMENT_SELECTED, this._onExperimentSelected);
-        webviewView.onDidDispose(
-            _event => {
-                signalManager().off(Signals.EXPERIMENT_SELECTED, this._onExperimentSelected);
-            },
-            undefined,
-            this._disposables
-        );
+    }
+
+    protected dispose() {
+        signalManager().off(Signals.EXPERIMENT_SELECTED, this._onExperimentSelected);
+        super.dispose();
     }
 
     private _onExperimentSelected = (experiment: Experiment | undefined): void =>

vscode-trace-extension/src/trace-explorer/opened-traces/trace-explorer-opened-traces-webview-provider.ts

@@ -65,13 +65,14 @@ export class TraceExplorerOpenedTracesViewProvider extends AbstractTraceExplorer
                 const data: any = message.data;
                 switch (command) {
                     case VSCODE_MESSAGES.CONNECTION_STATUS:
-                        if (data && data.status) {
-                            this._statusService.checkAndUpdateServerStatus();
+                        if (data?.status) {
+                            const status: boolean = JSON.parse(message.data.status);
+                            this._statusService.updateServerStatus(status);
                         }
                         return;
                     case VSCODE_MESSAGES.WEBVIEW_READY:
                         // Post the tspTypescriptClient
-                        this._view.webview.postMessage({
+                        this._view?.webview.postMessage({
                             command: VSCODE_MESSAGES.SET_TSP_CLIENT,
                             data: getTspClientUrl()
                         });
@@ -137,4 +138,10 @@ export class TraceExplorerOpenedTracesViewProvider extends AbstractTraceExplorer
             this._disposables
         );
     }
+    protected dispose() {
+        signalManager().off(Signals.TRACEVIEWERTAB_ACTIVATED, this._onOpenedTracesWidgetActivated);
+        signalManager().off(Signals.EXPERIMENT_SELECTED, this._onExperimentSelected);
+        signalManager().off(Signals.EXPERIMENT_OPENED, this._onExperimentOpened);
+        super.dispose();
+    }
 }

vscode-trace-extension/src/trace-explorer/time-range/trace-explorer-time-range-data-webview-provider.ts

@@ -41,12 +41,12 @@ export class TraceExplorerTimeRangeDataProvider extends AbstractTraceExplorerPro
         );
 
         webviewView.onDidChangeVisibility(() => {
-            if (this._view.visible) {
+            if (this._view?.visible) {
                 const data = {
                     mapArray: Array.from(this._experimentDataMap.experimentDataMap.values()),
                     activeData: this._experimentDataMap.activeData
                 };
-                this._view.webview.postMessage({
+                this._view?.webview.postMessage({
                     command: VSCODE_MESSAGES.RESTORE_VIEW,
                     data: JSONBig.stringify(data)
                 });
@@ -59,31 +59,28 @@ export class TraceExplorerTimeRangeDataProvider extends AbstractTraceExplorerPro
         signalManager().on(Signals.EXPERIMENT_UPDATED, this.onExperimentUpdated);
         signalManager().on(Signals.EXPERIMENT_CLOSED, this.onExperimentClosed);
         signalManager().on(Signals.CLOSE_TRACEVIEWERTAB, this.onExperimentTabClosed);
+    }
 
-        webviewView.onDidDispose(
-            _event => {
-                signalManager().off(Signals.VIEW_RANGE_UPDATED, this.onViewRangeUpdated);
-                signalManager().off(Signals.SELECTION_RANGE_UPDATED, this.onSelectionRangeUpdated);
-                signalManager().off(Signals.EXPERIMENT_SELECTED, this.onExperimentSelected);
-                signalManager().off(Signals.EXPERIMENT_UPDATED, this.onExperimentUpdated);
-                signalManager().off(Signals.EXPERIMENT_CLOSED, this.onExperimentClosed);
-                signalManager().off(Signals.CLOSE_TRACEVIEWERTAB, this.onExperimentTabClosed);
-            },
-            undefined,
-            this._disposables
-        );
+    protected dispose() {
+        signalManager().off(Signals.VIEW_RANGE_UPDATED, this.onViewRangeUpdated);
+        signalManager().off(Signals.SELECTION_RANGE_UPDATED, this.onSelectionRangeUpdated);
+        signalManager().off(Signals.EXPERIMENT_SELECTED, this.onExperimentSelected);
+        signalManager().off(Signals.EXPERIMENT_UPDATED, this.onExperimentUpdated);
+        signalManager().off(Signals.EXPERIMENT_CLOSED, this.onExperimentClosed);
+        signalManager().off(Signals.CLOSE_TRACEVIEWERTAB, this.onExperimentTabClosed);
+        super.dispose();
     }
 
     private onViewRangeUpdated = (update: TimeRangeUpdatePayload) => {
-        this._view.webview.postMessage({
+        this._view?.webview.postMessage({
             command: VSCODE_MESSAGES.VIEW_RANGE_UPDATED,
             data: JSONBig.stringify(update)
         });
         this._experimentDataMap.updateViewRange(update);
     };
 
     private onSelectionRangeUpdated = (update: TimeRangeUpdatePayload) => {
-        this._view.webview.postMessage({
+        this._view?.webview.postMessage({
             command: VSCODE_MESSAGES.SELECTION_RANGE_UPDATED,
             data: JSONBig.stringify(update)
         });
@@ -92,7 +89,7 @@ export class TraceExplorerTimeRangeDataProvider extends AbstractTraceExplorerPro
 
     private onExperimentSelected = (experiment: Experiment | undefined) => {
         const data = { wrapper: experiment ? JSONBig.stringify(experiment) : undefined };
-        this._view.webview.postMessage({ command: VSCODE_MESSAGES.EXPERIMENT_SELECTED, data });
+        this._view?.webview.postMessage({ command: VSCODE_MESSAGES.EXPERIMENT_SELECTED, data });
         if (experiment) {
             this._experimentDataMap.updateAbsoluteRange(experiment);
         }
@@ -101,18 +98,18 @@ export class TraceExplorerTimeRangeDataProvider extends AbstractTraceExplorerPro
 
     private onExperimentUpdated = (experiment: Experiment) => {
         const data = { wrapper: JSONBig.stringify(experiment) };
-        this._view.webview.postMessage({ command: VSCODE_MESSAGES.EXPERIMENT_UPDATED, data });
+        this._view?.webview.postMessage({ command: VSCODE_MESSAGES.EXPERIMENT_UPDATED, data });
         this._experimentDataMap.updateAbsoluteRange(experiment);
     };
 
     private onExperimentClosed = (experiment: Experiment) => {
         const data = { wrapper: JSONBig.stringify(experiment) };
-        this._view.webview.postMessage({ command: VSCODE_MESSAGES.EXPERIMENT_CLOSED, data });
+        this._view?.webview.postMessage({ command: VSCODE_MESSAGES.EXPERIMENT_CLOSED, data });
         this._experimentDataMap.delete(experiment);
     };
 
     private onExperimentTabClosed = (experimentUUID: string) => {
-        this._view.webview.postMessage({ command: VSCODE_MESSAGES.TRACE_VIEWER_TAB_CLOSED, data: experimentUUID });
+        this._view?.webview.postMessage({ command: VSCODE_MESSAGES.TRACE_VIEWER_TAB_CLOSED, data: experimentUUID });
         this._experimentDataMap.delete(experimentUUID);
     };
 }

vscode-trace-extension/src/trace-viewer-panel/trace-viewer-webview-panel.ts

@@ -255,7 +255,8 @@ export class TraceViewerPanel {
                         return;
                     case VSCODE_MESSAGES.CONNECTION_STATUS:
                         if (message.data?.status && this._statusService) {
-                            this._statusService.checkAndUpdateServerStatus();
+                            const status: boolean = JSON.parse(message.data.status);
+                            this._statusService.updateServerStatus(status);
                         }
                         return;
                     case VSCODE_MESSAGES.SHOW_MARKER_CATEGORIES:

vscode-trace-extension/src/utils/trace-server-status.ts

@@ -1,5 +1,6 @@
 import { ThemeColor, StatusBarItem } from 'vscode';
 import { isTraceServerUp } from './backend-tsp-client-provider';
+import * as vscode from 'vscode';
 
 export class TraceServerConnectionStatusService {
     private statusBarItem: StatusBarItem;
@@ -11,9 +12,14 @@ export class TraceServerConnectionStatusService {
 
     public checkAndUpdateServerStatus = async (): Promise<void> => {
         const isUp = await isTraceServerUp();
-        this.render(isUp);
+        await this.updateServerStatus(isUp);
     };
 
+    public async updateServerStatus(status: boolean): Promise<void> {
+        await vscode.commands.executeCommand('setContext', 'traceViewer.serverUp', status);
+        this.render(status);
+    }
+
     private render = (status: boolean): void => {
         if (status) {
             this.statusBarItem.backgroundColor = new ThemeColor('statusBarItem.warningBackground');