Skip to content

Commit

Permalink
Merge pull request #709 from cloudflare/malonso/hibernatable-ws-traces
Browse files Browse the repository at this point in the history 
Add support for hibernatable websocket tracing
  • Loading branch information
@MellowYarker
MellowYarker authored Jan 24, 2024
2 parents f4f29e8 + cee21cb commit bce7d42
Show file tree
Hide file tree
Showing 7 changed files with 280 additions and 23 deletions.
32 changes: 31 additions & 1 deletion src/workerd/api/hibernatable-web-socket.c++
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,9 +71,39 @@ kj::Promise<WorkerInterface::CustomEvent::Result> HibernatableWebSocketCustomEve
a.setHibernationManager(kj::addRef(KJ_REQUIRE_NONNULL(manager)));
}

auto eventParameters = consumeParams();

KJ_IF_MAYBE(t, incomingRequest->getWorkerTracer()) {
Trace::HibernatableWebSocketEventInfo::Type type = [&]()
-> Trace::HibernatableWebSocketEventInfo::Type {
KJ_SWITCH_ONEOF(eventParameters.eventType) {
KJ_CASE_ONEOF(_, HibernatableSocketParams::Text) {
return Trace::HibernatableWebSocketEventInfo::Message{};
}
KJ_CASE_ONEOF(data, HibernatableSocketParams::Data) {
return Trace::HibernatableWebSocketEventInfo::Message{};
}
KJ_CASE_ONEOF(close, HibernatableSocketParams::Close) {
return Trace::HibernatableWebSocketEventInfo::Close{
.code = close.code,
.wasClean = close.wasClean
};
}
KJ_CASE_ONEOF(_, HibernatableSocketParams::Error) {
return Trace::HibernatableWebSocketEventInfo::Error{};
}
}
KJ_UNREACHABLE;
}();

t->setEventInfo(context.now(),
Trace::HibernatableWebSocketEventInfo(kj::mv(type))
);
}

