Add contentType support for sending Queues messages (#641)

* Add contentType support for Queues messages

* Address PR comments

* Address more PR comments

Author: kewbish

packages/miniflare/src/plugins/queues/gateway.ts

@@ -4,7 +4,7 @@ import v8 from "v8";
 import { stringify } from "devalue";
 import { Colorize, bold, green, grey, red, reset, yellow } from "kleur/colors";
 import { z } from "zod";
-import { Log, Timers } from "../../shared";
+import { Base64DataSchema, Log, Timers, viewToBuffer } from "../../shared";
 import { Storage } from "../../storage";
 import { CoreHeaders, structuredSerializableReducers } from "../../workers";
 import { DispatchFetch, QueueConsumer } from "../shared";
@@ -44,6 +44,15 @@ const exceptionQueueResponse: QueueResponse = {
   explicitAcks: [],
 };
 
+export const QueueContentTypeSchema = z
+  .enum(["text", "json", "bytes", "v8"])
+  .default("v8");
+export const GatewayMessageSchema = z.object({
+  body: Base64DataSchema,
+  contentType: QueueContentTypeSchema,
+});
+export type GatewayMessage = z.infer<typeof GatewayMessageSchema>;
+
 export class Message {
   #failedAttempts = 0;
 
@@ -82,7 +91,7 @@ interface PendingFlush {
 export type QueueEnqueueOn = (
   queueName: string,
   consumer: QueueConsumer,
-  messages: (Message | Buffer)[]
+  messages: (Message | GatewayMessage)[]
 ) => void;
 
 // `QueuesGateway` slightly misrepresents what this class does. Each queue will
@@ -225,15 +234,24 @@ export class QueuesGateway {
   enqueue(
     enqueueOn: QueueEnqueueOn,
     consumer: QueueConsumer,
-    messages: (Message | Buffer)[]
+    messages: (Message | GatewayMessage)[]
   ) {
     for (const message of messages) {
       if (message instanceof Message) {
         this.#messages.push(message);
       } else {
         const id = crypto.randomBytes(16).toString("hex");
         const timestamp = this.timers.now();
-        const body = v8.deserialize(message);
+        let body: unknown;
+        if (message.contentType === "text") {
+          body = message.body.toString();
+        } else if (message.contentType === "json") {
+          body = JSON.parse(message.body.toString());
+        } else if (message.contentType === "bytes") {
+          body = viewToBuffer(message.body);
+        } else {
+          body = v8.deserialize(message.body);
+        }
         this.#messages.push(new Message(id, timestamp, body));
       }
     }

packages/miniflare/src/plugins/queues/router.ts

@@ -1,15 +1,20 @@
 import { parse } from "devalue";
 import { z } from "zod";
 import { Headers, Response } from "../../http";
-import { Base64DataSchema, HttpError } from "../../shared";
+import { HttpError } from "../../shared";
 import {
   HEADER_PERSIST,
   POST,
   QueueConsumer,
   RouteHandler,
   Router,
 } from "../shared";
-import { QueueEnqueueOn, QueuesGateway } from "./gateway";
+import {
+  GatewayMessageSchema,
+  QueueContentTypeSchema,
+  QueueEnqueueOn,
+  QueuesGateway,
+} from "./gateway";
 
 const MAX_MESSAGE_SIZE_BYTES = 128 * 1000;
 const MAX_MESSAGE_BATCH_COUNT = 100;
@@ -30,6 +35,21 @@ function validateMessageSize(headers: Headers) {
   }
 }
 
+function validateContentType(
+  headers: Headers
+): z.infer<typeof QueueContentTypeSchema> {
+  const format = headers.get("X-Msg-Fmt") ?? undefined; // zod will throw if null
+  const result = QueueContentTypeSchema.safeParse(format);
+  if (!result.success) {
+    throw new HttpError(
+      400,
+      `message content type ${format} is invalid; if specified, must be one of 'text', 'json', 'bytes', or 'v8'`
+    );
+  } else {
+    return result.data;
+  }
+}
+
 function validateBatchSize(headers: Headers) {
   const count = headers.get("CF-Queue-Batch-Count");
   if (count !== null && parseInt(count) > MAX_MESSAGE_BATCH_COUNT) {
@@ -63,7 +83,7 @@ async function decodeQueueConsumer(
 }
 
 const QueuesBatchRequestSchema = z.object({
-  messages: z.array(z.object({ body: Base64DataSchema })),
+  messages: z.array(GatewayMessageSchema),
 });
 
 export interface QueuesParams {
@@ -78,15 +98,20 @@ export class QueuesRouter extends Router<QueuesGateway> {
   @POST("/:queue/message")
   message: RouteHandler<QueuesParams> = async (req, params) => {
     validateMessageSize(req.headers);
+    const contentType = validateContentType(req.headers);
 
     // Get consumer from persistence header, if we don't have a consumer,
     // drop the message
     const consumer = await decodeQueueConsumer(req.headers);
     if (consumer === undefined) return new Response();
 
     const queue = decodeURIComponent(params.queue);
-    const serialisedBody = Buffer.from(await req.arrayBuffer());
-    this.#enqueueOn(queue, consumer, [serialisedBody]);
+    this.#enqueueOn(queue, consumer, [
+      {
+        body: Buffer.from(await req.arrayBuffer()),
+        contentType,
+      },
+    ]);
     return new Response();
   };
 
@@ -101,8 +126,7 @@ export class QueuesRouter extends Router<QueuesGateway> {
 
     const queue = decodeURIComponent(params.queue);
     const body = QueuesBatchRequestSchema.parse(await req.json());
-    const messages = body.messages.map(({ body }) => body);
-    this.#enqueueOn(queue, consumer, messages);
+    this.#enqueueOn(queue, consumer, body.messages);
     return new Response();
   };
 }

packages/miniflare/test/plugins/queues/index.spec.ts

@@ -657,3 +657,83 @@ test("operations permit strange queue names", async (t) => {
     { queue: id, id: batch[1].id, body: "msg2" },
   ]);
 });
+
+test("supports message contentTypes", async (t) => {
+  const MessageContentTypeTestSchema = z
+    .object({ queue: z.string(), id: z.string(), body: z.any() })
+    .array();
+  const promise = new DeferredPromise<
+    z.infer<typeof MessageContentTypeTestSchema>
+  >();
+  const timers = new TestTimers();
+  const id = "my/ Queue";
+  const log = new TestLog(t);
+  const mf = new Miniflare({
+    log,
+    timers,
+    verbose: true,
+    queueProducers: { QUEUE: id },
+    queueConsumers: [id],
+    serviceBindings: {
+      async REPORTER(request) {
+        promise.resolve(
+          MessageContentTypeTestSchema.parse(await request.json())
+        );
+        return new Response();
+      },
+    },
+    modules: true,
+    script: `export default {
+  async fetch(request, env, ctx) {
+    await env.QUEUE.send("msg1", { contentType: "text" });
+    await env.QUEUE.send([{ message: "msg2" }], { contentType: "json" });
+    const arrayBuffer = new Uint8Array([0, 1, 2, 3, 4, 5, 6, 7]);
+    await env.QUEUE.send(arrayBuffer, { contentType: "bytes" });
+    await env.QUEUE.send(new Date(1600000000000), { contentType: "v8" });
+    return new Response();
+  },
+  async queue(batch, env, ctx) {
+    delete Date.prototype.toJSON; // JSON.stringify calls .toJSON before the replacer
+    await env.REPORTER.fetch("http://localhost", {
+      method: "POST",
+      body: JSON.stringify(
+        batch.messages.map(({ id, body }) => ({
+          queue: batch.queue,
+          id,
+          body,
+        })),
+        (_, value) => {
+          if (value instanceof ArrayBuffer) {
+            return {
+              $type: "ArrayBuffer",
+              value: Array.from(new Uint8Array(value)),
+            };
+          } else if (value instanceof Date) {
+            return { $type: "Date", value: value.getTime() };
+          }
+          return value;
+        },
+      ),
+    });
+  },
+};`,
+  });
+  await mf.dispatchFetch("http://localhost");
+  timers.timestamp += 1000;
+  await timers.waitForTasks();
+  const batch = await promise;
+  t.deepEqual(batch, [
+    { queue: id, id: batch[0].id, body: "msg1" },
+    { queue: id, id: batch[1].id, body: [{ message: "msg2" }] },
+    {
+      queue: id,
+      id: batch[2].id,
+      body: { $type: "ArrayBuffer", value: [0, 1, 2, 3, 4, 5, 6, 7] },
+    },
+    {
+      queue: id,
+      id: batch[3].id,
+      body: { $type: "Date", value: 1600000000000 },
+    },
+  ]);
+});