From 7ea7fe783ce77be2c0342d03be61aa9f96724822 Mon Sep 17 00:00:00 2001 From: Nicolas Meienberger Date: Tue, 31 Mar 2026 18:50:05 +0200 Subject: [PATCH] refactor: context as effectful callbacks --- .../agents/controller-agent-session.ts | 76 ++++++------- app/server/modules/backups/backup-executor.ts | 36 +++--- apps/agent/src/commands/backup-cancel.ts | 4 +- apps/agent/src/commands/backup-run.ts | 106 ++++++++---------- apps/agent/src/commands/heartbeat-ping.ts | 4 +- apps/agent/src/context.ts | 9 +- apps/agent/src/controller-session.ts | 66 +++++------ 7 files changed, 141 insertions(+), 160 deletions(-) diff --git a/app/server/modules/agents/controller-agent-session.ts b/app/server/modules/agents/controller-agent-session.ts index 50ead28b..1d22e511 100644 --- a/app/server/modules/agents/controller-agent-session.ts +++ b/app/server/modules/agents/controller-agent-session.ts @@ -24,12 +24,6 @@ export type AgentConnectionData = { type AgentSocket = Bun.ServerWebSocket; -type SessionState = { - isReady: boolean; - lastSeenAt: number | null; - lastPongAt: number | null; -}; - type ControllerAgentSessionHandlers = { onBackupStarted?: (payload: BackupStartedPayload) => void; onBackupProgress?: (payload: BackupProgressPayload) => void; @@ -55,13 +49,7 @@ export const createControllerAgentSession = ( const outboundQueue = Effect.runSync(Queue.bounded(64)); const activeBackupJobs = Effect.runSync(Ref.make>(new Map())); const pendingBackupJobs = Effect.runSync(Ref.make>(new Map())); - const state = Effect.runSync( - Ref.make({ - isReady: false, - lastSeenAt: null, - lastPongAt: null, - }), - ); + const isReadyRef = Effect.runSync(Ref.make(false)); const offerOutbound = (message: ControllerWireMessage) => { void Effect.runPromise(Queue.offer(outboundQueue, message)).catch((error) => { @@ -69,8 +57,8 @@ export const createControllerAgentSession = ( }); }; - const updateState = (update: (current: SessionState) => SessionState) => { - Effect.runSync(Ref.update(state, update)); + const setIsReady = (isReady: boolean) => { + Effect.runSync(Ref.set(isReadyRef, isReady)); }; const setActiveBackupJob = (jobId: string, scheduleId: string) => { @@ -113,35 +101,41 @@ export const createControllerAgentSession = ( ); }; + const clearTrackedBackupJob = (jobId: string) => { + deletePendingBackupJob(jobId); + deleteActiveBackupJob(jobId); + }; + + const cancelTrackedJobs = (jobs: Map, message: string) => { + for (const [jobId, scheduleId] of jobs) { + handlers.onBackupCancelled?.({ + jobId, + scheduleId, + message, + }); + } + }; + const closeSession = () => { if (isClosed) { return; } isClosed = true; - updateState((current) => ({ ...current, isReady: false })); + setIsReady(false); const pendingJobs = Effect.runSync(Ref.get(pendingBackupJobs)); Effect.runSync(Ref.set(pendingBackupJobs, new Map())); - - for (const [jobId, scheduleId] of pendingJobs) { - handlers.onBackupCancelled?.({ - jobId, - scheduleId, - message: - "The connection to the backup agent was lost before this backup started. Restart the backup to ensure it completes.", - }); - } + cancelTrackedJobs( + pendingJobs, + "The connection to the backup agent was lost before this backup started. Restart the backup to ensure it completes.", + ); const activeJobs = Effect.runSync(Ref.get(activeBackupJobs)); Effect.runSync(Ref.set(activeBackupJobs, new Map())); - for (const [jobId, scheduleId] of activeJobs) { - handlers.onBackupCancelled?.({ - jobId, - scheduleId, - message: - "The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes.", - }); - } + cancelTrackedJobs( + activeJobs, + "The connection to the backup agent was lost while this backup was running. Restart the backup to ensure it completes.", + ); void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {}); void Effect.runPromise(Fiber.interrupt(heartbeatFiber)).catch(() => {}); void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {}); @@ -190,11 +184,9 @@ export const createControllerAgentSession = ( ); const handleAgentMessage = (message: AgentMessage) => { - updateState((current) => ({ ...current, lastSeenAt: Date.now() })); - switch (message.type) { case "agent.ready": { - updateState((current) => ({ ...current, isReady: true })); + setIsReady(true); logger.info(`Agent "${socket.data.agentName}" (${socket.data.agentId}) is ready`); break; } @@ -212,25 +204,21 @@ export const createControllerAgentSession = ( break; } case "backup.completed": { - deletePendingBackupJob(message.payload.jobId); - deleteActiveBackupJob(message.payload.jobId); + clearTrackedBackupJob(message.payload.jobId); handlers.onBackupCompleted?.(message.payload); break; } case "backup.failed": { - deletePendingBackupJob(message.payload.jobId); - deleteActiveBackupJob(message.payload.jobId); + clearTrackedBackupJob(message.payload.jobId); handlers.onBackupFailed?.(message.payload); break; } case "backup.cancelled": { - deletePendingBackupJob(message.payload.jobId); - deleteActiveBackupJob(message.payload.jobId); + clearTrackedBackupJob(message.payload.jobId); handlers.onBackupCancelled?.(message.payload); break; } case "heartbeat.pong": { - updateState((current) => ({ ...current, lastPongAt: message.payload.sentAt })); break; } } @@ -261,7 +249,7 @@ export const createControllerAgentSession = ( sendBackupCancel: (payload) => { offerOutbound(createControllerMessage("backup.cancel", payload)); }, - isReady: () => Effect.runSync(Ref.get(state)).isReady, + isReady: () => Effect.runSync(Ref.get(isReadyRef)), close: closeSession, }; }; diff --git a/app/server/modules/backups/backup-executor.ts b/app/server/modules/backups/backup-executor.ts index cf1418b8..5743f9c0 100644 --- a/app/server/modules/backups/backup-executor.ts +++ b/app/server/modules/backups/backup-executor.ts @@ -99,17 +99,25 @@ const clearActiveExecution = (jobId: string) => { return activeExecution; }; -const getActiveExecution = (jobId: string, scheduleId: string, eventName: string, executorId: string) => { +const clearExecutionState = (jobId: string, scheduleId: number) => { + requestedCancellationsByScheduleId.delete(scheduleId); + clearActiveExecution(jobId); +}; + +const resolveExecution = (jobId: string, activeExecution: ActiveBackupExecution, result: BackupExecutionResult) => { + clearExecutionState(jobId, activeExecution.scheduleId); + activeExecution.resolve(result); +}; + +const getActiveExecution = (jobId: string, scheduleId: string, eventName: string, agentId: string) => { const activeExecution = activeExecutionsByJobId.get(jobId); if (!activeExecution) { - logger.warn(`Received ${eventName} for unknown job ${jobId} from executor ${executorId}`); + logger.warn(`Received ${eventName} for unknown job ${jobId} from agent ${agentId}`); return null; } if (activeExecution.scheduleShortId !== scheduleId) { - logger.warn( - `Ignoring ${eventName} for job ${jobId} due to schedule mismatch ${scheduleId} from executor ${executorId}`, - ); + logger.warn(`Ignoring ${eventName} for job ${jobId} due to schedule mismatch ${scheduleId} from agent ${agentId}`); return null; } @@ -134,9 +142,7 @@ agentManager.setBackupEventHandlers({ return; } - requestedCancellationsByScheduleId.delete(activeExecution.scheduleId); - clearActiveExecution(payload.jobId); - activeExecution.resolve({ + resolveExecution(payload.jobId, activeExecution, { status: "completed", exitCode: payload.exitCode, result: payload.result, @@ -149,9 +155,7 @@ agentManager.setBackupEventHandlers({ return; } - requestedCancellationsByScheduleId.delete(activeExecution.scheduleId); - clearActiveExecution(payload.jobId); - activeExecution.resolve({ + resolveExecution(payload.jobId, activeExecution, { status: "failed", error: payload.errorDetails ?? payload.error, }); @@ -163,9 +167,7 @@ agentManager.setBackupEventHandlers({ } const wasRequested = requestedCancellationsByScheduleId.has(activeExecution.scheduleId); - requestedCancellationsByScheduleId.delete(activeExecution.scheduleId); - clearActiveExecution(payload.jobId); - activeExecution.resolve({ + resolveExecution(payload.jobId, activeExecution, { status: "cancelled", message: wasRequested ? undefined : payload.message, }); @@ -207,8 +209,7 @@ export const backupExecutor = { } if (!agentManager.sendBackup(LOCAL_AGENT_ID, payload)) { - requestedCancellationsByScheduleId.delete(request.scheduleId); - clearActiveExecution(jobId); + clearExecutionState(jobId, request.scheduleId); return { status: "unavailable", error: new Error("Local backup agent is not connected"), @@ -217,8 +218,7 @@ export const backupExecutor = { return completion; } catch (error) { - requestedCancellationsByScheduleId.delete(request.scheduleId); - clearActiveExecution(jobId); + clearExecutionState(jobId, request.scheduleId); throw error; } }, diff --git a/apps/agent/src/commands/backup-cancel.ts b/apps/agent/src/commands/backup-cancel.ts index 2dec65dd..c36c21e9 100644 --- a/apps/agent/src/commands/backup-cancel.ts +++ b/apps/agent/src/commands/backup-cancel.ts @@ -4,8 +4,8 @@ import { logger } from "@zerobyte/core/node"; import type { ControllerCommandContext } from "../context"; export const handleBackupCancelCommand = (context: ControllerCommandContext, payload: BackupCancelPayload) => - Effect.sync(() => { - const running = context.getRunningJob(payload.jobId); + Effect.gen(function* () { + const running = yield* context.getRunningJob(payload.jobId); if (!running) { logger.warn(`Backup ${payload.jobId} is not running`); return; diff --git a/apps/agent/src/commands/backup-run.ts b/apps/agent/src/commands/backup-run.ts index c77e1647..5342c004 100644 --- a/apps/agent/src/commands/backup-run.ts +++ b/apps/agent/src/commands/backup-run.ts @@ -9,9 +9,9 @@ import type { ControllerCommandContext } from "../context"; export const handleBackupRunCommand = (context: ControllerCommandContext, payload: BackupRunPayload) => Effect.fork( Effect.gen(function* () { - const existing = context.getRunningJob(payload.jobId); + const existing = yield* context.getRunningJob(payload.jobId); if (existing) { - context.offerOutbound( + yield* context.offerOutbound( createAgentMessage("backup.failed", { jobId: payload.jobId, scheduleId: payload.scheduleId, @@ -23,9 +23,19 @@ export const handleBackupRunCommand = (context: ControllerCommandContext, payloa logger.info(`Starting backup ${payload.jobId} for schedule ${payload.scheduleId}`); const abortController = new AbortController(); - context.setRunningJob(payload.jobId, { scheduleId: payload.scheduleId, abortController }); + yield* context.setRunningJob(payload.jobId, { scheduleId: payload.scheduleId, abortController }); - context.offerOutbound( + const sendCancelled = () => { + return context.offerOutbound( + createAgentMessage("backup.cancelled", { + jobId: payload.jobId, + scheduleId: payload.scheduleId, + message: "Backup was cancelled", + }), + ); + }; + + yield* context.offerOutbound( createAgentMessage("backup.started", { jobId: payload.jobId, scheduleId: payload.scheduleId, @@ -49,70 +59,52 @@ export const handleBackupRunCommand = (context: ControllerCommandContext, payloa ...payload.options, signal: abortController.signal, onProgress: (progress) => { - context.offerOutbound( - createAgentMessage("backup.progress", { - jobId: payload.jobId, - scheduleId: payload.scheduleId, - progress, - }), - ); + void Effect.runPromise( + context.offerOutbound( + createAgentMessage("backup.progress", { + jobId: payload.jobId, + scheduleId: payload.scheduleId, + progress, + }), + ), + ).catch((error) => { + logger.error(`Failed to send backup progress update: ${toMessage(error)}`); + }); }, }) .pipe( Effect.matchEffect({ onSuccess: (result) => { - return Effect.sync(() => { - if (abortController.signal.aborted) { - context.offerOutbound( - createAgentMessage("backup.cancelled", { - jobId: payload.jobId, - scheduleId: payload.scheduleId, - message: "Backup was cancelled", - }), - ); - return; - } + if (abortController.signal.aborted) { + return sendCancelled(); + } - context.offerOutbound( - createAgentMessage("backup.completed", { - jobId: payload.jobId, - scheduleId: payload.scheduleId, - exitCode: result.exitCode, - result: result.result, - warningDetails: result.warningDetails ?? undefined, - }), - ); - }); + return context.offerOutbound( + createAgentMessage("backup.completed", { + jobId: payload.jobId, + scheduleId: payload.scheduleId, + exitCode: result.exitCode, + result: result.result, + warningDetails: result.warningDetails ?? undefined, + }), + ); }, onFailure: (error) => { - return Effect.sync(() => { - if (abortController.signal.aborted) { - context.offerOutbound( - createAgentMessage("backup.cancelled", { - jobId: payload.jobId, - scheduleId: payload.scheduleId, - message: "Backup was cancelled", - }), - ); - return; - } + if (abortController.signal.aborted) { + return sendCancelled(); + } - context.offerOutbound( - createAgentMessage("backup.failed", { - jobId: payload.jobId, - scheduleId: payload.scheduleId, - error: toMessage(error), - errorDetails: toErrorDetails(error), - }), - ); - }); + return context.offerOutbound( + createAgentMessage("backup.failed", { + jobId: payload.jobId, + scheduleId: payload.scheduleId, + error: toMessage(error), + errorDetails: toErrorDetails(error), + }), + ); }, }), - Effect.ensuring( - Effect.sync(() => { - context.deleteRunningJob(payload.jobId); - }), - ), + Effect.ensuring(context.deleteRunningJob(payload.jobId)), ); }), ).pipe(Effect.asVoid); diff --git a/apps/agent/src/commands/heartbeat-ping.ts b/apps/agent/src/commands/heartbeat-ping.ts index 1e7d34f8..4587f549 100644 --- a/apps/agent/src/commands/heartbeat-ping.ts +++ b/apps/agent/src/commands/heartbeat-ping.ts @@ -5,8 +5,8 @@ import type { ControllerCommandContext } from "../context"; type HeartbeatPingPayload = Extract["payload"]; export const handleHeartbeatPingCommand = (context: ControllerCommandContext, payload: HeartbeatPingPayload) => - Effect.sync(() => { - context.offerOutbound( + Effect.gen(function* () { + yield* context.offerOutbound( createAgentMessage("heartbeat.pong", { sentAt: payload.sentAt, }), diff --git a/apps/agent/src/context.ts b/apps/agent/src/context.ts index 87b76506..4daa68bc 100644 --- a/apps/agent/src/context.ts +++ b/apps/agent/src/context.ts @@ -1,4 +1,5 @@ import type { AgentWireMessage } from "@zerobyte/contracts/agent-protocol"; +import type { Effect } from "effect"; export type RunningJob = { scheduleId: string; @@ -6,8 +7,8 @@ export type RunningJob = { }; export type ControllerCommandContext = { - getRunningJob: (jobId: string) => RunningJob | undefined; - setRunningJob: (jobId: string, job: RunningJob) => void; - deleteRunningJob: (jobId: string) => void; - offerOutbound: (message: AgentWireMessage) => void; + getRunningJob: (jobId: string) => Effect.Effect; + setRunningJob: (jobId: string, job: RunningJob) => Effect.Effect; + deleteRunningJob: (jobId: string) => Effect.Effect; + offerOutbound: (message: AgentWireMessage) => Effect.Effect; }; diff --git a/apps/agent/src/controller-session.ts b/apps/agent/src/controller-session.ts index b2e43163..c1b45b9d 100644 --- a/apps/agent/src/controller-session.ts +++ b/apps/agent/src/controller-session.ts @@ -8,6 +8,7 @@ import { import { logger } from "@zerobyte/core/node"; import { toMessage } from "@zerobyte/core/utils"; import { handleControllerCommand } from "./commands"; +import type { ControllerCommandContext, RunningJob } from "./context"; export type ControllerSession = { onOpen: () => void; @@ -18,45 +19,44 @@ export type ControllerSession = { export const createControllerSession = (ws: WebSocket): ControllerSession => { const outboundQueue = Effect.runSync(Queue.bounded(64)); const inboundQueue = Effect.runSync(Queue.bounded(64)); - const runningJobsRef = Effect.runSync( - Ref.make>(new Map()), - ); + const runningJobsRef = Effect.runSync(Ref.make>(new Map())); - const getRunningJob = (jobId: string) => Effect.runSync(Ref.get(runningJobsRef)).get(jobId); + const getRunningJob = (jobId: string) => Ref.get(runningJobsRef).pipe(Effect.map((map) => map.get(jobId))); - const setRunningJob = (jobId: string, job: { scheduleId: string; abortController: AbortController }) => { - Effect.runSync( - Ref.update(runningJobsRef, (current) => { - const next = new Map(current); - next.set(jobId, job); - return next; - }), - ); + const setRunningJob = (jobId: string, job: RunningJob) => { + return Ref.update(runningJobsRef, (current) => { + const next = new Map(current); + next.set(jobId, job); + return next; + }); }; + const abortRunningJobs = Effect.gen(function* () { + const runningJobs = yield* Ref.modify(runningJobsRef, (current) => [current, new Map()]); + yield* Effect.sync(() => { + for (const runningJob of runningJobs.values()) { + runningJob.abortController.abort(); + } + }); + }); + const deleteRunningJob = (jobId: string) => { - Effect.runSync( - Ref.update(runningJobsRef, (current) => { - const next = new Map(current); - next.delete(jobId); - return next; - }), - ); + return Ref.update(runningJobsRef, (current) => { + const next = new Map(current); + next.delete(jobId); + return next; + }); }; const offerOutbound = (message: AgentWireMessage) => { - void Effect.runPromise(Queue.offer(outboundQueue, message)).catch((error) => { - logger.error(`Failed to queue outbound controller message: ${toMessage(error)}`); - }); + return Queue.offer(outboundQueue, message); }; const offerInbound = (message: ControllerWireMessage) => { - void Effect.runPromise(Queue.offer(inboundQueue, message)).catch((error) => { - logger.error(`Failed to queue inbound controller message: ${toMessage(error)}`); - }); + return Queue.offer(inboundQueue, message); }; - const commandContext = { + const commandContext: ControllerCommandContext = { getRunningJob, setRunningJob, deleteRunningJob, @@ -101,7 +101,9 @@ export const createControllerSession = (ws: WebSocket): ControllerSession => { return { onOpen: () => { - offerOutbound(createAgentMessage("agent.ready", { agentId: "" })); + void Effect.runPromise(offerOutbound(createAgentMessage("agent.ready", { agentId: "" }))).catch((error) => { + logger.error(`Failed to queue ready message: ${toMessage(error)}`); + }); }, onMessage: (data) => { if (typeof data !== "string") { @@ -109,14 +111,12 @@ export const createControllerSession = (ws: WebSocket): ControllerSession => { return; } - offerInbound(data as ControllerWireMessage); + void Effect.runPromise(offerInbound(data as ControllerWireMessage)).catch((error) => { + logger.error(`Failed to queue inbound message: ${toMessage(error)}`); + }); }, close: () => { - const runningJobs = Effect.runSync(Ref.get(runningJobsRef)); - for (const running of runningJobs.values()) { - running.abortController.abort(); - } - Effect.runSync(Ref.set(runningJobsRef, new Map())); + void Effect.runPromise(abortRunningJobs).catch(() => {}); void Effect.runPromise(Fiber.interrupt(writerFiber)).catch(() => {}); void Effect.runPromise(Fiber.interrupt(processorFiber)).catch(() => {}); void Effect.runPromise(Queue.shutdown(outboundQueue)).catch(() => {});