From 67b78b845462ed213615a84174a67bd5748a5fd3 Mon Sep 17 00:00:00 2001 From: Coolfan Date: Sat, 8 Aug 2026 16:22:30 +0800 Subject: [PATCH 1/2] docs: plan console session recovery --- README/CONSOLE_SESSION_RECOVERY_PLAN.zh-CN.md | 395 ++++++++++++++++++ 1 file changed, 395 insertions(+) create mode 100644 README/CONSOLE_SESSION_RECOVERY_PLAN.zh-CN.md diff --git a/README/CONSOLE_SESSION_RECOVERY_PLAN.zh-CN.md b/README/CONSOLE_SESSION_RECOVERY_PLAN.zh-CN.md new file mode 100644 index 0000000..a5ac340 --- /dev/null +++ b/README/CONSOLE_SESSION_RECOVERY_PLAN.zh-CN.md @@ -0,0 +1,395 @@ +# Console 重启后的终端会话恢复实施计划 + +## 1. 目标 + +Console 正常重启或进程崩溃后,仍能恢复尚未过期的 terminal session route,使 Worker 重连时继续通过现有恢复握手核对后端资源,调用方随后可以使用原有 `session_id` 和沙箱文件系统。 + +本阶段建立以下完整链路: + +```text +terminalExec 成功 + -> Console 持久化 scoped session route 与绝对 lease + -> Console 重启并加载有效 route + -> route 以 unavailable 状态进入内存 + -> 原 Worker 重连并收到 recovery candidates + -> Worker 核对后端资源并上报结果 + -> Console 确认 route 或持久化删除失效 route + -> session 恢复可用 +``` + +范围: + +- 持久化已确认 terminal session 的 Worker 归属和绝对 lease。 +- 支持 Console 正常退出、进程崩溃和主机重启后的恢复。 +- 复用现有 Worker recovery candidate、report 和 ack 协议。 +- 同等支持 `worker-docker`、`worker-boxlite` 和 `worker-bridge-e2b`。 +- 保持 owner-scoped session ID 的隔离语义。 +- 恢复后继续支持 terminalExec、terminalResource 和新建 public preview route。 + +不包含: + +- 不恢复 Console 重启时正在执行的命令、输出流或非终态 task;这些 task 继续以 `console_restarted` 失败。 +- 不持久化 provisional route、reservation ID 或 provisional uses。 +- 不恢复已存在的 public preview URL。Preview route 仍是独立的内存状态;Console 重启后可基于恢复的 terminal session 重新创建。 +- 不改变 Worker 恢复协议、后端资源标记或 Worker session manager。 +- 不支持多 Console 实例共享同一 SQLite 数据库或同时调度同一 Worker。 + +## 2. 当前基础 + +当前分支已经实现 Worker 重启恢复: + +- `terminalSessionRoute` 在内存中保存 `NodeID`、`LeaseExpiresUnixMs` 和恢复状态。 +- Worker 断连时,已确认 route 从 `ready` 进入 `unavailable`,不会立即删除。 +- Worker 重连时,Console 发送 recovery candidates;Worker 返回 `RECOVERED`、`MISSING` 或 `INVALID`。 +- Worker 完成核对前不可调度;恢复失败或 lease 过期的 route 会被删除。 +- terminal session 使用 owner-scoped session ID,Worker 不接触外部账号身份。 + +当前缺口是 `terminalSessionToNode` 和 `terminalNodeToSessionIDIndex` 仅存在于进程内存。Console 重启后,新进程无法生成任何 recovery candidate,即使 Worker 后端资源和 lease 仍然有效。 + +Console 已使用 SQLite、Goose 和 sqlc,并在启动时完成 migration、将非终态 task 标记为失败、清空旧 Worker connection session。Terminal route 持久化应接入同一数据库和启动顺序。 + +## 3. 状态归属 + +### 3.1 持久化稳定事实 + +数据库只保存跨 Console 进程仍成立的事实: + +| 字段 | 含义 | +| --- | --- | +| `scoped_session_id` | owner-scoped terminal session ID,唯一主键 | +| `node_id` | session 固定归属的 Worker node ID | +| `lease_expires_unix_ms` | Worker 已确认的绝对 lease 到期时间 | +| `last_used_unix_ms` | 最近一次确认或 lease 更新的时间,用于审计和诊断 | +| `created_at_unix_ms` | route 首次持久化时间 | +| `updated_at_unix_ms` | route 最近一次持久化更新时间 | + +### 3.2 仅保留在内存的瞬态状态 + +以下状态不得持久化: + +- `RecoveryState`:`ready`、`unavailable` 和 `reconciling` 只描述当前 Console 进程与 Worker connection 的关系。每次 Console 启动时,所有加载的 route 都必须从 `unavailable` 开始。 +- `ReservationID` 与 `ProvisionalUses`:它们只保护当前进程中的首次并发 dispatch。Console 重启会中断对应 task,不能把不确定的创建结果恢复为已确认 route。 +- active Worker connection session ID、inflight capability 数和 recovery report 的临时幂等缓存。 + +### 3.3 权威边界 + +- SQLite 是已确认 route 和 lease 的持久化权威来源。 +- `terminalSessionToNode` 是运行时索引,不得包含数据库尚未确认的 durable route。 +- Worker 后端是沙箱及文件系统的权威来源。 +- Worker recovery report 决定持久化 route 对应的后端资源是否仍然存在且有效。 +- Console 不持久化 Docker container ID、Boxlite Box ID、E2B sandbox ID、envd token 或其他后端凭据。 + +## 4. 数据库设计 + +新增 migration: + +```sql +CREATE TABLE terminal_session_routes ( + scoped_session_id TEXT PRIMARY KEY, + node_id TEXT NOT NULL, + lease_expires_unix_ms INTEGER NOT NULL CHECK (lease_expires_unix_ms > 0), + last_used_unix_ms INTEGER NOT NULL, + created_at_unix_ms INTEGER NOT NULL, + updated_at_unix_ms INTEGER NOT NULL +); + +CREATE INDEX idx_terminal_session_routes_node + ON terminal_session_routes(node_id); + +CREATE INDEX idx_terminal_session_routes_lease + ON terminal_session_routes(lease_expires_unix_ms); +``` + +`node_id` 不声明指向 `worker_nodes` 的外键,原因如下: + +- runtime Worker 记录可能被 offline pruner 删除,但有效 terminal lease 必须继续等待同一 node ID 重连。 +- Console 启动时会清空 `worker_nodes.session_id`,这不应影响 terminal route。 +- 删除 Worker 凭据与删除 terminal route 是不同的业务动作,需要在 Worker 删除流程中显式处理,而不是依赖 registry 表的级联副作用。 + +新增 sqlc queries: + +- `UpsertTerminalSessionRoute` +- `DeleteTerminalSessionRouteBySessionAndNode` +- `DeleteTerminalSessionRoutesByNode` +- `DeleteExpiredTerminalSessionRoutes` +- `ListActiveTerminalSessionRoutes` +- 测试和诊断需要的精确查询或计数查询 + +`UpsertTerminalSessionRoute` 必须满足: + +- 同一 `scoped_session_id` 只能绑定一个 `node_id`。 +- 已确认的同 node lease 可以更新,但不得因迟到结果缩短现有 lease。 +- 不允许迟到的旧 Worker 结果覆盖已经属于其他 node 的 route。 +- `created_at_unix_ms` 在更新时保持不变。 + +## 5. 持久化接口 + +在 gRPC registry 与 SQLite 之间增加窄接口,避免 route 状态机直接依赖 sqlc 类型: + +```text +LoadActive(ctx, now_unix_ms) -> []PersistedTerminalSessionRoute +UpsertConfirmed(ctx, route) -> persisted | node_conflict +Delete(ctx, scoped_session_id, expected_node_id) -> deleted | not_owned +DeleteByNode(ctx, node_id) +DeleteExpired(ctx, now_unix_ms) -> deleted_count +``` + +实现要求: + +- 生产实现使用 `persistence.DB` 和 sqlc。 +- 单元测试使用内存 fake,能够注入写入、删除和加载失败。 +- RegistryService 初始化时显式注入该接口,不允许生产环境静默退化为纯内存 route。 +- 只有不需要重启恢复的局部测试可以使用 no-op store。 + +## 6. 写入一致性 + +### 6.1 已确认 route 的唯一提交点 + +合并当前分离的 route 确认与 lease 更新路径,形成单一 durable commit 操作: + +```text +commitConfirmedTerminalRoute( + scoped_session_id, + expected_node_id, + reservation_id, + lease_expires_unix_ms, + now +) +``` + +执行顺序: + +1. 在 route 锁下验证 reservation、node 归属和 ABA 条件。 +2. 要求 Worker 成功结果包含正数 `lease_expires_unix_ms`。 +3. 使用独立、有限时长的数据库 context 持久化 route,不复用可能刚好到期的请求 context。 +4. 数据库成功后,才把内存 route 确认为 durable、清除 reservation 并更新 lease。 +5. route durable commit 成功后,才允许成功结果返回调用方并进入 task 成功状态。 + +数据库写入失败时: + +- 不把 provisional route 暴露为已确认 route。 +- 当前命令返回内部错误;Worker 已创建但未确认的资源由远端 lease 或后续明确清理回收。 +- 记录不包含原始命令、文件内容或凭据的结构化错误。 + +### 6.2 现有 session 的 lease 延长 + +复用现有 session 的命令成功后,Worker 可能已经延长后端 lease。Console 必须先持久化不缩短的绝对 lease,再更新内存并返回成功。 + +若数据库失败: + +- 不得只更新内存 lease。 +- route 保持原 durable lease;请求返回内部错误。 +- 后端 lease 可能更长,但不会比 Console 记录更早删除资源,最多形成由后端 timeout 最终回收的孤儿,不会错误恢复过期 session。 + +### 6.3 删除 + +以下事件必须同步删除持久化 route: + +- Worker 返回 `session_not_found`。 +- recovery report 返回 `MISSING` 或 `INVALID`。 +- recovery 时发现 lease 已过期。 +- route janitor 清理到期 route。 +- 管理端明确删除对应 Worker。 + +删除使用 `scoped_session_id + expected_node_id` 条件,防止迟到事件删除 ABA 后的新归属。 + +Recovery report 的删除必须在发送 recovery ack 前完成。数据库失败时关闭本次 Worker connection,不发送 ack;Worker 重连后重新核对,不能以部分持久化状态开始接单。 + +### 6.4 不持久化 provisional route + +首次 dispatch 前的 route reservation 保持纯内存: + +- Console 重启时对应 task 已被标记为 `console_restarted`。 +- 未收到成功结果前,Console 无法确认 Worker 是否创建了资源。 +- 把 provisional route 写入数据库会把不确定状态错误升级为可恢复 session。 + +只有收到合法 terminalExec 成功结果并取得绝对 lease 后,route 才进入 `terminal_session_routes`。 + +## 7. Console 启动恢复 + +启动顺序调整为: + +1. 打开 SQLite 并执行 Goose migrations。 +2. 在启动事务中保持现有 task 和 Worker connection 清理: + - 非终态 task 标记为 `console_restarted`。 + - `worker_nodes.session_id` 清空。 +3. 删除 `lease_expires_unix_ms <= now` 的持久化 terminal route。 +4. 读取剩余 route,并校验 session ID、node ID 和绝对 lease。 +5. 创建 RegistryService,把持久化 route 装载到两个内存索引: + - `terminalSessionToNode` + - `terminalNodeToSessionIDIndex` +6. 所有装载 route 的内存状态固定为: + - `RecoveryState=unavailable` + - `ReservationID=0` + - `ProvisionalUses=0` +7. 完成装载后才启动 gRPC 和 HTTP listener。 +8. Worker 重连后,复用现有 `beginTerminalSessionRecovery` 和 `applyTerminalSessionRecoveryReport` 完成资源核对。 + +启动采用 fail-closed: + +- migration、过期清理或 route 加载失败时,Console 启动失败。 +- 不允许记录警告后以空 route 集合启动,否则会静默丢失恢复能力并允许同名 session 被重新分配。 +- 单条不合法记录视为数据库完整性错误,不跳过后继续启动。 + +## 8. 并发与崩溃语义 + +### 8.1 锁和数据库顺序 + +Route mutation 必须遵循统一顺序: + +```text +terminalRoutesMu -> SQLite write -> in-memory mutation -> unlock +``` + +Console 当前是单实例、SQLite 单 writer 模型。持锁执行短 SQL 写入可保持现有 ABA 与 reservation 判断在数据库提交期间不失效,也避免引入第二套 revision 状态机。 + +约束: + +- route SQL 不得调用网络服务或 Worker。 +- 数据库操作必须有固定超时。 +- 不得在持有 route 锁时等待 recovery report、command result 或 proxy 请求。 +- 批量 recovery report 使用单个数据库事务。 + +### 8.2 崩溃点 + +| 崩溃点 | 重启后的结果 | +| --- | --- | +| Worker 成功前 | provisional route 未持久化,不恢复 | +| Worker 成功后、数据库提交前 | route 不恢复;后端资源按 lease 最终回收 | +| 数据库提交后、返回调用方前 | route 会恢复;调用方可用原 request/task 查询结果或重试原 session | +| route 删除提交前 | 重启后再次向 Worker核对,结果最终收敛 | +| route 删除提交后、内存删除前 | 重启后 route 不再出现 | +| recovery report 事务中途 | 事务回滚且不发送 ack,Worker 重连重试 | + +系统目标是 session route 的 at-least-once reconciliation,不提供中断命令的 exactly-once 执行。 + +## 9. Worker 生命周期与清理 + +### 9.1 Worker 重连 + +Console 重启后,持久化 route 不依赖旧 Worker connection session。相同 `node_id` 的 Worker 建立新连接时会收到所有未过期 candidate,恢复协议不需要新增字段。 + +### 9.2 Worker 删除 + +管理端删除 Worker 时: + +1. 关闭当前 connection。 +2. 删除该 node 的持久化 terminal route。 +3. 删除内存 route 和 node 索引。 +4. 删除 Worker registry、credential 和 owner claim。 + +删除 Worker 后不再允许其使用旧凭据重连。后端遗留资源由 Docker/Boxlite 的受限 orphan 清理或 E2B timeout 回收。 + +Offline pruner 删除非 provisioned Worker registry 行时,不删除 terminal route。有效 route 保持到 lease 到期,以允许同一 node ID 在凭据仍有效时重连。 + +## 10. Owner 隔离与安全 + +- 数据库主键保存完整 owner-scoped session ID,不另存一份外部 session ID。 +- 所有 API 输入仍先通过 `scopeTerminalSessionID`,输出仍通过 `unscopeTerminalSessionID`。 +- route 查询、更新和删除不得接受未 scoped 的外部 session ID。 +- 日志不记录 scoped session ID,因为其中包含 owner 标识;使用 session hash、node ID、数量和状态。 +- 不新增后端凭据、Worker secret 或 proxy traffic token 的持久化。 +- 数据库文件沿用 Console 现有访问控制和备份策略。 + +## 11. Public preview 边界 + +Terminal route 恢复后: + +- 用户可以对原 session 创建新的 public preview route。 +- 已有 preview route key 仍不会跨 Console 重启恢复。 +- Preview 解析必须继续要求 terminal route 已完成 Worker reconciliation,不能仅因数据库存在 route 就转发流量。 + +持久化 preview URL 涉及匿名访问凭据生命周期、撤销、账号删除和独立过期索引,应在单独 PR 中设计,不能与 terminal route 恢复隐式绑定。 + +## 12. 实施阶段 + +### 阶段 A:Schema 与持久化层 + +- 新增 `terminal_session_routes` migration。 +- 新增 sqlc queries 并重新生成代码。 +- 实现 persistence adapter 和 fake。 +- 覆盖 upsert 不缩短 lease、node 冲突、条件删除、批量过期删除和事务回滚。 + +### 阶段 B:RegistryService durable mutation + +- 注入 terminal route store。 +- 合并 confirm 与 lease update 为 durable commit。 +- 把 session-not-found、recovery failure、janitor 和 Worker 删除接入持久化删除。 +- 为数据库失败定义统一内部错误和结构化日志。 + +### 阶段 C:启动恢复 + +- 在 listener 启动前加载有效 route。 +- 重建两个内存索引并统一标记为 unavailable。 +- 保持 task startup recovery 和 Worker session 清理行为。 +- 增加启动失败和非法数据测试。 + +### 阶段 D:重启故障矩阵 + +- 使用临时 SQLite 文件创建第一个 Console service,确认 route 后关闭。 +- 使用同一数据库创建第二个 service,模拟 Worker 重连和 recovery report。 +- Docker、Boxlite 和 E2B 分别执行真实后端重启场景。 +- 覆盖数据库写失败、事务回滚和各个崩溃边界。 + +## 13. 测试矩阵 + +### 13.1 Persistence + +- migration 可从当前 schema 正常升级和回滚。 +- confirmed route 写入后可重新加载。 +- upsert 不缩短 lease。 +- 不同 node 不能覆盖已有 route。 +- 条件删除不删除其他 node 的 route。 +- 过期清理只删除到期 route。 +- malformed route 使启动失败。 + +### 13.2 Registry 状态机 + +- provisional route 永不持久化。 +- terminalExec 成功必须在 durable commit 后返回。 +- 复用 session 的 lease 延长同步持久化。 +- 持久化失败不会产生仅内存 confirmed route。 +- `session_not_found` 同时删除数据库和内存 route。 +- recovery `MISSING`、`INVALID` 和过期结果同步删除。 +- recovery 批量写失败不发送 ack,且不应用部分内存结果。 +- stale reservation、迟到 command result 和迟到 recovery report 不能覆盖或删除新 route。 + +### 13.3 Console 重启 + +- 正常重启后有效 route 被加载为 unavailable。 +- 崩溃重启后行为与正常重启一致。 +- Worker 重连收到准确、稳定排序的 candidate 和原绝对 lease。 +- Worker 报告 `RECOVERED` 后原 session 可继续执行且 `created=false`。 +- Console 离线期间到期的 lease 不进入 candidate。 +- Worker 尚未重连时,terminalExec、terminalResource 和 preview 解析返回 unavailable,不发生重新分配。 +- 重启时的 provisional route 不恢复。 +- 非终态 task 仍以 `console_restarted` 失败。 + +### 13.4 后端验收 + +Docker、Boxlite 和 E2B 分别验证: + +1. 创建 session 并写入文件。 +2. 停止 Console,但保持 Worker 后端资源。 +3. 重启 Console。 +4. Worker 自动重连并完成 recovery handshake。 +5. 使用原 external `session_id` 读取文件。 +6. 确认 lease 未重置、session 未被重新创建、容量计数正确。 + +## 14. 验收标准 + +- Console 重启后,所有未过期 confirmed route 都能进入 Worker recovery candidate。 +- Worker 恢复成功后,原 session ID、文件系统、Worker 归属和绝对 lease 保持不变。 +- Console 离线期间到期的 route 不恢复,后端资源最终清理。 +- Worker 未重连或正在核对时返回 `session_unavailable`,不改派到其他 Worker。 +- 数据库故障不会产生仅内存成功、部分 recovery ack 或静默空状态启动。 +- Provisional route 和中断命令不会被误恢复。 +- Owner 隔离在持久化、恢复、命令和资源访问路径中保持不变。 +- 三个 Worker 后端通过同一 Console 重启故障矩阵,无需修改现有 recovery proto。 +- 现有 Worker 重启恢复、capacity routing、terminalResource 和 public preview 测试继续通过。 + +## 15. 发布说明 + +首次部署本功能时,旧 Console 进程中的 route 尚未写入新表,因此部署前已经存在的 terminal session 无法跨这次升级重启恢复。功能生效后新建或成功续租的 confirmed route 才具备 Console 重启恢复能力。 + +发布时应明确这一单次边界;不通过猜测 Worker 资源或账号级扫描来补建旧 route。 From 76b568d303469d290892d2aa0e4da0987cdc82f6 Mon Sep 17 00:00:00 2001 From: Coolfan Date: Sat, 8 Aug 2026 18:47:24 +0800 Subject: [PATCH 2/2] feat(console): persist terminal session routes for restart recovery MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Persist confirmed terminal session routes to SQLite so that Console restarts no longer lose Worker→session ownership. On startup, active routes are loaded as unavailable and offered as recovery candidates when the original Worker reconnects. Expired routes are pruned both at startup and lazily during dispatch. Key changes: - Add terminal_session_routes table with scoped_session_id primary key, node_id, and absolute lease_expires_unix_ms. - Merge route confirmation and lease update into a single durable commitConfirmedTerminalSessionRoute that persists before updating in-memory state. - Wire session_not_found, recovery report, Worker deletion, and route janitor deletions through the persistent store with conditional scoped_session_id + node_id matching to prevent ABA overwrites. - DeleteProvisionedWorker now returns an error when persistence fails. - Add comprehensive tests for persistence faults, ABA scenarios, and full Console restart recovery round-trips. --- console/cmd/console/main.go | 6 + .../00006_terminal_session_routes.sql | 20 + .../db/queries/terminal_session_routes.sql | 50 ++ .../grpcserver/connect_service_test.go | 6 +- console/internal/grpcserver/service.go | 3 + .../internal/grpcserver/service_connect.go | 5 +- .../grpcserver/service_credentials_runtime.go | 11 +- .../internal/grpcserver/service_dispatch.go | 42 +- .../terminal_capacity_dispatch_test.go | 10 +- .../grpcserver/terminal_session_routes.go | 273 +++++++- .../terminal_session_routes_test.go | 651 ++++++++++++++++++ console/internal/httpapi/worker_handler.go | 9 +- .../internal/httpapi/worker_handler_test.go | 8 +- console/internal/persistence/db.go | 3 + console/internal/persistence/sqlc/models.go | 9 + .../sqlc/terminal_session_routes.sql.go | 164 +++++ .../registry/store_terminal_session_routes.go | 119 ++++ console/sqlc.yaml | 2 + 18 files changed, 1334 insertions(+), 57 deletions(-) create mode 100644 console/db/migrations/00006_terminal_session_routes.sql create mode 100644 console/db/queries/terminal_session_routes.sql create mode 100644 console/internal/grpcserver/terminal_session_routes_test.go create mode 100644 console/internal/persistence/sqlc/terminal_session_routes.sql.go create mode 100644 console/internal/registry/store_terminal_session_routes.go diff --git a/console/cmd/console/main.go b/console/cmd/console/main.go index fd5db99..2fb8b97 100644 --- a/console/cmd/console/main.go +++ b/console/cmd/console/main.go @@ -95,6 +95,12 @@ func main() { registryService.SetHasher(db.Hasher) registryService.SetTaskRetention(time.Duration(cfg.TaskRetentionDays) * 24 * time.Hour) registryService.ConfigureProxy(cfg.ProxyEnabled, cfg.ProxyAllowedWorkerCIDRs, cfg.ProxyAllowedWorkerPorts, cfg.ProxyAllowedDirectDomains) + restoreCtx, restoreCancel := context.WithTimeout(context.Background(), 10*time.Second) + if err := registryService.RestoreTerminalSessionRoutes(restoreCtx, time.Now()); err != nil { + restoreCancel() + fatal("failed to restore terminal session routes", "error", err) + } + restoreCancel() grpcSrv := grpcserver.NewServer(registryService) httpHandler := httpapi.NewWorkerHandler( store, diff --git a/console/db/migrations/00006_terminal_session_routes.sql b/console/db/migrations/00006_terminal_session_routes.sql new file mode 100644 index 0000000..908ba93 --- /dev/null +++ b/console/db/migrations/00006_terminal_session_routes.sql @@ -0,0 +1,20 @@ +-- +goose Up +CREATE TABLE terminal_session_routes ( + scoped_session_id TEXT PRIMARY KEY, + node_id TEXT NOT NULL, + lease_expires_unix_ms INTEGER NOT NULL CHECK (lease_expires_unix_ms > 0), + last_used_unix_ms INTEGER NOT NULL, + created_at_unix_ms INTEGER NOT NULL, + updated_at_unix_ms INTEGER NOT NULL +); + +CREATE INDEX idx_terminal_session_routes_node + ON terminal_session_routes(node_id); + +CREATE INDEX idx_terminal_session_routes_lease + ON terminal_session_routes(lease_expires_unix_ms); + +-- +goose Down +DROP INDEX IF EXISTS idx_terminal_session_routes_lease; +DROP INDEX IF EXISTS idx_terminal_session_routes_node; +DROP TABLE IF EXISTS terminal_session_routes; diff --git a/console/db/queries/terminal_session_routes.sql b/console/db/queries/terminal_session_routes.sql new file mode 100644 index 0000000..b6814b8 --- /dev/null +++ b/console/db/queries/terminal_session_routes.sql @@ -0,0 +1,50 @@ +-- name: UpsertTerminalSessionRoute :execrows +INSERT INTO terminal_session_routes ( + scoped_session_id, + node_id, + lease_expires_unix_ms, + last_used_unix_ms, + created_at_unix_ms, + updated_at_unix_ms +) VALUES (?, ?, ?, ?, ?, ?) +ON CONFLICT(scoped_session_id) DO UPDATE SET + lease_expires_unix_ms = MAX(terminal_session_routes.lease_expires_unix_ms, excluded.lease_expires_unix_ms), + last_used_unix_ms = MAX(terminal_session_routes.last_used_unix_ms, excluded.last_used_unix_ms), + updated_at_unix_ms = MAX(terminal_session_routes.updated_at_unix_ms, excluded.updated_at_unix_ms) +WHERE terminal_session_routes.node_id = excluded.node_id; + +-- name: DeleteTerminalSessionRouteBySessionAndNode :execrows +DELETE FROM terminal_session_routes +WHERE scoped_session_id = ? AND node_id = ?; + +-- name: DeleteTerminalSessionRoutesByNode :execrows +DELETE FROM terminal_session_routes +WHERE node_id = ?; + +-- name: DeleteExpiredTerminalSessionRoutes :execrows +DELETE FROM terminal_session_routes +WHERE lease_expires_unix_ms <= ?; + +-- name: ListActiveTerminalSessionRoutes :many +SELECT + scoped_session_id, + node_id, + lease_expires_unix_ms, + last_used_unix_ms, + created_at_unix_ms, + updated_at_unix_ms +FROM terminal_session_routes +WHERE lease_expires_unix_ms > ? +ORDER BY node_id ASC, scoped_session_id ASC; + +-- name: GetTerminalSessionRouteBySession :one +SELECT + scoped_session_id, + node_id, + lease_expires_unix_ms, + last_used_unix_ms, + created_at_unix_ms, + updated_at_unix_ms +FROM terminal_session_routes +WHERE scoped_session_id = ? +LIMIT 1; diff --git a/console/internal/grpcserver/connect_service_test.go b/console/internal/grpcserver/connect_service_test.go index 59f8515..63e738e 100644 --- a/console/internal/grpcserver/connect_service_test.go +++ b/console/internal/grpcserver/connect_service_test.go @@ -214,8 +214,8 @@ func TestDeleteProvisionedWorkerDisconnectsSessionAndRevokesCredential(t *testin t.Fatalf("connect worker failed: %v", err) } - if removed := svc.DeleteProvisionedWorker(workerID); !removed { - t.Fatalf("expected delete to return true") + if removed, err := svc.DeleteProvisionedWorker(workerID); err != nil || !removed { + t.Fatalf("expected delete to return true, removed=%t err=%v", removed, err) } if _, ok := svc.GetWorkerSecret(workerID); ok { t.Fatalf("expected credential to be revoked") @@ -1272,7 +1272,7 @@ func TestDispatchCommandTerminalSessionCapacityDoesNotClearConcurrentProvisional session.resolvePending(®istryv1.CommandResult{ CommandId: secondDispatch.GetCommandId(), - PayloadJson: []byte(`{"session_id":"session-shared"}`), + PayloadJson: []byte(`{"session_id":"session-shared","lease_expires_unix_ms":4102444800000}`), CompletedUnixMs: now.UnixMilli(), }) second := <-secondDone diff --git a/console/internal/grpcserver/service.go b/console/internal/grpcserver/service.go index 2f14701..d8a516b 100644 --- a/console/internal/grpcserver/service.go +++ b/console/internal/grpcserver/service.go @@ -27,6 +27,7 @@ const ( defaultCommandDispatchTimeout = 60 * time.Second defaultTerminalRouteTTL = 30 * time.Minute terminalRoutePruneMinInterval = 1 * time.Minute + terminalRouteStoreTimeout = 5 * time.Second computerUseCapabilityName = "computeruse" computerUseCapabilityDeclared = "computerUse" readImageCapabilityName = "readimage" @@ -69,6 +70,7 @@ type RegistryService struct { terminalNodeToSessionIDIndex map[string]map[string]struct{} terminalRouteReservationSeq uint64 terminalRouteTTL time.Duration + terminalRouteStore terminalSessionRouteStore lastTerminalRoutePruneUnixMs atomic.Int64 tasksMu sync.RWMutex @@ -109,6 +111,7 @@ func NewRegistryService( terminalSessionToNode: make(map[string]terminalSessionRoute), terminalNodeToSessionIDIndex: make(map[string]map[string]struct{}), terminalRouteTTL: defaultTerminalRouteTTL, + terminalRouteStore: store, tasks: make(map[string]*taskRecord), taskRequestReservations: make(map[string]struct{}), criticalPersistenceFailureFn: func(error) {}, diff --git a/console/internal/grpcserver/service_connect.go b/console/internal/grpcserver/service_connect.go index c1d5bcb..52da916 100644 --- a/console/internal/grpcserver/service_connect.go +++ b/console/internal/grpcserver/service_connect.go @@ -71,7 +71,10 @@ func (s *RegistryService) Connect(stream grpc.BidiStreamingServer[registryv1.Con } session := newActiveSessionAt(hello.GetNodeId(), sessionID, hello, now) - recoveryCandidates := s.beginTerminalSessionRecovery(session.nodeID, now) + recoveryCandidates, err := s.beginTerminalSessionRecoveryWithError(session.nodeID, now) + if err != nil { + return status.Errorf(codes.Internal, "prepare terminal session recovery: %v", err) + } session.setRecoveryCandidates(recoveryCandidates) if err := s.configureSessionProxy(session, hello, workerSecret); err != nil { return status.Errorf(codes.InvalidArgument, "invalid proxy endpoint: %v", err) diff --git a/console/internal/grpcserver/service_credentials_runtime.go b/console/internal/grpcserver/service_credentials_runtime.go index 3f70d54..a822ea7 100644 --- a/console/internal/grpcserver/service_credentials_runtime.go +++ b/console/internal/grpcserver/service_credentials_runtime.go @@ -139,21 +139,24 @@ func normalizeProvisioningWorkerType(workerType string) string { } } -func (s *RegistryService) DeleteProvisionedWorker(nodeID string) bool { +func (s *RegistryService) DeleteProvisionedWorker(nodeID string) (bool, error) { trimmedNodeID := strings.TrimSpace(nodeID) if trimmedNodeID == "" { - return false + return false, nil + } + if _, err := s.deleteTerminalSessionRoutesByNode(trimmedNodeID); err != nil { + return false, err } deletedCredentialInMemory := s.deleteCredential(trimmedNodeID) deletedCredentialInDB := s.store.DeleteCredential(trimmedNodeID) deletedNode := s.store.Delete(trimmedNodeID) if !deletedCredentialInMemory && !deletedCredentialInDB && !deletedNode { - return false + return false, nil } s.disconnectWorker(trimmedNodeID, "worker credential revoked") - return true + return true, nil } func (s *RegistryService) getCredential(nodeID string) (string, bool) { diff --git a/console/internal/grpcserver/service_dispatch.go b/console/internal/grpcserver/service_dispatch.go index 3b360ea..6748554 100644 --- a/console/internal/grpcserver/service_dispatch.go +++ b/console/internal/grpcserver/service_dispatch.go @@ -236,7 +236,7 @@ func (s *RegistryService) dispatchCommandAttempt( terminalRouteReservationID, ) } - confirmTerminalRoute := func(resultPayload []byte) { + confirmTerminalRouteInMemory := func() { if terminalSessionID != "" { s.confirmTerminalSessionRoute( terminalSessionID, @@ -244,11 +244,27 @@ func (s *RegistryService) dispatchCommandAttempt( terminalRouteReservationID, s.nowFn(), ) - if leaseExpiresUnixMs := terminalSessionLeaseExpiresUnixMs(resultPayload); leaseExpiresUnixMs > 0 { - s.updateTerminalSessionRouteLease(terminalSessionID, session.nodeID, leaseExpiresUnixMs, s.nowFn()) - } } } + commitTerminalRoute := func(resultPayload []byte) error { + if terminalSessionID == "" { + return nil + } + confirmed, err := s.commitConfirmedTerminalSessionRoute( + terminalSessionID, + session.nodeID, + terminalRouteReservationID, + terminalSessionLeaseExpiresUnixMs(resultPayload), + s.nowFn(), + ) + if err != nil { + return err + } + if !confirmed { + return errors.New("terminal session route changed before persistence") + } + return nil + } commandID, err := s.newCommandIDFn() if err != nil { @@ -297,7 +313,7 @@ func (s *RegistryService) dispatchCommandAttempt( if onDispatched != nil { if err := onDispatched(commandID); err != nil { if terminalRouteReservationID != 0 { - confirmTerminalRoute(nil) + confirmTerminalRouteInMemory() } return dispatchAttemptResult{}, err } @@ -308,7 +324,7 @@ func (s *RegistryService) dispatchCommandAttempt( if terminalRouteReservationID != 0 && terminalSessionID != "" { // The dispatch reached the worker stream, so cancellation does not // prove that session creation failed. - confirmTerminalRoute(nil) + confirmTerminalRouteInMemory() } if errors.Is(commandCtx.Err(), context.DeadlineExceeded) { return dispatchAttemptResult{}, context.DeadlineExceeded @@ -320,7 +336,13 @@ func (s *RegistryService) dispatchCommandAttempt( return dispatchAttemptResult{}, status.Error(codes.Unavailable, "worker session closed before command result") } if outcome.err == nil && terminalSessionID != "" { - confirmTerminalRoute(outcome.payloadJSON) + if capability == taskCapabilityTerminalExec { + if err := commitTerminalRoute(outcome.payloadJSON); err != nil { + return dispatchAttemptResult{}, status.Errorf(codes.Internal, "persist terminal session route: %v", err) + } + } else { + confirmTerminalRouteInMemory() + } return dispatchAttemptResult{outcome: outcome}, nil } if outcome.err == nil || terminalSessionID == "" { @@ -331,8 +353,8 @@ func (s *RegistryService) dispatchCommandAttempt( case isSessionNotFoundCommandError(outcome.err): if terminalRouteReservationID != 0 { rollbackTerminalRouteReservation() - } else { - s.clearTerminalSessionRoute(terminalSessionID, session.nodeID) + } else if err := s.clearTerminalSessionRoute(terminalSessionID, session.nodeID); err != nil { + return dispatchAttemptResult{}, status.Errorf(codes.Internal, "delete terminal session route: %v", err) } case isSessionCapacityCommandError(outcome.err): releaseResult := rollbackTerminalRouteReservation() @@ -344,7 +366,7 @@ func (s *RegistryService) dispatchCommandAttempt( }, nil case terminalRouteReservationID != 0: // Other execution errors do not prove that session creation failed. - confirmTerminalRoute(nil) + confirmTerminalRouteInMemory() } return dispatchAttemptResult{outcome: outcome}, nil } diff --git a/console/internal/grpcserver/terminal_capacity_dispatch_test.go b/console/internal/grpcserver/terminal_capacity_dispatch_test.go index d7065c1..df94672 100644 --- a/console/internal/grpcserver/terminal_capacity_dispatch_test.go +++ b/console/internal/grpcserver/terminal_capacity_dispatch_test.go @@ -270,7 +270,7 @@ func TestDispatchCommandRetriesTerminalCapacityOnAnotherWorker(t *testing.T) { } workerB.resolvePending(®istryv1.CommandResult{ CommandId: dispatchB.GetCommandId(), - PayloadJson: []byte(`{"session_id":"session-retry","stdout":"ok"}`), + PayloadJson: []byte(`{"session_id":"session-retry","stdout":"ok","lease_expires_unix_ms":4102444800000}`), CompletedUnixMs: now.UnixMilli(), }) @@ -371,7 +371,7 @@ func TestSubmitTaskSkipsReportedFullConnectedWorker(t *testing.T) { Payload: ®istryv1.ConnectRequest_CommandResult{ CommandResult: ®istryv1.CommandResult{ CommandId: dispatchB.GetCommandId(), - PayloadJson: []byte(`{"session_id":"obx:owner-a:session-connected-skip","stdout":"ok"}`), + PayloadJson: []byte(`{"session_id":"obx:owner-a:session-connected-skip","stdout":"ok","lease_expires_unix_ms":4102444800000}`), CompletedUnixMs: now.UnixMilli(), }, }, @@ -531,7 +531,7 @@ func TestSubmitTaskRetriesCapacityAcrossConnectedWorkers(t *testing.T) { Payload: ®istryv1.ConnectRequest_CommandResult{ CommandResult: ®istryv1.CommandResult{ CommandId: dispatchB.GetCommandId(), - PayloadJson: []byte(`{"session_id":"obx:owner-a:external-session","stdout":"ok"}`), + PayloadJson: []byte(`{"session_id":"obx:owner-a:external-session","stdout":"ok","lease_expires_unix_ms":4102444800000}`), CompletedUnixMs: now.UnixMilli(), }, }, @@ -618,7 +618,7 @@ func TestTerminalCapacityRetryWorksForAllTaskModes(t *testing.T) { workerB.resolvePending(®istryv1.CommandResult{ CommandId: dispatchB.GetCommandId(), PayloadJson: []byte(`{"session_id":"obx:owner-a:session-mode-` + - string(mode) + `","stdout":"ok"}`), + string(mode) + `","stdout":"ok","lease_expires_unix_ms":4102444800000}`), CompletedUnixMs: now.UnixMilli(), }) @@ -778,7 +778,7 @@ func TestConcurrentProvisionalCapacityOnlyLastRollbackRetries(t *testing.T) { retryDispatch := receiveCommandDispatch(t, workerB) workerB.resolvePending(®istryv1.CommandResult{ CommandId: retryDispatch.GetCommandId(), - PayloadJson: []byte(`{"session_id":"session-shared-retry","stdout":"ok"}`), + PayloadJson: []byte(`{"session_id":"session-shared-retry","stdout":"ok","lease_expires_unix_ms":4102444800000}`), CompletedUnixMs: now.UnixMilli(), }) second := <-secondDone diff --git a/console/internal/grpcserver/terminal_session_routes.go b/console/internal/grpcserver/terminal_session_routes.go index 5d86152..d291057 100644 --- a/console/internal/grpcserver/terminal_session_routes.go +++ b/console/internal/grpcserver/terminal_session_routes.go @@ -1,14 +1,27 @@ package grpcserver import ( + "context" + "errors" + "fmt" "log/slog" "sort" "strings" "time" registryv1 "github.com/onlyboxes/onlyboxes/api/gen/go/registry/v1" + "github.com/onlyboxes/onlyboxes/console/internal/registry" ) +type terminalSessionRouteStore interface { + LoadActiveTerminalSessionRoutes(context.Context, int64) ([]registry.TerminalSessionRoute, error) + UpsertConfirmedTerminalSessionRoute(context.Context, registry.TerminalSessionRoute) error + DeleteTerminalSessionRoute(context.Context, string, string) (bool, error) + DeleteTerminalSessionRoutes(context.Context, []registry.TerminalSessionRouteRef) error + DeleteTerminalSessionRoutesByNode(context.Context, string) (int64, error) + DeleteExpiredTerminalSessionRoutes(context.Context, int64) (int64, error) +} + type routeReservationReleaseResult int const ( @@ -24,8 +37,9 @@ type terminalSessionRoute struct { RecoveryState terminalSessionRecoveryState // ReservationID is non-zero only while the first dispatch is provisional. // A successful result confirms the route by clearing it. - ReservationID uint64 - ProvisionalUses uint64 + ReservationID uint64 + ConfirmedReservationID uint64 + ProvisionalUses uint64 } type terminalSessionRecoveryState uint8 @@ -36,6 +50,56 @@ const ( terminalSessionRecoveryReconciling ) +func (s *RegistryService) RestoreTerminalSessionRoutes(ctx context.Context, now time.Time) error { + if s == nil || s.terminalRouteStore == nil { + return errors.New("terminal session route store is required") + } + nowUnixMs := routeNowUnixMs(now) + if _, err := s.terminalRouteStore.DeleteExpiredTerminalSessionRoutes(ctx, nowUnixMs); err != nil { + return fmt.Errorf("delete expired terminal session routes: %w", err) + } + persisted, err := s.terminalRouteStore.LoadActiveTerminalSessionRoutes(ctx, nowUnixMs) + if err != nil { + return fmt.Errorf("load terminal session routes: %w", err) + } + routes := make(map[string]terminalSessionRoute, len(persisted)) + index := make(map[string]map[string]struct{}) + for _, route := range persisted { + sessionID := strings.TrimSpace(route.ScopedSessionID) + nodeID := strings.TrimSpace(route.NodeID) + if !isValidPersistedTerminalSessionID(sessionID) || nodeID == "" || route.LeaseExpiresUnixMs <= nowUnixMs { + return errors.New("persisted terminal session route is invalid") + } + if _, duplicate := routes[sessionID]; duplicate { + return errors.New("persisted terminal session route is duplicated") + } + routes[sessionID] = terminalSessionRoute{ + NodeID: nodeID, + LastUsedUnixMs: route.LastUsedUnixMs, + LeaseExpiresUnixMs: route.LeaseExpiresUnixMs, + RecoveryState: terminalSessionRecoveryUnavailable, + } + if index[nodeID] == nil { + index[nodeID] = make(map[string]struct{}) + } + index[nodeID][sessionID] = struct{}{} + } + + s.terminalRoutesMu.Lock() + defer s.terminalRoutesMu.Unlock() + if len(s.terminalSessionToNode) != 0 || len(s.terminalNodeToSessionIDIndex) != 0 { + return errors.New("terminal session routes are already initialized") + } + s.terminalSessionToNode = routes + s.terminalNodeToSessionIDIndex = index + return nil +} + +func isValidPersistedTerminalSessionID(sessionID string) bool { + parts := strings.SplitN(strings.TrimSpace(sessionID), taskOwnerScopeSeparator, 3) + return len(parts) == 3 && parts[0] == taskOwnerScopePrefix && strings.TrimSpace(parts[1]) != "" && strings.TrimSpace(parts[2]) != "" +} + func (s *RegistryService) bindTerminalSessionRoute(sessionID string, nodeID string, now time.Time) { if s == nil { return @@ -175,12 +239,13 @@ func (s *RegistryService) confirmTerminalSessionRoute( defer s.terminalRoutesMu.Unlock() route, ok := s.terminalSessionToNode[normalizedSessionID] - if !ok || route.NodeID != normalizedNodeID || route.ReservationID != reservationID { + if !ok || route.NodeID != normalizedNodeID || (route.ReservationID != reservationID && route.ConfirmedReservationID != reservationID) { return false } route.LastUsedUnixMs = nowUnixMs if reservationID != 0 { route.ReservationID = 0 + route.ConfirmedReservationID = reservationID route.ProvisionalUses = 0 } s.terminalSessionToNode[normalizedSessionID] = route @@ -226,13 +291,67 @@ func (s *RegistryService) updateTerminalSessionRouteLease(sessionID string, expe if !ok || route.NodeID != normalizedNodeID || route.ReservationID != 0 { return false } - route.LeaseExpiresUnixMs = leaseExpiresUnixMs + if route.LeaseExpiresUnixMs < leaseExpiresUnixMs { + route.LeaseExpiresUnixMs = leaseExpiresUnixMs + } route.LastUsedUnixMs = routeNowUnixMs(now) route.RecoveryState = terminalSessionRecoveryReady s.terminalSessionToNode[normalizedSessionID] = route return true } +func (s *RegistryService) commitConfirmedTerminalSessionRoute( + sessionID string, + expectedNodeID string, + reservationID uint64, + leaseExpiresUnixMs int64, + now time.Time, +) (bool, error) { + if s == nil || leaseExpiresUnixMs <= 0 { + return false, errors.New("terminal session result is missing a valid lease") + } + normalizedSessionID := strings.TrimSpace(sessionID) + normalizedNodeID := strings.TrimSpace(expectedNodeID) + if normalizedSessionID == "" || normalizedNodeID == "" { + return false, errors.New("terminal session route identity is required") + } + nowUnixMs := routeNowUnixMs(now) + + s.terminalRoutesMu.Lock() + defer s.terminalRoutesMu.Unlock() + route, ok := s.terminalSessionToNode[normalizedSessionID] + if !ok || route.NodeID != normalizedNodeID || (route.ReservationID != reservationID && route.ConfirmedReservationID != reservationID) { + return false, nil + } + if route.LeaseExpiresUnixMs < leaseExpiresUnixMs { + route.LeaseExpiresUnixMs = leaseExpiresUnixMs + } + if s.terminalRouteStore != nil { + ctx, cancel := context.WithTimeout(context.Background(), terminalRouteStoreTimeout) + err := s.terminalRouteStore.UpsertConfirmedTerminalSessionRoute(ctx, registry.TerminalSessionRoute{ + ScopedSessionID: normalizedSessionID, + NodeID: normalizedNodeID, + LeaseExpiresUnixMs: route.LeaseExpiresUnixMs, + LastUsedUnixMs: nowUnixMs, + CreatedAtUnixMs: nowUnixMs, + UpdatedAtUnixMs: nowUnixMs, + }) + cancel() + if err != nil { + return false, fmt.Errorf("persist terminal session route: %w", err) + } + } + route.LastUsedUnixMs = nowUnixMs + route.RecoveryState = terminalSessionRecoveryReady + route.ReservationID = 0 + if reservationID != 0 { + route.ConfirmedReservationID = reservationID + } + route.ProvisionalUses = 0 + s.terminalSessionToNode[normalizedSessionID] = route + return true, nil +} + func (s *RegistryService) terminalSessionRouteSnapshot(sessionID string, now time.Time) (terminalSessionRoute, bool) { if s == nil { return terminalSessionRoute{}, false @@ -249,6 +368,15 @@ func (s *RegistryService) terminalSessionRouteSnapshot(sessionID string, now tim return terminalSessionRoute{}, false } if route.LeaseExpiresUnixMs > 0 && route.LeaseExpiresUnixMs <= nowUnixMs { + if s.terminalRouteStore != nil { + ctx, cancel := context.WithTimeout(context.Background(), terminalRouteStoreTimeout) + _, err := s.terminalRouteStore.DeleteTerminalSessionRoute(ctx, normalizedSessionID, route.NodeID) + cancel() + if err != nil { + slog.Error("failed to delete expired terminal session route", "node_id", route.NodeID, "error", err) + return terminalSessionRoute{}, false + } + } s.deleteTerminalSessionRouteLocked(normalizedSessionID, route) return terminalSessionRoute{}, false } @@ -256,18 +384,24 @@ func (s *RegistryService) terminalSessionRouteSnapshot(sessionID string, now tim } func (s *RegistryService) beginTerminalSessionRecovery(nodeID string, now time.Time) []*registryv1.TerminalSessionRecoveryCandidate { + candidates, _ := s.beginTerminalSessionRecoveryWithError(nodeID, now) + return candidates +} + +func (s *RegistryService) beginTerminalSessionRecoveryWithError(nodeID string, now time.Time) ([]*registryv1.TerminalSessionRecoveryCandidate, error) { if s == nil { - return nil + return nil, nil } normalizedNodeID := strings.TrimSpace(nodeID) if normalizedNodeID == "" { - return nil + return nil, nil } nowUnixMs := routeNowUnixMs(now) s.terminalRoutesMu.Lock() defer s.terminalRoutesMu.Unlock() index := s.terminalNodeToSessionIDIndex[normalizedNodeID] candidates := make([]*registryv1.TerminalSessionRecoveryCandidate, 0, len(index)) + expired := make([]registry.TerminalSessionRouteRef, 0) for sessionID := range index { route, ok := s.terminalSessionToNode[sessionID] if !ok || route.NodeID != normalizedNodeID { @@ -278,19 +412,40 @@ func (s *RegistryService) beginTerminalSessionRecovery(nodeID string, now time.T leaseExpiresUnixMs = route.LastUsedUnixMs + s.terminalRouteTTL.Milliseconds() } if route.ReservationID != 0 || leaseExpiresUnixMs <= nowUnixMs { - s.deleteTerminalSessionRouteLocked(sessionID, route) + if route.LeaseExpiresUnixMs > 0 { + expired = append(expired, registry.TerminalSessionRouteRef{ScopedSessionID: sessionID, NodeID: route.NodeID}) + } continue } - route.LeaseExpiresUnixMs = leaseExpiresUnixMs - route.RecoveryState = terminalSessionRecoveryReconciling - s.terminalSessionToNode[sessionID] = route candidates = append(candidates, ®istryv1.TerminalSessionRecoveryCandidate{ SessionId: sessionID, LeaseExpiresUnixMs: leaseExpiresUnixMs, }) } + if len(expired) > 0 && s.terminalRouteStore != nil { + ctx, cancel := context.WithTimeout(context.Background(), terminalRouteStoreTimeout) + err := s.terminalRouteStore.DeleteTerminalSessionRoutes(ctx, expired) + cancel() + if err != nil { + return nil, fmt.Errorf("delete expired recovery candidates: %w", err) + } + } + for _, ref := range expired { + if route, ok := s.terminalSessionToNode[ref.ScopedSessionID]; ok && route.NodeID == ref.NodeID { + s.deleteTerminalSessionRouteLocked(ref.ScopedSessionID, route) + } + } + for _, candidate := range candidates { + route, ok := s.terminalSessionToNode[candidate.GetSessionId()] + if !ok || route.NodeID != normalizedNodeID { + continue + } + route.LeaseExpiresUnixMs = candidate.GetLeaseExpiresUnixMs() + route.RecoveryState = terminalSessionRecoveryReconciling + s.terminalSessionToNode[candidate.GetSessionId()] = route + } sort.Slice(candidates, func(i, j int) bool { return candidates[i].GetSessionId() < candidates[j].GetSessionId() }) - return candidates + return candidates, nil } func (s *RegistryService) markTerminalSessionRoutesUnavailable(nodeID string) int { @@ -357,13 +512,33 @@ func (s *RegistryService) applyTerminalSessionRecoveryReport(session *activeSess return &terminalRecoveryValidationError{message: "recovery candidate no longer belongs to worker"} } } + deleteRefs := make([]registry.TerminalSessionRouteRef, 0) + for sessionID, leaseExpiresUnixMs := range candidates { + route, ok := s.terminalSessionToNode[sessionID] + if !ok || route.NodeID != session.nodeID { + continue + } + if results[sessionID] != registryv1.TerminalSessionRecoveryResult_RECOVERED || (leaseExpiresUnixMs > 0 && leaseExpiresUnixMs <= nowUnixMs) { + if route.LeaseExpiresUnixMs > 0 { + deleteRefs = append(deleteRefs, registry.TerminalSessionRouteRef{ScopedSessionID: sessionID, NodeID: route.NodeID}) + } + } + } + if len(deleteRefs) > 0 && s.terminalRouteStore != nil { + ctx, cancel := context.WithTimeout(context.Background(), terminalRouteStoreTimeout) + err := s.terminalRouteStore.DeleteTerminalSessionRoutes(ctx, deleteRefs) + cancel() + if err != nil { + s.terminalRoutesMu.Unlock() + return fmt.Errorf("persist terminal session recovery report: %w", err) + } + } for sessionID, leaseExpiresUnixMs := range candidates { route, ok := s.terminalSessionToNode[sessionID] if !ok || route.NodeID != session.nodeID { continue } - status := results[sessionID] - if status != registryv1.TerminalSessionRecoveryResult_RECOVERED || (leaseExpiresUnixMs > 0 && leaseExpiresUnixMs <= nowUnixMs) { + if results[sessionID] != registryv1.TerminalSessionRecoveryResult_RECOVERED || (leaseExpiresUnixMs > 0 && leaseExpiresUnixMs <= nowUnixMs) { s.deleteTerminalSessionRouteLocked(sessionID, route) failures++ continue @@ -429,13 +604,13 @@ func (s *RegistryService) clearTerminalSessionRouteReservation(sessionID string, return routeReservationRemoved } -func (s *RegistryService) clearTerminalSessionRoute(sessionID string, expectedNodeID string) { +func (s *RegistryService) clearTerminalSessionRoute(sessionID string, expectedNodeID string) error { if s == nil { - return + return nil } normalizedSessionID := strings.TrimSpace(sessionID) if normalizedSessionID == "" { - return + return nil } normalizedExpectedNodeID := strings.TrimSpace(expectedNodeID) @@ -444,13 +619,52 @@ func (s *RegistryService) clearTerminalSessionRoute(sessionID string, expectedNo route, ok := s.terminalSessionToNode[normalizedSessionID] if !ok { - return + return nil } if normalizedExpectedNodeID != "" && route.NodeID != normalizedExpectedNodeID { - return + return nil + } + if route.LeaseExpiresUnixMs > 0 && s.terminalRouteStore != nil { + ctx, cancel := context.WithTimeout(context.Background(), terminalRouteStoreTimeout) + _, err := s.terminalRouteStore.DeleteTerminalSessionRoute(ctx, normalizedSessionID, route.NodeID) + cancel() + if err != nil { + return fmt.Errorf("delete terminal session route: %w", err) + } } s.deleteTerminalSessionRouteLocked(normalizedSessionID, route) + return nil +} + +func (s *RegistryService) deleteTerminalSessionRoutesByNode(nodeID string) (int, error) { + if s == nil { + return 0, nil + } + normalizedNodeID := strings.TrimSpace(nodeID) + if normalizedNodeID == "" { + return 0, nil + } + s.terminalRoutesMu.Lock() + defer s.terminalRoutesMu.Unlock() + if s.terminalRouteStore != nil { + ctx, cancel := context.WithTimeout(context.Background(), terminalRouteStoreTimeout) + _, err := s.terminalRouteStore.DeleteTerminalSessionRoutesByNode(ctx, normalizedNodeID) + cancel() + if err != nil { + return 0, fmt.Errorf("delete terminal session routes for worker: %w", err) + } + } + removed := 0 + for sessionID := range s.terminalNodeToSessionIDIndex[normalizedNodeID] { + route, ok := s.terminalSessionToNode[sessionID] + if !ok || route.NodeID != normalizedNodeID { + continue + } + s.deleteTerminalSessionRouteLocked(sessionID, route) + removed++ + } + return removed, nil } func (s *RegistryService) deleteTerminalSessionRouteLocked(sessionID string, route terminalSessionRoute) { @@ -473,26 +687,29 @@ func (s *RegistryService) pruneExpiredTerminalSessionRoutes(now time.Time) int { nowUnixMs := routeNowUnixMs(now) expireBefore := nowUnixMs - ttl.Milliseconds() - removed := 0 s.terminalRoutesMu.Lock() defer s.terminalRoutesMu.Unlock() + persistedDeleteSucceeded := true + if s.terminalRouteStore != nil { + ctx, cancel := context.WithTimeout(context.Background(), terminalRouteStoreTimeout) + _, err := s.terminalRouteStore.DeleteExpiredTerminalSessionRoutes(ctx, nowUnixMs) + cancel() + if err != nil { + persistedDeleteSucceeded = false + slog.Error("failed to prune persisted terminal session routes", "error", err) + } + } + removed := 0 for sessionID, route := range s.terminalSessionToNode { if route.LeaseExpiresUnixMs > 0 { - if route.LeaseExpiresUnixMs > nowUnixMs { + if route.LeaseExpiresUnixMs > nowUnixMs || !persistedDeleteSucceeded { continue } } else if ttl <= 0 || route.LastUsedUnixMs > expireBefore { continue } - delete(s.terminalSessionToNode, sessionID) - index := s.terminalNodeToSessionIDIndex[route.NodeID] - if index != nil { - delete(index, sessionID) - if len(index) == 0 { - delete(s.terminalNodeToSessionIDIndex, route.NodeID) - } - } + s.deleteTerminalSessionRouteLocked(sessionID, route) removed++ } return removed diff --git a/console/internal/grpcserver/terminal_session_routes_test.go b/console/internal/grpcserver/terminal_session_routes_test.go new file mode 100644 index 0000000..01a6341 --- /dev/null +++ b/console/internal/grpcserver/terminal_session_routes_test.go @@ -0,0 +1,651 @@ +package grpcserver + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + registryv1 "github.com/onlyboxes/onlyboxes/api/gen/go/registry/v1" + "github.com/onlyboxes/onlyboxes/console/internal/registry" + "github.com/onlyboxes/onlyboxes/console/internal/testutil/registrytest" +) + +// fakeTerminalSessionRouteStore wraps a real store and allows injecting +// failures for specific operations. It implements terminalSessionRouteStore. +type fakeTerminalSessionRouteStore struct { + mu sync.Mutex + inner *registry.Store + upsertErr error + deleteErr error + deleteByNodeErr error + deleteExpiredErr error + loadErr error + upsertCallCount int + deleteCallCount int + deleteByNodeCount int + deleteExpiredCount int + loadCallCount int + lastUpsertRoute registry.TerminalSessionRoute + lastDeleteRefs []registry.TerminalSessionRouteRef + lastDeleteNodeID string +} + +func (f *fakeTerminalSessionRouteStore) LoadActiveTerminalSessionRoutes(ctx context.Context, nowUnixMs int64) ([]registry.TerminalSessionRoute, error) { + f.mu.Lock() + f.loadCallCount++ + err := f.loadErr + f.mu.Unlock() + if err != nil { + return nil, err + } + return f.inner.LoadActiveTerminalSessionRoutes(ctx, nowUnixMs) +} + +func (f *fakeTerminalSessionRouteStore) UpsertConfirmedTerminalSessionRoute(ctx context.Context, route registry.TerminalSessionRoute) error { + f.mu.Lock() + f.upsertCallCount++ + f.lastUpsertRoute = route + err := f.upsertErr + f.mu.Unlock() + if err != nil { + return err + } + return f.inner.UpsertConfirmedTerminalSessionRoute(ctx, route) +} + +func (f *fakeTerminalSessionRouteStore) DeleteTerminalSessionRoute(ctx context.Context, scopedSessionID string, expectedNodeID string) (bool, error) { + f.mu.Lock() + f.deleteCallCount++ + err := f.deleteErr + f.mu.Unlock() + if err != nil { + return false, err + } + return f.inner.DeleteTerminalSessionRoute(ctx, scopedSessionID, expectedNodeID) +} + +func (f *fakeTerminalSessionRouteStore) DeleteTerminalSessionRoutes(ctx context.Context, routes []registry.TerminalSessionRouteRef) error { + f.mu.Lock() + f.deleteCallCount++ + f.lastDeleteRefs = routes + err := f.deleteErr + f.mu.Unlock() + if err != nil { + return err + } + return f.inner.DeleteTerminalSessionRoutes(ctx, routes) +} + +func (f *fakeTerminalSessionRouteStore) DeleteTerminalSessionRoutesByNode(ctx context.Context, nodeID string) (int64, error) { + f.mu.Lock() + f.deleteByNodeCount++ + f.lastDeleteNodeID = nodeID + err := f.deleteByNodeErr + f.mu.Unlock() + if err != nil { + return 0, err + } + return f.inner.DeleteTerminalSessionRoutesByNode(ctx, nodeID) +} + +func (f *fakeTerminalSessionRouteStore) DeleteExpiredTerminalSessionRoutes(ctx context.Context, nowUnixMs int64) (int64, error) { + f.mu.Lock() + f.deleteExpiredCount++ + err := f.deleteExpiredErr + f.mu.Unlock() + if err != nil { + return 0, err + } + return f.inner.DeleteExpiredTerminalSessionRoutes(ctx, nowUnixMs) +} + +func (f *fakeTerminalSessionRouteStore) setUpsertErr(err error) { + f.mu.Lock() + defer f.mu.Unlock() + f.upsertErr = err +} + +func (f *fakeTerminalSessionRouteStore) setDeleteErr(err error) { + f.mu.Lock() + defer f.mu.Unlock() + f.deleteErr = err +} + +func (f *fakeTerminalSessionRouteStore) setUpsertCallCount() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.upsertCallCount +} + +func (f *fakeTerminalSessionRouteStore) setDeleteCallCount() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.deleteCallCount +} + +func newFakeRouteStore(t testing.TB) *fakeTerminalSessionRouteStore { + t.Helper() + return &fakeTerminalSessionRouteStore{ + inner: registrytest.NewStore(t), + } +} + +// ---- Persistence fault tests ---- + +func TestCommitConfirmedRouteFailsOnStoreError(t *testing.T) { + store := newFakeRouteStore(t) + svc := NewRegistryService(store.inner, nil, 5, 15, time.Minute) + svc.terminalRouteStore = store + base := time.Unix(1_700_700_000, 0) + svc.nowFn = func() time.Time { return base } + + nodeID, reservationID := svc.reserveTerminalSessionRoute("session-fault", "node-a", base) + if reservationID == 0 { + t.Fatal("expected non-zero reservation ID") + } + + store.setUpsertErr(errors.New("disk full")) + confirmed, err := svc.commitConfirmedTerminalSessionRoute( + "session-fault", + nodeID, + reservationID, + base.Add(10*time.Minute).UnixMilli(), + base, + ) + if err == nil { + t.Fatal("expected persistence error from commit") + } + if confirmed { + t.Fatal("route must not be confirmed when persistence fails") + } + + svc.terminalRoutesMu.RLock() + route := svc.terminalSessionToNode["session-fault"] + svc.terminalRoutesMu.RUnlock() + if route.ReservationID == 0 { + t.Fatal("reservation must not be cleared when persistence fails") + } + if route.ConfirmedReservationID != 0 { + t.Fatal("route must not be marked confirmed when persistence fails") + } + if route.LeaseExpiresUnixMs != 0 { + t.Fatal("lease must not be set when persistence fails") + } +} + +func TestRecoveryReportFailsOnStoreError(t *testing.T) { + store := newFakeRouteStore(t) + svc := NewRegistryService(store.inner, nil, 5, 15, time.Minute) + svc.terminalRouteStore = store + base := time.Unix(1_700_710_000, 0) + svc.nowFn = func() time.Time { return base } + + svc.bindTerminalSessionRoute("session-recovery-fault", "node-a", base) + svc.updateTerminalSessionRouteLease("session-recovery-fault", "node-a", base.Add(time.Minute).UnixMilli(), base) + candidates := svc.beginTerminalSessionRecovery("node-a", base) + session := newActiveSessionAt("node-a", "worker-session-a", ®istryv1.ConnectHello{ + Capabilities: []*registryv1.CapabilityDeclaration{{Name: taskCapabilityTerminalExec, MaxInflight: 1}}, + }, base) + session.setRecoveryCandidates(candidates) + + store.setDeleteErr(errors.New("disk full")) + err := svc.applyTerminalSessionRecoveryReport(session, ®istryv1.TerminalSessionRecoveryReport{ + Results: []*registryv1.TerminalSessionRecoveryResult{ + {SessionId: "session-recovery-fault", Status: registryv1.TerminalSessionRecoveryResult_MISSING}, + }, + }, base.Add(time.Second)) + if err == nil { + t.Fatal("expected persistence error from recovery report") + } + + svc.terminalRoutesMu.RLock() + _, stillExists := svc.terminalSessionToNode["session-recovery-fault"] + svc.terminalRoutesMu.RUnlock() + if !stillExists { + t.Fatal("in-memory route must not be deleted when persistence fails") + } +} + +// ---- ABA tests ---- + +func TestLateCommandResultDoesNotConfirmReassignedRoute(t *testing.T) { + svc := NewRegistryService(registrytest.NewStore(t), nil, 5, 15, time.Minute) + base := time.Unix(1_700_720_000, 0) + svc.nowFn = func() time.Time { return base } + + nodeID, reservationID := svc.reserveTerminalSessionRoute("session-aba", "node-a", base) + if reservationID == 0 { + t.Fatal("expected non-zero reservation ID") + } + + // Simulate the route being reassigned to a different worker before the + // late result arrives. + svc.bindTerminalSessionRoute("session-aba", "node-b", base.Add(time.Second)) + + confirmed, err := svc.commitConfirmedTerminalSessionRoute( + "session-aba", + nodeID, + reservationID, + base.Add(10*time.Minute).UnixMilli(), + base.Add(2*time.Second), + ) + if err != nil { + t.Fatalf("commit should not error on ABA mismatch: %v", err) + } + if confirmed { + t.Fatal("late result must not confirm a route reassigned to another worker") + } + + svc.terminalRoutesMu.RLock() + route := svc.terminalSessionToNode["session-aba"] + svc.terminalRoutesMu.RUnlock() + if route.NodeID != "node-b" { + t.Fatalf("route should still belong to node-b, got %s", route.NodeID) + } +} + +func TestStaleReservationClearDoesNotDeleteConfirmedRoute(t *testing.T) { + svc := NewRegistryService(registrytest.NewStore(t), nil, 5, 15, time.Minute) + base := time.Unix(1_700_730_000, 0) + svc.nowFn = func() time.Time { return base } + + nodeID, reservationID := svc.reserveTerminalSessionRoute("session-stale", "node-a", base) + if reservationID == 0 { + t.Fatal("expected non-zero reservation ID") + } + + // Confirm the route via durable commit. + confirmed, err := svc.commitConfirmedTerminalSessionRoute( + "session-stale", + nodeID, + reservationID, + base.Add(10*time.Minute).UnixMilli(), + base, + ) + if err != nil || !confirmed { + t.Fatalf("durable commit failed: confirmed=%v err=%v", confirmed, err) + } + + // A late dispatch sharing the old reservation tries to roll back. + result := svc.clearTerminalSessionRouteReservation("session-stale", nodeID, reservationID) + if result != routeReservationNotOwned { + t.Fatalf("stale reservation clear should be not-owned, got %v", result) + } + + svc.terminalRoutesMu.RLock() + _, ok := svc.terminalSessionToNode["session-stale"] + svc.terminalRoutesMu.RUnlock() + if !ok { + t.Fatal("confirmed route must not be deleted by stale reservation clear") + } +} + +func TestLateRecoveryReportDoesNotDeleteReassignedRoute(t *testing.T) { + svc := NewRegistryService(registrytest.NewStore(t), nil, 5, 15, time.Minute) + base := time.Unix(1_700_740_000, 0) + svc.nowFn = func() time.Time { return base } + + svc.bindTerminalSessionRoute("session-late-recovery", "node-a", base) + svc.updateTerminalSessionRouteLease("session-late-recovery", "node-a", base.Add(time.Minute).UnixMilli(), base) + candidates := svc.beginTerminalSessionRecovery("node-a", base) + session := newActiveSessionAt("node-a", "worker-session-a", ®istryv1.ConnectHello{ + Capabilities: []*registryv1.CapabilityDeclaration{{Name: taskCapabilityTerminalExec, MaxInflight: 1}}, + }, base) + session.setRecoveryCandidates(candidates) + + // Route gets reassigned to node-b while node-a's recovery is in flight. + svc.bindTerminalSessionRoute("session-late-recovery", "node-b", base.Add(time.Second)) + + err := svc.applyTerminalSessionRecoveryReport(session, ®istryv1.TerminalSessionRecoveryReport{ + Results: []*registryv1.TerminalSessionRecoveryResult{ + {SessionId: "session-late-recovery", Status: registryv1.TerminalSessionRecoveryResult_MISSING}, + }, + }, base.Add(2*time.Second)) + if err == nil { + t.Fatal("late recovery report should be rejected for reassigned route") + } + + svc.terminalRoutesMu.RLock() + route := svc.terminalSessionToNode["session-late-recovery"] + svc.terminalRoutesMu.RUnlock() + if route.NodeID != "node-b" { + t.Fatalf("reassigned route should still belong to node-b, got %s", route.NodeID) + } +} + +func TestSessionNotFoundDeletesOnlyMatchingNode(t *testing.T) { + svc := NewRegistryService(registrytest.NewStore(t), nil, 5, 15, time.Minute) + base := time.Unix(1_700_750_000, 0) + svc.nowFn = func() time.Time { return base } + + svc.bindTerminalSessionRoute("session-node-a", "node-a", base) + svc.updateTerminalSessionRouteLease("session-node-a", "node-a", base.Add(time.Minute).UnixMilli(), base) + + // A late session_not_found from a different node must not delete the route. + if err := svc.clearTerminalSessionRoute("session-node-a", "node-b"); err != nil { + t.Fatalf("clear from wrong node should be noop: %v", err) + } + svc.terminalRoutesMu.RLock() + _, ok := svc.terminalSessionToNode["session-node-a"] + svc.terminalRoutesMu.RUnlock() + if !ok { + t.Fatal("route must not be deleted by a different node's session_not_found") + } + + // The correct node's session_not_found should delete it. + if err := svc.clearTerminalSessionRoute("session-node-a", "node-a"); err != nil { + t.Fatalf("clear from correct node failed: %v", err) + } + svc.terminalRoutesMu.RLock() + _, ok = svc.terminalSessionToNode["session-node-a"] + svc.terminalRoutesMu.RUnlock() + if ok { + t.Fatal("route must be deleted by the correct node's session_not_found") + } +} + +// ---- Console restart tests ---- + +func TestRestoreTerminalSessionRoutesLoadsActiveRoutes(t *testing.T) { + store := registrytest.NewStore(t) + base := time.Unix(1_700_800_000, 0) + lease := base.Add(30 * time.Minute).UnixMilli() + + // First service: persist a confirmed route. + svc1 := NewRegistryService(store, nil, 5, 15, time.Minute) + svc1.nowFn = func() time.Time { return base } + ctx := context.Background() + if err := store.UpsertConfirmedTerminalSessionRoute(ctx, registry.TerminalSessionRoute{ + ScopedSessionID: "obx:owner-a:session-restart", + NodeID: "node-a", + LeaseExpiresUnixMs: lease, + LastUsedUnixMs: base.UnixMilli(), + CreatedAtUnixMs: base.UnixMilli(), + UpdatedAtUnixMs: base.UnixMilli(), + }); err != nil { + t.Fatalf("persist route: %v", err) + } + + // Second service: simulate restart. + svc2 := NewRegistryService(store, nil, 5, 15, time.Minute) + svc2.nowFn = func() time.Time { return base.Add(time.Second) } + if err := svc2.RestoreTerminalSessionRoutes(ctx, base.Add(time.Second)); err != nil { + t.Fatalf("restore routes: %v", err) + } + + svc2.terminalRoutesMu.RLock() + route, ok := svc2.terminalSessionToNode["obx:owner-a:session-restart"] + svc2.terminalRoutesMu.RUnlock() + if !ok { + t.Fatal("restored route not found in memory") + } + if route.NodeID != "node-a" { + t.Fatalf("restored route node=%s, want node-a", route.NodeID) + } + if route.RecoveryState != terminalSessionRecoveryUnavailable { + t.Fatalf("restored route recovery state=%v, want unavailable", route.RecoveryState) + } + if route.ReservationID != 0 { + t.Fatal("restored route must not carry a reservation") + } + if route.LeaseExpiresUnixMs != lease { + t.Fatalf("restored route lease=%d, want %d", route.LeaseExpiresUnixMs, lease) + } +} + +func TestRestoreTerminalSessionRoutesDeletesExpiredRoutes(t *testing.T) { + store := registrytest.NewStore(t) + base := time.Unix(1_700_810_000, 0) + ctx := context.Background() + + // Persist an already-expired route. + if err := store.UpsertConfirmedTerminalSessionRoute(ctx, registry.TerminalSessionRoute{ + ScopedSessionID: "obx:owner-a:session-expired", + NodeID: "node-a", + LeaseExpiresUnixMs: base.Add(-time.Minute).UnixMilli(), + LastUsedUnixMs: base.Add(-2 * time.Minute).UnixMilli(), + CreatedAtUnixMs: base.Add(-2 * time.Minute).UnixMilli(), + UpdatedAtUnixMs: base.Add(-2 * time.Minute).UnixMilli(), + }); err != nil { + t.Fatalf("persist expired route: %v", err) + } + + svc := NewRegistryService(store, nil, 5, 15, time.Minute) + svc.nowFn = func() time.Time { return base } + if err := svc.RestoreTerminalSessionRoutes(ctx, base); err != nil { + t.Fatalf("restore routes: %v", err) + } + + svc.terminalRoutesMu.RLock() + _, ok := svc.terminalSessionToNode["obx:owner-a:session-expired"] + svc.terminalRoutesMu.RUnlock() + if ok { + t.Fatal("expired route must not be loaded into memory") + } + + // Verify it was deleted from the database. + routes, err := store.LoadActiveTerminalSessionRoutes(ctx, base.UnixMilli()) + if err != nil { + t.Fatalf("load active routes: %v", err) + } + for _, r := range routes { + if r.ScopedSessionID == "obx:owner-a:session-expired" { + t.Fatal("expired route must be deleted from database") + } + } +} + +func TestRestoreTerminalSessionRoutesRejectsAlreadyInitialized(t *testing.T) { + svc := NewRegistryService(registrytest.NewStore(t), nil, 5, 15, time.Minute) + base := time.Unix(1_700_820_000, 0) + svc.nowFn = func() time.Time { return base } + + // Manually populate the in-memory map to simulate already-initialized state. + svc.terminalRoutesMu.Lock() + svc.terminalSessionToNode["session-existing"] = terminalSessionRoute{NodeID: "node-a"} + svc.terminalRoutesMu.Unlock() + + if err := svc.RestoreTerminalSessionRoutes(context.Background(), base); err == nil { + t.Fatal("restore should fail when routes are already initialized") + } +} + +func TestConsoleRestartRecoveryRoundTrip(t *testing.T) { + store := registrytest.NewStore(t) + base := time.Unix(1_700_830_000, 0) + lease := base.Add(30 * time.Minute).UnixMilli() + ctx := context.Background() + + // First service: persist a confirmed route, then simulate crash. + if err := store.UpsertConfirmedTerminalSessionRoute(ctx, registry.TerminalSessionRoute{ + ScopedSessionID: "obx:owner-a:session-restart-rt", + NodeID: "node-a", + LeaseExpiresUnixMs: lease, + LastUsedUnixMs: base.UnixMilli(), + CreatedAtUnixMs: base.UnixMilli(), + UpdatedAtUnixMs: base.UnixMilli(), + }); err != nil { + t.Fatalf("persist route: %v", err) + } + + // Second service: restart and restore. + svc := NewRegistryService(store, nil, 5, 15, time.Minute) + restartTime := base.Add(5 * time.Second) + svc.nowFn = func() time.Time { return restartTime } + if err := svc.RestoreTerminalSessionRoutes(ctx, restartTime); err != nil { + t.Fatalf("restore routes: %v", err) + } + + // Route should be unavailable until worker reconnects and recovers. + _, _, err := svc.pickSessionForDispatch( + taskCapabilityTerminalExec, + "owner-a", + "obx:owner-a:session-restart-rt", + sessionPickOptions{terminalSessionIntent: terminalSessionIntentKnownNew}, + ) + var commandErr *CommandExecutionError + if !errors.As(err, &commandErr) || commandErr.Code != terminalSessionUnavailableCode { + t.Fatalf("expected session_unavailable after restart, got %v", err) + } + + // Worker reconnects and recovery begins. + candidates := svc.beginTerminalSessionRecovery("node-a", restartTime) + if len(candidates) != 1 || candidates[0].GetSessionId() != "obx:owner-a:session-restart-rt" { + t.Fatalf("unexpected recovery candidates: %#v", candidates) + } + if candidates[0].GetLeaseExpiresUnixMs() != lease { + t.Fatalf("recovery candidate lease=%d, want %d", candidates[0].GetLeaseExpiresUnixMs(), lease) + } + + session := newActiveSessionAt("node-a", "worker-session-a", ®istryv1.ConnectHello{ + Capabilities: []*registryv1.CapabilityDeclaration{{Name: taskCapabilityTerminalExec, MaxInflight: 1}}, + }, restartTime) + session.setRecoveryCandidates(candidates) + + // Worker reports successful recovery. + err = svc.applyTerminalSessionRecoveryReport(session, ®istryv1.TerminalSessionRecoveryReport{ + Results: []*registryv1.TerminalSessionRecoveryResult{ + {SessionId: "obx:owner-a:session-restart-rt", Status: registryv1.TerminalSessionRecoveryResult_RECOVERED}, + }, + }, restartTime.Add(time.Second)) + if err != nil { + t.Fatalf("apply recovery report: %v", err) + } + + // Route should now be ready. + route, ok := svc.terminalSessionRouteSnapshot("obx:owner-a:session-restart-rt", restartTime.Add(time.Second)) + if !ok || route.RecoveryState != terminalSessionRecoveryReady { + t.Fatalf("route should be ready after recovery: %#v ok=%v", route, ok) + } + if route.LeaseExpiresUnixMs != lease { + t.Fatalf("recovered route lease=%d, want %d", route.LeaseExpiresUnixMs, lease) + } +} + +func TestConsoleRestartExpiredLeaseNotOfferedForRecovery(t *testing.T) { + store := registrytest.NewStore(t) + base := time.Unix(1_700_840_000, 0) + ctx := context.Background() + + // Persist a route whose lease expires during the console offline period. + if err := store.UpsertConfirmedTerminalSessionRoute(ctx, registry.TerminalSessionRoute{ + ScopedSessionID: "obx:owner-a:session-offline-expired", + NodeID: "node-a", + LeaseExpiresUnixMs: base.Add(time.Minute).UnixMilli(), + LastUsedUnixMs: base.UnixMilli(), + CreatedAtUnixMs: base.UnixMilli(), + UpdatedAtUnixMs: base.UnixMilli(), + }); err != nil { + t.Fatalf("persist route: %v", err) + } + + // Restart after the lease has expired. + restartTime := base.Add(2 * time.Minute) + svc := NewRegistryService(store, nil, 5, 15, time.Minute) + svc.nowFn = func() time.Time { return restartTime } + if err := svc.RestoreTerminalSessionRoutes(ctx, restartTime); err != nil { + t.Fatalf("restore routes: %v", err) + } + + // No candidates should be offered. + candidates := svc.beginTerminalSessionRecovery("node-a", restartTime) + if len(candidates) != 0 { + t.Fatalf("expired route should not be offered for recovery: %#v", candidates) + } +} + +func TestDeleteProvisionedWorkerRemovesPersistedRoutes(t *testing.T) { + store := registrytest.NewStore(t) + base := time.Unix(1_700_850_000, 0) + ctx := context.Background() + + // Persist a route for the worker that will be deleted. + if err := store.UpsertConfirmedTerminalSessionRoute(ctx, registry.TerminalSessionRoute{ + ScopedSessionID: "obx:owner-a:session-delete-worker", + NodeID: "node-delete", + LeaseExpiresUnixMs: base.Add(30 * time.Minute).UnixMilli(), + LastUsedUnixMs: base.UnixMilli(), + CreatedAtUnixMs: base.UnixMilli(), + UpdatedAtUnixMs: base.UnixMilli(), + }); err != nil { + t.Fatalf("persist route: %v", err) + } + + svc := NewRegistryService(store, nil, 5, 15, time.Minute) + svc.nowFn = func() time.Time { return base } + + // Manually bind the route in memory so deleteTerminalSessionRoutesByNode has something to remove. + svc.bindTerminalSessionRoute("obx:owner-a:session-delete-worker", "node-delete", base) + svc.updateTerminalSessionRouteLease("obx:owner-a:session-delete-worker", "node-delete", base.Add(30*time.Minute).UnixMilli(), base) + + removed, err := svc.deleteTerminalSessionRoutesByNode("node-delete") + if err != nil { + t.Fatalf("delete routes by node: %v", err) + } + if removed != 1 { + t.Fatalf("removed=%d, want 1", removed) + } + + // Verify the route is gone from the database. + routes, err := store.LoadActiveTerminalSessionRoutes(ctx, base.UnixMilli()) + if err != nil { + t.Fatalf("load active routes: %v", err) + } + for _, r := range routes { + if r.NodeID == "node-delete" { + t.Fatal("persisted route should be deleted when worker is deleted") + } + } + + // Verify the route is gone from memory. + svc.terminalRoutesMu.RLock() + _, ok := svc.terminalSessionToNode["obx:owner-a:session-delete-worker"] + svc.terminalRoutesMu.RUnlock() + if ok { + t.Fatal("in-memory route should be deleted when worker is deleted") + } +} + +func TestPruneExpiredTerminalSessionRoutesDeletesFromPersistence(t *testing.T) { + store := registrytest.NewStore(t) + svc := NewRegistryService(store, nil, 5, 15, time.Minute) + base := time.Unix(1_700_860_000, 0) + svc.nowFn = func() time.Time { return base } + ctx := context.Background() + + // Persist a route with a short lease. + if err := store.UpsertConfirmedTerminalSessionRoute(ctx, registry.TerminalSessionRoute{ + ScopedSessionID: "obx:owner-a:session-prune", + NodeID: "node-a", + LeaseExpiresUnixMs: base.Add(time.Minute).UnixMilli(), + LastUsedUnixMs: base.UnixMilli(), + CreatedAtUnixMs: base.UnixMilli(), + UpdatedAtUnixMs: base.UnixMilli(), + }); err != nil { + t.Fatalf("persist route: %v", err) + } + + // Load it into memory via restore. + if err := svc.RestoreTerminalSessionRoutes(ctx, base); err != nil { + t.Fatalf("restore routes: %v", err) + } + + // Prune after the lease has expired. + removed := svc.pruneExpiredTerminalSessionRoutes(base.Add(2 * time.Minute)) + if removed != 1 { + t.Fatalf("removed=%d, want 1", removed) + } + + // Verify it's gone from the database. + routes, err := store.LoadActiveTerminalSessionRoutes(ctx, base.Add(2*time.Minute).UnixMilli()) + if err != nil { + t.Fatalf("load active routes: %v", err) + } + for _, r := range routes { + if r.ScopedSessionID == "obx:owner-a:session-prune" { + t.Fatal("expired route should be pruned from database") + } + } +} diff --git a/console/internal/httpapi/worker_handler.go b/console/internal/httpapi/worker_handler.go index 19deb93..8084be7 100644 --- a/console/internal/httpapi/worker_handler.go +++ b/console/internal/httpapi/worker_handler.go @@ -37,7 +37,7 @@ type WorkerHandler struct { type WorkerProvisioning interface { CreateProvisionedWorkerForOwner(ownerID string, workerType string, now time.Time, offlineTTL time.Duration) (string, string, error) - DeleteProvisionedWorker(nodeID string) bool + DeleteProvisionedWorker(nodeID string) (bool, error) } type workerItem struct { @@ -330,7 +330,12 @@ func (h *WorkerHandler) DeleteWorker(c *gin.Context) { return } } - if !h.provisioning.DeleteProvisionedWorker(nodeID) { + deleted, err := h.provisioning.DeleteProvisionedWorker(nodeID) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to delete worker"}) + return + } + if !deleted { c.JSON(http.StatusNotFound, gin.H{"error": "worker not found"}) return } diff --git a/console/internal/httpapi/worker_handler_test.go b/console/internal/httpapi/worker_handler_test.go index 46b5385..0dbbf87 100644 --- a/console/internal/httpapi/worker_handler_test.go +++ b/console/internal/httpapi/worker_handler_test.go @@ -52,15 +52,15 @@ func (p *fakeWorkerProvisioning) CreateProvisionedWorkerForOwner(ownerID string, return p.createNodeID, p.createSecret, nil } -func (p *fakeWorkerProvisioning) DeleteProvisionedWorker(nodeID string) bool { +func (p *fakeWorkerProvisioning) DeleteProvisionedWorker(nodeID string) (bool, error) { if p == nil || p.secrets == nil { - return false + return false, nil } if _, ok := p.secrets[nodeID]; !ok { - return false + return false, nil } delete(p.secrets, nodeID) - return true + return true, nil } func TestListWorkersEmpty(t *testing.T) { diff --git a/console/internal/persistence/db.go b/console/internal/persistence/db.go index d5b857b..79e3847 100644 --- a/console/internal/persistence/db.go +++ b/console/internal/persistence/db.go @@ -85,6 +85,9 @@ func Open(ctx context.Context, opts Options) (*DB, error) { queries := sqlc.New(db) nowMS := time.Now().UnixMilli() + if _, err := queries.DeleteExpiredTerminalSessionRoutes(ctx, nowMS); err != nil { + return nil, fmt.Errorf("startup recovery terminal session routes: %w", err) + } retentionMS := int64((time.Duration(taskRetentionDays) * 24 * time.Hour).Milliseconds()) expiresMS := nowMS + retentionMS if _, err := queries.MarkNonTerminalTasksFailedOnStartup(ctx, sqlc.MarkNonTerminalTasksFailedOnStartupParams{ diff --git a/console/internal/persistence/sqlc/models.go b/console/internal/persistence/sqlc/models.go index 0d949ba..ee80fd6 100644 --- a/console/internal/persistence/sqlc/models.go +++ b/console/internal/persistence/sqlc/models.go @@ -44,6 +44,15 @@ type Task struct { ExpiresAtUnixMs int64 `json:"expires_at_unix_ms"` } +type TerminalSessionRoute struct { + ScopedSessionID string `json:"scoped_session_id"` + NodeID string `json:"node_id"` + LeaseExpiresUnixMs int64 `json:"lease_expires_unix_ms"` + LastUsedUnixMs int64 `json:"last_used_unix_ms"` + CreatedAtUnixMs int64 `json:"created_at_unix_ms"` + UpdatedAtUnixMs int64 `json:"updated_at_unix_ms"` +} + type TrustedToken struct { TokenID string `json:"token_id"` AccountID string `json:"account_id"` diff --git a/console/internal/persistence/sqlc/terminal_session_routes.sql.go b/console/internal/persistence/sqlc/terminal_session_routes.sql.go new file mode 100644 index 0000000..63b3e28 --- /dev/null +++ b/console/internal/persistence/sqlc/terminal_session_routes.sql.go @@ -0,0 +1,164 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 +// source: terminal_session_routes.sql + +package sqlc + +import ( + "context" +) + +const deleteExpiredTerminalSessionRoutes = `-- name: DeleteExpiredTerminalSessionRoutes :execrows +DELETE FROM terminal_session_routes +WHERE lease_expires_unix_ms <= ? +` + +func (q *Queries) DeleteExpiredTerminalSessionRoutes(ctx context.Context, leaseExpiresUnixMs int64) (int64, error) { + result, err := q.db.ExecContext(ctx, deleteExpiredTerminalSessionRoutes, leaseExpiresUnixMs) + if err != nil { + return 0, err + } + return result.RowsAffected() +} + +const deleteTerminalSessionRouteBySessionAndNode = `-- name: DeleteTerminalSessionRouteBySessionAndNode :execrows +DELETE FROM terminal_session_routes +WHERE scoped_session_id = ? AND node_id = ? +` + +type DeleteTerminalSessionRouteBySessionAndNodeParams struct { + ScopedSessionID string `json:"scoped_session_id"` + NodeID string `json:"node_id"` +} + +func (q *Queries) DeleteTerminalSessionRouteBySessionAndNode(ctx context.Context, arg DeleteTerminalSessionRouteBySessionAndNodeParams) (int64, error) { + result, err := q.db.ExecContext(ctx, deleteTerminalSessionRouteBySessionAndNode, arg.ScopedSessionID, arg.NodeID) + if err != nil { + return 0, err + } + return result.RowsAffected() +} + +const deleteTerminalSessionRoutesByNode = `-- name: DeleteTerminalSessionRoutesByNode :execrows +DELETE FROM terminal_session_routes +WHERE node_id = ? +` + +func (q *Queries) DeleteTerminalSessionRoutesByNode(ctx context.Context, nodeID string) (int64, error) { + result, err := q.db.ExecContext(ctx, deleteTerminalSessionRoutesByNode, nodeID) + if err != nil { + return 0, err + } + return result.RowsAffected() +} + +const getTerminalSessionRouteBySession = `-- name: GetTerminalSessionRouteBySession :one +SELECT + scoped_session_id, + node_id, + lease_expires_unix_ms, + last_used_unix_ms, + created_at_unix_ms, + updated_at_unix_ms +FROM terminal_session_routes +WHERE scoped_session_id = ? +LIMIT 1 +` + +func (q *Queries) GetTerminalSessionRouteBySession(ctx context.Context, scopedSessionID string) (TerminalSessionRoute, error) { + row := q.db.QueryRowContext(ctx, getTerminalSessionRouteBySession, scopedSessionID) + var i TerminalSessionRoute + err := row.Scan( + &i.ScopedSessionID, + &i.NodeID, + &i.LeaseExpiresUnixMs, + &i.LastUsedUnixMs, + &i.CreatedAtUnixMs, + &i.UpdatedAtUnixMs, + ) + return i, err +} + +const listActiveTerminalSessionRoutes = `-- name: ListActiveTerminalSessionRoutes :many +SELECT + scoped_session_id, + node_id, + lease_expires_unix_ms, + last_used_unix_ms, + created_at_unix_ms, + updated_at_unix_ms +FROM terminal_session_routes +WHERE lease_expires_unix_ms > ? +ORDER BY node_id ASC, scoped_session_id ASC +` + +func (q *Queries) ListActiveTerminalSessionRoutes(ctx context.Context, leaseExpiresUnixMs int64) ([]TerminalSessionRoute, error) { + rows, err := q.db.QueryContext(ctx, listActiveTerminalSessionRoutes, leaseExpiresUnixMs) + if err != nil { + return nil, err + } + defer rows.Close() + var items []TerminalSessionRoute + for rows.Next() { + var i TerminalSessionRoute + if err := rows.Scan( + &i.ScopedSessionID, + &i.NodeID, + &i.LeaseExpiresUnixMs, + &i.LastUsedUnixMs, + &i.CreatedAtUnixMs, + &i.UpdatedAtUnixMs, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const upsertTerminalSessionRoute = `-- name: UpsertTerminalSessionRoute :execrows +INSERT INTO terminal_session_routes ( + scoped_session_id, + node_id, + lease_expires_unix_ms, + last_used_unix_ms, + created_at_unix_ms, + updated_at_unix_ms +) VALUES (?, ?, ?, ?, ?, ?) +ON CONFLICT(scoped_session_id) DO UPDATE SET + lease_expires_unix_ms = MAX(terminal_session_routes.lease_expires_unix_ms, excluded.lease_expires_unix_ms), + last_used_unix_ms = MAX(terminal_session_routes.last_used_unix_ms, excluded.last_used_unix_ms), + updated_at_unix_ms = MAX(terminal_session_routes.updated_at_unix_ms, excluded.updated_at_unix_ms) +WHERE terminal_session_routes.node_id = excluded.node_id +` + +type UpsertTerminalSessionRouteParams struct { + ScopedSessionID string `json:"scoped_session_id"` + NodeID string `json:"node_id"` + LeaseExpiresUnixMs int64 `json:"lease_expires_unix_ms"` + LastUsedUnixMs int64 `json:"last_used_unix_ms"` + CreatedAtUnixMs int64 `json:"created_at_unix_ms"` + UpdatedAtUnixMs int64 `json:"updated_at_unix_ms"` +} + +func (q *Queries) UpsertTerminalSessionRoute(ctx context.Context, arg UpsertTerminalSessionRouteParams) (int64, error) { + result, err := q.db.ExecContext(ctx, upsertTerminalSessionRoute, + arg.ScopedSessionID, + arg.NodeID, + arg.LeaseExpiresUnixMs, + arg.LastUsedUnixMs, + arg.CreatedAtUnixMs, + arg.UpdatedAtUnixMs, + ) + if err != nil { + return 0, err + } + return result.RowsAffected() +} diff --git a/console/internal/registry/store_terminal_session_routes.go b/console/internal/registry/store_terminal_session_routes.go new file mode 100644 index 0000000..c682053 --- /dev/null +++ b/console/internal/registry/store_terminal_session_routes.go @@ -0,0 +1,119 @@ +package registry + +import ( + "context" + "errors" + "strings" + + "github.com/onlyboxes/onlyboxes/console/internal/persistence/sqlc" +) + +var ErrTerminalSessionRouteNodeConflict = errors.New("terminal session route belongs to another worker") + +type TerminalSessionRoute struct { + ScopedSessionID string + NodeID string + LeaseExpiresUnixMs int64 + LastUsedUnixMs int64 + CreatedAtUnixMs int64 + UpdatedAtUnixMs int64 +} + +type TerminalSessionRouteRef struct { + ScopedSessionID string + NodeID string +} + +func (s *Store) LoadActiveTerminalSessionRoutes(ctx context.Context, nowUnixMs int64) ([]TerminalSessionRoute, error) { + if s == nil || s.queries == nil { + return nil, ErrPersistenceDBRequired + } + rows, err := s.queries.ListActiveTerminalSessionRoutes(ctx, nowUnixMs) + if err != nil { + return nil, err + } + routes := make([]TerminalSessionRoute, 0, len(rows)) + for _, row := range rows { + routes = append(routes, terminalSessionRouteFromSQL(row)) + } + return routes, nil +} + +func (s *Store) UpsertConfirmedTerminalSessionRoute(ctx context.Context, route TerminalSessionRoute) error { + if s == nil || s.queries == nil { + return ErrPersistenceDBRequired + } + route.ScopedSessionID = strings.TrimSpace(route.ScopedSessionID) + route.NodeID = strings.TrimSpace(route.NodeID) + if route.ScopedSessionID == "" || route.NodeID == "" || route.LeaseExpiresUnixMs <= 0 { + return errors.New("invalid terminal session route") + } + rows, err := s.queries.UpsertTerminalSessionRoute(ctx, sqlc.UpsertTerminalSessionRouteParams{ + ScopedSessionID: route.ScopedSessionID, + NodeID: route.NodeID, + LeaseExpiresUnixMs: route.LeaseExpiresUnixMs, + LastUsedUnixMs: route.LastUsedUnixMs, + CreatedAtUnixMs: route.CreatedAtUnixMs, + UpdatedAtUnixMs: route.UpdatedAtUnixMs, + }) + if err != nil { + return err + } + if rows == 0 { + return ErrTerminalSessionRouteNodeConflict + } + return nil +} + +func (s *Store) DeleteTerminalSessionRoute(ctx context.Context, scopedSessionID string, expectedNodeID string) (bool, error) { + if s == nil || s.queries == nil { + return false, ErrPersistenceDBRequired + } + rows, err := s.queries.DeleteTerminalSessionRouteBySessionAndNode(ctx, sqlc.DeleteTerminalSessionRouteBySessionAndNodeParams{ + ScopedSessionID: strings.TrimSpace(scopedSessionID), + NodeID: strings.TrimSpace(expectedNodeID), + }) + return rows > 0, err +} + +func (s *Store) DeleteTerminalSessionRoutes(ctx context.Context, routes []TerminalSessionRouteRef) error { + if s == nil || s.db == nil { + return ErrPersistenceDBRequired + } + return s.db.WithTx(ctx, func(q *sqlc.Queries) error { + for _, route := range routes { + if _, err := q.DeleteTerminalSessionRouteBySessionAndNode(ctx, sqlc.DeleteTerminalSessionRouteBySessionAndNodeParams{ + ScopedSessionID: strings.TrimSpace(route.ScopedSessionID), + NodeID: strings.TrimSpace(route.NodeID), + }); err != nil { + return err + } + } + return nil + }) +} + +func (s *Store) DeleteTerminalSessionRoutesByNode(ctx context.Context, nodeID string) (int64, error) { + if s == nil || s.queries == nil { + return 0, ErrPersistenceDBRequired + } + return s.queries.DeleteTerminalSessionRoutesByNode(ctx, strings.TrimSpace(nodeID)) +} + +func (s *Store) DeleteExpiredTerminalSessionRoutes(ctx context.Context, nowUnixMs int64) (int64, error) { + if s == nil || s.queries == nil { + return 0, ErrPersistenceDBRequired + } + return s.queries.DeleteExpiredTerminalSessionRoutes(ctx, nowUnixMs) +} + +func terminalSessionRouteFromSQL(row sqlc.TerminalSessionRoute) TerminalSessionRoute { + return TerminalSessionRoute{ + ScopedSessionID: row.ScopedSessionID, + NodeID: row.NodeID, + LeaseExpiresUnixMs: row.LeaseExpiresUnixMs, + LastUsedUnixMs: row.LastUsedUnixMs, + CreatedAtUnixMs: row.CreatedAtUnixMs, + UpdatedAtUnixMs: row.UpdatedAtUnixMs, + } +} diff --git a/console/sqlc.yaml b/console/sqlc.yaml index 5017a3a..d1a32ff 100644 --- a/console/sqlc.yaml +++ b/console/sqlc.yaml @@ -7,6 +7,7 @@ sql: - "db/migrations/00003_accounts_and_token_binding.sql" - "db/migrations/00004_worker_sys_owner_claims.sql" - "db/migrations/00005_api_keys.sql" + - "db/migrations/00006_terminal_session_routes.sql" queries: - "db/queries/accounts.sql" - "db/queries/workers.sql" @@ -14,6 +15,7 @@ sql: - "db/queries/api_keys.sql" - "db/queries/tasks.sql" - "db/queries/maintenance.sql" + - "db/queries/terminal_session_routes.sql" gen: go: package: "sqlc"