direct batch/last for tests

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

jetstream/examples/util.ts

@@ -14,8 +14,8 @@
  */
 
 import { createConsumer, fill, initStream } from "../tests/jstest_util.ts";
-import type { NatsConnection } from "jsr:@nats-io/nats-core@3.0.0-31";
-import { nuid } from "jsr:@nats-io/nats-core@3.0.0-31";
+import type { NatsConnection } from "@nats-io/nats-core";
+import { nuid } from "@nats-io/nats-core";
 
 export async function setupStreamAndConsumer(
   nc: NatsConnection,

jetstream/src/jsapi_types.ts

@@ -479,33 +479,50 @@ export type StreamInfoRequestOptions = {
    */
   "subjects_filter": string;
 } & ApiPagedRequest;
+
+export type MsgRequest = SeqMsgRequest | LastForMsgRequest;
+
 /**
- * Request the next stream message by sequence for the specified subject.
+ * Request the first message matching specified subject found on or after the specified optional sequence.
  */
-export type NextMsgRequest = {
-  /**
-   * The seq to start looking. If the message under the specified sequence
-   * matches, it will be returned.
-   */
-  seq: number;
-  /**
-   * The subject to look for
-   */
-  next_by_subj: string;
-};
+export type NextMsgRequest = { seq?: number; next_by_subj: string };
+
+/**
+ * Retrieves the last message with the given subject
+ */
+export type LastForMsgRequest = { "last_by_subj": string };
+
+/**
+ * Stream sequence number of the message to retrieve
+ */
+export type SeqMsgRequest = { seq: number };
+
+/**
+ * Start time for the message
+ */
+export type StartTimeMsgRequest = { start_time: Date | string };
+
 export type DirectMsgRequest =
   | SeqMsgRequest
   | LastForMsgRequest
-  | NextMsgRequest;
+  | NextMsgRequest
+  | StartTimeMsgRequest;
 
-// FIXME: new options on the server
-export type DirectBatchOptions = {
+export type StartSeq = { seq?: number };
+export type StartTime = { start_time?: Date | string };
+export type DirectBatchLimits = {
   batch?: number;
   max_bytes?: number;
+};
+export type DirectBatchStartSeq = StartSeq & DirectBatchLimits;
+export type DirectBatchStartTime = StartTime & DirectBatchLimits;
+export type DirectBatchOptions = DirectBatchStartSeq | DirectBatchStartTime;
+
+export type DirectLastFor = {
   multi_last: string[];
-  up_to_seq?: number;
   up_to_time?: Date | string;
-};
+  up_to_seq?: number;
+} & DirectBatchLimits;
 
 export interface StreamState {
   /**
@@ -789,30 +806,13 @@ export interface Success {
 
 export type SuccessResponse = ApiResponse & Success;
 
-export interface LastForMsgRequest {
-  /**
-   * Retrieves the last message for the given subject
-   */
-  "last_by_subj": string;
-}
-
-export interface SeqMsgRequest {
-  /**
-   * Stream sequence number of the message to retrieve
-   */
-  seq: number;
-}
-
-// FIXME: remove number as it is deprecated
-export type MsgRequest = SeqMsgRequest | LastForMsgRequest | number;
-
-export interface MsgDeleteRequest extends SeqMsgRequest {
+export type MsgDeleteRequest = SeqMsgRequest & {
   /**
    * Default will securely remove a message and rewrite the data with random data,
    * set this to true to only remove the message
    */
   "no_erase"?: boolean;
-}
+};
 
 export interface AccountLimits {
   /**

jetstream/src/jsclient.ts

@@ -53,8 +53,8 @@ import type {
   ApiResponse,
   JetStreamAccountStats,
 } from "./jsapi_types.ts";
-import { DirectStreamAPIImpl } from "./jsm.ts";
 import { JetStreamError, JetStreamNotEnabled } from "./jserrors.ts";
+import { DirectStreamAPIImpl } from "./jsm_direct.ts";
 
 export function toJetStreamClient(
   nc: NatsConnection | JetStreamClient,

jetstream/src/jserrors.ts

@@ -172,6 +172,10 @@ export class JetStreamStatus {
   isMessageNotFound(): boolean {
     return this.code === 404 && this.description === "message not found";
   }
+
+  isEndOfBatch(): boolean {
+    return this.code === 204 && this.description === "eob";
+  }
 }
 
 export enum JetStreamApiCodes {

jetstream/src/jsm.ts → jetstream/src/jsm_direct.ts

@@ -30,14 +30,14 @@ import type {
   ReviverFn,
 } from "@nats-io/nats-core";
 import {
+  createInbox,
   Empty,
-  errors,
   QueuedIteratorImpl,
-  RequestStrategy,
   TD,
 } from "@nats-io/nats-core/internal";
 import type {
   DirectBatchOptions,
+  DirectLastFor,
   DirectMsgRequest,
   LastForMsgRequest,
 } from "./jsapi_types.ts";
@@ -101,71 +101,75 @@ export class DirectStreamAPIImpl extends BaseApiClientImpl
     return Promise.resolve(dm);
   }
 
-  async getBatch(
+  getBatch(
     stream: string,
     opts: DirectBatchOptions,
+  ): Promise<QueuedIterator<StoredMsg>> {
+    opts.batch = opts.batch || 1024;
+    return this.get(stream, opts);
+  }
+
+  getLastMessagesFor(
+    stream: string,
+    opts: DirectLastFor,
+  ): Promise<QueuedIterator<StoredMsg>> {
+    return this.get(stream, opts);
+  }
+
+  get(
+    stream: string,
+    opts: Record<string, unknown>,
   ): Promise<QueuedIterator<StoredMsg>> {
     validateStreamName(stream);
-    const pre = this.opts.apiPrefix || "$JS.API";
-    const subj = `${pre}.DIRECT.GET.${stream}`;
-    if (!Array.isArray(opts.multi_last) || opts.multi_last.length === 0) {
-      return Promise.reject(
-        errors.InvalidArgumentError.format("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 inbox = createInbox(this.nc.options.inboxPrefix);
+    let batchSupported = false;
+    const sub = this.nc.subscribe(inbox, {
+      timeout: 10_000,
+      callback: (err, msg) => {
+        if (err) {
+          iter.push(() => {
+            iter.stop(err);
+          });
+          sub.unsubscribe();
+          return;
+        }
+        const status = JetStreamStatus.maybeParseStatus(msg);
+        if (status) {
+          iter.push(() => {
+            status.isEndOfBatch() ? iter.stop() : iter.stop(status.toError());
+          });
+          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(() => {
+              iter.stop();
+            });
+          } else {
+            batchSupported = true;
+          }
+        }
 
-    const raw = await this.nc.requestMany(
-      subj,
-      payload,
-      {
-        strategy: RequestStrategy.SentinelMsg,
+        iter.push(new DirectMsgImpl(msg));
       },
-    );
+    });
 
-    (async () => {
-      let gotFirst = false;
-      let badServer = false;
-      let status: JetStreamStatus | null = null;
-      for await (const m of raw) {
-        if (!gotFirst) {
-          gotFirst = true;
-          status = JetStreamStatus.maybeParseStatus(m);
-          if (status) {
-            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));
+    const pre = this.opts.apiPrefix || "$JS.API";
+    const subj = `${pre}.DIRECT.GET.${stream}`;
+
+    const payload = JSON.stringify(opts, (key, value) => {
+      if (
+        (key === "up_to_time" || key === "start_time") && value instanceof Date
+      ) {
+        return value.toISOString();
       }
-      //@ts-ignore: term function
-      iter.push((): void => {
-        if (badServer) {
-          throw new Error("batch direct get not supported by the server");
-        }
-        if (status) {
-          throw status.toError();
-        }
-        iter.stop();
-      });
-    })().catch((err) => {
-      iter.stop(err);
+      return value;
     });
+    this.nc.publish(subj, payload, { reply: inbox });
 
     return Promise.resolve(iter);
   }
@@ -205,6 +209,11 @@ export class DirectMsgImpl implements DirectMsg {
     return this.header.last(DirectMsgHeaders.Stream);
   }
 
+  get lastSequence(): number {
+    const v = this.header.last(DirectMsgHeaders.LastSequence);
+    return typeof v === "string" ? parseInt(v) : 0;
+  }
+
   json<T = unknown>(reviver?: ReviverFn): T {
     return JSON.parse(new TextDecoder().decode(this.data), reviver);
   }

jetstream/src/types.ts

@@ -21,7 +21,11 @@ import type {
   ReviverFn,
 } from "@nats-io/nats-core/internal";
 
-import type { DeliverPolicy, ReplayPolicy } from "./jsapi_types.ts";
+import type {
+  DeliverPolicy,
+  DirectLastFor,
+  ReplayPolicy,
+} from "./jsapi_types.ts";
 
 import type {
   ConsumerConfig,
@@ -800,6 +804,20 @@ export interface DirectStreamAPI {
     stream: string,
     opts: DirectBatchOptions,
   ): Promise<QueuedIterator<StoredMsg>>;
+
+  /**
+   * Retrieves the last message for each subject in the filter.
+   * If no filter is specified, a maximum of 1024 subjects are returned.
+   * Care should be given on the specified filters to ensure that
+   * the results match what the client is expecting and to avoid missing
+   * expected data.
+   * @param stream
+   * @param opts
+   */
+  getLastMessagesFor(
+    stream: string,
+    opts: DirectLastFor,
+  ): Promise<QueuedIterator<StoredMsg>>;
 }
 
 /**
@@ -851,6 +869,11 @@ export interface DirectMsg extends StoredMsg {
    * The name of the Stream storing message
    */
   stream: string;
+
+  /**
+   * Previous sequence delivered to the client
+   */
+  lastSequence: number;
 }
 
 /**
@@ -966,6 +989,7 @@ export enum DirectMsgHeaders {
   Sequence = "Nats-Sequence",
   TimeStamp = "Nats-Time-Stamp",
   Subject = "Nats-Subject",
+  LastSequence = "Nats-Last-Sequence",
 }
 
 export enum RepublishHeaders {