@comprehend/telemetry-node 0.1.4 → 0.2.1

package/src/WebSocketConnection.ts

@@ -1,33 +1,95 @@
 import WebSocket from 'ws';
+import { randomUUID } from 'crypto';
 import {
+    AttributeType,
+    CustomMetricSpecification,
     InitMessage,
     NewObservedEntityMessage,
     NewObservedInteractionMessage,
     ObservationInputMessage,
-    ObservationMessage,
     ObservationOutputMessage,
+    StartProcessContextMessage,
+    InitAck,
 } from './wire-protocol';
+import { hrTimeNow } from './util';
 
-const INGESTION_ENDPOINT = 'wss://ingestion.comprehend.dev';
+const INGESTION_ENDPOINT = process.env.COMPREHEND_INGESTION_ENDPOINT || 'wss://ingestion.comprehend.dev';
+
+/** A queue of unacknowledged messages keyed by sequence number, with an associated ack type. */
+class SeqQueue<T extends { seq: number }> {
+    readonly ackType: string;
+    private readonly pending = new Map<number, T>();
+
+    constructor(ackType: string) {
+        this.ackType = ackType;
+    }
+
+    add(message: T): void {
+        this.pending.set(message.seq, message);
+    }
+
+    ack(seq: number): void {
+        this.pending.delete(seq);
+    }
+
+    values(): IterableIterator<T> {
+        return this.pending.values();
+    }
+}
 
 export class WebSocketConnection {
     private readonly organization: string;
     private readonly token: string;
     private readonly logger?: (message: string) => void;
+    private readonly onAuthorized?: (ack: InitAck) => void;
+    private readonly onCustomMetricChange?: (specs: CustomMetricSpecification[]) => void;
     private readonly unacknowledgedObserved = new Map<string, NewObservedEntityMessage | NewObservedInteractionMessage>();
-    private readonly unacknowledgedObservations = new Map<number, ObservationMessage>();
+
+    private readonly seqQueues: Record<string, SeqQueue<any>>;
+    private readonly observationsQueue = new SeqQueue<any>('ack-observations');
+    private readonly timeseriesQueue = new SeqQueue<any>('ack-timeseries');
+    private readonly cumulativeQueue = new SeqQueue<any>('ack-cumulative');
+    private readonly traceSpansQueue = new SeqQueue<any>('ack-tracespans');
+    private readonly dbQueryQueue = new SeqQueue<any>('ack-db-query');
+    private readonly contextQueue = new SeqQueue<StartProcessContextMessage>('ack-context');
+
     private socket: WebSocket | null = null;
     private reconnectDelay = 1000;
     private shouldReconnect = true;
     private authorized = false;
 
-    constructor(organization: string, token: string, logger?: (message: string) => void) {
-        this.organization = organization;
-        this.token = token;
-        this.logger = logger;
+    private processContext: { serviceEntityHash: string, resources: Record<string, AttributeType> } | null = null;
+    private ingestionId: string = randomUUID();
+
+    private _seq = 1;
+
+    constructor(options: {
+        organization: string,
+        token: string,
+        logger?: (message: string) => void,
+        onAuthorized?: (ack: InitAck) => void,
+        onCustomMetricChange?: (specs: CustomMetricSpecification[]) => void,
+    }) {
+        this.organization = options.organization;
+        this.token = options.token;
+        this.logger = options.logger;
+        this.onAuthorized = options.onAuthorized;
+        this.onCustomMetricChange = options.onCustomMetricChange;
+
+        // Build a lookup from ack type to queue for dispatch in onMessage
+        this.seqQueues = {};
+        for (const q of [this.observationsQueue, this.timeseriesQueue, this.cumulativeQueue,
+                         this.traceSpansQueue, this.dbQueryQueue, this.contextQueue]) {
+            this.seqQueues[q.ackType] = q;
+        }
+
         this.connect();
     }
 
+    public nextSeq(): number {
+        return this._seq++;
+    }
+
     private log(message: string) {
         if (this.logger) {
             this.logger(message);
@@ -47,9 +109,10 @@ export class WebSocketConnection {
 
     private onOpen(): void {
         this.log('WebSocket connected. Sending init/auth message.');
+        this.ingestionId = randomUUID();
         const init: InitMessage = {
             event: 'init',
-            protocolVersion: 1,
+            protocolVersion: 2,
             token: this.token,
         };
         this.sendRaw(init);
@@ -63,18 +126,42 @@ export class WebSocketConnection {
                 this.authorized = true;
                 this.log('Authorization acknowledged by server.');
 
+                if (this.onAuthorized) {
+                    this.onAuthorized(msg);
+                }
+
+                // Send context first if we have one
+                if (this.processContext) {
+                    this.sendContextStart();
+                }
+
+                // Replay entities and interactions
                 for (const message of this.unacknowledgedObserved.values()) {
                     this.sendRaw(message);
                 }
-                for (const message of this.unacknowledgedObservations.values()) {
-                    this.sendRaw(message);
+
+                // Replay all seq-based queues
+                for (const q of Object.values(this.seqQueues)) {
+                    if (q === this.contextQueue) continue; // already sent above
+                    for (const message of q.values()) {
+                        this.sendRaw(message);
+                    }
                 }
             }
             else if (msg.type === 'ack-observed') {
                 this.unacknowledgedObserved.delete(msg.hash);
             }
-            else if (msg.type === 'ack-observations') {
-                this.unacknowledgedObservations.delete(msg.seq);
+            else if (msg.type === 'custom-metric-change') {
+                if (this.onCustomMetricChange) {
+                    this.onCustomMetricChange(msg.customMetrics);
+                }
+            }
+            else if ('seq' in msg) {
+                // Dispatch seq-based acks to the appropriate queue
+                const queue = this.seqQueues[msg.type];
+                if (queue) {
+                    queue.ack(msg.seq);
+                }
             }
         } catch (e) {
             this.log('Error parsing message from server: ' + (e instanceof Error ? e.message : String(e)));
@@ -99,12 +186,47 @@ export class WebSocketConnection {
         }
     }
 
+    private sendContextStart(): void {
+        if (!this.processContext) return;
+        const seq = this.nextSeq();
+        const contextMsg: StartProcessContextMessage = {
+            event: 'context-start',
+            seq,
+            timestamp: hrTimeNow(),
+            ingestionId: this.ingestionId,
+            type: 'process',
+            serviceEntityHash: this.processContext.serviceEntityHash,
+            resources: this.processContext.resources,
+        };
+        this.contextQueue.add(contextMsg);
+        this.sendRaw(contextMsg);
+    }
+
+    public setProcessContext(serviceEntityHash: string, resources: Record<string, AttributeType>): void {
+        this.processContext = { serviceEntityHash, resources };
+        if (this.authorized) {
+            this.sendContextStart();
+        }
+    }
+
     public sendMessage(message: ObservationInputMessage): void {
         if (message.event === 'new-entity' || message.event === 'new-interaction') {
             this.unacknowledgedObserved.set(message.hash, message);
         }
         else if (message.event === 'observations') {
-            this.unacknowledgedObservations.set(message.seq, message);
+            this.observationsQueue.add(message);
+        }
+        else if (message.event === 'timeseries') {
+            this.timeseriesQueue.add(message);
+        }
+        else if (message.event === 'cumulative') {
+            this.cumulativeQueue.add(message);
+        }
+        else if (message.event === 'tracespans') {
+            this.traceSpansQueue.add(message);
+        }
+        else if (message.event === 'db-query') {
+            this.dbQueryQueue.add(message);
         }
 
         if (this.authorized) {

package/src/index.ts

@@ -1,2 +1,3 @@
-export {ComprehendDevSpanProcessor} from "./ComprehendDevSpanProcessor";
-
+export { ComprehendSDK } from './ComprehendSDK';
+export { ComprehendDevSpanProcessor } from './ComprehendDevSpanProcessor';
+export { ComprehendMetricsExporter } from './ComprehendMetricsExporter';

package/src/util.ts

@@ -0,0 +1,6 @@
+import { HrTime } from './wire-protocol';
+
+export function hrTimeNow(): HrTime {
+    const now = Date.now();
+    return [Math.floor(now / 1000), (now % 1000) * 1_000_000];
+}

package/src/wire-protocol.ts

@@ -1,20 +1,61 @@
-import {HrTime} from "@opentelemetry/api";
+export type HrTime = [number, number]; // [seconds, nanoseconds]
 
-export type ObservationInputMessage = InitMessage | NewObservedEntityMessage | NewObservedInteractionMessage | ObservationMessage;
-export type ObservationOutputMessage = InitAck | ObservedAck | ObservationsAck;
+export type ObservationInputMessage = InitMessage | StartProcessContextMessage | EndContextMessage
+                                      | NewObservedEntityMessage | NewObservedInteractionMessage
+                                      | ObservationMessage | TimeSeriesMetricsMessage | CumulativeMetricsMessage
+                                      | TraceSpansMessage | DatabaseQueryMessage;
+export type ObservationOutputMessage = InitAck | ContextAck | CustomMetricChange | ObservedAck | ObservationsAck
+                                       | TimeSeriesMetricsAck | CumulativeMetricsAck | TraceSpansAck
+                                       | DatabaseQueryAck;
 
 
+// Initialization
+
 export interface InitMessage {
     event: "init";
-    protocolVersion: 1;
+    protocolVersion: 2;
     token: string;
 }
 
+export interface InitAck {
+    type: "ack-authorized";
+    customMetrics: CustomMetricSpecification[];
+}
+
+
+// Contexts
+
+export interface StartProcessContextMessage {
+    event: "context-start";
+    seq: number;
+    timestamp: HrTime;
+    ingestionId: string;
+    type: "process";
+    serviceEntityHash: string;
+    resources: {
+        [key: string]: AttributeType
+    };
+}
+
+export interface EndContextMessage {
+    event: "context-end";
+    seq: number;
+    timestamp: HrTime;
+    ingestionId: string;
+}
+
+export interface ContextAck {
+    type: "ack-context";
+    seq: number;
+}
+
+
+// Observed entities
 
 export interface NewObservedEntityMessage {
     event: "new-entity";
-    type: string;
     hash: string;
+    type: string;
 }
 
 export interface NewObservedServiceMessage extends NewObservedEntityMessage {
@@ -46,13 +87,20 @@ export interface NewObservedHttpServiceMessage extends NewObservedEntityMessage
     port: number;
 }
 
+export interface ObservedAck {
+    type: "ack-observed";
+    hash: string;
+}
+
+
+// Observed interactions
 
 export interface NewObservedInteractionMessage {
     event: "new-interaction";
-    type: string;
     hash: string;
     from: string;
     to: string;
+    type: string;
 }
 
 export interface NewObservedHttpRequestMessage extends NewObservedInteractionMessage {
@@ -65,15 +113,8 @@ export interface NewObservedDatabaseConnectionMessage extends NewObservedInterac
     user?: string;
 }
 
-export interface NewObservedDatabaseQueryMessage extends NewObservedInteractionMessage {
-    type: "db-query";
-    query: string;
-    selects?: string[];
-    inserts?: string[];
-    updates?: string[];
-    deletes?: string[];
-}
 
+// Observations
 
 export interface ObservationMessage {
     event: "observations";
@@ -81,16 +122,20 @@ export interface ObservationMessage {
     observations: Array<Observation>;
 }
 
-export interface Observation {
-    type: string;
-    subject: string; // Hash of the entity or interaction the observation relates to
+export type Observation = HttpClientObservation | HttpServerObservation | CustomObservation;
+
+interface BaseObservation {
+    subject: string;
+    spanId: string;
+    traceId: string;
     timestamp: HrTime;
     errorMessage?: string;
     errorType?: string;
     stack?: string;
+    type: string;
 }
 
-export interface HttpClientObservation extends Observation {
+export interface HttpClientObservation extends BaseObservation {
     type: "http-client";
     path: string;
     method: string;
@@ -101,7 +146,7 @@ export interface HttpClientObservation extends Observation {
     responseBytes?: number;
 }
 
-export interface HttpServerObservation extends Observation {
+export interface HttpServerObservation extends BaseObservation {
     type: "http-server";
     path: string;
     status: number;
@@ -112,23 +157,153 @@ export interface HttpServerObservation extends Observation {
     userAgent?: string;
 }
 
-export interface DatabaseQueryObservation extends Observation {
-    type: "db-query";
+export type AttributeType = string | number | boolean;
+
+export interface CustomObservation extends BaseObservation {
+    type: "custom";
+    id: string;
+    attributes: {
+        [key: string]: AttributeType
+    };
+}
+
+export interface ObservationsAck {
+    type: "ack-observations";
+    seq: number;
+}
+
+
+// Time series metrics
+
+export interface TimeSeriesMetricsMessage {
+    event: "timeseries";
+    seq: number;
+    data: TimeSeriesDataPoint[];
+}
+
+export interface TimeSeriesDataPoint {
+    subject: string;
+    type: string;
+    timestamp: HrTime;
+    value: number;
+    unit: string;
+    attributes: {
+        [key: string]: AttributeType
+    };
+}
+
+export interface TimeSeriesMetricsAck {
+    type: "ack-timeseries";
+    seq: number;
+}
+
+
+// Cumulative metrics
+
+export interface CumulativeMetricsMessage {
+    event: "cumulative";
+    seq: number;
+    data: CumulativeDataPoint[];
+}
+
+export interface CumulativeDataPoint {
+    subject: string;
+    type: string;
+    timestamp: HrTime;
+    value: number;
+    unit: string;
+    attributes: {
+        [key: string]: AttributeType
+    };
+}
+
+export interface CumulativeMetricsAck {
+    type: "ack-cumulative";
+    seq: number;
+}
+
+
+// Trace spans
+
+export interface TraceSpansMessage {
+    event: "tracespans";
+    seq: number;
+    data: TraceSpan[];
+}
+
+export interface TraceSpan {
+    trace: string;
+    span: string;
+    parent: string;
+    name: string;
+    timestamp: HrTime;
+}
+
+export interface TraceSpansAck {
+    type: "ack-tracespans";
+    seq: number;
+}
+
+
+// Database query
+
+export interface DatabaseQueryMessage {
+    event: "db-query";
+    seq: number;
+    query: string;
+    from: string;
+    to: string;
+    timestamp: HrTime;
     duration: HrTime;
+    traceId?: string;
+    spanId?: string;
+    errorMessage?: string;
+    errorType?: string;
+    stack?: string;
     returnedRows?: number;
 }
 
+export interface DatabaseQueryAck {
+    type: "ack-db-query";
+    seq: number;
+}
 
-export interface InitAck {
-    type: "ack-authorized";
+
+// Custom metrics
+
+export interface CustomMetricChange {
+    type: "custom-metric-change";
+    customMetrics: CustomMetricSpecification[];
 }
 
-export interface ObservedAck {
-    type: "ack-observed";
-    hash: string;
+export type CustomMetricSpecification = CustomCumulativeMetricSpecification | CustomTimeSeriesMetricSpecification |
+                                        CustomSpanObservationSpecification;
+
+export interface CustomCumulativeMetricSpecification {
+    type: "cumulative";
+    id: string;
+    attributes: string[];
+    subject: string;
 }
 
-export interface ObservationsAck {
-    type: "ack-observations";
-    seq: number;
+export interface CustomTimeSeriesMetricSpecification {
+    type: "timeseries";
+    id: string;
+    attributes: string[];
+    subject: string;
 }
+
+export interface CustomSpanObservationSpecification {
+    type: "span";
+    rule: SpanMatcherRule;
+    subject: string;
+}
+
+export type SpanMatcherRule =
+    | { kind: "type", value: "client" | "server" | "internal" }
+    | { kind: "attribute-present", key: string }
+    | { kind: "attribute-absent", key: string }
+    | { kind: "attribute-equals", key: string, value: string }
+    | { kind: "attribute-not-equals", key: string, value: string }
+    | { kind: "all", rules: SpanMatcherRule[] }
+    | { kind: "any", rules: SpanMatcherRule[] };