DX-1614: Parallelism and Calls per second parameters

src/client/client.test.ts

@@ -4,6 +4,7 @@ import { afterAll, beforeAll, describe, expect, test } from "bun:test";
 import { nanoid } from "nanoid";
 import { Client } from "./client";
 import type { PublishToUrlResponse } from "../../dist";
+import { MOCK_QSTASH_SERVER_URL, mockQStashServer } from "./workflow/test-utils";
 
 export const clearQueues = async (client: Client) => {
   const queueDetails = await client.queue().list();
@@ -415,3 +416,117 @@ describe("E2E Queue", () => {
     { timeout: 35_000 }
   );
 });
+
+describe("flow control", () => {
+  const token = nanoid();
+  const client = new Client({
+    baseUrl: MOCK_QSTASH_SERVER_URL,
+    token,
+  });
+
+  const flowControlKey = nanoid();
+
+  test("should throw if key is passed but no ratePerSec or parallelism", () => {
+    // eslint-disable-next-line unicorn/consistent-function-scoping
+    const throws = async () => {
+      await client.publishJSON({
+        url: "https://example.com/",
+        // @ts-expect-error missing ratePerSecond or parallelism for test purposes
+        flowControl: {
+          key: flowControlKey,
+        },
+      });
+    };
+    expect(throws).toThrow("Provide at least one of parallelism or ratePerSecond for flowControl");
+  });
+
+  test("should publish a message with flow control", async () => {
+    await mockQStashServer({
+      execute: async () => {
+        await client.publishJSON({
+          urlGroup: "my-group",
+          flowControl: {
+            key: flowControlKey,
+            parallelism: 3,
+            ratePerSecond: 5,
+          },
+        });
+      },
+      responseFields: {
+        body: { messageId: "msgId" },
+        status: 200,
+      },
+      receivesRequest: {
+        method: "POST",
+        token,
+        url: "http://localhost:8080/v2/publish/my-group",
+        body: undefined,
+        headers: {
+          "Upstash-Flow-Control-Key": flowControlKey,
+          "Upstash-Flow-Control-Value": "parallelism=3, rate=5",
+        },
+      },
+    });
+  });
+
+  test("should batch messages with flow control", async () => {
+    const flowControlKeyOne = nanoid();
+    const flowControlKeyTwo = nanoid();
+    await mockQStashServer({
+      execute: async () => {
+        await client.batch([
+          {
+            url: "https://example.com/one",
+            flowControl: {
+              key: flowControlKeyOne,
+              ratePerSecond: 10,
+            },
+            body: "some-body",
+          },
+          {
+            url: "https://example.com/two",
+            flowControl: {
+              key: flowControlKeyTwo,
+              parallelism: 5,
+            },
+            method: "GET",
+          },
+        ]);
+      },
+      responseFields: {
+        body: { messageId: "msgId" },
+        status: 200,
+      },
+      receivesRequest: {
+        method: "POST",
+        token,
+        url: "http://localhost:8080/v2/batch",
+        body: [
+          {
+            body: "some-body",
+            destination: "https://example.com/one",
+            headers: {
+              "upstash-flow-control-key": flowControlKeyOne,
+              "upstash-flow-control-value": "rate=10",
+              "upstash-method": "POST",
+            },
+          },
+          {
+            destination: "https://example.com/two",
+            headers: {
+              "upstash-flow-control-key": flowControlKeyTwo,
+              "upstash-flow-control-value": "parallelism=5",
+              "upstash-method": "GET",
+            },
+          },
+        ],
+        headers: {
+          // eslint-disable-next-line unicorn/no-null
+          "Upstash-Flow-Control-Key": null,
+          // eslint-disable-next-line unicorn/no-null
+          "Upstash-Flow-Control-Value": null,
+        },
+      },
+    });
+  });
+});

src/client/client.ts

@@ -5,7 +5,15 @@ import { Chat } from "./llm/chat";
 import { Messages } from "./messages";
 import { Queue } from "./queue";
 import { Schedules } from "./schedules";
-import type { BodyInit, Event, GetEventsPayload, HeadersInit, HTTPMethods, State } from "./types";
+import type {
+  BodyInit,
+  Event,
+  FlowControl,
+  GetEventsPayload,
+  HeadersInit,
+  HTTPMethods,
+  State,
+} from "./types";
 import { UrlGroups } from "./url-groups";
 import { getRequestPath, prefixHeaders, processHeaders, wrapWithGlobalHeaders } from "./utils";
 import { Workflow } from "./workflow";
