From 4be5671dffabab4d4d154e91206a180bfabe0266 Mon Sep 17 00:00:00 2001
From: Grzegorz Godlewski <gg@gitgis.com>
Date: Thu, 23 Nov 2023 21:14:29 +0100
Subject: [PATCH] Add driveFetch instrumentation

---
 apps/ui/src/components/ZipkinViewer.vue   |  2 +-
 src/containers/job/JobManagerContainer.ts |  2 --
 src/google/driveFetch.ts                  |  7 ++--
 src/telemetry.ts                          | 41 +++++++++++++++++++++++
 4 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/apps/ui/src/components/ZipkinViewer.vue b/apps/ui/src/components/ZipkinViewer.vue
index 9662eadd..e533b508 100644
--- a/apps/ui/src/components/ZipkinViewer.vue
+++ b/apps/ui/src/components/ZipkinViewer.vue
@@ -52,7 +52,7 @@ export default {
   data() {
     return {
       traces: [],
-      ZIPKIN_URL
+      ZIPKIN_URL: ZIPKIN_URL.startsWith('http://localhost:9411') ? ZIPKIN_URL + '/zipkin' : ZIPKIN_URL
     };
   },
   created() {
diff --git a/src/containers/job/JobManagerContainer.ts b/src/containers/job/JobManagerContainer.ts
index b3eb147c..64606412 100644
--- a/src/containers/job/JobManagerContainer.ts
+++ b/src/containers/job/JobManagerContainer.ts
@@ -4,8 +4,6 @@ import {GoogleFolderContainer} from '../google_folder/GoogleFolderContainer';
 import {TransformContainer} from '../transform/TransformContainer';
 
 import {fileURLToPath} from 'url';
-import {WatchChangesContainer} from '../changes/WatchChangesContainer';
-import {GoogleFile} from '../../model/GoogleFile';
 import {UserConfigService} from '../google_folder/UserConfigService';
 import {MarkdownTreeProcessor} from '../transform/MarkdownTreeProcessor';
 import {WorkerPool} from './WorkerPool';
diff --git a/src/google/driveFetch.ts b/src/google/driveFetch.ts
index f51733ca..1184d966 100644
--- a/src/google/driveFetch.ts
+++ b/src/google/driveFetch.ts
@@ -2,6 +2,7 @@ import {Readable} from 'stream';
 import {SimpleFile} from '../model/GoogleFile';
 import opentelemetry from '@opentelemetry/api';
 import {QuotaLimiter} from './QuotaLimiter';
+import {instrumentFunction} from '../telemetry';
 
 async function handleReadable(obj): Promise<string> {
   if (obj instanceof Readable) {
@@ -141,7 +142,8 @@ async function driveRequest(quotaLimiter: QuotaLimiter, accessToken: string, met
   }
 
   if (!quotaLimiter) {
-    const response = await fetch(url, {
+    const fetchInstrumented = instrumentFunction(fetch, 1);
+    const response = await fetchInstrumented(url, {
       method,
       headers: {
         Authorization: 'Bearer ' + accessToken,
@@ -161,7 +163,8 @@ async function driveRequest(quotaLimiter: QuotaLimiter, accessToken: string, met
   return await new Promise<Response>(async (resolve, reject) => { /* eslint-disable-line no-async-promise-executor */
     const job = async () => {
       try {
-        const response = await fetch(url, {
+        const fetchInstrumented = instrumentFunction(fetch, 1);
+        const response = await fetchInstrumented(url, {
           method,
           headers: {
             Authorization: 'Bearer ' + accessToken,
diff --git a/src/telemetry.ts b/src/telemetry.ts
index 7ea96ce9..24aa1452 100644
--- a/src/telemetry.ts
+++ b/src/telemetry.ts
@@ -155,6 +155,47 @@ export class ClassInstrumentation extends InstrumentationBase {
   }
 }
 
+export function instrumentFunction<K>(func, telemetryParamCount = 0, telemetryParamOffset = 0) {
+  return async function (...args) {
+    if (!process.env.ZIPKIN_URL) {
+      return await func(...args);
+    }
+
+    const stackTrace = new Error().stack.split('\n').splice(2);
+    const tracer = opentelemetry.trace.getTracer(
+      provider.resource.attributes[SemanticResourceAttributes.SERVICE_NAME].toString(),
+      '1.0'
+    );
+
+    let spanName = func.name;
+    if (telemetryParamCount) {
+      spanName += '(';
+      for (let i = telemetryParamOffset; i < telemetryParamCount + telemetryParamOffset; i++) {
+        if (i > 0) {
+          spanName += ', ';
+        }
+        spanName += args[0];
+      }
+      spanName += ')';
+    }
+
+    return tracer.startActiveSpan(
+      spanName,
+      { kind: SpanKind.INTERNAL },
+      async (span) => {
+        try {
+          return await func(...args);
+        } catch (err) {
+          err.stack = [err.message].concat(stackTrace).join('\n');
+          span.recordException(err);
+          throw err;
+        } finally {
+          span.end();
+        }
+      });
+  };
+}
+
 export async function addTelemetry(serviceName: string, mainDir: string) {
   if (!process.env.ZIPKIN_URL) {
     return;