feat(core): tracing (#265)

Introduce `TraceOptions` in `PublishOptions` to enable message tracing with NATS. Messages can now be traced using `traceDestination` and optionally flagged as `traceOnly` to avoid delivery and generate trace data only. Updated `publish` implementation and added corresponding tests.

Signed-off-by: Alberto Ricart <[email protected]>

Author: aricart

core/src/core.ts

@@ -288,7 +288,7 @@ export interface MsgHdrs extends Iterable<[string, string[]]> {
   last(k: string, match?: Match): string;
 }
 
-export interface RequestOptions {
+export interface RequestOptions extends TraceOptions {
   /**
    * number of milliseconds before the request will timeout.
    */
@@ -311,7 +311,7 @@ export interface RequestOptions {
 
 export type RequestStrategy = "timer" | "count" | "stall" | "sentinel";
 
-export interface RequestManyOptions {
+export interface RequestManyOptions extends TraceOptions {
   strategy: RequestStrategy;
   maxWait: number;
   headers?: MsgHdrs;
@@ -653,7 +653,25 @@ export interface Subscription extends AsyncIterable<Msg> {
   getMax(): number | undefined;
 }
 
-export interface PublishOptions {
+/**
+ * These options enable message tracing through NATS.
+ */
+export interface TraceOptions {
+  /**
+   * If set, the server will send events representing the flow of the
+   * message as it moves through the system to this subject.
+   */
+  traceDestination?: string;
+  /**
+   * If true, the message will NOT be delivered, and instead will just
+   * generate trace information. Note that in the context of requests,
+   * this means the service will not be triggered so the operation will
+   * timeout or return no responses.
+   */
+  traceOnly?: boolean;
+}
+
+export interface PublishOptions extends TraceOptions {
   /**
    * An optional subject where a response should be sent.
    * Note you must have a subscription listening on this subject

core/src/nats.ts

@@ -16,6 +16,7 @@
 import { deferred } from "./util.ts";
 import { ProtocolHandler, SubscriptionImpl } from "./protocol.ts";
 import { Empty } from "./encoders.ts";
+import { headers } from "./headers.ts";
 
 import type { Features, SemVer } from "./semver.ts";
 import { parseSemVer } from "./semver.ts";
@@ -103,6 +104,18 @@ export class NatsConnectionImpl implements NatsConnection {
     if (options?.reply) {
       this._check(options.reply, false, true);
     }
+
+    if (typeof options?.traceOnly === "boolean") {
+      const hdrs = options.headers || headers();
+      hdrs.set("Nats-Trace-Only", "true");
+      options.headers = hdrs;
+    }
+    if (typeof options?.traceDestination === "string") {
+      const hdrs = options.headers || headers();
+      hdrs.set("Nats-Trace-Dest", options.traceDestination);
+      options.headers = hdrs;
+    }
+
     this.protocol.publish(subject, data, options);
   }
 
@@ -293,8 +306,14 @@ export class NatsConnectionImpl implements NatsConnection {
           sub?.unsubscribe();
         });
 
+      const { headers, traceDestination, traceOnly } = opts;
       try {
-        this.publish(subject, data, { reply: sub.getSubject() });
+        this.publish(subject, data, {
+          reply: sub.getSubject(),
+          headers,
+          traceDestination,
+          traceOnly,
+        });
       } catch (err) {
         cancel(err as Error);
       }
@@ -322,13 +341,17 @@ export class NatsConnectionImpl implements NatsConnection {
       const r = new RequestMany(this.protocol.muxSubscriptions, subject, rmo);
       this.protocol.request(r);
 
+      const { headers, traceDestination, traceOnly } = opts;
+
       try {
         this.publish(
           subject,
           data,
           {
             reply: `${this.protocol.muxSubscriptions.baseInbox}${r.token}`,
-            headers: opts.headers,
+            headers,
+            traceDestination,
+            traceOnly,
           },
         );
       } catch (err) {
@@ -415,13 +438,17 @@ export class NatsConnectionImpl implements NatsConnection {
       );
       this.protocol.request(r);
 
+      const { headers, traceDestination, traceOnly } = opts;
+
       try {
         this.publish(
           subject,
           data,
           {
             reply: `${this.protocol.muxSubscriptions.baseInbox}${r.token}`,
-            headers: opts.headers,
+            headers,
+            traceDestination,
+            traceOnly,
           },
         );
       } catch (err) {

core/tests/basics_test.ts

@@ -1726,3 +1726,56 @@ Deno.test("basics - internal close listener", async () => {
 
   await ns.stop();
 });
+
+Deno.test("basics - publish tracing", async () => {
+  const { ns, nc } = await setup();
+  const sub = nc.subscribe("foo", { callback: () => {} });
+
+  const traces = nc.subscribe("traces", {
+    callback: () => {},
+    max: 2,
+  });
+  nc.flush();
+
+  nc.publish("foo", Empty, { traceDestination: "traces" });
+  nc.publish("foo", Empty, { traceDestination: "traces", traceOnly: true });
+
+  await traces.closed;
+  assertEquals(sub.getReceived(), 1);
+  assertEquals(traces.getReceived(), 2);
+
+  await cleanup(ns, nc);
+});
+
+Deno.test("basics - request tracing", async () => {
+  const { ns, nc } = await setup();
+  const sub = nc.subscribe("foo", {
+    callback: (_, m) => {
+      m.respond();
+    },
+  });
+
+  const traces = nc.subscribe("traces", {
+    callback: () => {},
+    max: 2,
+  });
+  nc.flush();
+
+  await nc.request("foo", Empty, {
+    timeout: 2_000,
+    traceDestination: "traces",
+  });
+  await assertRejects(() => {
+    return nc.request("foo", Empty, {
+      timeout: 2_000,
+      traceDestination: "traces",
+      traceOnly: true,
+    });
+  });
+
+  await traces.closed;
+  assertEquals(sub.getReceived(), 1);
+  assertEquals(traces.getReceived(), 2);
+
+  await cleanup(ns, nc);
+});

core/tests/mrequest_test.ts

@@ -511,3 +511,49 @@ Deno.test("mreq - no mux request no perms doesn't leak subs", async () => {
 
   await cleanup(ns, nc);
 });
+
+Deno.test("basics - request many tracing", async () => {
+  const { ns, nc } = await setup();
+  const sub = nc.subscribe("foo", {
+    callback: (_, m) => {
+      m.respond();
+      m.respond();
+    },
+  });
+
+  const traces = nc.subscribe("traces", {
+    callback: () => {},
+    max: 2,
+  });
+  nc.flush();
+
+  let iter = await nc.requestMany("foo", Empty, {
+    strategy: "stall",
+    maxWait: 2_000,
+    traceDestination: "traces",
+  });
+  let count = 0;
+  for await (const _ of iter) {
+    count++;
+  }
+  assertEquals(count, 2);
+
+  iter = await nc.requestMany("foo", Empty, {
+    strategy: "stall",
+    maxWait: 2_000,
+    traceDestination: "traces",
+    traceOnly: true,
+  });
+
+  count = 0;
+  for await (const _ of iter) {
+    count++;
+  }
+  assertEquals(count, 0);
+
+  await traces.closed;
+  assertEquals(sub.getReceived(), 1);
+  assertEquals(traces.getReceived(), 2);
+
+  await cleanup(ns, nc);
+});