Merge pull request #654 from nats-io/batch-direct

[FEAT] direct batch

jetstream/jsapi_types.ts

@@ -483,6 +483,14 @@ export type DirectMsgRequest =
   | LastForMsgRequest
   | NextMsgRequest;
 
+export type DirectBatchOptions = {
+  batch?: number;
+  max_bytes?: number;
+  multi_last: string[];
+  up_to_seq?: number;
+  up_to_time?: Date | string;
+};
+
 export interface StreamState {
   /**
    * Number of messages stored in the Stream

jetstream/jsm.ts

@@ -33,11 +33,14 @@ import {
   Msg,
   MsgHdrs,
   NatsConnection,
+  QueuedIterator,
+  RequestStrategy,
   ReviverFn,
 } from "../nats-base-client/core.ts";
 import {
   AccountInfoResponse,
   ApiResponse,
+  DirectBatchOptions,
   DirectMsgRequest,
   JetStreamAccountStats,
   LastForMsgRequest,
@@ -84,6 +87,73 @@ export class DirectStreamAPIImpl extends BaseApiClient
     const dm = new DirectMsgImpl(r);
     return Promise.resolve(dm);
   }
+
+  async getBatch(
+    stream: string,
+    opts: DirectBatchOptions,
+  ): Promise<QueuedIterator<StoredMsg>> {
+    validateStreamName(stream);
+    const pre = this.opts.apiPrefix || "$JS.API";
+    const subj = `${pre}.DIRECT.GET.${stream}`;
+    opts = opts || {};
+    if (!Array.isArray(opts.multi_last) || opts.multi_last.length === 0) {
+      return Promise.reject("multi_last is required");
+    }
+    const payload = JSON.stringify(opts, (key, value) => {
+      if (key === "up_to_time" && value instanceof Date) {
+        return value.toISOString();
+      }
+      return value;
+    });
+
+    const iter = new QueuedIteratorImpl<StoredMsg>();
+
+    const raw = await this.nc.requestMany(
+      subj,
+      payload,
+      {
+        strategy: RequestStrategy.SentinelMsg,
+      },
+    );
+
+    (async () => {
+      let gotFirst = false;
+      let badServer = false;
+      let badRequest: string | undefined;
+      for await (const m of raw) {
+        if (!gotFirst) {
+          gotFirst = true;
+          const code = m.headers?.code || 0;
+          if (code !== 0 && code < 200 || code > 299) {
+            badRequest = m.headers?.description.toLowerCase();
+            break;
+          }
+          // inspect the message and make sure that we have a supported server
+          const v = m.headers?.get("Nats-Num-Pending");
+          if (v === "") {
+            badServer = true;
+            break;
+          }
+        }
+        if (m.data.length === 0) {
+          break;
+        }
+        iter.push(new DirectMsgImpl(m));
+      }
+      //@ts-ignore: term function
+      iter.push((): void => {
+        if (badServer) {
+          throw new Error("batch direct get not supported by the server");
+        }
+        if (badRequest) {
+          throw new Error(`bad request: ${badRequest}`);
+        }
+        iter.stop();
+      });
+    })();
+
+    return Promise.resolve(iter);
+  }
 }
 
 export class DirectMsgImpl implements DirectMsg {

jetstream/tests/jsm_test.ts

@@ -22,7 +22,7 @@ import {
   fail,
 } from "https://deno.land/[email protected]/assert/mod.ts";
 
-import { NatsConnectionImpl } from "../../nats-base-client/internal_mod.ts";
+import { NatsConnectionImpl } from "../../nats-base-client/nats.ts";
 import {
   deferred,
   Empty,
@@ -2785,3 +2785,72 @@ Deno.test("jsm - pause/unpause", async () => {
 
   await cleanup(ns, nc);
 });
+
+Deno.test("jsm - batch direct get multi_last", async () => {
+  const { ns, nc } = await setup(jetstreamServerConf());
+  if (await notCompatible(ns, nc, "2.11.0")) {
+    return;
+  }
+  const jsm = await nc.jetstreamManager() as JetStreamManagerImpl;
+  await jsm.streams.add({
+    name: "A",
+    subjects: ["a.>"],
+    storage: StorageType.Memory,
+    allow_direct: true,
+  });
+
+  const js = nc.jetstream();
+  await Promise.all([
+    js.publish("a.foo", "foo"),
+    js.publish("a.bar", "bar"),
+    js.publish("a.baz", "baz"),
+  ]);
+
+  const iter = await jsm.direct.getBatch("A", {
+    multi_last: ["a.foo", "a.baz"],
+  });
+
+  const keys = [];
+  for await (const m of iter) {
+    keys.push(m.subject);
+  }
+  assertEquals(keys.length, 2);
+  assertArrayIncludes(keys, ["a.foo", "a.baz"]);
+
+  await cleanup(ns, nc);
+});
+
+Deno.test("jsm - batch direct get batch", async () => {
+  const { ns, nc } = await setup(jetstreamServerConf());
+  if (await notCompatible(ns, nc, "2.11.0")) {
+    return;
+  }
+  const jsm = await nc.jetstreamManager() as JetStreamManagerImpl;
+  await jsm.streams.add({
+    name: "A",
+    subjects: ["a.>"],
+    storage: StorageType.Memory,
+    allow_direct: true,
+  });
+
+  const js = nc.jetstream();
+  await Promise.all([
+    js.publish("a.foo", "foo"),
+    js.publish("a.bar", "bar"),
+    js.publish("a.baz", "baz"),
+    js.publish("a.foobar", "foobar"),
+  ]);
+
+  const iter = await jsm.direct.getBatch("A", {
+    batch: 3,
+    multi_last: [">"],
+  });
+
+  const buf = [];
+  for await (const m of iter) {
+    buf.push(m);
+  }
+  assertEquals(buf.length, 3);
+
+  await cleanup(ns, nc);
+});

jetstream/types.ts

@@ -37,6 +37,7 @@ import {
   ConsumerInfo,
   defaultConsumer,
   DeliverPolicy,
+  DirectBatchOptions,
   DirectMsgRequest,
   JetStreamAccountStats,
   MsgRequest,
@@ -764,6 +765,16 @@ export interface DirectStreamAPI {
    * @param query
    */
   getMessage(stream: string, query: DirectMsgRequest): Promise<StoredMsg>;
+
+  /**
+   * Retrieves all last subject messages for the specified subjects
+   * @param stream
+   * @param opts
+   */
+  getBatch(
+    stream: string,
+    opts: DirectBatchOptions,
+  ): Promise<QueuedIterator<StoredMsg>>;
 }
 
 /**

nats-base-client/semver.ts

@@ -52,6 +52,7 @@ export enum Feature {
   JS_STREAM_SOURCE_SUBJECT_TRANSFORM = "js_stream_source_subject_transform",
   JS_STREAM_COMPRESSION = "js_stream_compression",
   JS_DEFAULT_CONSUMER_LIMITS = "js_default_consumer_limits",
+  JS_BATCH_DIRECT_GET = "js_batch_direct_get",
 }
 
 type FeatureVersion = {
@@ -109,6 +110,7 @@ export class Features {
     this.set(Feature.JS_STREAM_SOURCE_SUBJECT_TRANSFORM, "2.10.0");
     this.set(Feature.JS_STREAM_COMPRESSION, "2.10.0");
     this.set(Feature.JS_DEFAULT_CONSUMER_LIMITS, "2.10.0");
+    this.set(Feature.JS_BATCH_DIRECT_GET, "2.11.0");
 
     this.disabled.forEach((f) => {
       this.features.delete(f);