@@ -156,6 +164,12 @@ export type PublishRequest<TBody = BodyInit> = {
    * @default undefined
    */
   timeout?: Duration | number;
+
+  /**
+   * Settings for controlling the number of active requests
+   * and number of requests per second with the same key.
+   */
+  flowControl?: FlowControl;
 } & (
   | {
       /**

src/client/dlq.test.ts

@@ -115,4 +115,39 @@ describe("DLQ", () => {
     },
     { timeout: 20_000 }
   );
+
+  test(
+    "should get flow control",
+    async () => {
+      const parallelism = 10;
+      const ratePerSecond = 5;
+      const { messageId } = await client.publish({
+        url: "https://httpstat.us/400",
+        body: "hello",
+        retries: 0,
+        flowControl: {
+          key: "flow-key",
+          parallelism,
+          ratePerSecond,
+        },
+      });
+
+      await sleep(5000);
+
+      const result = await client.dlq.listMessages({
+        filter: {
+          messageId,
+        },
+      });
+      expect(result.messages.length).toBe(1);
+      const message = result.messages[0];
+
+      expect(message.flowControlKey).toBe("flow-key");
+      expect(message.parallelism).toBe(parallelism);
+      expect(message.ratePerSecond).toBe(ratePerSecond);
+    },
+    {
+      timeout: 10_000,
+    }
+  );
 });

src/client/dlq.ts

@@ -128,6 +128,7 @@ export class DLQ {
         return {
           ...message,
           urlGroup: message.topicName,
+          ratePerSecond: "rate" in message ? (message.rate as number) : undefined,
         };
       }),
       cursor: messagesPayload.cursor,

src/client/llm/chat.test.ts

@@ -117,7 +117,7 @@ describe("Test QStash chat", () => {
     { timeout: 30_000, retry: 3 }
   );
 
-  test("should publish with llm api", async () => {
+  test.skip("should publish with llm api", async () => {
     const result = await client.publishJSON({
       api: { name: "llm", provider: upstash() },
       body: {
@@ -290,7 +290,7 @@ describe("Test QStash chat with third party LLMs", () => {
     { timeout: 30_000, retry: 3 }
   );
 
-  test("should publish with llm api", async () => {
+  test.skip("should publish with llm api", async () => {
     const result = await client.publishJSON({
       api: {
         name: "llm",
@@ -321,7 +321,7 @@ describe("Test QStash chat with third party LLMs", () => {
     expect(deliveredEvent).not.toBeUndefined();
   });
 
-  test("should publish with llm api", () => {
+  test("should not be able to without callback", () => {
     //@ts-expect-error We intentionally omit the callback to ensure the function fails as expected
     const resultPromise = client.publishJSON({
       api: {

src/client/messages.test.ts

@@ -89,4 +89,23 @@ describe("Messages", () => {
     },
     { timeout: 20_000 }
   );
+
+  test("should create message with flow control", async () => {
+    const parallelism = 10;
+    const ratePerSecond = 5;
+    const { messageId } = await client.publish({
+      url: "https://httpstat.us/200?sleep=30000",
+      body: "hello",
+      flowControl: {
+        key: "flow-key",
+        parallelism,
+        ratePerSecond,
+      },
+    });
+
+    const message = await client.messages.get(messageId);
+    expect(message.flowControlKey).toBe("flow-key");
+    expect(message.parallelism).toBe(parallelism);
+    expect(message.ratePerSecond).toBe(ratePerSecond);
+  });
 });

src/client/messages.ts

@@ -93,6 +93,19 @@ export type Message = {
    * IP address of the publisher of this message
    */
   callerIp?: string;
+
+  /**
+   * flow control key
+   */
+  flowControlKey: string;
+  /**
+   * number of requests which can be active with the same flow control key
+   */
+  parallelism?: number;
+  /**
+   * number of requests to activate per second with the same flow control key
+   */
+  ratePerSecond?: number;
 };
 
 export type MessagePayload = Omit<Message, "urlGroup"> & { topicName: string };
@@ -115,6 +128,7 @@ export class Messages {
     const message: Message = {
       ...messagePayload,
       urlGroup: messagePayload.topicName,
+      ratePerSecond: "rate" in messagePayload ? (messagePayload.rate as number) : undefined,
     };
     return message;
   }

src/client/schedules.test.ts

@@ -218,4 +218,26 @@ describe("Schedules", () => {
       },
     });
   });
+
+  test("should create schedule with flow control", async () => {
+    const parallelism = 10;
+    const ratePerSecond = 5;
+    const scheduleId = nanoid();
+    await client.schedules.create({
+      destination: "https://www.initial.com",
+      cron: "*/5 * * * *",
+      scheduleId,
+      body: "my-payload",
+      flowControl: {
+        key: "flow-key",
+        parallelism,
+        ratePerSecond,
+      },
+    });
+
+    const schedule = await client.schedules.get(scheduleId);
+    expect(schedule.flowControlKey).toBe("flow-key");
+    expect(schedule.parallelism).toBe(parallelism);
+    expect(schedule.ratePerSecond).toBe(ratePerSecond);
+  });
 });