From cbe15a42bdea292bd2bf07e8e8515675cb2aa9ff Mon Sep 17 00:00:00 2001
From: KernelDeimos <eric.alex.dube@gmail.com>
Date: Wed, 11 Dec 2024 16:36:20 -0500
Subject: [PATCH] dev: update TelemetryService

---
 .../src/filesystem/ll_operations/ll_copy.js   |  4 +--
 .../src/modules/perfmon/TelemetryService.js   | 26 +++++++++++++++++++
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/src/backend/src/filesystem/ll_operations/ll_copy.js b/src/backend/src/filesystem/ll_operations/ll_copy.js
index 79ff1d45a..47592d678 100644
--- a/src/backend/src/filesystem/ll_operations/ll_copy.js
+++ b/src/backend/src/filesystem/ll_operations/ll_copy.js
@@ -165,7 +165,7 @@ class LLCopy extends LLFilesystemOperation {
 
         this.checkpoint('before parallel tasks');
         const tasks = new ParallelTasks({ tracer, max: 4 });
-        await tracer.startActiveSpan(`fs:cp:parallel-portion`, async span => {
+        await Context.arun(`fs:cp:parallel-portion`, async () => {
             this.checkpoint('starting parallel tasks');
             // Add child copy tasks if this is a directory
             if ( source.entry.is_dir ) {
@@ -216,8 +216,6 @@ class LLCopy extends LLFilesystemOperation {
             this.checkpoint('waiting for parallel tasks');
             await tasks.awaitAll();
             this.checkpoint('finishing up');
-
-            span.end();
         });
 
         node = node || await fs.node(new NodeUIDSelector(uuid));
diff --git a/src/backend/src/modules/perfmon/TelemetryService.js b/src/backend/src/modules/perfmon/TelemetryService.js
index 97515e151..6ba040ad7 100644
--- a/src/backend/src/modules/perfmon/TelemetryService.js
+++ b/src/backend/src/modules/perfmon/TelemetryService.js
@@ -16,6 +16,7 @@
  * You should have received a copy of the GNU Affero General Public License
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
+const opentelemetry = require("@opentelemetry/api");
 const { NodeSDK } = require('@opentelemetry/sdk-node');
 const { getNodeAutoInstrumentations } = require('@opentelemetry/auto-instrumentations-node');
 const { PeriodicExportingMetricReader, ConsoleMetricExporter } = require('@opentelemetry/sdk-metrics');
@@ -59,6 +60,31 @@ class TelemetryService extends BaseService {
         this.sdk = sdk;
 
         this.sdk.start();
+
+        this.tracer_ = opentelemetry.trace.getTracer(
+            'puter-tracer'
+        );
+    }
+    
+    _init () {
+        const svc_context = this.services.get('context');
+        svc_context.register_context_hook('pre_arun', ({ hints, trace_name, callback, replace_callback }) => {
+            if ( ! trace_name ) return;
+            if ( ! hints.trace ) return;
+            console.log('APPLYING TRACE NAME', trace_name);
+            replace_callback(async () => {
+                return await this.tracer_.startActiveSpan(trace_name, async span => {
+                    try {
+                        return await callback();
+                    } catch (error) {
+                        span.setStatus({ code: opentelemetry.SpanStatusCode.ERROR, message: error.message });
+                        throw error;
+                    } finally {
+                        span.end();
+                    }
+                });
+            });
+        });
     }
 
     getConfiguredExporter_() {