Merge pull request #3071 from codefori/cleanDisconnect

Dispose SSH connection when disconnecting & fixed disconnection events handling

Author: sebjulliand

package-lock.json

@@ -3247,6 +3247,7 @@
 		"@types/glob": "^7.1.3",
 		"@types/node": "22.19.7",
 		"@types/source-map-support": "^0.5.6",
+		"@types/ssh2": "^1.15.5",
 		"@types/tar": "6.1.13",
 		"@types/tmp": "0.2.5",
 		"@types/vscode": "^1.90.0",
@@ -3281,4 +3282,4 @@
 		"halcyontechltd.vscode-ibmi-walkthroughs",
 		"vscode.git"
 	]
-}
+}

package.json

@@ -1,7 +1,7 @@
 import { EventEmitter } from "stream";
 import * as vscode from "vscode";
 import { ILELibrarySettings } from "./api/CompileTools";
-import IBMi, { ConnectionResult } from "./api/IBMi";
+import IBMi, { ConnectionResult, DisconnectedCallback } from "./api/IBMi";
 import { BaseStorage } from "./api/configuration/storage/BaseStorage";
 import { CodeForIStorage } from "./api/configuration/storage/CodeForIStorage";
 import { ConnectionStorage } from "./api/configuration/storage/ConnectionStorage";
