internal(direct): handle with optional callback (#153)

* internal(direct): handle with optional callback

Removed the unnecessary error parameter from the callback function in the get method to streamline the code. Corrected the handling of errors by ensuring the check is done on the 'done' object directly, improving code clarity and maintainability.

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

---------

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

Author: aricart

Diff for: jetstream/src/jsapi_types.ts

@@ -15,6 +15,7 @@
 
 import type { Nanos } from "@nats-io/nats-core";
 import { nanos } from "@nats-io/nats-core";
+import type { StoredMsg } from "./types.ts";
 
 export interface ApiPaged {
   total: number;
@@ -509,11 +510,14 @@ export type DirectMsgRequest =
   | NextMsgRequest
   | StartTimeMsgRequest;
 
+export type CompletionResult = { err?: Error };
+export type BatchCallback<T> = (done: CompletionResult | null, d: T) => void;
 export type StartSeq = { seq?: number };
 export type StartTime = { start_time?: Date | string };
 export type DirectBatchLimits = {
   batch?: number;
   max_bytes?: number;
+  callback?: BatchCallback<StoredMsg>;
 };
 export type DirectBatchStartSeq = StartSeq & DirectBatchLimits;
 export type DirectBatchStartTime = StartTime & DirectBatchLimits;

Diff for: jetstream/src/jsm_direct.ts

@@ -22,13 +22,14 @@ import type {
 } from "./types.ts";
 import { DirectMsgHeaders } from "./types.ts";
 import type {
+  CallbackFn,
   Codec,
   Msg,
   MsgHdrs,
   NatsConnection,
   QueuedIterator,
   ReviverFn,
-} from "@nats-io/nats-core";
+} from "@nats-io/nats-core/internal";
 import {
   createInbox,
   Empty,
@@ -37,6 +38,7 @@ import {
   TD,
 } from "@nats-io/nats-core/internal";
 import type {
+  CompletionResult,
   DirectBatchOptions,
   DirectLastFor,
   DirectMsgRequest,
@@ -121,40 +123,80 @@ export class DirectStreamAPIImpl extends BaseApiClientImpl
       throw new Error(`batch direct require server ${min}`);
     }
     validateStreamName(stream);
-
+    const callback = typeof opts.callback === "function" ? opts.callback : null;
     const iter = new QueuedIteratorImpl<StoredMsg>();
+
+    function pushIter(
+      done: CompletionResult | null,
+      d: StoredMsg | CallbackFn,
+    ) {
+      if (done) {
+        iter.push(() => {
+          done.err ? iter.stop(done.err) : iter.stop();
+        });
+        return;
+      }
+      iter.push(d);
+    }
+
+    function pushCb(
+      done: CompletionResult | null,
+      m: StoredMsg | CallbackFn,
+    ) {
+      const cb = callback!;
+      if (typeof m === "function") {
+        m();
+        return;
+      }
+      cb(done, m);
+    }
+
+    if (callback) {
+      iter.iterClosed.then((err) => {
+        push({ err: err ? err : undefined }, {} as StoredMsg);
+        sub.unsubscribe();
+      });
+    }
+
+    const push = callback ? pushCb : pushIter;
+
     const inbox = createInbox(this.nc.options.inboxPrefix);
     let batchSupported = false;
     const sub = this.nc.subscribe(inbox, {
-      timeout: 10_000,
+      timeout: 5000,
       callback: (err, msg) => {
         if (err) {
-          iter.push(() => {
-            iter.stop(err);
-          });
+          iter.stop(err);
           sub.unsubscribe();
           return;
         }
         const status = JetStreamStatus.maybeParseStatus(msg);
         if (status) {
-          iter.push(() => {
-            status.isEndOfBatch() ? iter.stop() : iter.stop(status.toError());
-          });
+          if (status.isEndOfBatch()) {
+            push({}, () => {
+              iter.stop();
+            });
+          } else {
+            const err = status.toError();
+            push({ err }, () => {
+              iter.stop(err);
+            });
+          }
           return;
         }
         if (!batchSupported) {
           if (typeof msg.headers?.get("Nats-Num-Pending") !== "string") {
             // no batch/max_bytes option was provided, so single response
             sub.unsubscribe();
-            iter.push(() => {
+            push({}, () => {
               iter.stop();
             });
           } else {
             batchSupported = true;
           }
         }
 
-        iter.push(new DirectMsgImpl(msg));
+        push(null, new DirectMsgImpl(msg));
       },
     });
 

Diff for: jetstream/tests/jsm_direct_test.ts

@@ -17,7 +17,9 @@ import {
   assertArrayIncludes,
   assertEquals,
   assertExists,
+  assertIsError,
   assertRejects,
+  fail,
 } from "jsr:@std/assert";
 
 import { deferred, delay } from "@nats-io/nats-core";
@@ -27,6 +29,7 @@ import {
   JetStreamError,
   jetstreamManager,
   StorageType,
+  type StoredMsg,
 } from "../src/mod.ts";
 import {
   cleanup,
@@ -36,9 +39,15 @@ import {
   setup,
 } from "test_helpers";
 
+import { JetStreamStatusError } from "../src/jserrors.ts";
+
 import type { JetStreamManagerImpl } from "../src/jsclient.ts";
 import type { DirectBatchOptions, DirectLastFor } from "../src/jsapi_types.ts";
-import type { NatsConnectionImpl } from "@nats-io/nats-core/internal";
+import {
+  type NatsConnectionImpl,
+  type QueuedIteratorImpl,
+  TimeoutError,
+} from "@nats-io/nats-core/internal";
 
 Deno.test("direct - decoder", async (t) => {
   const { ns, nc } = await setup(jetstreamServerConf({}));
@@ -174,6 +183,96 @@ Deno.test("direct - get", async (t) => {
   await cleanup(ns, nc);
 });
 
+Deno.test("direct callback", async (t) => {
+  const { ns, nc } = await setup(jetstreamServerConf());
+  if (await notCompatible(ns, nc, "2.11.0")) {
+    return;
+  }
+  const nci = nc as NatsConnectionImpl;
+  const jsm = await jetstreamManager(nci) as JetStreamManagerImpl;
+
+  await t.step("no stream", async () => {
+    const d = deferred();
+    const iter = await jsm.direct.getBatch("hello", {
+      //@ts-ignore: test
+      seq: 1,
+      callback: (done, _) => {
+        assertExists(done);
+        assertIsError(done.err, TimeoutError);
+        d.resolve();
+      },
+    }) as QueuedIteratorImpl<StoredMsg>;
+
+    const err = await iter.iterClosed;
+    assertIsError(err, TimeoutError);
+  });
+
+  await t.step("message not found", async () => {
+    await jsm.streams.add({
+      name: "empty",
+      subjects: ["empty"],
+      storage: StorageType.Memory,
+      allow_direct: true,
+    });
+
+    const iter = await jsm.direct.getBatch("empty", {
+      //@ts-ignore: test
+      seq: 1,
+      callback: (done, _) => {
+        assertExists(done);
+        assertIsError(done.err, JetStreamStatusError, "message not found");
+      },
+    }) as QueuedIteratorImpl<StoredMsg>;
+
+    const err = await iter.iterClosed;
+    assertIsError(err, JetStreamStatusError, "message not found");
+  });
+
+  await t.step("6 messages", async () => {
+    await jsm.streams.add({
+      name: "A",
+      subjects: ["a.*"],
+      storage: StorageType.Memory,
+      allow_direct: true,
+    });
+
+    const js = jetstream(nc);
+    await Promise.all([
+      js.publish("a.a"),
+      js.publish("a.b"),
+      js.publish("a.c"),
+    ]);
+
+    const buf: StoredMsg[] = [];
+
+    const iter = await jsm.direct.getBatch("A", {
+      batch: 10,
+      seq: 1,
+      callback: (done, sm) => {
+        if (done) {
+          if (done.err) {
+            console.log(done.err);
+            fail(done.err.message);
+          }
+          return;
+        }
+        buf.push(sm);
+      },
+    }) as QueuedIteratorImpl<StoredMsg>;
+
+    const err = await iter.iterClosed;
+    assertEquals(err, undefined);
+
+    const subj = buf.map((m) => m.subject);
+    assertEquals(subj.length, 3);
+    assertEquals(subj[0], "a.a");
+    assertEquals(subj[1], "a.b");
+    assertEquals(subj[2], "a.c");
+  });
+
+  await cleanup(ns, nc);
+});
+
 Deno.test("direct - batch", async (t) => {
   const { ns, nc } = await setup(jetstreamServerConf());
   if (await notCompatible(ns, nc, "2.11.0")) {