From 23186c9132606243a28eb3cf1cf3e2d6317ca03c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Serkan=20=C3=96ZAL?= <serkan@thundra.io>
Date: Sun, 27 Mar 2022 09:48:08 +0300
Subject: [PATCH] Added flag to enable/disable (enable by default for now) load
 time instrumentation by `require-in-the-middle`

---
 src/config/ConfigMetadata.ts |  4 ++++
 src/config/ConfigNames.ts    |  2 ++
 src/init/InitManager.ts      |  7 +++++++
 src/utils/ModuleUtils.ts     | 33 ++++++++++++++++++++++++++-------
 4 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/src/config/ConfigMetadata.ts b/src/config/ConfigMetadata.ts
index e81ce84c..5912f7b2 100644
--- a/src/config/ConfigMetadata.ts
+++ b/src/config/ConfigMetadata.ts
@@ -136,6 +136,10 @@ export const ConfigMetadata: {[key: string]: { type: string, defaultValue?: any
     [ConfigNames.THUNDRA_TRACE_INSTRUMENT_FILE_PREFIX]: {
         type: 'string',
     },
+    [ConfigNames.THUNDRA_TRACE_INSTRUMENT_ONLOAD]: {
+        type: 'boolean',
+        defaultValue: true,
+    },
     [ConfigNames.THUNDRA_TRACE_SPAN_LISTENERCONFIG]: {
         type: 'string',
     },
diff --git a/src/config/ConfigNames.ts b/src/config/ConfigNames.ts
index dcd0e40c..9bcc27e0 100644
--- a/src/config/ConfigNames.ts
+++ b/src/config/ConfigNames.ts
@@ -100,6 +100,8 @@ class ConfigNames {
         'thundra.agent.trace.instrument.traceableconfig';
     public static readonly THUNDRA_TRACE_INSTRUMENT_FILE_PREFIX: string =
         'thundra.agent.trace.instrument.file.prefix';
+    public static readonly THUNDRA_TRACE_INSTRUMENT_ONLOAD: string =
+        'thundra.agent.trace.instrument.onload';
 
     /////////////////////////////////////////////////////////////////////////////
 
diff --git a/src/init/InitManager.ts b/src/init/InitManager.ts
index 3e26f691..85be5200 100644
--- a/src/init/InitManager.ts
+++ b/src/init/InitManager.ts
@@ -3,6 +3,9 @@
  */
 import { INITIALIZERS } from './Initializers';
 import ThundraLogger from '../ThundraLogger';
+import ModuleUtils from '../utils/ModuleUtils';
+import ConfigProvider from '../config/ConfigProvider';
+import ConfigNames from '../config/ConfigNames';
 
 /**
  * Manages initialization of initializers
@@ -20,6 +23,10 @@ class InitManager {
     static init(): void {
         ThundraLogger.debug(`<InitManager> Initializing initializers ...`);
         if (!InitManager.initialized) {
+            const instrumentOnLoad: boolean =
+                ConfigProvider.get<boolean>(ConfigNames.THUNDRA_TRACE_INSTRUMENT_ONLOAD);
+            ModuleUtils.setInstrumentOnLoad(instrumentOnLoad);
+
             INITIALIZERS.forEach((initializer: any) => {
                 ThundraLogger.debug(`<InitManager> Initializing ${initializer.name} ...`);
                 if (!initializer.initialized) {
diff --git a/src/utils/ModuleUtils.ts b/src/utils/ModuleUtils.ts
index a7683dc2..5bb9a5b5 100644
--- a/src/utils/ModuleUtils.ts
+++ b/src/utils/ModuleUtils.ts
@@ -21,10 +21,27 @@ class ModuleUtils {
 
     private static readonly instrumenters: any = [];
     private static readonly pendingModulesToInstrument: any = [];
+    private static instrumentOnLoad: boolean = true;
 
     private constructor() {
     }
 
+    /**
+     * Gets the load time instrumentation mode flag
+     * @return the load time instrumentation mode flag
+     */
+    static isInstrumentOnLoad(): boolean {
+        return ModuleUtils.instrumentOnLoad;
+    }
+
+    /**
+     * Gets the load time instrumentation mode flag
+     * @param instrumentOnLoad the load time instrumentation mode flag to be set
+     */
+    static setInstrumentOnLoad(instrumentOnLoad: boolean): void {
+        ModuleUtils.instrumentOnLoad = instrumentOnLoad;
+    }
+
     /**
      * Tries to require given module by its name and paths
      * @param {string} name the module name
@@ -145,13 +162,15 @@ class ModuleUtils {
                     ModuleUtils.doInstrument(requiredLib, libs, null, moduleName, null, wrapper, config);
                 }
             }
-            const hook = Hook(moduleName, { internals: true }, (lib: any, name: string, basedir: string) => {
-                if (name === moduleName) {
-                    ModuleUtils.doInstrument(lib, libs, basedir, moduleName, version, wrapper, config);
-                }
-                return lib;
-            });
-            hooks.push(hook);
+            if (ModuleUtils.instrumentOnLoad) {
+                const hook = Hook(moduleName, {internals: true}, (lib: any, name: string, basedir: string) => {
+                    if (name === moduleName) {
+                        ModuleUtils.doInstrument(lib, libs, basedir, moduleName, version, wrapper, config);
+                    }
+                    return lib;
+                });
+                hooks.push(hook);
+            }
         }
         return {
             uninstrument: () => {