@@ -79,70 +79,71 @@ export default class Instance {
 
     let result: ConnectionResult;
 
-    const timeoutHandler = async (conn: IBMi) => {
-      if (conn) {
-        const choice = await vscode.window.showWarningMessage(`Connection lost`, {
-          modal: true,
-          detail: `Connection to ${conn.currentConnectionName} has dropped. Would you like to reconnect?`
-        }, `Yes`, `No, get logs`);
+    const onDisconnected: DisconnectedCallback = async (connection, error) => {
+      if (connection.connectionSuccessful) {
+        this.fire(`disconnected`);
 
-        let reconnect = choice === `Yes`;
-        let collectLogs = choice === `No, get logs`;
+        if (error) {
+          const choice = await vscode.window.showWarningMessage(`Connection lost: ${error.description || error.message || error.level}`, {
+            modal: true,
+            detail: `Connection to ${connection.currentConnectionName} has dropped. Would you like to reconnect?`
+          }, `Yes`, `No, get logs`);
 
-        if (collectLogs) {
-          const logs = this.output.content;
-          vscode.workspace.openTextDocument({ content: logs, language: `plaintext` }).then(doc => {
-            vscode.window.showTextDocument(doc);
-          });
-        }
+          const reconnect = choice === `Yes`;
+          const collectLogs = choice === `No, get logs`;
 
-        this.disconnect();
+          if (collectLogs) {
+            const logs = this.output.content;
+            vscode.workspace.openTextDocument({ content: logs, language: `plaintext` }).then(doc => {
+              vscode.window.showTextDocument(doc);
+            });
+          }
+
+          this.disconnect();
 
-        if (reconnect) {
-          await this.connect({ ...options, reconnecting: true });
+          if (reconnect) {
+            await this.connect({ ...options, reconnecting: true });
+          }
         }
       }
+      else {
+        this.disconnect();
+      }
     };
 
     return VscodeTools.withContext("code-for-ibmi:connecting", async () => {
       while (true) {
         let customError: string | undefined;
         await vscode.window.withProgress({ location: vscode.ProgressLocation.Notification, title: options.data.name, cancellable: true }, async (p, cancelToken) => {
-          try {
-            const cancelEmitter = new EventEmitter();
+          const cancelEmitter = new EventEmitter();
 
-            cancelToken.onCancellationRequested(() => {
-              cancelEmitter.emit(`cancel`);
-            });
+          cancelToken.onCancellationRequested(() => {
+            cancelEmitter.emit(`cancel`);
+          });
 
-            result = await connection.connect(
-              options.data,
-              {
-                callbacks: {
-                  timeoutCallback: timeoutHandler,
-                  onConnectedOperations: options.onConnectedOperations || [],
-                  uiErrorHandler: handleConnectionResults,
-                  progress: (message) => { p.report(message) },
-                  message: messageCallback,
-                  inputBox: async (prompt: string, placeHolder: string, ignoreFocusOut: boolean) => {
-                    return await inputBoxCallback(prompt, placeHolder, ignoreFocusOut, cancelToken)
-                  },
-                  cancelEmitter
+          result = await connection.connect(
+            options.data,
+            {
+              callbacks: {
+                onDisconnected,
+                onConnectedOperations: options.onConnectedOperations,
+                uiErrorHandler: handleConnectionResults,
+                progress: (message) => { p.report(message) },
+                message: messageCallback,
+                inputBox: async (prompt: string, placeHolder: string, ignoreFocusOut: boolean) => {
+                  return await inputBoxCallback(prompt, placeHolder, ignoreFocusOut, cancelToken)
                 },
-                reconnecting: options.reconnecting,
-                reloadServerSettings: options.reloadServerSettings,
+                cancelEmitter
               },
-            );
-          } catch (e: any) {
-            customError = e.message;
-            result = { success: false };
-          }
+              reconnecting: options.reconnecting,
+              reloadServerSettings: options.reloadServerSettings,
+            },
+          )
         });
 
         if (result.success) {
           await this.setConnection(connection);
           break;
-
         } else {
           await this.disconnect();
           if (options.reconnecting && await vscode.window.showWarningMessage(`Could not reconnect`, {
@@ -159,10 +160,6 @@ export default class Instance {
         }
       }
 
-      if (result.success === false) {
-        connection.dispose();
-      }
-
       return result;
     });
   }
@@ -179,15 +176,10 @@ export default class Instance {
 
   private async setConnection(connection?: IBMi) {
     if (this.connection) {
-      await this.connection.dispose();
+      this.connection.disconnect();
     }
 
     if (connection) {
-      connection.setDisconnectedCallback(async () => {
-        this.setConnection();
-        this.fire(`disconnected`);
-      });
-
       this.connection = connection;
       this.storage.setConnectionName(connection.currentConnectionName);
       await IBMi.GlobalStorage.setLastConnection(connection.currentConnectionName);

src/Instance.ts

@@ -25,6 +25,11 @@ export namespace CompileTools {
     updateProgress?: (message: string) => void
   }
 
+  export function reset(){
+    ileQueue.clear();
+    jobLogOrdinal = 0;
+  }
+
   /**
    * Execute a command
    */

src/api/CompileTools.ts

@@ -1,6 +1,7 @@
 import { BindingValue } from "@ibm/mapepire-js";
 import * as node_ssh from "node-ssh";
 import path, { parse as parsePath } from 'path';
+import { ClientErrorExtensions } from "ssh2";
 import { EventEmitter } from 'stream';
 import { CompileTools } from "./CompileTools";
 import IBMiContent from "./IBMiContent";
@@ -30,6 +31,8 @@ export interface ConnectionResult {
   errorCodes?: ConnectionErrorCode[]
 }
 
+export type DisconnectedCallback = (conn: IBMi, error?: Error & ClientErrorExtensions) => Promise<void>;
+
 const remoteApps = [ // All names MUST also be defined as key in 'remoteFeatures' below!!
   {
     path: `/usr/bin/`,
@@ -46,11 +49,9 @@ const remoteApps = [ // All names MUST also be defined as key in 'remoteFeatures
   }
 ];
 
-type DisconnectCallback = (conn: IBMi) => Promise<void>;
-
 interface ConnectionCallbacks {
   onConnectedOperations?: Function[],
-  timeoutCallback?: (conn: IBMi) => Promise<void>,
+  onDisconnected?: DisconnectedCallback,
   uiErrorHandler: (connection: IBMi, error: ConnectionErrorCode, data?: any) => Promise<boolean>,
   progress: (detail: { message: string }) => void,
   message: (type: ConnectionMessageType, message: string) => void,
@@ -117,6 +118,8 @@ export default class IBMi {
   private currentAsp: string | undefined;
   private libraryAsps = new Map<string, number>();
 
+  connectionSuccessful = false;
+
   /**
    * @deprecated Will be replaced with {@link IBMi.getAllIAsps} in v3.0.0
    */
@@ -148,15 +151,6 @@ export default class IBMi {
     process.stdout.write(text);
   };
 
-  private disconnectedCallback: (DisconnectCallback) | undefined;
-
-  /**
-   * Will only be called once per connection.
-   */
-  setDisconnectedCallback(callback: DisconnectCallback) {
-    this.disconnectedCallback = callback;
-  }
-
   /**
    * getConfigFile can return pre-defined configuration files,
    * but can lazy load new configuration files as well.
@@ -268,7 +262,7 @@ export default class IBMi {
       if (callbacks.cancelEmitter) {
         callbacks.cancelEmitter.once('cancel', () => {
           wasCancelled = true;
-          this.dispose();
+          this.disconnect();
         });
       }
 
@@ -297,7 +291,8 @@ export default class IBMi {
         privateKeyPath: connectionObject.privateKeyPath ? Tools.resolvePath(connectionObject.privateKeyPath) : undefined,
         passphrase: connectionObject.privateKeyPath ? connectionObject.passphrase : undefined,
         debug: connectionObject.sshDebug ? (message: string) => this.appendOutput(`\n[SSH debug] ${message}`) : undefined
-      } as node_ssh.Config);
+      });
+      this.connectionSuccessful = true;
 
       this.currentConnectionName = connectionObject.name;
       this.currentHost = connectionObject.host;
@@ -338,18 +333,18 @@ export default class IBMi {
         };
       }
 
-      if (callbacks.timeoutCallback) {
-        const timeoutCallbackWrapper = () => {
-          // Don't call the callback function if it was based on a user cancellation request.
-          if (!wasCancelled) {
-            callbacks.timeoutCallback!(this);
-          }
-        }
+      // Trigger callbacks unless the connection is cancelled
+      const onDisconnected = async (error?: Error & ClientErrorExtensions) => {
+        await this.dispose();
+        callbacks.onDisconnected?.(this, error);
+      };
 
-        // Register handlers after we might have to abort due to bad configuration.
-        this.client.connection!.once(`timeout`, timeoutCallbackWrapper);
-        this.client.connection!.once(`end`, timeoutCallbackWrapper);
-        this.client.connection!.once(`error`, timeoutCallbackWrapper);
+      if (this.client.connection) {
+        //end: Disconnected by the user
+        this.client.connection.once(`end`, onDisconnected);
+        //error/tiemout: connection dropped for some reason (details given in the SSHError type) 
+        this.client.connection.once(`error`, onDisconnected);
+        this.client.connection.once(`timeout`, onDisconnected);
       }
 
       callbacks.progress({
@@ -958,8 +953,7 @@ export default class IBMi {
       }
 
       if (!options.reconnecting) {
-        const delayedOperations: Function[] = callbacks.onConnectedOperations ? [...callbacks.onConnectedOperations] : [];
-        for (const operation of delayedOperations) {
+        for (const operation of callbacks.onConnectedOperations || []) {
           await operation();
         }
       }
@@ -984,7 +978,7 @@ export default class IBMi {
       };
 
     } catch (e: any) {
-      this.disconnect(true);
+      this.disconnect();
 
       let error = e.message;
       if (wasCancelled) {
@@ -1166,24 +1160,27 @@ export default class IBMi {
     };
   }
 
-  private disconnect(failedToConnect = false) {
-    if (this.sqlJob) {
-      this.sqlJob.close();
-      this.sqlJob = undefined;
-      this.splfUserData = undefined;
+  disconnect() {    
+    if (this.client?.connection) {
+      //Close the connection and triggers its 'end' event
+      this.client.dispose();
     }
-
-    if (this.client) {
-      this.client = undefined;
-
-      if (failedToConnect === false && this.disconnectedCallback) {
-        this.disconnectedCallback(this);
-      }
+    else{
+      //There is no connection: dispose directly
+      this.dispose();
     }
   }
 
-  async dispose() {
-    this.disconnect();
+  private async dispose() {
+    CompileTools.reset();
+
+    //Clear connected resources
+    if (this.sqlJob) {
+      delete this.sqlJob;
+      delete this.splfUserData;
+    }
+    await this.getComponent<Mapepire>(Mapepire.ID)?.endJobs();
+    delete this.client;
   }
 
   /**