From c65efbc21e08a687294cc2276fa214bfb6e9e300 Mon Sep 17 00:00:00 2001
From: Ines Fazlic <fazlic.ines@gmail.com>
Date: Wed, 24 Jan 2024 14:21:58 +0000
Subject: [PATCH] fix(publish-metrics): otel tracing bug (#2439)

* fix: wait for pending playwright spans

* feat: adjust batching defaults and expose their configuration

* refactor: implement timeout for waiting for traces
---
 .../lib/open-telemetry/index.js               |  4 +-
 .../lib/open-telemetry/tracing/base.js        | 45 +++++++++++++++----
 .../lib/open-telemetry/tracing/playwright.js  |  7 +++
 3 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/packages/artillery-plugin-publish-metrics/lib/open-telemetry/index.js b/packages/artillery-plugin-publish-metrics/lib/open-telemetry/index.js
index 30e73f358d..f2ff3ce3d1 100644
--- a/packages/artillery-plugin-publish-metrics/lib/open-telemetry/index.js
+++ b/packages/artillery-plugin-publish-metrics/lib/open-telemetry/index.js
@@ -101,10 +101,10 @@ class OTelReporter {
       return done();
     }
     if (this.httpReporter) {
-      await this.httpReporter.cleanup();
+      await this.httpReporter.cleanup('http');
     }
     if (this.playwrightReporter) {
-      await this.playwrightReporter.cleanup();
+      await this.playwrightReporter.cleanup('playwright');
     }
     await this.traceConfig.shutDown();
     return done();
diff --git a/packages/artillery-plugin-publish-metrics/lib/open-telemetry/tracing/base.js b/packages/artillery-plugin-publish-metrics/lib/open-telemetry/tracing/base.js
index 08c580aacb..38d5927061 100644
--- a/packages/artillery-plugin-publish-metrics/lib/open-telemetry/tracing/base.js
+++ b/packages/artillery-plugin-publish-metrics/lib/open-telemetry/tracing/base.js
@@ -60,15 +60,16 @@ class OTelTraceConfig {
       this.exporterOpts
     );
 
+    this.processorOpts = {
+      scheduledDelayMillis: this.config.scheduledDelayMillis || 5000,
+      maxExportBatchSize: this.config.maxExportBatchSize || 1000,
+      maxQueueSize: this.config.maxQueueSize || 2000
+    };
     const Processor = this.config.smartSampling
       ? OutlierDetectionBatchSpanProcessor
       : BatchSpanProcessor;
 
-    this.tracerProvider.addSpanProcessor(
-      new Processor(this.exporter, {
-        scheduledDelayMillis: 1000
-      })
-    );
+    this.tracerProvider.addSpanProcessor(new Processor(this.exporter));
     this.tracerProvider.register();
   }
 
@@ -89,6 +90,9 @@ class OTelTraceBase {
     this.script = script;
     this.pendingRequestSpans = 0;
     this.pendingScenarioSpans = 0;
+    this.pendingPageSpans = 0;
+    this.pendingStepSpans = 0;
+    this.pendingPlaywrightScenarioSpans = 0;
   }
   setTracer(engine) {
     // Get and set the tracer by engine
@@ -141,13 +145,38 @@ class OTelTraceBase {
   otelTraceOnError(scenarioErr, req, userContext, ee, done) {
     done();
   }
-
-  async cleanup() {
-    while (this.pendingRequestSpans > 0 || this.pendingScenarioSpans > 0) {
+  async waitOnPendingSpans(pendingRequests, pendingScenarios, maxWaitTime) {
+    let waitedTime = 0;
+    while (
+      (pendingRequests > 0 || pendingScenarios > 0) &&
+      waitedTime < maxWaitTime
+    ) {
       debug('Waiting for pending traces ...');
       await new Promise((resolve) => setTimeout(resolve, 500));
+      waitedTime += 500;
+    }
+    return true;
+  }
+
+  async cleanup(engines) {
+    if (engines.includes('http')) {
+      await this.waitOnPendingSpans(
+        this.pendingRequestSpans,
+        this.pendingScenarioSpans,
+        5000
+      );
     }
+    if (engines.includes('playwright')) {
+      await this.waitOnPendingSpans(
+        this.pendingPlaywrightScenarioSpans,
+        this.pendingPlaywrightScenarioSpans,
+        5000
+      );
+    }
+
     debug('Pending traces done');
+    debug('Waiting for flush period to complete');
+    await new Promise((resolve) => setTimeout(resolve, 5000));
   }
 }
 
diff --git a/packages/artillery-plugin-publish-metrics/lib/open-telemetry/tracing/playwright.js b/packages/artillery-plugin-publish-metrics/lib/open-telemetry/tracing/playwright.js
index 1e34969593..f45af9c373 100644
--- a/packages/artillery-plugin-publish-metrics/lib/open-telemetry/tracing/playwright.js
+++ b/packages/artillery-plugin-publish-metrics/lib/open-telemetry/tracing/playwright.js
@@ -44,6 +44,7 @@ class OTelPlaywrightTraceReporter extends OTelTraceBase {
           'vu.uuid': vuContext.vars.$uuid,
           ...(this.config.attributes || {})
         });
+        this.pendingPlaywrightScenarioSpans++;
         // Set variables to track state and context
         const ctx = context.active();
         let lastPageUrl;
@@ -104,6 +105,7 @@ class OTelPlaywrightTraceReporter extends OTelTraceBase {
             scenarioSpan.addEvent(`navigated to ${page.url()}`);
             if (pageSpan) {
               pageSpan.end();
+              this.pendingPlaywrightSpans--;
             }
 
             pageSpan = this.playwrightTracer.startSpan(
@@ -116,6 +118,7 @@ class OTelPlaywrightTraceReporter extends OTelTraceBase {
               ...(this.config.attributes || {})
             });
             lastPageUrl = pageUrl;
+            this.pendingPlaywrightSpans++;
           }
         });
 
@@ -143,8 +146,10 @@ class OTelPlaywrightTraceReporter extends OTelTraceBase {
         } finally {
           if (pageSpan && !pageSpan.endTime[0]) {
             pageSpan.end();
+            this.pendingPlaywrightSpans--;
           }
           scenarioSpan.end();
+          this.pendingPlaywrightScenarioSpans--;
         }
       }
     );
@@ -159,6 +164,7 @@ class OTelPlaywrightTraceReporter extends OTelTraceBase {
           { kind: SpanKind.CLIENT },
           context.active()
         );
+        this.pendingPlaywrightSpans++;
         const startTime = Date.now();
 
         try {
@@ -180,6 +186,7 @@ class OTelPlaywrightTraceReporter extends OTelTraceBase {
           const difference = Date.now() - startTime;
           events.emit('histogram', `browser.step.${stepName}`, difference);
           span.end();
+          this.pendingPlaywrightSpans--;
         }
       });
     };