try {
co_await context.run(
[entrypointName=entrypointName, &context, eventParameters=consumeParams()]
[entrypointName=entrypointName, &context, eventParameters = kj::mv(eventParameters)]
(Worker::Lock& lock) mutable {
KJ_SWITCH_ONEOF(eventParameters.eventType) {
KJ_CASE_ONEOF(text, HibernatableSocketParams::Text) {
Expand Down
5 changes: 3 additions & 2 deletions src/workerd/api/hibernation-event-params.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

#pragma once

#include <cstdint>
#include <kj/common.h>
#include <kj/string.h>
#include <kj/debug.h>
Expand All @@ -21,7 +22,7 @@ namespace workerd::api {
};

struct Close {
int code;
uint16_t code;
kj::String reason;
bool wasClean;
};
Expand All @@ -38,7 +39,7 @@ namespace workerd::api {
: eventType(Text { kj::mv(message) }), websocketId(kj::mv(id)) {}
explicit HibernatableSocketParams(kj::Array<kj::byte> message, kj::String id)
: eventType(Data { kj::mv(message) }), websocketId(kj::mv(id)) {}
explicit HibernatableSocketParams(int code, kj::String reason, bool wasClean, kj::String id)
explicit HibernatableSocketParams(uint16_t code, kj::String reason, bool wasClean, kj::String id)
: eventType(Close { code, kj::mv(reason), wasClean }), websocketId(kj::mv(id)) {}
explicit HibernatableSocketParams(kj::Exception e, kj::String id)
: eventType(Error { kj::mv(e) }), websocketId(kj::mv(id)) {}
Expand Down
57 changes: 57 additions & 0 deletions src/workerd/api/trace.c++
View file
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,20 @@ kj::Maybe<TraceItem::EventInfo> getTraceEvent(jsg::Lock& js, const Trace& trace)
KJ_CASE_ONEOF(tracedTrace, Trace::TraceEventInfo) {
return kj::Maybe(jsg::alloc<TraceItem::TailEventInfo>(trace, tracedTrace));
}
KJ_CASE_ONEOF(hibWs, Trace::HibernatableWebSocketEventInfo) {
KJ_SWITCH_ONEOF(hibWs.type) {
KJ_CASE_ONEOF(message, Trace::HibernatableWebSocketEventInfo::Message) {
return kj::Maybe(jsg::alloc<TraceItem::HibernatableWebSocketEventInfo>(trace, message));
}
KJ_CASE_ONEOF(close, Trace::HibernatableWebSocketEventInfo::Close) {
return kj::Maybe(jsg::alloc<TraceItem::HibernatableWebSocketEventInfo>(trace, close));
}
KJ_CASE_ONEOF(error, Trace::HibernatableWebSocketEventInfo::Error) {
return kj::Maybe(jsg::alloc<TraceItem::HibernatableWebSocketEventInfo>(trace, error));
}
}
KJ_UNREACHABLE;
}
KJ_CASE_ONEOF(custom, Trace::CustomEventInfo) {
return kj::Maybe(jsg::alloc<TraceItem::CustomEventInfo>(trace, custom));
}
Expand Down Expand Up @@ -194,6 +208,7 @@ kj::Maybe<TraceItem::EventInfo> TraceItem::getEvent(jsg::Lock& js) {
KJ_CASE_ONEOF(info, jsg::Ref<QueueEventInfo>) { return info.addRef(); }
KJ_CASE_ONEOF(info, jsg::Ref<EmailEventInfo>) { return info.addRef(); }
KJ_CASE_ONEOF(info, jsg::Ref<TailEventInfo>) { return info.addRef(); }
KJ_CASE_ONEOF(info, jsg::Ref<HibernatableWebSocketEventInfo>) { return info.addRef(); }
KJ_CASE_ONEOF(info, jsg::Ref<CustomEventInfo>) { return info.addRef(); }
}
KJ_UNREACHABLE;
Expand Down Expand Up @@ -453,6 +468,48 @@ TraceItem::CustomEventInfo::CustomEventInfo(const Trace& trace,
const Trace::CustomEventInfo& eventInfo)
: eventInfo(eventInfo) {}

TraceItem::HibernatableWebSocketEventInfo::HibernatableWebSocketEventInfo(
const Trace& trace,
const Trace::HibernatableWebSocketEventInfo::Message& eventInfo)
: eventType(
jsg::alloc<TraceItem::HibernatableWebSocketEventInfo::Message>(trace, eventInfo)) {}

TraceItem::HibernatableWebSocketEventInfo::HibernatableWebSocketEventInfo(
const Trace& trace,
const Trace::HibernatableWebSocketEventInfo::Close& eventInfo)
: eventType(
jsg::alloc<TraceItem::HibernatableWebSocketEventInfo::Close>(trace, eventInfo)) {}

TraceItem::HibernatableWebSocketEventInfo::HibernatableWebSocketEventInfo(
const Trace& trace,
const Trace::HibernatableWebSocketEventInfo::Error& eventInfo)
: eventType(
jsg::alloc<TraceItem::HibernatableWebSocketEventInfo::Error>(trace, eventInfo)) {}

TraceItem::HibernatableWebSocketEventInfo::Type
TraceItem::HibernatableWebSocketEventInfo::getEvent() {
KJ_SWITCH_ONEOF(eventType) {
KJ_CASE_ONEOF(m, jsg::Ref<TraceItem::HibernatableWebSocketEventInfo::Message>) {
return m.addRef();
}
KJ_CASE_ONEOF(c, jsg::Ref<TraceItem::HibernatableWebSocketEventInfo::Close>) {
return c.addRef();
}
KJ_CASE_ONEOF(e, jsg::Ref<TraceItem::HibernatableWebSocketEventInfo::Error>) {
return e.addRef();
}
}
KJ_UNREACHABLE;
}

uint16_t TraceItem::HibernatableWebSocketEventInfo::Close::getCode() {
return eventInfo.code;
}

bool TraceItem::HibernatableWebSocketEventInfo::Close::getWasClean() {
return eventInfo.wasClean;
}

TraceLog::TraceLog(jsg::Lock& js, const Trace& trace, const Trace::Log& log)
: timestamp(getTraceLogTimestamp(log)),
level(getTraceLogLevel(log)),
Expand Down
122 changes: 103 additions & 19 deletions src/workerd/api/trace.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ class TraceItem final: public jsg::Object {
class QueueEventInfo;
class EmailEventInfo;
class TailEventInfo;
class HibernatableWebSocketEventInfo;
class CustomEventInfo;

explicit TraceItem(jsg::Lock& js, const Trace& trace);
Expand All @@ -74,7 +75,8 @@ class TraceItem final: public jsg::Object {
jsg::Ref<QueueEventInfo>,
jsg::Ref<EmailEventInfo>,
jsg::Ref<TailEventInfo>,
jsg::Ref<CustomEventInfo>> EventInfo;
jsg::Ref<CustomEventInfo>,
jsg::Ref<HibernatableWebSocketEventInfo>> EventInfo;
kj::Maybe<EventInfo> getEvent(jsg::Lock& js);
kj::Maybe<double> getEventTimestamp();

Expand Down Expand Up @@ -312,6 +314,84 @@ class TraceItem::TailEventInfo::TailItem final: public jsg::Object {
kj::Maybe<kj::String> scriptName;
};

class TraceItem::HibernatableWebSocketEventInfo final: public jsg::Object {
public:
class Message;
class Close;
class Error;

explicit HibernatableWebSocketEventInfo(const Trace& trace,
const Trace::HibernatableWebSocketEventInfo::Message& eventInfo);
explicit HibernatableWebSocketEventInfo(const Trace& trace,
const Trace::HibernatableWebSocketEventInfo::Close& eventInfo);
explicit HibernatableWebSocketEventInfo(const Trace& trace,
const Trace::HibernatableWebSocketEventInfo::Error& eventInfo);

using Type = kj::OneOf<jsg::Ref<Message>, jsg::Ref<Close>, jsg::Ref<Error>>;

Type getEvent();

JSG_RESOURCE_TYPE(HibernatableWebSocketEventInfo) {
JSG_LAZY_READONLY_INSTANCE_PROPERTY(getWebSocketEvent, getEvent);
}

private:
Type eventType;
};

class TraceItem::HibernatableWebSocketEventInfo::Message final: public jsg::Object {
public:
explicit Message(const Trace& trace,
const Trace::HibernatableWebSocketEventInfo::Message& eventInfo): eventInfo(eventInfo) {}

static constexpr kj::StringPtr webSocketEventType = "message"_kj;
kj::StringPtr getWebSocketEventType() { return webSocketEventType; }

JSG_RESOURCE_TYPE(Message) {
JSG_READONLY_INSTANCE_PROPERTY(webSocketEventType, getWebSocketEventType);
}

private:
const Trace::HibernatableWebSocketEventInfo::Message& eventInfo;
};

class TraceItem::HibernatableWebSocketEventInfo::Close final: public jsg::Object {
public:
explicit Close(const Trace& trace,
const Trace::HibernatableWebSocketEventInfo::Close& eventInfo): eventInfo(eventInfo) {}

static constexpr kj::StringPtr webSocketEventType = "close"_kj;
kj::StringPtr getWebSocketEventType() { return webSocketEventType; }

uint16_t getCode();
bool getWasClean();

JSG_RESOURCE_TYPE(Close) {
JSG_READONLY_INSTANCE_PROPERTY(webSocketEventType, getWebSocketEventType);
JSG_READONLY_INSTANCE_PROPERTY(code, getCode);
JSG_READONLY_INSTANCE_PROPERTY(wasClean, getWasClean);
}

private:
const Trace::HibernatableWebSocketEventInfo::Close& eventInfo;
};

class TraceItem::HibernatableWebSocketEventInfo::Error final: public jsg::Object {
public:
explicit Error(const Trace& trace,
const Trace::HibernatableWebSocketEventInfo::Error& eventInfo): eventInfo(eventInfo) {}

static constexpr kj::StringPtr webSocketEventType = "error"_kj;
kj::StringPtr getWebSocketEventType() { return webSocketEventType; }

JSG_RESOURCE_TYPE(Error) {
JSG_READONLY_INSTANCE_PROPERTY(webSocketEventType, getWebSocketEventType);
}

private:
const Trace::HibernatableWebSocketEventInfo::Error& eventInfo;
};

class TraceItem::CustomEventInfo final: public jsg::Object {
public:
explicit CustomEventInfo(const Trace& trace, const Trace::CustomEventInfo& eventInfo);
Expand Down Expand Up @@ -439,24 +519,28 @@ class TraceCustomEventImpl final: public WorkerInterface::CustomEvent {
kj::Array<kj::Own<workerd::Trace>> traces;
};

#define EW_TRACE_ISOLATE_TYPES \
api::ScriptVersion, \
api::TailEvent, \
api::TraceItem, \
api::TraceItem::AlarmEventInfo, \
api::TraceItem::CustomEventInfo, \
api::TraceItem::ScheduledEventInfo, \
api::TraceItem::QueueEventInfo, \
api::TraceItem::EmailEventInfo, \
api::TraceItem::TailEventInfo, \
api::TraceItem::TailEventInfo::TailItem, \
api::TraceItem::FetchEventInfo, \
api::TraceItem::FetchEventInfo::Request, \
api::TraceItem::FetchEventInfo::Response, \
api::TraceLog, \
api::TraceException, \
api::TraceDiagnosticChannelEvent, \
api::TraceMetrics, \
#define EW_TRACE_ISOLATE_TYPES \
api::ScriptVersion, \
api::TailEvent, \
api::TraceItem, \
api::TraceItem::AlarmEventInfo, \
api::TraceItem::CustomEventInfo, \
api::TraceItem::ScheduledEventInfo, \
api::TraceItem::QueueEventInfo, \
api::TraceItem::EmailEventInfo, \
api::TraceItem::TailEventInfo, \
api::TraceItem::TailEventInfo::TailItem, \
api::TraceItem::FetchEventInfo, \
api::TraceItem::FetchEventInfo::Request, \
api::TraceItem::FetchEventInfo::Response, \
api::TraceItem::HibernatableWebSocketEventInfo, \
api::TraceItem::HibernatableWebSocketEventInfo::Message, \
api::TraceItem::HibernatableWebSocketEventInfo::Close, \
api::TraceItem::HibernatableWebSocketEventInfo::Error, \
api::TraceLog, \
api::TraceException, \
api::TraceDiagnosticChannelEvent, \
api::TraceMetrics, \
api::UnsafeTraceMetrics
// The list of trace.h types that are added to worker.c++'s JSG_DECLARE_ISOLATE_TYPE

Expand Down
53 changes: 53 additions & 0 deletions src/workerd/io/trace.c++
View file
Original file line number Diff line number Diff line change
Expand Up @@ -167,6 +167,51 @@ void Trace::DiagnosticChannelEvent::copyTo(
builder.setMessage(message);
}

Trace::HibernatableWebSocketEventInfo::HibernatableWebSocketEventInfo(Type type)
: type(type) {}

Trace::HibernatableWebSocketEventInfo::HibernatableWebSocketEventInfo(
rpc::Trace::HibernatableWebSocketEventInfo::Reader reader)
: type(readFrom(reader)) {}

void Trace::HibernatableWebSocketEventInfo::copyTo(
rpc::Trace::HibernatableWebSocketEventInfo::Builder builder) {
auto typeBuilder = builder.initType();
KJ_SWITCH_ONEOF(type) {
KJ_CASE_ONEOF(_, Message) {
typeBuilder.setMessage();
}
KJ_CASE_ONEOF(close, Close) {
auto closeBuilder = typeBuilder.initClose();
closeBuilder.setCode(close.code);
closeBuilder.setWasClean(close.wasClean);
}
KJ_CASE_ONEOF(_, Error) {
typeBuilder.setError();
}
}
}

Trace::HibernatableWebSocketEventInfo::Type Trace::HibernatableWebSocketEventInfo::readFrom(
rpc::Trace::HibernatableWebSocketEventInfo::Reader reader) {
auto type = reader.getType();
switch(type.which()) {
case rpc::Trace::HibernatableWebSocketEventInfo::Type::MESSAGE: {
return Message{};
}
case rpc::Trace::HibernatableWebSocketEventInfo::Type::CLOSE: {
auto close = type.getClose();
return Close {
.code = close.getCode(),
.wasClean = close.getWasClean(),
};
}
case rpc::Trace::HibernatableWebSocketEventInfo::Type::ERROR: {
return Error{};
}
}
}

Trace::FetchResponseInfo::FetchResponseInfo(uint16_t statusCode)
: statusCode(statusCode) {}

Expand Down Expand Up @@ -262,6 +307,10 @@ void Trace::copyTo(rpc::Trace::Builder builder) {
auto traceBuilder = eventInfoBuilder.initTrace();
trace.copyTo(traceBuilder);
}
KJ_CASE_ONEOF(hibWs, HibernatableWebSocketEventInfo) {
auto hibWsBuilder = eventInfoBuilder.initHibernatableWebSocket();
hibWs.copyTo(hibWsBuilder);
}
KJ_CASE_ONEOF(custom, CustomEventInfo) {
eventInfoBuilder.initCustom();
}
Expand Down Expand Up @@ -347,6 +396,9 @@ void Trace::mergeFrom(rpc::Trace::Reader reader, PipelineLogLevel pipelineLogLev
case rpc::Trace::EventInfo::Which::TRACE:
eventInfo = TraceEventInfo(e.getTrace());
break;
case rpc::Trace::EventInfo::Which::HIBERNATABLE_WEB_SOCKET:
eventInfo = HibernatableWebSocketEventInfo(e.getHibernatableWebSocket());
break;
case rpc::Trace::EventInfo::Which::CUSTOM:
eventInfo = CustomEventInfo(e.getCustom());
break;
Expand Down Expand Up @@ -563,6 +615,7 @@ void WorkerTracer::setEventInfo(kj::Date timestamp, Trace::EventInfo&& info) {
KJ_CASE_ONEOF(_, Trace::QueueEventInfo) {}
KJ_CASE_ONEOF(_, Trace::EmailEventInfo) {}
KJ_CASE_ONEOF(_, Trace::TraceEventInfo) {}
KJ_CASE_ONEOF(_, Trace::HibernatableWebSocketEventInfo) {}
KJ_CASE_ONEOF(_, Trace::CustomEventInfo) {}
}
trace->bytesUsed = newSize;
Expand Down
Loading

0 comments on commit bce7d42

Please sign in to comment.