From a17ffcc3273f65f881bbc65ca6e7ec683275b92c Mon Sep 17 00:00:00 2001
From: Robin Townsend <robin@robin.town>
Date: Wed, 12 Apr 2023 17:07:18 -0400
Subject: [PATCH] Include unended spans in rageshakes

By turning the RageshakeSpanExporter into a SpanProcessor, it can now be notified of spans as soon as they're started.
---
 ...nExporter.ts => RageshakeSpanProcessor.ts} | 35 +++++++++++--------
 src/otel/otel.ts                              | 11 +++---
 src/settings/submit-rageshake.ts              |  2 +-
 3 files changed, 26 insertions(+), 22 deletions(-)
 rename src/analytics/{RageshakeSpanExporter.ts => RageshakeSpanProcessor.ts} (77%)

diff --git a/src/analytics/RageshakeSpanExporter.ts b/src/analytics/RageshakeSpanProcessor.ts
similarity index 77%
rename from src/analytics/RageshakeSpanExporter.ts
rename to src/analytics/RageshakeSpanProcessor.ts
index 249ff25..428d2d8 100644
--- a/src/analytics/RageshakeSpanExporter.ts
+++ b/src/analytics/RageshakeSpanProcessor.ts
@@ -1,10 +1,10 @@
 import { Attributes } from "@opentelemetry/api";
+import { hrTimeToMilliseconds } from "@opentelemetry/core";
 import {
-  ExportResult,
-  ExportResultCode,
-  hrTimeToMilliseconds,
-} from "@opentelemetry/core";
-import { SpanExporter, ReadableSpan } from "@opentelemetry/sdk-trace-base";
+  SpanProcessor,
+  ReadableSpan,
+  Span,
+} from "@opentelemetry/sdk-trace-base";
 
 const dumpAttributes = (attr: Attributes) =>
   Object.entries(attr).map(([key, value]) => ({
@@ -17,21 +17,22 @@ const dumpAttributes = (attr: Attributes) =>
  * Exports spans on demand to the Jaeger JSON format, which can be attached to
  * rageshakes and loaded into analysis tools like Jaeger and Stalk.
  */
-export class RageshakeSpanExporter implements SpanExporter {
+export class RageshakeSpanProcessor implements SpanProcessor {
   private readonly spans: ReadableSpan[] = [];
 
-  export(
-    spans: ReadableSpan[],
-    resultCallback: (result: ExportResult) => void
-  ): void {
-    this.spans.push(...spans);
-    resultCallback({ code: ExportResultCode.SUCCESS });
+  async forceFlush(): Promise<void> {}
+
+  onStart(span: Span): void {
+    this.spans.push(span);
   }
 
+  onEnd(): void {}
+
   /**
    * Dumps the spans collected so far as Jaeger-compatible JSON.
    */
   public dump(): string {
+    const now = Date.now();
     const traces = new Map<string, ReadableSpan[]>();
 
     // Organize spans by their trace IDs
@@ -69,6 +70,12 @@ export class RageshakeSpanExporter implements SpanExporter {
         processes,
         spans: spans.map((span) => {
           const ctx = span.spanContext();
+          const startTime = hrTimeToMilliseconds(span.startTime);
+          // If the span has not yet ended, pretend that it ends now
+          const duration =
+            span.duration[0] === -1
+              ? now - startTime
+              : hrTimeToMilliseconds(span.duration);
 
           return {
             traceID: traceId,
@@ -76,8 +83,8 @@ export class RageshakeSpanExporter implements SpanExporter {
             operationName: span.name,
             processID: processId,
             warnings: null,
-            startTime: hrTimeToMilliseconds(span.startTime),
-            duration: hrTimeToMilliseconds(span.duration),
+            startTime,
+            duration,
             references:
               span.parentSpanId === undefined
                 ? []
diff --git a/src/otel/otel.ts b/src/otel/otel.ts
index a73c6a0..312d85b 100644
--- a/src/otel/otel.ts
+++ b/src/otel/otel.ts
@@ -28,7 +28,7 @@ import { logger } from "matrix-js-sdk/src/logger";
 import { PosthogSpanExporter } from "../analytics/OtelPosthogExporter";
 import { Anonymity } from "../analytics/PosthogAnalytics";
 import { Config } from "../config/Config";
-import { RageshakeSpanExporter } from "../analytics/RageshakeSpanExporter";
+import { RageshakeSpanProcessor } from "../analytics/RageshakeSpanProcessor";
 
 const SERVICE_NAME = "element-call";
 
@@ -39,7 +39,7 @@ export class ElementCallOpenTelemetry {
   private _tracer: Tracer;
   private _anonymity: Anonymity;
   private otlpExporter: OTLPTraceExporter;
-  public readonly rageshakeExporter?: RageshakeSpanExporter;
+  public readonly rageshakeProcessor?: RageshakeSpanProcessor;
 
   static globalInit(): void {
     const config = Config.get();
@@ -89,11 +89,8 @@ export class ElementCallOpenTelemetry {
     }
 
     if (rageshakeUrl) {
-      logger.info("Enabling rageshake collector");
-      this.rageshakeExporter = new RageshakeSpanExporter();
-      this._provider.addSpanProcessor(
-        new SimpleSpanProcessor(this.rageshakeExporter)
-      );
+      this.rageshakeProcessor = new RageshakeSpanProcessor();
+      this._provider.addSpanProcessor(this.rageshakeProcessor);
     }
 
     const consoleExporter = new ConsoleSpanExporter();
diff --git a/src/settings/submit-rageshake.ts b/src/settings/submit-rageshake.ts
index a41b886..8080714 100644
--- a/src/settings/submit-rageshake.ts
+++ b/src/settings/submit-rageshake.ts
@@ -248,7 +248,7 @@ export function useSubmitRageshake(): {
 
           body.append(
             "file",
-            gzip(ElementCallOpenTelemetry.instance.rageshakeExporter!.dump()),
+            gzip(ElementCallOpenTelemetry.instance.rageshakeProcessor!.dump()),
             "traces.json"
           );