Skip to content
rachel/workerd
File

Blob: src/workerd/api/queue.c++

34.0 KB
1// Copyright (c) 2023 Cloudflare, Inc.
2// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3//     https://opensource.org/licenses/Apache-2.0
4 
5#include "queue.h"
6 
7#include "util.h"
8 
9#include <workerd/api/global-scope.h>
10#include <workerd/io/features.h>
11#include <workerd/io/tracer.h>
12#include <workerd/jsg/jsg.h>
13#include <workerd/jsg/ser.h>
14#include <workerd/util/mimetype.h>
15#include <workerd/util/strings.h>
16 
17#include <kj/encoding.h>
18 
19namespace workerd::api {
20 
21namespace {
22 
23// Header for the message format.
24static constexpr kj::StringPtr HDR_MSG_FORMAT = "X-Msg-Fmt"_kj;
25 
26// The upstream service sends 0 when there is "no data" available on a timestamp field (e.g. no `oldestMessageTimestamp`).
27// This method converts it to kj::none so users see `undefined`.
28void clearEpochSentinel(jsg::Optional<kj::Date>& ts) {
29  KJ_IF_SOME(date, ts) {
30    if (date == kj::UNIX_EPOCH) {
31      ts = kj::none;
32    }
33  }
34}
35 
36// Returns a callback suitable for IoContext::awaitIo() that parses a JSON response string into
37// a typed struct via the given TypeHandler, then clears the epoch sentinel on
38// oldestMessageTimestamp.
39//
40// The returned callback captures `handler` by reference. TypeHandler instances are managed by
41// the JSG type registration system and live for the lifetime of the isolate, so this is safe.
42//
43// getOldestMessageTimestamp: (T&) -> jsg::Optional<kj::Date>&
44template <typename T>
45auto parseQueueResponse(
46    const jsg::TypeHandler<T>& handler, kj::StringPtr errorMsg, auto getOldestMessageTimestamp) {
47  return [&handler, errorMsg, getOldestMessageTimestamp](jsg::Lock& js, kj::String text) -> T {
48    auto parsed = jsg::JsValue::fromJson(js, text);
49    auto result = JSG_REQUIRE_NONNULL(handler.tryUnwrap(js, parsed), Error, errorMsg, text);
50    clearEpochSentinel(getOldestMessageTimestamp(result));
51    return kj::mv(result);
52  };
53}
54 
55// Header for the message delivery delay.
56static constexpr kj::StringPtr HDR_MSG_DELAY = "X-Msg-Delay-Secs"_kj;
57 
58auto buildQueueErrorMessage(
59    const kj::HttpClient::Response& response, const ThreadContext::HeaderIdBundle& headerIds) {
60  auto errorCode = response.headers->get(headerIds.cfQueuesErrorCode).orDefault("15000"_kj);
61  auto errorCause =
62      response.headers->get(headerIds.cfQueuesErrorCause).orDefault("Unknown Internal Error"_kj);
63 
64  return kj::str(errorCause, " (", errorCode, ")");
65}
66 
67kj::StringPtr validateContentType(kj::StringPtr contentType) {
68  auto lowerCase = toLower(contentType);
69  if (lowerCase == IncomingQueueMessage::ContentType::TEXT) {
70    return IncomingQueueMessage::ContentType::TEXT;
71  } else if (lowerCase == IncomingQueueMessage::ContentType::BYTES) {
72    return IncomingQueueMessage::ContentType::BYTES;
73  } else if (lowerCase == IncomingQueueMessage::ContentType::JSON) {
74    return IncomingQueueMessage::ContentType::JSON;
75  } else if (lowerCase == IncomingQueueMessage::ContentType::V8) {
76    return IncomingQueueMessage::ContentType::V8;
77  } else {
78    JSG_FAIL_REQUIRE(TypeError, kj::str("Unsupported queue message content type: ", contentType));
79  }
80}
81 
82struct Serialized {
83  kj::Maybe<kj::OneOf<kj::String, kj::Array<kj::byte>, jsg::BufferSource, jsg::BackingStore>> own;
84  // Holds onto the owner of a given array of serialized data.
85  kj::ArrayPtr<kj::byte> data;
86  // A pointer into that data that can be directly written into an outgoing queue send, regardless
87  // of its holder.
88};
89 
90Serialized serializeV8(jsg::Lock& js, const jsg::JsValue& body) {
91  // Use a specific serialization version to avoid sending messages using a new version before all
92  // runtimes at the edge know how to read it.
93  jsg::Serializer serializer(js,
94      jsg::Serializer::Options{
95        .version = 15,
96        .omitHeader = false,
97      });
98  serializer.write(js, jsg::JsValue(body));
99  kj::Array<kj::byte> bytes = serializer.release().data;
100  Serialized result;
101  result.data = bytes;
102  result.own = kj::mv(bytes);
103  return kj::mv(result);
104}
105 
106// Control whether the serialize() method makes a deep copy of provided ArrayBuffer types or if it
107// just returns a shallow reference that is only valid until the given method returns.
108enum class SerializeArrayBufferBehavior {
109  DEEP_COPY,
110  SHALLOW_REFERENCE,
111};
112 
113Serialized serialize(jsg::Lock& js,
114    const jsg::JsValue& body,
115    kj::StringPtr contentType,
116    SerializeArrayBufferBehavior bufferBehavior) {
117  if (contentType == IncomingQueueMessage::ContentType::TEXT) {
118    JSG_REQUIRE(body.isString(), TypeError,
119        kj::str("Content Type \"", IncomingQueueMessage::ContentType::TEXT,
120            "\" requires a value of type string, but received: ", body.typeOf(js)));
121 
122    kj::String s = body.toString(js);
123    Serialized result;
124    result.data = s.asBytes();
125    result.own = kj::mv(s);
126    return kj::mv(result);
127  } else if (contentType == IncomingQueueMessage::ContentType::BYTES) {
128    JSG_REQUIRE(body.isArrayBufferView(), TypeError,
129        kj::str("Content Type \"", IncomingQueueMessage::ContentType::BYTES,
130            "\" requires a value of type ArrayBufferView, but received: ", body.typeOf(js)));
131 
132    jsg::BufferSource source(js, body);
133    if (bufferBehavior == SerializeArrayBufferBehavior::SHALLOW_REFERENCE) {
134      // If we know the data will be consumed synchronously, we can avoid copying it.
135      Serialized result;
136      result.data = source.asArrayPtr();
137      result.own = kj::mv(source);
138      return kj::mv(result);
139    } else if (source.canDetach(js)) {
140      // Prefer detaching the input ArrayBuffer whenever possible to avoid needing to copy it.
141      auto backingSource = source.detach(js);
142      Serialized result;
143      result.data = backingSource.asArrayPtr();
144      result.own = kj::mv(backingSource);
145      return kj::mv(result);
146    } else {
147      kj::Array<kj::byte> bytes = kj::heapArray(source.asArrayPtr());
148      Serialized result;
149      result.data = bytes;
150      result.own = kj::mv(bytes);
151      return kj::mv(result);
152    }
153  } else if (contentType == IncomingQueueMessage::ContentType::JSON) {
154    kj::String s = body.toJson(js);
155    Serialized result;
156    result.data = s.asBytes();
157    result.own = kj::mv(s);
158    return kj::mv(result);
159  } else if (contentType == IncomingQueueMessage::ContentType::V8) {
160    return serializeV8(js, body);
161  } else {
162    JSG_FAIL_REQUIRE(TypeError, kj::str("Unsupported queue message content type: ", contentType));
163  }
164}
165 
166struct SerializedWithOptions {
167  Serialized body;
168  kj::Maybe<kj::StringPtr> contentType;
169  kj::Maybe<int> delaySeconds;
170};
171 
172jsg::JsValue deserialize(
173    jsg::Lock& js, kj::Array<kj::byte> body, kj::Maybe<kj::StringPtr> contentType) {
174  auto type = contentType.orDefault(IncomingQueueMessage::ContentType::V8);
175 
176  if (type == IncomingQueueMessage::ContentType::TEXT) {
177    return js.str(body);
178  } else if (type == IncomingQueueMessage::ContentType::BYTES) {
179    return jsg::JsValue(js.bytes(kj::mv(body)).getHandle(js));
180  } else if (type == IncomingQueueMessage::ContentType::JSON) {
181    return jsg::JsValue::fromJson(js, body.asChars());
182  } else if (type == IncomingQueueMessage::ContentType::V8) {
183    return jsg::JsValue(jsg::Deserializer(js, body.asPtr()).readValue(js));
184  } else {
185    JSG_FAIL_REQUIRE(TypeError, kj::str("Unsupported queue message content type: ", type));
186  }
187}
188 
189jsg::JsValue deserialize(jsg::Lock& js, rpc::QueueMessage::Reader message) {
190  kj::StringPtr type = message.getContentType();
191  if (type == "") {
192    // default to v8 format
193    type = IncomingQueueMessage::ContentType::V8;
194  }
195 
196  if (type == IncomingQueueMessage::ContentType::TEXT) {
197    return js.str(message.getData().asChars());
198  } else if (type == IncomingQueueMessage::ContentType::BYTES) {
199    kj::Array<kj::byte> bytes = kj::heapArray(message.getData().asBytes());
200    return jsg::JsValue(js.bytes(kj::mv(bytes)).getHandle(js));
201  } else if (type == IncomingQueueMessage::ContentType::JSON) {
202    return jsg::JsValue::fromJson(js, message.getData().asChars());
203  } else if (type == IncomingQueueMessage::ContentType::V8) {
204    return jsg::JsValue(jsg::Deserializer(js, message.getData()).readValue(js));
205  } else {
206    JSG_FAIL_REQUIRE(TypeError, kj::str("Unsupported queue message content type: ", type));
207  }
208}
209}  // namespace
210 
211jsg::Promise<WorkerQueue::SendResponse> WorkerQueue::send(jsg::Lock& js,
212    jsg::JsValue body,
213    jsg::Optional<SendOptions> options,
214    const jsg::TypeHandler<SendResponse>& responseHandler) {
215  auto& context = IoContext::current();
216 
217  JSG_REQUIRE(!body.isUndefined(), TypeError, "Message body cannot be undefined");
218 
219  auto headers = kj::HttpHeaders(context.getHeaderTable());
220  headers.set(kj::HttpHeaderId::CONTENT_TYPE, MimeType::OCTET_STREAM.toString());
221 
222  kj::Maybe<kj::StringPtr> contentType;
223  KJ_IF_SOME(opts, options) {
224    KJ_IF_SOME(type, opts.contentType) {
225      auto validatedType = validateContentType(type);
226      headers.addPtrPtr(HDR_MSG_FORMAT, validatedType);
227      contentType = validatedType;
228    }
229    KJ_IF_SOME(secs, opts.delaySeconds) {
230      headers.addPtr(HDR_MSG_DELAY, kj::str(secs));
231    }
232  }
233 
234  Serialized serialized;
235  KJ_IF_SOME(type, contentType) {
236    serialized = serialize(js, body, type, SerializeArrayBufferBehavior::DEEP_COPY);
237  } else if (workerd::FeatureFlags::get(js).getQueuesJsonMessages()) {
238    headers.addPtrPtr("X-Msg-Fmt", IncomingQueueMessage::ContentType::JSON);
239    serialized = serialize(
240        js, body, IncomingQueueMessage::ContentType::JSON, SerializeArrayBufferBehavior::DEEP_COPY);
241  } else {
242    serialized = serializeV8(js, body);
243  }
244 
245  auto client = context.getHttpClient(subrequestChannel, true, kj::none, "queue_send"_kjc);
246  auto req = client->request(
247      kj::HttpMethod::POST, "https://fake-host/message"_kjc, headers, serialized.data.size());
248 
249  const auto& headerIds = context.getHeaderIds();
250  const auto exposeErrorCodes = workerd::FeatureFlags::get(js).getQueueExposeErrorCodes();
251 
252  static constexpr auto handleSend = [](auto req, auto serialized, auto client, auto& headerIds,
253                                         bool exposeErrorCodes) -> kj::Promise<kj::String> {
254    co_await req.body->write(serialized.data);
255    auto response = co_await req.response;
256 
257    if (exposeErrorCodes) {
258      JSG_REQUIRE(response.statusCode == 200, Error, buildQueueErrorMessage(response, headerIds));
259    } else {
260      JSG_REQUIRE(
261          response.statusCode == 200, Error, kj::str("Queue send failed: ", response.statusText));
262    }
263 
264    auto responseBody = co_await response.body->readAllBytes();
265    co_return kj::str(responseBody.asChars());
266  };
267 
268  auto promise =
269      handleSend(kj::mv(req), kj::mv(serialized), kj::mv(client), headerIds, exposeErrorCodes);
270 
271  return context.awaitIo(js, kj::mv(promise),
272      parseQueueResponse(responseHandler, "Failed to parse queue send response"_kj,
273          [](SendResponse& r) -> auto& { return r.metadata.metrics.oldestMessageTimestamp; }));
274}
275 
276jsg::Promise<WorkerQueue::Metrics> WorkerQueue::metrics(
277    jsg::Lock& js, const jsg::TypeHandler<Metrics>& metricsHandler) {
278  auto& context = IoContext::current();
279 
280  auto headers = kj::HttpHeaders(context.getHeaderTable());
281 
282  auto client = context.getHttpClient(subrequestChannel, true, kj::none, "queue_metrics"_kjc);
283  auto req = client->request(
284      kj::HttpMethod::GET, "https://fake-host/metrics"_kjc, headers, static_cast<uint64_t>(0));
285  const auto& headerIds = context.getHeaderIds();
286 
287  static constexpr auto handleMetrics = [](auto req, auto client,
288                                            auto& headerIds) -> kj::Promise<kj::String> {
289    auto response = co_await req.response;
290 
291    JSG_REQUIRE(response.statusCode == 200, Error, buildQueueErrorMessage(response, headerIds));
292 
293    co_return co_await response.body->readAllText();
294  };
295 
296  auto promise = handleMetrics(kj::mv(req), kj::mv(client), headerIds);
297 
298  return context.awaitIo(js, kj::mv(promise),
299      parseQueueResponse(metricsHandler, "Failed to parse queue metrics response"_kj,
300          [](Metrics& m) -> auto& { return m.oldestMessageTimestamp; }));
301}
302 
303jsg::Promise<WorkerQueue::SendBatchResponse> WorkerQueue::sendBatch(jsg::Lock& js,
304    jsg::Sequence<MessageSendRequest> batch,
305    jsg::Optional<SendBatchOptions> options,
306    const jsg::TypeHandler<SendBatchResponse>& responseHandler) {
307  auto& context = IoContext::current();
308 
309  JSG_REQUIRE(batch.size() > 0, TypeError, "sendBatch() requires at least one message");
310 
311  size_t totalSize = 0;
312  size_t largestMessage = 0;
313  auto messageCount = batch.size();
314  auto builder = kj::heapArrayBuilder<SerializedWithOptions>(messageCount);
315  for (auto& message: batch) {
316    auto body = message.body.getHandle(js);
317    JSG_REQUIRE(!body.isUndefined(), TypeError, "Message body cannot be undefined");
318 
319    SerializedWithOptions item;
320    KJ_IF_SOME(secs, message.delaySeconds) {
321      item.delaySeconds = secs;
322    }
323 
324    KJ_IF_SOME(contentType, message.contentType) {
325      item.contentType = validateContentType(contentType);
326      item.body = serialize(js, body, contentType, SerializeArrayBufferBehavior::SHALLOW_REFERENCE);
327    } else if (workerd::FeatureFlags::get(js).getQueuesJsonMessages()) {
328      item.contentType = IncomingQueueMessage::ContentType::JSON;
329      item.body = serialize(js, body, IncomingQueueMessage::ContentType::JSON,
330          SerializeArrayBufferBehavior::SHALLOW_REFERENCE);
331    } else {
332      item.body = serializeV8(js, body);
333    }
334 
335    builder.add(kj::mv(item));
336    totalSize += builder.back().body.data.size();
337    largestMessage = kj::max(largestMessage, builder.back().body.data.size());
338  }
339  auto serializedBodies = builder.finish();
340 
341  auto estimatedSize = (totalSize + 2) / 3 * 4 + messageCount * 64 + 32;
342  kj::Vector<char> bodyBuilder(estimatedSize);
343  bodyBuilder.addAll("{\"messages\":["_kj);
344  for (size_t i = 0; i < messageCount; ++i) {
345    bodyBuilder.addAll("{\"body\":\""_kj);
346    bodyBuilder.addAll(kj::encodeBase64(serializedBodies[i].body.data));
347    bodyBuilder.add('"');
348 
349    KJ_IF_SOME(contentType, serializedBodies[i].contentType) {
350      bodyBuilder.addAll(",\"contentType\":\""_kj);
351      bodyBuilder.addAll(contentType);
352      bodyBuilder.add('"');
353    }
354 
355    KJ_IF_SOME(delaySecs, serializedBodies[i].delaySeconds) {
356      bodyBuilder.addAll(",\"delaySecs\": "_kj);
357      bodyBuilder.addAll(kj::str(delaySecs));
358    }
359 
360    bodyBuilder.addAll("}"_kj);
361    if (i < messageCount - 1) {
362      bodyBuilder.add(',');
363    }
364  }
365  bodyBuilder.addAll("]}"_kj);
366  bodyBuilder.add('\0');
367  KJ_DASSERT(bodyBuilder.size() <= estimatedSize);
368  kj::String body(bodyBuilder.releaseAsArray());
369  KJ_DASSERT(jsg::JsValue::fromJson(js, body).isObject());
370 
371  auto client = context.getHttpClient(subrequestChannel, true, kj::none, "queue_send"_kjc);
372 
373  auto headers = kj::HttpHeaders(context.getHeaderTable());
374  headers.addPtr("CF-Queue-Batch-Count"_kj, kj::str(messageCount));
375  headers.addPtr("CF-Queue-Batch-Bytes"_kj, kj::str(totalSize));
376  headers.addPtr("CF-Queue-Largest-Msg"_kj, kj::str(largestMessage));
377  headers.set(kj::HttpHeaderId::CONTENT_TYPE, MimeType::JSON.toString());
378 
379  KJ_IF_SOME(opts, options) {
380    KJ_IF_SOME(secs, opts.delaySeconds) {
381      headers.addPtr(HDR_MSG_DELAY, kj::str(secs));
382    }
383  }
384 
385  auto req =
386      client->request(kj::HttpMethod::POST, "https://fake-host/batch"_kjc, headers, body.size());
387 
388  const auto& headerIds = context.getHeaderIds();
389  const auto exposeErrorCodes = workerd::FeatureFlags::get(js).getQueueExposeErrorCodes();
390  static constexpr auto handleWrite = [](auto req, auto body, auto client, auto& headerIds,
391                                          bool exposeErrorCodes) -> kj::Promise<kj::String> {
392    co_await req.body->write(body.asBytes());
393    auto response = co_await req.response;
394 
395    if (exposeErrorCodes) {
396      JSG_REQUIRE(response.statusCode == 200, Error, buildQueueErrorMessage(response, headerIds));
397    } else {
398      JSG_REQUIRE(response.statusCode == 200, Error,
399          kj::str("Queue sendBatch failed: ", response.statusText));
400    }
401 
402    auto responseBody = co_await response.body->readAllBytes();
403    co_return kj::str(responseBody.asChars());
404  };
405 
406  auto promise =
407      handleWrite(kj::mv(req), kj::mv(body), kj::mv(client), headerIds, exposeErrorCodes);
408 
409  return context.awaitIo(js, kj::mv(promise),
410      parseQueueResponse(responseHandler, "Failed to parse queue send response"_kj,
411          [](SendBatchResponse& r) -> auto& { return r.metadata.metrics.oldestMessageTimestamp; }));
412}
413 
414QueueMessage::QueueMessage(
415    jsg::Lock& js, rpc::QueueMessage::Reader message, IoPtr<QueueEventResult> result)
416    : id(kj::str(message.getId())),
417      timestamp(message.getTimestampNs() * kj::NANOSECONDS + kj::UNIX_EPOCH),
418      body(deserialize(js, message).addRef(js)),
419      attempts(message.getAttempts()),
420      result(result) {}
421// Note that we must make deep copies of all data here since the incoming Reader may be
422// deallocated while JS's GC wrappers still exist.
423 
424QueueMessage::QueueMessage(
425    jsg::Lock& js, IncomingQueueMessage message, IoPtr<QueueEventResult> result)
426    : id(kj::mv(message.id)),
427      timestamp(message.timestamp),
428      body(deserialize(js, kj::mv(message.body), message.contentType).addRef(js)),
429      attempts(message.attempts),
430      result(result) {}
431 
432jsg::JsValue QueueMessage::getBody(jsg::Lock& js) {
433  return body.getHandle(js);
434}
435 
436void QueueMessage::retry(jsg::Optional<QueueRetryOptions> options) {
437  if (result->ackAll) {
438    auto msg = kj::str("Received a call to retry() on message ", id,
439        " after ackAll() was already called. "
440        "Calling retry() on a message after calling ackAll() has no effect.");
441    IoContext::current().logWarning(msg);
442    return;
443  }
444 
445  if (result->explicitAcks.contains(id)) {
446    auto msg = kj::str("Received a call to retry() on message ", id,
447        " after ack() was already called. "
448        "Calling retry() on a message after calling ack() has no effect.");
449    IoContext::current().logWarning(msg);
450    return;
451  }
452 
453  auto& entry = result->retries.upsert(kj::heapString(id), {});
454  KJ_IF_SOME(opts, options) {
455    KJ_IF_SOME(secs, opts.delaySeconds) {
456      entry.value.delaySeconds = secs;
457    }
458  }
459}
460 
461void QueueMessage::ack() {
462  if (result->ackAll) {
463    return;
464  }
465 
466  if (result->retryBatch.retry) {
467    auto msg = kj::str("Received a call to ack() on message ", id,
468        " after retryAll() was already called. "
469        "Calling ack() on a message after calling retryAll() has no effect.");
470    IoContext::current().logWarning(msg);
471    return;
472  }
473 
474  if (result->retries.find(id) != kj::none) {
475    auto msg = kj::str("Received a call to ack() on message ", id,
476        " after retry() was already called. "
477        "Calling ack() on a message after calling retry() has no effect.");
478    IoContext::current().logWarning(msg);
479    return;
480  }
481  result->explicitAcks.findOrCreate(id, [this]() { return kj::heapString(id); });
482}
483 
484QueueEvent::QueueEvent(
485    jsg::Lock& js, rpc::EventDispatcher::QueueParams::Reader params, IoPtr<QueueEventResult> result)
486    : ExtendableEvent("queue"),
487      queueName(kj::heapString(params.getQueueName())),
488      result(result) {
489  // Note that we must make deep copies of all data here since the incoming Reader may be
490  // deallocated while JS's GC wrappers still exist.
491  auto incoming = params.getMessages();
492  auto messagesBuilder = kj::heapArrayBuilder<jsg::Ref<QueueMessage>>(incoming.size());
493  for (auto i: kj::indices(incoming)) {
494    messagesBuilder.add(js.alloc<QueueMessage>(js, incoming[i], result));
495  }
496  messages = messagesBuilder.finish();
497 
498  // Extract metadata. If the sender didn't set the field, capnp defaults all to the zero values.
499  auto m = params.getMetadata().getMetrics();
500  jsg::Optional<kj::Date> oldestTimestamp;
501  if (m.getOldestMessageTimestamp() != 0) {
502    oldestTimestamp =
503        kj::UNIX_EPOCH + static_cast<int64_t>(m.getOldestMessageTimestamp()) * kj::MILLISECONDS;
504  }
505  metadata = MessageBatchMetadata{
506    .metrics =
507        MessageBatchMetrics{
508          .backlogCount = m.getBacklogCount(),
509          .backlogBytes = m.getBacklogBytes(),
510          .oldestMessageTimestamp = oldestTimestamp,
511        },
512  };
513}
514 
515QueueEvent::QueueEvent(jsg::Lock& js, Params params, IoPtr<QueueEventResult> result)
516    : ExtendableEvent("queue"),
517      queueName(kj::mv(params.queueName)),
518      metadata(kj::mv(params.metadata)),
519      result(result) {
520  clearEpochSentinel(metadata.metrics.oldestMessageTimestamp);
521 
522  auto messagesBuilder = kj::heapArrayBuilder<jsg::Ref<QueueMessage>>(params.messages.size());
523  for (auto i: kj::indices(params.messages)) {
524    messagesBuilder.add(js.alloc<QueueMessage>(js, kj::mv(params.messages[i]), result));
525  }
526  messages = messagesBuilder.finish();
527}
528 
529void QueueEvent::retryAll(jsg::Optional<QueueRetryOptions> options) {
530  if (result->ackAll) {
531    IoContext::current().logWarning(
532        "Received a call to retryAll() after ackAll() was already called. "
533        "Calling retryAll() after calling ackAll() has no effect.");
534    return;
535  }
536 
537  result->retryBatch.retry = true;
538  KJ_IF_SOME(opts, options) {
539    KJ_IF_SOME(secs, opts.delaySeconds) {
540      result->retryBatch.delaySeconds = secs;
541    }
542  }
543}
544 
545void QueueEvent::ackAll() {
546  if (result->retryBatch.retry) {
547    IoContext::current().logWarning(
548        "Received a call to ackAll() after retryAll() was already called. "
549        "Calling ackAll() after calling retryAll() has no effect.");
550    return;
551  }
552  result->ackAll = true;
553}
554 
555namespace {
556 
557struct StartQueueEventResponse {
558  jsg::Ref<QueueEvent> event = nullptr;
559  kj::Maybe<kj::Promise<void>> exportedHandlerProm;
560  bool isServiceWorkerHandler = false;
561};
562 
563StartQueueEventResponse startQueueEvent(EventTarget& globalEventTarget,
564    IoContext& context,
565    kj::OneOf<rpc::EventDispatcher::QueueParams::Reader, QueueEvent::Params> params,
566    IoPtr<QueueEventResult> result,
567    Worker::Lock& lock,
568    kj::Maybe<ExportedHandler&> exportedHandler,
569    const jsg::TypeHandler<QueueExportedHandler>& handlerHandler) {
570  jsg::Lock& js = lock;
571  jsg::Ref<QueueEvent> event(nullptr);
572  KJ_SWITCH_ONEOF(params) {
573    KJ_CASE_ONEOF(p, rpc::EventDispatcher::QueueParams::Reader) {
574      event = js.alloc<QueueEvent>(js, p, result);
575    }
576    KJ_CASE_ONEOF(p, QueueEvent::Params) {
577      event = js.alloc<QueueEvent>(js, kj::mv(p), result);
578    }
579  }
580 
581  kj::Maybe<kj::Promise<void>> exportedHandlerProm;
582  bool isServiceWorkerHandler = false;
583  KJ_IF_SOME(h, exportedHandler) {
584    auto queueHandler = KJ_ASSERT_NONNULL(handlerHandler.tryUnwrap(lock, h.self.getHandle(lock)));
585    KJ_IF_SOME(f, queueHandler.queue) {
586      auto promise = f(lock, js.alloc<QueueController>(event.addRef()),
587          jsg::JsValue(h.env.getHandle(js)).addRef(js), h.getCtx())
588                         .then([event = event.addRef(), &context]() mutable {
589        event->setCompletionStatus(QueueEvent::CompletedSuccessfully{});
590        KJ_IF_SOME(t, context.getWorkerTracer()) {
591          t.setReturn(context.now());
592        }
593      }, [event = event.addRef()](kj::Exception&& e) mutable {
594        event->setCompletionStatus(QueueEvent::CompletedWithError{e.clone()});
595        return kj::mv(e);
596      });
597      if (FeatureFlags::get(js).getQueueConsumerNoWaitForWaitUntil()) {
598        exportedHandlerProm = kj::mv(promise);
599      } else {
600        event->waitUntil(kj::mv(promise));
601      }
602    } else {
603      lock.logWarningOnce("Received a QueueEvent but we lack a handler for QueueEvents. "
604                          "Did you remember to export a queue() function?");
605      JSG_FAIL_REQUIRE(Error, "Handler does not export a queue() function.");
606    }
607  } else {
608    isServiceWorkerHandler = true;
609    if (globalEventTarget.getHandlerCount("queue") == 0) {
610      lock.logWarningOnce("Received a QueueEvent but we lack an event listener for queue events. "
611                          "Did you remember to call addEventListener(\"queue\", ...)?");
612      JSG_FAIL_REQUIRE(Error, "No event listener registered for queue messages.");
613    }
614    globalEventTarget.dispatchEventImpl(lock, event.addRef());
615    event->setCompletionStatus(QueueEvent::CompletedSuccessfully{});
616  }
617 
618  return StartQueueEventResponse{
619    kj::mv(event), kj::mv(exportedHandlerProm), isServiceWorkerHandler};
620}
621 
622}  // namespace
623 
624tracing::EventInfo QueueCustomEvent::getEventInfo() const {
625  kj::String queueName;
626  uint32_t batchSize;
627  KJ_SWITCH_ONEOF(params) {
628    KJ_CASE_ONEOF(p, rpc::EventDispatcher::QueueParams::Reader) {
629      queueName = kj::heapString(p.getQueueName());
630      batchSize = p.getMessages().size();
631    }
632    KJ_CASE_ONEOF(p, QueueEvent::Params) {
633      queueName = kj::heapString(p.queueName);
634      batchSize = p.messages.size();
635    }
636  }
637 
638  return tracing::QueueEventInfo(kj::mv(queueName), batchSize);
639}
640 
641kj::Promise<WorkerInterface::CustomEvent::Result> QueueCustomEvent::run(
642    kj::Own<IoContext_IncomingRequest> incomingRequest,
643    kj::Maybe<kj::StringPtr> entrypointName,
644    kj::Maybe<Worker::VersionInfo> versionInfo,
645    Frankenvalue props,
646    kj::TaskSet& waitUntilTasks,
647    bool isDynamicDispatch) {
648  // This method has three main chunks of logic:
649  // 1. Do all necessary setup work. This starts right below this comment.
650  // 2. Call into the worker's queue event handler.
651  // 3. Wait on the necessary portions of the worker's code to complete.
652  incomingRequest->delivered();
653  auto& context = incomingRequest->getContext();
654 
655  // Create a custom refcounted type for holding the queueEvent so that we can pass it to the
656  // waitUntil'ed callback safely without worrying about whether this coroutine gets canceled.
657  struct QueueEventHolder: public kj::Refcounted {
658    jsg::Ref<QueueEvent> event = nullptr;
659    kj::Maybe<kj::Promise<void>> exportedHandlerProm;
660    bool isServiceWorkerHandler = false;
661  };
662  auto queueEventHolder = kj::refcounted<QueueEventHolder>();
663 
664  // 2. This is where we call into the worker's queue event handler
665  auto runProm = context.run(
666      [this, entrypointName = entrypointName, &context, queueEvent = kj::addRef(*queueEventHolder),
667          &metrics = incomingRequest->getMetrics(), versionInfo = kj::mv(versionInfo),
668          props = kj::mv(props), isDynamicDispatch](Worker::Lock& lock) mutable {
669    jsg::AsyncContextFrame::StorageScope traceScope = context.makeAsyncTraceScope(lock);
670    jsg::AsyncContextFrame::StorageScope userTraceScope = context.makeUserAsyncTraceScope(lock);
671 
672    auto& typeHandler = lock.getWorker().getIsolate().getApi().getQueueTypeHandler(lock);
673    auto startResp = startQueueEvent(lock.getGlobalScope(), context, kj::mv(params),
674        context.addObject(result), lock,
675        lock.getExportedHandler(entrypointName, kj::mv(versionInfo), kj::mv(props),
676            context.getActor(), isDynamicDispatch),
677        typeHandler);
678    queueEvent->event = kj::mv(startResp.event);
679    queueEvent->exportedHandlerProm = kj::mv(startResp.exportedHandlerProm);
680    queueEvent->isServiceWorkerHandler = startResp.isServiceWorkerHandler;
681  });
682 
683  // 3. Now that we've (asynchronously) called into the event handler, wait on all necessary async
684  // work to complete. This logic is split into two completely separate code paths depending on
685  // whether the queueConsumerNoWaitForWaitUntil compatibility flag is enabled.
686  // * In the enabled path, the queue event can be considered complete as soon as the event handler
687  //   returns and the promise that it returns (if any) has resolved.
688  // * In the disabled path, the queue event isn't complete until all waitUntil'ed promises resolve.
689  //   This was how Queues originally worked, but made for a poor user experience.
690  auto compatFlags = context.getWorker().getIsolate().getApi().getFeatureFlags();
691  if (compatFlags.getQueueConsumerNoWaitForWaitUntil()) {
692    // The user has opted in to only waiting on their event handler rather than all waitUntil'd
693    // promises.
694    auto timeoutPromise = context.getLimitEnforcer().limitScheduled();
695    // Start invoking the queue handler. The promise chain here is intended to mimic the behavior of
696    // finishScheduled, but only waiting on the promise returned by the event handler rather than on
697    // all waitUntil'ed promises.
698    auto outcome = co_await runProm
699                       .then([queueEvent = kj::addRef(
700                                  *queueEventHolder)]() mutable -> kj::Promise<EventOutcome> {
701      // If the queue handler returned a promise, wait on the promise.
702      KJ_IF_SOME(handlerProm, queueEvent->exportedHandlerProm) {
703        return handlerProm.then([]() { return EventOutcome::OK; });
704      }
705      // If not, we can consider the invocation complete.
706      return EventOutcome::OK;
707    })
708                       .catch_([](kj::Exception&& e) {
709      // If any exceptions were thrown, mark the outcome accordingly.
710      return EventOutcome::EXCEPTION;
711    })
712                       .exclusiveJoin(timeoutPromise.then([] {
713      // Join everything against a timeout to ensure queue handlers can't run forever.
714      return EventOutcome::EXCEEDED_CPU;
715    })).exclusiveJoin(context.onAbort().then([] {
716      // Also handle anything that might cause the worker to get aborted.
717      // This is a change from the outcome we returned on abort before the compat flag, but better
718      // matches the behavior of fetch() handlers and the semantics of what's actually happening.
719      return EventOutcome::EXCEPTION;
720    }, [](kj::Exception&&) { return EventOutcome::EXCEPTION; }));
721 
722    if (outcome == EventOutcome::OK && queueEventHolder->isServiceWorkerHandler) {
723      // HACK: For service-worker syntax, we effectively ignore the compatibility flag and wait
724      // for all waitUntil tasks anyway, since otherwise there's no way to do async work from an
725      // event listener callback.
726      // It'd be nicer if we could fall through to the code below for the non-compat-flag logic in
727      // this case, but we don't even know if the worker uses service worker syntax until after
728      // runProm resolves, so we just copy the bare essentials here.
729      auto scheduledResult = co_await incomingRequest->finishScheduled();
730      bool completed = scheduledResult == EventOutcome::OK;
731      outcome = completed ? context.waitUntilStatus() : scheduledResult;
732    } else {
733      // We're responsible for calling drain() on the incomingRequest to ensure that waitUntil tasks
734      // can continue to run in the backgound for a while even after we return a result to the
735      // caller of this event. But this is only needed in this code path because in all other code
736      // paths we call incomingRequest->finishScheduled(), which already takes care of waiting on
737      // waitUntil tasks.
738      waitUntilTasks.add(incomingRequest->drain().attach(
739          kj::mv(incomingRequest), kj::addRef(*queueEventHolder), kj::addRef(*this)));
740    }
741 
742    KJ_IF_SOME(status, context.getLimitEnforcer().getLimitsExceeded()) {
743      outcome = status;
744    }
745    co_return WorkerInterface::CustomEvent::Result{.outcome = outcome};
746  } else {
747    // The user has not opted in to the new waitUntil behavior, so we need to add the queue()
748    // handler's promise to the waitUntil promises and then wait on them all to finish.
749    context.addWaitUntil(kj::mv(runProm));
750 
751    // We reuse the finishScheduled() method for convenience, since queues use the same wall clock
752    // timeout as scheduled workers.
753    auto scheduledResult = co_await incomingRequest->finishScheduled();
754    bool completed = scheduledResult == EventOutcome::OK;
755 
756    co_return WorkerInterface::CustomEvent::Result{
757      .outcome = completed ? context.waitUntilStatus() : scheduledResult,
758    };
759  }
760}
761 
762kj::Promise<WorkerInterface::CustomEvent::Result> QueueCustomEvent::sendRpc(
763    capnp::HttpOverCapnpFactory& httpOverCapnpFactory,
764    capnp::ByteStreamFactory& byteStreamFactory,
765    rpc::EventDispatcher::Client dispatcher) {
766  auto req = dispatcher.castAs<rpc::EventDispatcher>().queueRequest();
767  KJ_SWITCH_ONEOF(params) {
768    KJ_CASE_ONEOF(p, rpc::EventDispatcher::QueueParams::Reader) {
769      req.setQueueName(p.getQueueName());
770      req.setMessages(p.getMessages());
771      req.setMetadata(p.getMetadata());
772    }
773    KJ_CASE_ONEOF(p, QueueEvent::Params) {
774      req.setQueueName(p.queueName);
775      auto messages = req.initMessages(p.messages.size());
776      for (auto i: kj::indices(p.messages)) {
777        messages[i].setId(p.messages[i].id);
778        messages[i].setTimestampNs((p.messages[i].timestamp - kj::UNIX_EPOCH) / kj::NANOSECONDS);
779        messages[i].setData(p.messages[i].body);
780        KJ_IF_SOME(contentType, p.messages[i].contentType) {
781          messages[i].setContentType(contentType);
782        }
783        messages[i].setAttempts(p.messages[i].attempts);
784      }
785      {
786        auto metadataBuilder = req.initMetadata();
787        auto metricsBuilder = metadataBuilder.initMetrics();
788        metricsBuilder.setBacklogCount(p.metadata.metrics.backlogCount);
789        metricsBuilder.setBacklogBytes(p.metadata.metrics.backlogBytes);
790        KJ_IF_SOME(ts, p.metadata.metrics.oldestMessageTimestamp) {
791          metricsBuilder.setOldestMessageTimestamp((ts - kj::UNIX_EPOCH) / kj::MILLISECONDS);
792        }
793      }
794    }
795  }
796 
797  return req.send().then([this](auto resp) {
798    auto respResult = resp.getResult();
799    this->result.ackAll = respResult.getAckAll();
800    auto retryBatch = respResult.getRetryBatch();
801    this->result.retryBatch.retry = retryBatch.getRetry();
802    if (retryBatch.isDelaySeconds()) {
803      this->result.retryBatch.delaySeconds = retryBatch.getDelaySeconds();
804    }
805 
806    this->result.explicitAcks.clear();
807    for (const auto& msgId: respResult.getExplicitAcks()) {
808      this->result.explicitAcks.insert(kj::heapString(msgId));
809    }
810    this->result.retries.clear();
811    for (const auto& retry: respResult.getRetryMessages()) {
812      auto& entry = this->result.retries.upsert(kj::heapString(retry.getMsgId()), {});
813      if (retry.isDelaySeconds()) {
814        entry.value.delaySeconds = retry.getDelaySeconds();
815      }
816    }
817 
818    return WorkerInterface::CustomEvent::Result{
819      .outcome = respResult.getOutcome(),
820    };
821  });
822}
823 
824kj::Array<QueueRetryMessage> QueueCustomEvent::getRetryMessages() const {
825  auto retryMsgs = kj::heapArrayBuilder<QueueRetryMessage>(result.retries.size());
826  for (const auto& entry: result.retries) {
827    retryMsgs.add(QueueRetryMessage{
828      .msgId = kj::heapString(entry.key), .delaySeconds = entry.value.delaySeconds});
829  }
830  return retryMsgs.finish();
831}
832 
833kj::Array<kj::String> QueueCustomEvent::getExplicitAcks() const {
834  auto ackArray = kj::heapArrayBuilder<kj::String>(result.explicitAcks.size());
835  for (const auto& msgId: result.explicitAcks) {
836    ackArray.add(kj::heapString(msgId));
837  }
838  return ackArray.finish();
839}
840 
841}  // namespace workerd::api