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。 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"