Create a new timestamped span to enable let us keep the proper endTime

Author: udnay

packages/opentelemetry-instrumentation-ws/src/index.ts

@@ -1,6 +1,6 @@
 /* eslint-disable @typescript-eslint/no-this-alias */
 /* eslint-disable @typescript-eslint/ban-types */
-import { context, Context, diag, propagation, ROOT_CONTEXT, Span, SpanKind, SpanStatusCode, trace } from "@opentelemetry/api";
+import { context, Context, diag, Exception, propagation, ROOT_CONTEXT, Span, SpanAttributes, SpanAttributeValue, SpanContext, SpanKind, SpanStatus, SpanStatusCode, TimeInput, trace } from "@opentelemetry/api";
 import { RPCMetadata, RPCType, setRPCMetadata } from "@opentelemetry/core";
 import {
   InstrumentationBase,
@@ -386,7 +386,7 @@ export class WSInstrumentation extends InstrumentationBase<WS> {
       upgradeHead: Buffer,
       callback: (client: WebSocket, request: IncomingMessage) => void
     ) {
-      const parentSpan = self._requestSpans.get(request);
+      const parentSpan = self._requestSpans.get(request) !== undefined ? new TimeStampedSpan(self._requestSpans.get(request)!) : undefined;
       const span = self.tracer.startSpan(`WS upgrade`, {
         kind: SpanKind.SERVER,
         attributes: getIncomingRequestAttributes(request, {
@@ -419,7 +419,9 @@ export class WSInstrumentation extends InstrumentationBase<WS> {
               parentSpan?.setAttributes({
                 [SemanticAttributes.HTTP_STATUS_CODE]: 101,
               });
-              parentSpan?.end();
+              if (parentSpan) {
+                parentSpan.endTime = new Date(); // Set the end time to the current time but don't end the span yet
+              }
               self._requestSpans.delete(request);
 
               return callback.call(this, websocket, request);
@@ -430,3 +432,49 @@ export class WSInstrumentation extends InstrumentationBase<WS> {
     };
   };
 }
+
+class TimeStampedSpan implements Span {
+  constructor(private readonly span: Span) {}
+  spanContext(): SpanContext {
+    return this.span.spanContext();
+  }
+  setAttribute(key: string, value: SpanAttributeValue): this {
+    this.span.setAttribute(key, value);
+    return this;
+  }
+  setAttributes(attributes: SpanAttributes): this {
+    this.span.setAttributes(attributes);
+    return this;
+  }
+  addEvent(name: string, attributesOrStartTime?: SpanAttributes | TimeInput, startTime?: TimeInput): this {
+    this.span.addEvent(name, attributesOrStartTime, startTime);
+    return this;
+  }
+  setStatus(status: SpanStatus): this {
+    this.span.setStatus(status);
+    return this;
+  }
+  updateName(name: string): this {
+    this.span.updateName(name);
+    return this;
+  }
+  isRecording(): boolean {
+    return this.span.isRecording();
+  }
+  recordException(exception: Exception, time?: TimeInput): void {
+    this.span.recordException(exception, time);
+  }
+  private _endTime: TimeInput | undefined;
+
+  set endTime(endTime: TimeInput) {
+    this._endTime = endTime;
+  }
+
+  end(endTime?: TimeInput): void {
+    if (this._endTime) {
+      this.span.end(this._endTime);
+    } else {
+      this.span.end(endTime);
+    }
+  }
+}