nicktrn
101 lines
2.9 KiB
TypeScript
101 lines
2.9 KiB
TypeScript
import {
|
|
MessageCatalogToSocketIoEvents,
|
|
StructuredLogger,
|
|
ZodMessageHandler,
|
|
ZodMessageSender,
|
|
clientWebsocketMessages,
|
|
serverWebsocketMessages,
|
|
} from "@trigger.dev/core/v3";
|
|
import { Evt } from "evt";
|
|
import { randomUUID } from "node:crypto";
|
|
import { logger } from "~/services/logger.server";
|
|
import { SharedQueueConsumer } from "./marqs/sharedQueueConsumer.server";
|
|
import { DisconnectReason, Namespace, Socket } from "socket.io";
|
|
|
|
interface SharedSocketConnectionOptions {
|
|
namespace: Namespace<
|
|
MessageCatalogToSocketIoEvents<typeof clientWebsocketMessages>,
|
|
MessageCatalogToSocketIoEvents<typeof serverWebsocketMessages>
|
|
>;
|
|
socket: Socket<
|
|
MessageCatalogToSocketIoEvents<typeof clientWebsocketMessages>,
|
|
MessageCatalogToSocketIoEvents<typeof serverWebsocketMessages>
|
|
>;
|
|
logger?: StructuredLogger;
|
|
}
|
|
|
|
export class SharedSocketConnection {
|
|
public id: string;
|
|
public onClose: Evt<DisconnectReason> = new Evt();
|
|
|
|
private _sender: ZodMessageSender<typeof serverWebsocketMessages>;
|
|
private _sharedConsumer: SharedQueueConsumer;
|
|
private _messageHandler: ZodMessageHandler<typeof clientWebsocketMessages>;
|
|
|
|
constructor(opts: SharedSocketConnectionOptions) {
|
|
this.id = randomUUID();
|
|
|
|
this._sender = new ZodMessageSender({
|
|
schema: serverWebsocketMessages,
|
|
sender: async (message) => {
|
|
return new Promise((resolve, reject) => {
|
|
try {
|
|
const { type, ...payload } = message;
|
|
opts.namespace.emit(type, payload as any);
|
|
resolve();
|
|
} catch (err) {
|
|
reject(err);
|
|
}
|
|
});
|
|
},
|
|
});
|
|
|
|
this._sharedConsumer = new SharedQueueConsumer(this._sender, {
|
|
interval: 100,
|
|
nextTickInterval: 1000,
|
|
});
|
|
|
|
opts.socket.on("disconnect", this.#handleClose.bind(this));
|
|
opts.socket.on("error", this.#handleError.bind(this));
|
|
|
|
this._messageHandler = new ZodMessageHandler({
|
|
schema: clientWebsocketMessages,
|
|
messages: {
|
|
READY_FOR_TASKS: async (payload) => {
|
|
this._sharedConsumer.start();
|
|
},
|
|
BACKGROUND_WORKER_DEPRECATED: async (payload) => {
|
|
// await this._sharedConsumer.deprecateBackgroundWorker(payload.backgroundWorkerId);
|
|
},
|
|
BACKGROUND_WORKER_MESSAGE: async (payload) => {
|
|
switch (payload.data.type) {
|
|
case "TASK_RUN_COMPLETED": {
|
|
// handled in coordinator namespace
|
|
break;
|
|
}
|
|
case "TASK_HEARTBEAT": {
|
|
// handled in coordinator namespace
|
|
break;
|
|
}
|
|
}
|
|
},
|
|
},
|
|
});
|
|
this._messageHandler.registerHandlers(opts.socket, opts.logger ?? logger);
|
|
}
|
|
|
|
async initialize() {
|
|
this._sender.send("SERVER_READY", { id: this.id });
|
|
}
|
|
|
|
async #handleClose(ev: DisconnectReason) {
|
|
await this._sharedConsumer.stop();
|
|
|
|
this.onClose.post(ev);
|
|
}
|
|
|
|
async #handleError(ev: Error) {
|
|
logger.error("Websocket error", { ev });
|
|
}
|
|
}
|