diff --git a/Dockerfile b/Dockerfile index 92f2c576d6..fca8de8ed7 100644 --- a/Dockerfile +++ b/Dockerfile @@ -83,13 +83,13 @@ RUN --mount=type=cache,target=/root/.cache/uv,id=uv-${TARGETPLATFORM} \ fi; \ case "${UV_LOCK_STRATEGY}" in \ locked) \ - uv sync --locked --no-editable --reinstall-package openviking --extra bot --extra gemini --extra opengauss \ + uv sync --locked --no-editable --reinstall-package openviking --extra bot --extra gemini --extra opengauss --extra expiry \ ;; \ auto) \ if ! uv lock --check; then \ uv lock; \ fi; \ - uv sync --locked --no-editable --reinstall-package openviking --extra bot --extra gemini --extra opengauss \ + uv sync --locked --no-editable --reinstall-package openviking --extra bot --extra gemini --extra opengauss --extra expiry \ ;; \ *) \ echo "Unsupported UV_LOCK_STRATEGY: ${UV_LOCK_STRATEGY}" >&2; \ diff --git a/MANIFEST.in b/MANIFEST.in index 4b07bc56c5..69d225fbc6 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -12,6 +12,7 @@ include Cargo.lock graft crates/ragfs graft crates/ragfs-python recursive-include openviking *.yaml +recursive-include openviking/scheduling/expiry *.lua recursive-include openviking_gateway *.json # sdist should be source-only: never ship runtime binaries from working tree diff --git a/docs/en/api/05-sessions.md b/docs/en/api/05-sessions.md index c779abd2e1..27edb74e03 100644 --- a/docs/en/api/05-sessions.md +++ b/docs/en/api/05-sessions.md @@ -48,15 +48,15 @@ Create a new session. Sessions are containers for conversations, storing message `auto_commit_policy` fields (all optional; omitted fields fall back to the defaults when a policy is present): -| Field | Type | Default | Max | Description | -|-------|------|---------|-----|-------------| -| `pending_token_threshold` | int | 150000 | 1000000 | When uncommitted pending tokens exceed this value (strictly greater-than), an auto commit is triggered after a message write. | -| `message_count_threshold` | int | 100 | 1000 | When the uncommitted live message count exceeds this value (strictly greater-than), an auto commit is triggered after a message write. | -| `idle_timeout_seconds` | int | 86400 | 604800 | After this many idle seconds, a session with uncommitted content becomes eligible for the server-side idle scheduler. An idle-timeout commit archives the full backlog and ignores `keep_recent_count`. | -| `keep_recent_count` | int | 0 | 500 | Number of recent live messages to keep (not archived) on a threshold-triggered auto commit. Idle-timeout commits ignore this and commit everything. | -| `min_commit_interval_seconds` | int | 0 | 604800 | Minimum seconds between two automatic commits (throttle). | - -All fields have a minimum of `0` and are clamped into `[0, max]`. Unknown keys are rejected with `InvalidArgumentError`. Idle commits also require `memory.session_auto_commit.enabled` on the server; setting `idle_timeout_seconds` alone does not start the scanner. +| Field | Type | Default | Min | Max | Description | +|-------|------|---------|-----|-----|-------------| +| `pending_token_threshold` | int | 150000 | 0 | 1000000 | When uncommitted pending tokens exceed this value (strictly greater-than), an auto commit is triggered after a message write. | +| `message_count_threshold` | int | 100 | 0 | 1000 | When the uncommitted live message count exceeds this value (strictly greater-than), an auto commit is triggered after a message write. | +| `idle_timeout_seconds` | int | 86400 | 0 (off) or 300 | 604800 | After this many idle seconds, a session with uncommitted content becomes eligible for the server-side idle scheduler. An idle-timeout commit archives the full backlog and ignores `keep_recent_count`. | +| `keep_recent_count` | int | 0 | 0 | 500 | Number of recent live messages to keep (not archived) on a threshold-triggered auto commit. Idle-timeout commits ignore this and commit everything. | +| `min_commit_interval_seconds` | int | 0 | 0 | 604800 | Minimum seconds between two automatic commits (throttle). | + +New `idle_timeout_seconds` inputs must be `0` (off) or at least `300`; values above `604800` are clamped to that maximum. For stored-data compatibility, existing positive values below 300 continue to be read unchanged. Other fields have a minimum of `0` and are clamped to their ranges. Unknown keys are rejected with `InvalidArgumentError`. Idle commits also require `memory.session_auto_commit.enabled` on the server; setting `idle_timeout_seconds` alone does not start the scanner. #### 3. Usage Examples diff --git a/docs/en/guides/01-configuration.md b/docs/en/guides/01-configuration.md index 7463545071..fc78e2fa39 100644 --- a/docs/en/guides/01-configuration.md +++ b/docs/en/guides/01-configuration.md @@ -1760,6 +1760,7 @@ For memory-related settings, add a `memory` section in `ov.conf`: "memory": { "session_auto_commit": { "enabled": false, + "idle_backend": "local", "check_interval_seconds": 600.0, "scan_rate_limit_files_per_second": 2.0 } @@ -1770,6 +1771,7 @@ For memory-related settings, add a `memory` section in `ov.conf`: | Parameter | Type | Description | Default | |-----------|------|-------------|---------| | `enabled` | bool | Master switch for automatic session commits. When enabled, newly created sessions without an explicit `auto_commit_policy` get a default policy, and the idle-timeout background scheduler is started. When disabled, neither happens | `false` | +| `idle_backend` | string | Idle-timeout discovery backend: `local` scans session metadata; `redis` uses distributed expiry scheduling | `local` | | `check_interval_seconds` | float | Minimum interval between two idle scan rounds in seconds. Must be greater than `0` | `600.0` | | `scan_rate_limit_files_per_second` | float | Maximum number of session `.meta.json` files read per second during an idle scan. Must be greater than `0`. Bounds background storage IO when the sessions tree is very large | `2.0` | @@ -1780,22 +1782,56 @@ Notes: - When `enabled=false`, sessions created without an explicit or `server.user_config_defaults.auto_commit_policy` policy keep auto commit disabled and return `auto_commit_policy: null`; `SessionAutoCommitScheduler` is not started. - When `enabled=true`: - Sessions without an explicit or deployment-default policy get the built-in policy below. - - `SessionAutoCommitScheduler` starts a scan round immediately, then waits at least `check_interval_seconds` before the next round. If a round already took longer than that interval, the next round starts immediately. - - Within a round, session `.meta.json` files are read serially at no more than `scan_rate_limit_files_per_second` to avoid IO spikes. - - It does not perform a dedicated startup recovery sweep; idle detection happens only on periodic scans. + - With `idle_backend="local"`, `SessionAutoCommitScheduler` starts a scan round immediately, then waits at least `check_interval_seconds` before the next round. Session `.meta.json` files are read serially at no more than `scan_rate_limit_files_per_second`. + - With `idle_backend="redis"`, the server enables expiry reporting, publishing, and callback consumption together using the top-level `expiry` configuration. - Token- and message-count auto commit run inline after message writes, do not depend on the scheduler, but still require the session to carry an `auto_commit_policy`. +#### Distributed Expiry Configuration + +All API, publisher, and callback instances that share expiry state must use the same stable `service_name`. It identifies the deployed document service, not an individual pod. + +```json +{ + "expiry": { + "service_name": "openviking-documents", + "publisher_separate_process": true, + "shard_count": 128, + "session_consumer_concurrency": 8, + "redis": { + "mode": "cluster", + "endpoints": ["redis://redis-proxy.example.com:6379"], + "username": "openviking", + "password_env": "OPENVIKING_EXPIRY_REDIS_PASSWORD" + } + } +} +``` + +| Parameter | Type | Description | Default | +|-----------|------|-------------|---------| +| `service_name` | string | Stable document service name shared by all instances using the same expiry data | `""` | +| `publisher_separate_process` | bool | Let `openviking-server` supervise publisher in a dedicated child process; when `false`, publisher runs in each API process | `true` | +| `shard_count` | int | Number of logical expiry shards; all instances sharing a namespace must use the same value, in the range `1..512` | `128` | +| `session_consumer_concurrency` | int | Concurrent Session expiry callbacks per API process; must be greater than `0` | `8` | +| `redis.mode` | string | Redis deployment mode: `standalone` or `cluster` | `standalone` | +| `redis.endpoints` | string array | Redis endpoints; a cluster may use one proxy endpoint | `[]` | +| `redis.username` | string | Redis ACL username | `""` | +| `redis.password_env` | string | Environment variable containing the Redis password; startup fails if it is not set | `""` | +| `redis.password` | string | Legacy plaintext password; prefer `password_env`; mutually exclusive with it | `""` | + +`service_name` and Redis endpoints are required when Redis idle auto commit is enabled. Credentials must use the dedicated fields above and cannot be embedded in endpoint URIs. With the default `publisher_separate_process=true`, `openviking-server` starts one internal publisher before serving traffic, stops it during shutdown, and exits if that child fails unexpectedly. This publisher is not exposed as a public CLI command. + ##### Per-session Auto Commit Policy -When a session carries an `auto_commit_policy`, any field you omit falls back to the built-in default below. Sessions without a stored policy keep auto commit disabled. Values are clamped into `[0, max]`, and unknown keys are rejected with `InvalidArgumentError`. See [Sessions API](../api/05-sessions.md#create-session) for how to set and view it. +When a session carries an `auto_commit_policy`, any field you omit falls back to the built-in default below. Sessions without a stored policy keep auto commit disabled. Except for the idle-timeout write minimum, values are clamped into `[0, max]`; unknown keys are rejected with `InvalidArgumentError`. See [Sessions API](../api/05-sessions.md#create-session) for how to set and view it. -| Field | Type | Default | Max | Description | -|-------|------|---------|-----|-------------| -| `pending_token_threshold` | int | 150000 | 1000000 | When uncommitted pending tokens exceed this value (strictly greater-than), an auto commit is triggered after a message write. | -| `message_count_threshold` | int | 100 | 1000 | When the uncommitted live message count exceeds this value (strictly greater-than), an auto commit is triggered after a message write. | -| `idle_timeout_seconds` | int | 86400 | 604800 | After this many idle seconds, a session with uncommitted content becomes eligible for the server-side idle scheduler. Idle-timeout commits archive the full backlog and ignore `keep_recent_count`. | -| `keep_recent_count` | int | 0 | 500 | Number of recent live messages to keep (not archived) on a threshold-triggered auto commit. Idle-timeout commits ignore this and commit everything. | -| `min_commit_interval_seconds` | int | 0 | 604800 | Minimum seconds between two automatic commits (throttle). | +| Field | Type | Default | Min | Max | Description | +|-------|------|---------|-----|-----|-------------| +| `pending_token_threshold` | int | 150000 | 0 | 1000000 | When uncommitted pending tokens exceed this value (strictly greater-than), an auto commit is triggered after a message write. | +| `message_count_threshold` | int | 100 | 0 | 1000 | When the uncommitted live message count exceeds this value (strictly greater-than), an auto commit is triggered after a message write. | +| `idle_timeout_seconds` | int | 86400 | 0 (off) or 300 | 604800 | New positive inputs have a five-minute minimum; existing stored positive values below 300 continue to be read unchanged. Idle-timeout commits archive the full backlog and ignore `keep_recent_count`. | +| `keep_recent_count` | int | 0 | 0 | 500 | Number of recent live messages to keep (not archived) on a threshold-triggered auto commit. Idle-timeout commits ignore this and commit everything. | +| `min_commit_interval_seconds` | int | 0 | 0 | 604800 | Minimum seconds between two automatic commits (throttle). | Code entry: `openviking/session/auto_commit_policy.py:AutoCommitPolicy`. diff --git a/docs/zh/api/05-sessions.md b/docs/zh/api/05-sessions.md index 2607865a51..196c25afb3 100644 --- a/docs/zh/api/05-sessions.md +++ b/docs/zh/api/05-sessions.md @@ -48,15 +48,15 @@ Session API 按认证用户作用域访问会话，并返回 canonical user sess `auto_commit_policy` 字段（均为可选；存在 policy 时，未传字段回退到默认值）： -| 字段 | 类型 | 默认值 | 上限 | 说明 | -|------|------|--------|------|------| -| `pending_token_threshold` | int | 150000 | 1000000 | 当未提交的 pending token 超过该值（严格大于）时，会在消息写入后触发一次自动 commit。 | -| `message_count_threshold` | int | 100 | 1000 | 当未提交的 live message 数量超过该值（严格大于）时，会在消息写入后触发一次自动 commit。 | -| `idle_timeout_seconds` | int | 86400 | 604800 | 有未提交内容的 session 在空闲这么多秒后，进入服务端 idle scheduler 的处理范围。idle 触发的 commit 会归档全部积压消息，并忽略 `keep_recent_count`。 | -| `keep_recent_count` | int | 0 | 500 | 阈值触发的自动 commit 后保留（不归档）的最近 live message 数量。idle 超时触发的 commit 会忽略该值并归档所有消息。 | -| `min_commit_interval_seconds` | int | 0 | 604800 | 两次自动 commit 之间的最小间隔秒数（节流）。 | - -所有字段最小值为 `0`，会被 clamp 到 `[0, 上限]`。未知字段会以 `InvalidArgumentError` 拒绝。idle 超时还要求服务端开启 `memory.session_auto_commit.enabled`；仅设置 `idle_timeout_seconds` 不会启动扫描器。 +| 字段 | 类型 | 默认值 | 最小值 | 上限 | 说明 | +|------|------|--------|--------|------|------| +| `pending_token_threshold` | int | 150000 | 0 | 1000000 | 当未提交的 pending token 超过该值（严格大于）时，会在消息写入后触发一次自动 commit。 | +| `message_count_threshold` | int | 100 | 0 | 1000 | 当未提交的 live message 数量超过该值（严格大于）时，会在消息写入后触发一次自动 commit。 | +| `idle_timeout_seconds` | int | 86400 | 0（关闭）或 300 | 604800 | 有未提交内容的 session 在空闲这么多秒后，进入服务端 idle scheduler 的处理范围。idle 触发的 commit 会归档全部积压消息，并忽略 `keep_recent_count`。 | +| `keep_recent_count` | int | 0 | 0 | 500 | 阈值触发的自动 commit 后保留（不归档）的最近 live message 数量。idle 超时触发的 commit 会忽略该值并归档所有消息。 | +| `min_commit_interval_seconds` | int | 0 | 0 | 604800 | 两次自动 commit 之间的最小间隔秒数（节流）。 | + +新传入的 `idle_timeout_seconds` 必须为 `0`（关闭）或不小于 `300`，超过 `604800` 时会被 clamp 到该上限；为兼容存量数据，已经保存的小于 300 的正值仍按原值读取。其他字段最小值为 `0`，并会被 clamp 到各自范围。未知字段会以 `InvalidArgumentError` 拒绝。idle 超时还要求服务端开启 `memory.session_auto_commit.enabled`；仅设置 `idle_timeout_seconds` 不会启动扫描器。 #### 3. 使用示例 diff --git a/docs/zh/guides/01-configuration.md b/docs/zh/guides/01-configuration.md index a546685082..7a219cc789 100644 --- a/docs/zh/guides/01-configuration.md +++ b/docs/zh/guides/01-configuration.md @@ -1747,6 +1747,7 @@ openviking-server --config /path/to/ov.conf "memory": { "session_auto_commit": { "enabled": false, + "idle_backend": "local", "check_interval_seconds": 600.0, "scan_rate_limit_files_per_second": 2.0 } @@ -1757,6 +1758,7 @@ openviking-server --config /path/to/ov.conf | 参数 | 类型 | 说明 | 默认值 | |------|------|------|--------| | `enabled` | bool | 自动 commit 总开关。开启后：未显式传入 `auto_commit_policy` 的新 session 会套用默认 policy，同时启动后台 idle 扫描器；关闭后两者都不发生 | `false` | +| `idle_backend` | string | idle 超时发现后端：`local` 扫描 session 元数据；`redis` 使用分布式 expiry 调度 | `local` | | `check_interval_seconds` | float | 两轮 idle 扫描之间的最小间隔，单位秒，必须大于 `0` | `600.0` | | `scan_rate_limit_files_per_second` | float | idle 扫描时每秒最多读取的 session `.meta.json` 文件数，必须大于 `0`，用于限制大量 session 扫描时的存储 IO 压力 | `2.0` | @@ -1767,22 +1769,56 @@ openviking-server --config /path/to/ov.conf - `enabled=false` 时：既无显式 policy、也无 `server.user_config_defaults.auto_commit_policy` 的新 Session 保持 auto commit 关闭，并返回 `auto_commit_policy: null`；同时不启动 `SessionAutoCommitScheduler`。 - `enabled=true` 时： - 既无显式 policy、也无部署级默认 policy 的新 Session 会带上下方内置 policy。 - - `SessionAutoCommitScheduler` 启动后立即开始一轮扫描，之后每轮结束后至少等待 `check_interval_seconds` 再开始下一轮；若一轮耗时已超过该值，则立即开始下一轮。 - - 扫描过程中以 `scan_rate_limit_files_per_second` 为上限串行读取 session `.meta.json`，避免瞬时 IO 洪峰。 - - 不会做单独的启动恢复扫描，idle 检查只发生在周期扫描时。 + - `idle_backend="local"` 时，`SessionAutoCommitScheduler` 启动后立即开始扫描；扫描过程以 `scan_rate_limit_files_per_second` 为上限串行读取 session `.meta.json`，每轮结束后至少等待 `check_interval_seconds`。 + - `idle_backend="redis"` 时，服务通过顶层 `expiry` 配置整体启用 expiry 上报、发布和 callback 消费。 - token 和 message-count 自动触发在消息写入后内联执行，不依赖 scheduler，但同样以 session 是否带有 `auto_commit_policy` 为前提。 +#### 分布式 Expiry 配置 + +共享同一份 expiry 状态的 API、publisher 和 callback 实例必须使用相同且稳定的 `service_name`。它表示部署的文档服务名，不是单个 Pod 名。 + +```json +{ + "expiry": { + "service_name": "openviking-documents", + "publisher_separate_process": true, + "shard_count": 128, + "session_consumer_concurrency": 8, + "redis": { + "mode": "cluster", + "endpoints": ["redis://redis-proxy.example.com:6379"], + "username": "openviking", + "password_env": "OPENVIKING_EXPIRY_REDIS_PASSWORD" + } + } +} +``` + +| 参数 | 类型 | 说明 | 默认值 | +|------|------|------|--------| +| `service_name` | string | 使用同一份 expiry 数据的所有实例共享的稳定文档服务名 | `""` | +| `publisher_separate_process` | bool | 由 `openviking-server` 在独立子进程中托管 publisher；为 `false` 时 publisher 在各 API 进程内运行 | `true` | +| `shard_count` | int | expiry 逻辑分片数；所有共享同一 namespace 的实例必须一致，取值范围为 `1..512` | `128` | +| `session_consumer_concurrency` | int | 每个 API 进程中 Session expiry callback 的并发数，必须大于 `0` | `8` | +| `redis.mode` | string | Redis 部署模式：`standalone` 或 `cluster` | `standalone` | +| `redis.endpoints` | string array | Redis 地址；集群模式可以只配置一个统一 proxy 地址 | `[]` | +| `redis.username` | string | Redis ACL 用户名 | `""` | +| `redis.password_env` | string | 保存 Redis 密码的环境变量名；变量不存在时启动失败 | `""` | +| `redis.password` | string | 兼容用明文密码；建议使用 `password_env`，两者互斥 | `""` | + +启用 Redis idle auto commit 时必须配置 `service_name` 和 Redis 地址。用户名与密码只能使用上述独立字段，不能写入 endpoint URI。默认 `publisher_separate_process=true` 时，`openviking-server` 会在开始服务前拉起一个内部 publisher 子进程，在服务关闭时停止它，并在子进程异常退出时一并退出。publisher 不再暴露为公共 CLI 命令。 + ##### 单 session 自动 commit 策略 -当 session 带有 `auto_commit_policy` 时，未传的字段会回退到下方内置默认值。没有存储 policy 的 session 保持 auto commit 关闭。取值会被 clamp 到 `[0, 上限]`，未知字段会以 `InvalidArgumentError` 拒绝。设置和查看方式见 [Sessions API](../api/05-sessions.md#create-session)。 +当 session 带有 `auto_commit_policy` 时，未传的字段会回退到下方内置默认值。没有存储 policy 的 session 保持 auto commit 关闭。除 idle 超时的写入下限外，取值会被 clamp 到 `[0, 上限]`；未知字段会以 `InvalidArgumentError` 拒绝。设置和查看方式见 [Sessions API](../api/05-sessions.md#create-session)。 -| 字段 | 类型 | 默认值 | 上限 | 说明 | -|------|------|--------|------|------| -| `pending_token_threshold` | int | 150000 | 1000000 | 当未提交的 pending token 超过该值（严格大于）时，会在消息写入后触发一次自动 commit。 | -| `message_count_threshold` | int | 100 | 1000 | 当未提交的 live message 数量超过该值（严格大于）时，会在消息写入后触发一次自动 commit。 | -| `idle_timeout_seconds` | int | 86400 | 604800 | 有未提交内容的 session 在空闲这么多秒后，进入服务端 idle scheduler 的处理范围。idle 触发的 commit 会归档全部积压消息，并忽略 `keep_recent_count`。 | -| `keep_recent_count` | int | 0 | 500 | 阈值触发的自动 commit 后保留（不归档）的最近 live message 数量。idle 超时触发的 commit 会忽略该值并归档所有消息。 | -| `min_commit_interval_seconds` | int | 0 | 604800 | 两次自动 commit 之间的最小间隔秒数（节流）。 | +| 字段 | 类型 | 默认值 | 最小值 | 上限 | 说明 | +|------|------|--------|--------|------|------| +| `pending_token_threshold` | int | 150000 | 0 | 1000000 | 当未提交的 pending token 超过该值（严格大于）时，会在消息写入后触发一次自动 commit。 | +| `message_count_threshold` | int | 100 | 0 | 1000 | 当未提交的 live message 数量超过该值（严格大于）时，会在消息写入后触发一次自动 commit。 | +| `idle_timeout_seconds` | int | 86400 | 0（关闭）或 300 | 604800 | 新传入的正数最短为 5 分钟；存量的小于 300 的正值继续按原值读取。idle 触发的 commit 会归档全部积压消息，并忽略 `keep_recent_count`。 | +| `keep_recent_count` | int | 0 | 0 | 500 | 阈值触发的自动 commit 后保留（不归档）的最近 live message 数量。idle 超时触发的 commit 会忽略该值并归档所有消息。 | +| `min_commit_interval_seconds` | int | 0 | 0 | 604800 | 两次自动 commit 之间的最小间隔秒数（节流）。 | 代码入口：`openviking/session/auto_commit_policy.py:AutoCommitPolicy`。 diff --git a/openviking/scheduling/__init__.py b/openviking/scheduling/__init__.py new file mode 100644 index 0000000000..e908d0001e --- /dev/null +++ b/openviking/scheduling/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Scheduling implementations.""" diff --git a/openviking/scheduling/expiry/__init__.py b/openviking/scheduling/expiry/__init__.py new file mode 100644 index 0000000000..fcb0117a66 --- /dev/null +++ b/openviking/scheduling/expiry/__init__.py @@ -0,0 +1,31 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +from .api import ( + SESSION_ITEM_TYPE, + CallbackRequest, + CallbackResult, + CallbackStatus, + ExpiryError, + ExpiryHandler, + ExpiryMode, + ExpiryRegistration, + ExpiryRegistrationResult, + ExpiryScheduler, + HandlerOptions, + ItemIdentity, +) + +__all__ = [ + "ExpiryScheduler", + "SESSION_ITEM_TYPE", + "ItemIdentity", + "ExpiryRegistration", + "ExpiryRegistrationResult", + "ExpiryMode", + "ExpiryError", + "HandlerOptions", + "ExpiryHandler", + "CallbackRequest", + "CallbackResult", + "CallbackStatus", +] diff --git a/openviking/scheduling/expiry/api.py b/openviking/scheduling/expiry/api.py new file mode 100644 index 0000000000..4b54cdf4ff --- /dev/null +++ b/openviking/scheduling/expiry/api.py @@ -0,0 +1,86 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Public contracts for distributed expiry scheduling.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Protocol + +SESSION_ITEM_TYPE = "session" +SUPPORTED_ITEM_TYPES: frozenset[str] = frozenset({SESSION_ITEM_TYPE}) + + +class ExpiryMode(str, Enum): + IDLE_ONESHOT = "idle_oneshot" + FIXED = "fixed" + + +class CallbackStatus(str, Enum): + COMPLETED = "completed" + RETRY = "retry" + + +class ExpiryError(RuntimeError): + def __init__(self, code: str, message: str, *, retryable: bool = False) -> None: + super().__init__(message) + self.code = code + self.retryable = retryable + + +@dataclass(frozen=True) +class ItemIdentity: + item_type: str + account_id: str + user_id: str + local_id: str + + +@dataclass(frozen=True) +class ExpiryRegistration: + identity: ItemIdentity + due_at_ms: int + timeout_seconds: int + mode: ExpiryMode = ExpiryMode.IDLE_ONESHOT + + +@dataclass(frozen=True) +class ExpiryRegistrationResult: + accepted_due_at_ms: int + + +@dataclass(frozen=True) +class CallbackRequest: + identity: ItemIdentity + + +@dataclass(frozen=True) +class CallbackResult: + status: CallbackStatus + reason: str + + +@dataclass(frozen=True) +class HandlerOptions: + concurrency: int = 4 + max_attempts: int = 3 + metadata_enabled: bool = True + + +class ExpiryHandler(Protocol): + async def handle(self, request: CallbackRequest) -> CallbackResult: ... + + +class ExpiryScheduler(Protocol): + async def report(self, task: ExpiryRegistration) -> ExpiryRegistrationResult: ... + + async def unregister(self, identity: ItemIdentity) -> None: ... + + def register_handler( + self, + item_type: str, + handler: ExpiryHandler, + *, + options: HandlerOptions | None = None, + ) -> None: ... diff --git a/openviking/scheduling/expiry/config.py b/openviking/scheduling/expiry/config.py new file mode 100644 index 0000000000..cb27eabc5c --- /dev/null +++ b/openviking/scheduling/expiry/config.py @@ -0,0 +1,107 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Configuration independent from the aggregate OpenViking config package.""" + +from __future__ import annotations + +import os +from typing import Literal +from urllib.parse import urlsplit + +from pydantic import BaseModel, ConfigDict, Field, model_validator + + +class ExpiryRedisConfig(BaseModel): + model_config = ConfigDict(extra="forbid", hide_input_in_errors=True) + + mode: Literal["standalone", "cluster"] = "standalone" + endpoints: list[str] = Field(default_factory=list) + username: str = "" + password_env: str = "" + password: str = Field(default="", repr=False) + tls: bool = False + connect_timeout_seconds: float = Field(default=1.0, gt=0) + socket_timeout_seconds: float = Field(default=2.0, gt=0) + request_budget_seconds: float = Field(default=5.0, gt=0) + max_connections: int = Field(default=64, ge=1) + max_inflight_requests: int = Field(default=128, ge=1) + + @model_validator(mode="after") + def validate_connection(self) -> "ExpiryRedisConfig": + if self.password and self.password_env: + raise ValueError("Redis password and password_env are mutually exclusive") + if self.mode == "standalone" and len(self.endpoints) > 1: + raise ValueError("standalone Redis accepts at most one endpoint") + schemes: set[str] = set() + for endpoint in self.endpoints: + normalized = ( + endpoint + if "://" in endpoint + else (("rediss" if self.tls else "redis") + f"://{endpoint}") + ) + try: + parsed = urlsplit(normalized) + valid_address = bool(parsed.hostname and parsed.port) + except ValueError: + # urlsplit/.port errors can echo untrusted endpoint fragments. + raise ValueError("invalid Redis endpoint") from None + if parsed.scheme not in {"redis", "rediss"} or not valid_address: + raise ValueError("invalid Redis endpoint") + if ( + parsed.username is not None + or parsed.password is not None + or parsed.path not in {"", "/"} + or parsed.query + or parsed.fragment + ): + raise ValueError( + "Redis endpoints must not contain credentials, paths, " + "query parameters, or fragments; use dedicated Redis fields" + ) + schemes.add(parsed.scheme) + if len(schemes) > 1: + raise ValueError("Redis cluster endpoints must use one transport scheme") + return self + + def resolve_password(self) -> str | None: + if self.password_env.strip(): + try: + return os.environ[self.password_env] + except KeyError as exc: + raise ValueError( + f"Redis secret environment variable {self.password_env} is not set" + ) from exc + return self.password if self.password.strip() else None + + +class ExpiryConfig(BaseModel): + model_config = ConfigDict(extra="forbid", hide_input_in_errors=True) + + service_name: str = "" + publisher_separate_process: bool = True + shard_count: int = Field(default=128, ge=1, le=512) + session_consumer_concurrency: int = Field(default=8, ge=1) + precision_seconds: Literal[60] = 60 + owner_ttl_seconds: Literal[120] = 120 + owner_renew_interval_seconds: Literal[20] = 20 + scan_batch_size: int = Field(default=128, ge=1, le=4096) + cleanup_batch_size: int = Field(default=128, ge=1, le=4096) + registry_page_size: int = Field(default=128, ge=1, le=4096) + processing_timeout_seconds: int = Field(default=120, ge=1) + callback_poll_seconds: float = Field(default=0.5, gt=0) + shutdown_timeout_seconds: float = Field(default=30.0, gt=0) + redis: ExpiryRedisConfig = Field(default_factory=ExpiryRedisConfig) + + def validate_for_runtime(self) -> None: + if not self.service_name.strip(): + raise ValueError( + "expiry.service_name is required when Redis session auto commit is enabled" + ) + if not self.redis.endpoints: + raise ValueError( + "expiry.redis.endpoints is required when Redis session auto commit is enabled" + ) + if self.redis.mode == "standalone" and len(self.redis.endpoints) != 1: + raise ValueError("standalone Redis requires exactly one endpoint") + if self.redis.request_budget_seconds >= self.owner_renew_interval_seconds: + raise ValueError("Redis request budget must be shorter than owner renew interval") diff --git a/openviking/scheduling/expiry/consumer.py b/openviking/scheduling/expiry/consumer.py new file mode 100644 index 0000000000..bdb4485ca7 --- /dev/null +++ b/openviking/scheduling/expiry/consumer.py @@ -0,0 +1,392 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Published resource callback consumption.""" + +from __future__ import annotations + +import asyncio +import logging +import secrets +from dataclasses import dataclass, field +from typing import Awaitable, Callable + +from .api import CallbackRequest, CallbackStatus, ExpiryError +from .config import ExpiryConfig +from .keys import decode_item_id +from .metrics import ExpiryMetrics +from .models import ( + BucketItem, + BucketRef, + Claim, + OwnerRole, + StoreOperationError, + decode_claim_score, +) +from .redis_store import RedisExpiryStore +from .registry import HandlerBinding +from .shard_runner import ShardExecution, ShardRunner + +logger = logging.getLogger(__name__) + + +@dataclass +class _ConsumerState: + reserved: int = 0 + tasks: set[asyncio.Task[None]] = field(default_factory=set) + + +@dataclass(frozen=True) +class _BucketScan: + next_poll_at_ms: int | None = None + conflicted: bool = False + + +class ExpiryConsumer: + def __init__( + self, + store: RedisExpiryStore, + runner: ShardRunner, + item_type: str, + binding: HandlerBinding, + config: ExpiryConfig, + metrics: ExpiryMetrics, + *, + start_shard: int | None = None, + ) -> None: + self._store = store + self._runner = runner + self._item_type = item_type + self._binding = binding + self._config = config + self._metrics = metrics + self._callback_tasks: set[asyncio.Task[None]] = set() + self._next_shard = ( + secrets.randbelow(config.shard_count) if start_shard is None else start_shard + ) + + async def run(self, stop: asyncio.Event) -> None: + """Traverse every shard from a per-consumer offset.""" + while not stop.is_set(): + start_shard = self._next_shard + for offset in range(self._config.shard_count): + if stop.is_set(): + break + shard_id = (start_shard + offset) % self._config.shard_count + try: + await self._runner.run_if_acquired( + OwnerRole.CALLBACK, + shard_id, + self.process_shard, + stop, + item_type=self._item_type, + ) + except StoreOperationError as exc: + if exc.code == "shard_not_stopped": + logger.warning( + "expiry consumer skipped unfinished shard execution " + "item_type=%s shard_id=%d", + self._item_type, + shard_id, + ) + self._metrics.count( + "timer_shard_round", + item_type=self._item_type, + outcome="orphaned", + ) + continue + self._metrics.count( + "timer_shard_round", + item_type=self._item_type, + outcome="failed", + ) + except Exception: + self._metrics.count( + "timer_shard_round", + item_type=self._item_type, + outcome="failed", + ) + self._next_shard = (start_shard + 1) % self._config.shard_count + if not stop.is_set(): + try: + await asyncio.wait_for(stop.wait(), self._config.callback_poll_seconds) + except asyncio.TimeoutError: + pass + + async def process_shard(self, lease: ShardExecution, stop: asyncio.Event) -> None: + state = _ConsumerState() + bucket = await self._discover_bucket(lease) + if bucket is None: + return + await self._process_bucket(lease, stop, bucket, state) + + async def _discover_bucket(self, lease: ShardExecution) -> BucketRef | None: + cursor = None + while True: + lease.ensure_active() + page = await self._store.read_registry_page( + lease.owner.shard_id, + after=cursor, + limit=self._config.registry_page_size, + ) + for bucket in page.buckets: + if bucket.item_type == self._item_type: + return bucket + cursor = page.next_cursor + if cursor is None: + return None + + async def _process_bucket( + self, + lease: ShardExecution, + stop: asyncio.Event, + bucket: BucketRef, + state: _ConsumerState, + ) -> None: + while not stop.is_set(): + lease.ensure_active() + available = self._binding.options.concurrency - state.reserved + scan = _BucketScan() + redis_now_ms = 0 + if available > 0: + redis_now_ms = await self._store.read_time_ms() + lease.ensure_active() + scan = await self._claim_ready_items( + lease, + bucket, + redis_now_ms=redis_now_ms, + available=available, + state=state, + ) + + active = {task for task in state.tasks if not task.done()} + if active: + await lease.wait(self._config.callback_poll_seconds, stop) + continue + if scan.conflicted: + continue + + lease.ensure_active() + empty = await self._store.remove_registration_if_empty(bucket) + lease.ensure_active() + if empty: + return + if scan.next_poll_at_ms is None or scan.next_poll_at_ms <= redis_now_ms: + continue + await self._store.defer_registration( + lease.owner, + bucket, + next_poll_at_ms=scan.next_poll_at_ms, + ) + return + + tasks = {task for task in state.tasks if not task.done()} + if tasks: + await asyncio.wait(tasks) + lease.ensure_active() + + async def _claim_ready_items( + self, + lease: ShardExecution, + bucket: BucketRef, + *, + redis_now_ms: int, + available: int, + state: _ConsumerState, + ) -> _BucketScan: + next_poll_at_ms: int | None = None + started_at_seconds = redis_now_ms // 1000 + for attempt in range(self._binding.options.max_attempts): + while available > 0: + limit = min(self._config.scan_batch_size, available) + page = await self._store.read_bucket_page( + bucket, + attempt=attempt, + limit=limit, + ) + lease.ensure_active() + if not page.items: + break + + claimed_from_full_page = len(page.items) == limit + for item in page.items: + ready_at_ms = self._ready_at_ms(item, attempt) + if ready_at_ms is not None and ready_at_ms > redis_now_ms: + next_poll_at_ms = ( + ready_at_ms + if next_poll_at_ms is None + else min(next_poll_at_ms, ready_at_ms) + ) + claimed_from_full_page = False + break + + claim = await self._store.claim( + bucket, + item, + attempt=attempt, + started_at_seconds=started_at_seconds, + ) + lease.ensure_active() + if claim is None: + return _BucketScan(next_poll_at_ms, conflicted=True) + self._start_callback(lease, claim, state) + available -= 1 + if available == 0: + break + + if available == 0 or not claimed_from_full_page: + break + return _BucketScan(next_poll_at_ms) + + def _ready_at_ms(self, item: BucketItem, expected_attempt: int) -> int | None: + if item.score == 0: + if expected_attempt != 0: + raise StoreOperationError( + "namespace_corrupt", "Unclaimed expiry item is outside attempt zero" + ) + return None + attempt, started_at_seconds = decode_claim_score(item.score) + if attempt != expected_attempt: + raise StoreOperationError( + "namespace_corrupt", "Expiry item is outside its expected attempt layer" + ) + return (started_at_seconds + self._config.processing_timeout_seconds) * 1000 + + def _start_callback( + self, + lease: ShardExecution, + claim: Claim, + state: _ConsumerState, + ) -> None: + task = lease.spawn(self._callback_operation(lease, claim, state)) + state.reserved += 1 + state.tasks.add(task) + self._callback_tasks.add(task) + task.add_done_callback(state.tasks.discard) + task.add_done_callback(self._callback_tasks.discard) + + def _callback_operation( + self, + lease: ShardExecution, + claim: Claim, + state: _ConsumerState, + ) -> Callable[[], Awaitable[None]]: + async def operation() -> None: + await self._execute(lease, claim, state) + + return operation + + async def _execute( + self, + lease: ShardExecution, + claim: Claim, + state: _ConsumerState, + ) -> None: + try: + try: + lease.ensure_active() + request = await self._resolve_request(claim) + lease.ensure_active() + result = await asyncio.wait_for( + self._binding.handler.handle(request), + timeout=self._config.processing_timeout_seconds, + ) + except asyncio.CancelledError: + raise + except asyncio.TimeoutError as exc: + self._record_callback_error(claim, exc) + lease.ensure_active() + await self._finish_failed_attempt(lease, claim, "callback_timeout") + return + except StoreOperationError as exc: + if exc.code == "owner_lost" or exc.outcome_unknown: + raise + self._record_callback_error(claim, exc) + lease.ensure_active() + await self._finish_failed_attempt(lease, claim, "callback_exception") + return + except Exception as exc: + self._record_callback_error(claim, exc) + lease.ensure_active() + await self._finish_failed_attempt(lease, claim, "callback_exception") + return + + lease.ensure_active() + if result.status is CallbackStatus.COMPLETED: + await self._store.ack(claim) + lease.ensure_active() + else: + await self._finish_failed_attempt(lease, claim, "callback_retry") + finally: + state.reserved = max(0, state.reserved - 1) + + async def _finish_failed_attempt( + self, lease: ShardExecution, claim: Claim, reason: str + ) -> None: + attempt, _ = decode_claim_score(claim.claim_score) + if attempt + 1 >= self._binding.options.max_attempts: + await self._log_exhausted_and_ack(lease, claim, reason) + return + lease.ensure_active() + await self._store.advance_retry(claim) + lease.ensure_active() + + async def _log_exhausted_and_ack( + self, lease: ShardExecution, claim: Claim, reason: str + ) -> None: + lease.ensure_active() + self._metrics.count( + "timer_callback_exhausted", + item_type=claim.bucket.item_type, + outcome=reason, + ) + try: + self._metrics.log_exhausted( + claim, + max_attempts=self._binding.options.max_attempts, + reason=reason, + ) + except Exception as exc: + self._record_callback_error(claim, exc) + + try: + lease.ensure_active() + await self._store.ack(claim) + except StoreOperationError as exc: + if exc.code == "owner_lost" or exc.outcome_unknown: + raise + self._record_callback_error(claim, exc) + except Exception as exc: + self._record_callback_error(claim, exc) + + async def _resolve_request(self, claim: Claim) -> CallbackRequest: + identity = decode_item_id(claim.item_id) + if identity.item_type != claim.bucket.item_type: + raise ExpiryError("invalid_identity", "Bucket type differs from item ID") + if not self._binding.options.metadata_enabled: + return CallbackRequest(identity) + metadata = await self._store.read_metadata(claim.item_id) + if metadata is not None and metadata.identity != identity: + raise ExpiryError("invalid_identity", "Metadata differs from item ID") + self._metrics.count( + "timer_metadata_lookup", + item_type=identity.item_type, + outcome="miss" if metadata is None else "hit", + ) + return CallbackRequest(identity) + + async def drain(self, timeout_seconds: float) -> None: + tasks = {task for task in self._callback_tasks if not task.done()} + if not tasks: + return + _, pending = await asyncio.wait(tasks, timeout=max(0.0, timeout_seconds)) + for task in pending: + task.cancel() + if pending: + raise StoreOperationError("shard_not_stopped", "Callbacks did not drain within budget") + + def _record_callback_error(self, claim: Claim, exc: Exception) -> None: + self._metrics.count( + "timer_callback_error", + item_type=claim.bucket.item_type, + outcome=type(exc).__name__, + ) diff --git a/openviking/scheduling/expiry/keys.py b/openviking/scheduling/expiry/keys.py new file mode 100644 index 0000000000..dbd685ee41 --- /dev/null +++ b/openviking/scheduling/expiry/keys.py @@ -0,0 +1,122 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Redis key construction and reversible expiry identities.""" + +from __future__ import annotations + +import base64 +import binascii +from dataclasses import dataclass + +import xxhash + +from .api import ExpiryError, ItemIdentity +from .models import MINUTE_MS, BucketRef, OwnerRole, validate_identity, validate_item_type + +HASH_SEED = 0 +HASH_NAME = "xxh64" +IDENTITY_ENCODING = "urlsafe_base64" + + +def encode_component(value: str) -> str: + if not isinstance(value, str) or not value: + raise ExpiryError("invalid_identity", "Identity components must be non-empty strings") + return base64.urlsafe_b64encode(value.encode("utf-8")).decode("ascii").rstrip("=") + + +def decode_component(value: str) -> str: + if not isinstance(value, str) or not value or "=" in value: + raise ExpiryError("invalid_identity", "Invalid encoded identity component") + try: + padded = value + "=" * (-len(value) % 4) + decoded = base64.b64decode(padded, altchars=b"-_", validate=True).decode("utf-8") + except (binascii.Error, UnicodeDecodeError, ValueError) as exc: + raise ExpiryError("invalid_identity", "Invalid encoded identity component") from exc + if not decoded or encode_component(decoded) != value: + raise ExpiryError("invalid_identity", "Non-canonical encoded identity component") + return decoded + + +def encode_item_id(identity: ItemIdentity) -> str: + validate_identity(identity) + return ":".join( + ( + identity.item_type, + encode_component(identity.account_id), + encode_component(identity.user_id), + encode_component(identity.local_id), + ) + ) + + +def decode_item_id(item_id: str) -> ItemIdentity: + if not isinstance(item_id, str): + raise ExpiryError("invalid_identity", "Invalid expiry item ID") + parts = item_id.split(":") + if len(parts) != 4: + raise ExpiryError("invalid_identity", "Invalid expiry item ID") + validate_item_type(parts[0]) + identity = ItemIdentity(parts[0], *(decode_component(part) for part in parts[1:])) + validate_identity(identity) + return identity + + +def shard_for(item_id: str, shard_count: int) -> int: + if type(shard_count) is not int or shard_count <= 0: + raise ValueError("shard_count must be positive") + return xxhash.xxh64(item_id.encode("utf-8"), seed=HASH_SEED).intdigest() % shard_count + + +@dataclass(frozen=True) +class ExpiryKeys: + service_key: str + + def namespace(self) -> str: + return f"expiry:namespace:{self.service_key}" + + def metadata(self, item_id: str) -> str: + return f"expiry:metadata:{self.service_key}:{item_id}" + + def source(self, shard_id: int) -> str: + return f"expiry:publisher:source:{self.service_key}:{shard_id}" + + def control_tag(self, shard_id: int) -> str: + return f"{{{self.service_key}:{shard_id}}}" + + def progress(self, shard_id: int) -> str: + return f"expiry:publisher:progress:{self.control_tag(shard_id)}" + + def registry(self, shard_id: int) -> str: + return f"expiry:consumer:registry:{self.control_tag(shard_id)}" + + def owner(self, role: OwnerRole, shard_id: int, *, item_type: str | None = None) -> str: + if role is OwnerRole.SCHEDULER: + if item_type is not None: + raise ValueError("Scheduler owner must not include an item type") + return f"expiry:publisher:owner:{self.control_tag(shard_id)}" + if item_type is None: + raise ValueError("Callback owner requires an item type") + validate_item_type(item_type) + return f"expiry:consumer:owner:{item_type}:{self.control_tag(shard_id)}" + + def bucket(self, ref: BucketRef) -> str: + return ( + f"expiry:consumer:bucket:{self.service_key}:" + f"{ref.item_type}:{ref.shard_id}:{ref.minute_ms}" + ) + + def registry_member(self, ref: BucketRef) -> str: + return f"{ref.item_type}:{ref.shard_id}:{ref.minute_ms}" + + def parse_registry_member(self, value: str, *, expected_shard: int) -> BucketRef: + parts = value.split(":") + if len(parts) != 3: + raise ExpiryError("invalid_identity", "Invalid expiry registry member") + validate_item_type(parts[0]) + try: + shard_id, minute_ms = int(parts[1]), int(parts[2]) + except ValueError as exc: + raise ExpiryError("invalid_identity", "Invalid expiry registry member") from exc + if shard_id != expected_shard or minute_ms < 0 or minute_ms % MINUTE_MS: + raise ExpiryError("invalid_identity", "Invalid expiry registry member") + return BucketRef(parts[0], shard_id, minute_ms) diff --git a/openviking/scheduling/expiry/leases.py b/openviking/scheduling/expiry/leases.py new file mode 100644 index 0000000000..e812cc7aea --- /dev/null +++ b/openviking/scheduling/expiry/leases.py @@ -0,0 +1,101 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Redis token operations for temporary per-shard ownership.""" + +from __future__ import annotations + +import asyncio +import time +import uuid + +from .config import ExpiryConfig +from .metrics import ExpiryMetrics +from .models import OwnerRole, ShardLease, StoreOperationError, task_is_cancelling +from .redis_store import RedisExpiryStore + + +class ShardLeaseCoordinator: + def __init__( + self, + store: RedisExpiryStore, + config: ExpiryConfig, + metrics: ExpiryMetrics, + ) -> None: + self._store = store + self._config = config + self._metrics = metrics + + async def try_acquire( + self, role: OwnerRole, shard_id: int, *, item_type: str | None = None + ) -> ShardLease | None: + if role is OwnerRole.CALLBACK and item_type is None: + raise ValueError("Callback lease requires an item type") + if role is OwnerRole.SCHEDULER and item_type is not None: + raise ValueError("Scheduler lease must not include an item type") + started = time.monotonic() + token = uuid.uuid4().hex + owner = await self._store.try_acquire_owner( + role, + shard_id, + token, + self._config.owner_ttl_seconds, + item_type=item_type, + ) + if owner is None: + return None + lease = ShardLease(owner, asyncio.Event(), started + self._config.owner_ttl_seconds) + lease.ensure_active() + return lease + + async def renew(self, lease: ShardLease) -> bool: + """Never revive lost ownership, including a successful but late response.""" + try: + lease.ensure_active() + except StoreOperationError: + return False + started = time.monotonic() + previous_deadline = lease.valid_until_monotonic + try: + renewed = await self._store.renew_owner(lease.owner, self._config.owner_ttl_seconds) + except asyncio.CancelledError: + # Cancelling a request does not prove Redis did not execute it. + lease.lost.set() + raise + except Exception: + lease.lost.set() + return False + task = asyncio.current_task() + if ( + renewed is None + or lease.lost.is_set() + or task_is_cancelling(task) + or time.monotonic() >= previous_deadline + ): + lease.lost.set() + return False + if lease.revoked.is_set(): + return False + lease.owner = renewed + lease.valid_until_monotonic = started + self._config.owner_ttl_seconds + return True + + async def release(self, lease: ShardLease) -> None: + """Compare the token only after local work/renewal have terminated. + + Local revocation is expected during cleanup and is not ownership loss. + Failed or unknown release leaves the owner to expire; it never deletes a + successor's token. + """ + if time.monotonic() >= lease.valid_until_monotonic: + lease.lost.set() + if lease.lost.is_set(): + return + try: + await self._store.release_owner(lease.owner) + except asyncio.CancelledError: + lease.lost.set() + raise + except Exception as exc: + if isinstance(exc, StoreOperationError) and exc.outcome_unknown: + lease.lost.set() + self._metrics.count("timer_owner_release", outcome="failed") diff --git a/openviking/scheduling/expiry/lua/_common.lua b/openviking/scheduling/expiry/lua/_common.lua new file mode 100644 index 0000000000..69712a358a --- /dev/null +++ b/openviking/scheduling/expiry/lua/_common.lua @@ -0,0 +1,87 @@ +local BASE = 1099511627776 +local MAX_INT = 9007199254740991 + +local function decimal(n) + return string.format('%.0f', n) +end + +local function now_ms() + local t = redis.call('TIME') + return tonumber(t[1]) * 1000 + math.floor(tonumber(t[2]) / 1000) +end + +local function integer_arg(value, min_value, max_value) + if type(value) ~= 'string' or not string.match(value, '^%-?%d+$') then + error('invalid_integer') + end + local n = tonumber(value) + if not n or n ~= math.floor(n) or n < min_value or n > max_value or decimal(n) ~= value then + error('invalid_integer') + end + return n +end + +local function expect_type(key, allowed) + local actual = redis.call('TYPE', key)['ok'] + if actual ~= 'none' and actual ~= allowed then + error('wrong_key_type') + end +end + +local function require_owner(key, token) + if redis.call('GET', key) ~= token then + error('owner_lost') + end +end + +local function tuple_after(score, member, lower, last_score, last_member) + if score <= lower then return false end + if last_score == nil then return true end + return score > last_score or (score == last_score and member > last_member) +end + +local function page_after(key, lower, upper, last_score, last_member, limit) + local lo = nil + if last_score ~= nil then + local current = redis.call('ZSCORE', key, last_member) + if current and tonumber(current) == last_score and last_score > lower then + local rank = redis.call('ZRANK', key, last_member) + if rank then lo = rank + 1 end + end + end + if lo == nil then + local size = redis.call('ZCARD', key) + lo = 0 + local hi = size + while lo < hi do + local mid = math.floor((lo + hi) / 2) + local entry = redis.call('ZRANGE', key, mid, mid, 'WITHSCORES') + if #entry == 0 then + hi = mid + else + local score = tonumber(entry[2]) + if tuple_after(score, entry[1], lower, last_score, last_member) then + hi = mid + else + lo = mid + 1 + end + end + end + end + local raw = redis.call('ZRANGE', key, lo, lo + limit, 'WITHSCORES') + local items = {} + local count = 0 + for i = 1, #raw, 2 do + local score = tonumber(raw[i + 1]) + if score > upper then break end + count = count + 1 + items[count] = {raw[i], decimal(score)} + end + local next_value = false + if count > limit then + items[count] = nil + count = count - 1 + next_value = {items[count][2], items[count][1]} + end + return items, next_value +end diff --git a/openviking/scheduling/expiry/lua/cas_zset_member.lua b/openviking/scheduling/expiry/lua/cas_zset_member.lua new file mode 100644 index 0000000000..b08318743a --- /dev/null +++ b/openviking/scheduling/expiry/lua/cas_zset_member.lua @@ -0,0 +1,16 @@ +local action = ARGV[1] +local expected = integer_arg(ARGV[3], 0, 256 * BASE - 1) +local replacement = nil +if action == 'set' then + replacement = integer_arg(ARGV[4], 1, 256 * BASE - 1) +elseif action ~= 'delete' or ARGV[4] ~= '' then + return redis.error_reply('invalid_cas_action') +end + +local current = redis.call('ZSCORE', KEYS[1], ARGV[2]) +if not current or tonumber(current) ~= expected then return 0 end +if action == 'set' then + redis.call('ZADD', KEYS[1], decimal(replacement), ARGV[2]) + return 1 +end +return redis.call('ZREM', KEYS[1], ARGV[2]) diff --git a/openviking/scheduling/expiry/lua/cleanup_source_page.lua b/openviking/scheduling/expiry/lua/cleanup_source_page.lua new file mode 100644 index 0000000000..babba0ee09 --- /dev/null +++ b/openviking/scheduling/expiry/lua/cleanup_source_page.lua @@ -0,0 +1,8 @@ +local through = integer_arg(ARGV[1], 0, MAX_INT) +local limit = integer_arg(ARGV[2], 1, 4096) +local due = redis.call('ZCOUNT', KEYS[1], '-inf', decimal(through)) +local removed = math.min(due, limit) +if removed > 0 then + redis.call('ZREMRANGEBYRANK', KEYS[1], 0, removed - 1) +end +return removed diff --git a/openviking/scheduling/expiry/lua/compare_owner.lua b/openviking/scheduling/expiry/lua/compare_owner.lua new file mode 100644 index 0000000000..5164312a52 --- /dev/null +++ b/openviking/scheduling/expiry/lua/compare_owner.lua @@ -0,0 +1,13 @@ +local action = ARGV[1] +local ttl = nil +if action == 'renew' then + ttl = integer_arg(ARGV[3], 1, MAX_INT) +elseif action ~= 'release' or ARGV[3] ~= '' then + return redis.error_reply('invalid_owner_action') +end + +if redis.call('GET', KEYS[1]) ~= ARGV[2] then return 0 end +if action == 'renew' then + return redis.call('PEXPIRE', KEYS[1], ttl) +end +return redis.call('DEL', KEYS[1]) diff --git a/openviking/scheduling/expiry/lua/defer_registration.lua b/openviking/scheduling/expiry/lua/defer_registration.lua new file mode 100644 index 0000000000..c09fdb6430 --- /dev/null +++ b/openviking/scheduling/expiry/lua/defer_registration.lua @@ -0,0 +1,3 @@ +require_owner(KEYS[2], ARGV[1]) +local next_poll_at_ms = integer_arg(ARGV[3], 0, MAX_INT) +return redis.call('ZADD', KEYS[1], 'XX', decimal(next_poll_at_ms), ARGV[2]) diff --git a/openviking/scheduling/expiry/lua/ensure_metadata.lua b/openviking/scheduling/expiry/lua/ensure_metadata.lua new file mode 100644 index 0000000000..d770d14761 --- /dev/null +++ b/openviking/scheduling/expiry/lua/ensure_metadata.lua @@ -0,0 +1,19 @@ +expect_type(KEYS[1], 'hash') +if #ARGV[1] == 0 or #ARGV[2] == 0 or #ARGV[3] == 0 or #ARGV[4] == 0 then + return redis.error_reply('invalid_identity') +end +local ttl = integer_arg(ARGV[5], 1, MAX_INT) +local fields = {'item_type', 'local_id', 'account_id', 'user_id'} +local exists = redis.call('EXISTS', KEYS[1]) == 1 +if exists then + for i = 1, 4 do + local current = redis.call('HGET', KEYS[1], fields[i]) + if current ~= ARGV[i] then return redis.error_reply('metadata_conflict') end + end +else + redis.call('HSET', KEYS[1], + fields[1], ARGV[1], fields[2], ARGV[2], + fields[3], ARGV[3], fields[4], ARGV[4]) +end +redis.call('PEXPIRE', KEYS[1], ttl) +return exists and 'refreshed' or 'created' diff --git a/openviking/scheduling/expiry/lua/publish_minute.lua b/openviking/scheduling/expiry/lua/publish_minute.lua new file mode 100644 index 0000000000..360246a232 --- /dev/null +++ b/openviking/scheduling/expiry/lua/publish_minute.lua @@ -0,0 +1,15 @@ +require_owner(KEYS[3], ARGV[1]) +local expected = integer_arg(ARGV[2], 0, MAX_INT) +local minute = integer_arg(ARGV[3], 0, MAX_INT) +local count = integer_arg(ARGV[4], 1, 1024) +if minute % 60000 ~= 0 or #ARGV ~= 4 + count then return redis.error_reply('invalid_publish') end +local raw_current = redis.call('GET', KEYS[2]) +if not raw_current then return redis.error_reply('namespace_corrupt') end +local current = tonumber(raw_current) +if current >= minute then return {'already_published', decimal(current)} end +if current ~= expected or minute <= current then return redis.error_reply('progress_conflict') end +for i = 1, count do + redis.call('ZADD', KEYS[1], decimal(minute), ARGV[4 + i]) +end +redis.call('SET', KEYS[2], decimal(minute)) +return {'published', decimal(minute)} diff --git a/openviking/scheduling/expiry/lua/read_zset_page.lua b/openviking/scheduling/expiry/lua/read_zset_page.lua new file mode 100644 index 0000000000..8bc105695b --- /dev/null +++ b/openviking/scheduling/expiry/lua/read_zset_page.lua @@ -0,0 +1,13 @@ +local lower = integer_arg(ARGV[1], -1, MAX_INT) +local upper = nil +if ARGV[2] == '' then + upper = now_ms() +else + upper = integer_arg(ARGV[2], 0, MAX_INT) +end +local after_score = integer_arg(ARGV[3], 0, MAX_INT) +local limit = integer_arg(ARGV[5], 1, 4096) +local items, next_value = page_after( + KEYS[1], lower, upper, after_score, ARGV[4], limit +) +return cjson.encode({items=items, next=next_value}) diff --git a/openviking/scheduling/expiry/lua/upsert_resource.lua b/openviking/scheduling/expiry/lua/upsert_resource.lua new file mode 100644 index 0000000000..1bdcac33ea --- /dev/null +++ b/openviking/scheduling/expiry/lua/upsert_resource.lua @@ -0,0 +1,7 @@ +if #ARGV[1] == 0 then return redis.error_reply('invalid_identity') end +local requested = integer_arg(ARGV[2], 0, MAX_INT) +local accepted = math.max(requested, now_ms() + 1) +local max_supported_due = math.floor(MAX_INT / 60000) * 60000 +if accepted > max_supported_due then return redis.error_reply('invalid_report') end +redis.call('ZADD', KEYS[1], decimal(accepted), ARGV[1]) +return decimal(accepted) diff --git a/openviking/scheduling/expiry/metrics.py b/openviking/scheduling/expiry/metrics.py new file mode 100644 index 0000000000..cbbe16df5f --- /dev/null +++ b/openviking/scheduling/expiry/metrics.py @@ -0,0 +1,75 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Low-cardinality metrics and structured timer logging.""" + +from __future__ import annotations + +import logging +from collections.abc import Callable, Mapping + +from .keys import decode_item_id +from .models import Claim + +logger = logging.getLogger(__name__) + + +class ExpiryMetrics: + def __init__(self, sink: Callable[[str, float, Mapping[str, str]], None] | None = None) -> None: + self._sink = sink + + def _emit(self, name: str, value: float, labels: Mapping[str, str]) -> None: + if self._sink is None: + return + try: + self._sink(name, value, labels) + except Exception: + logger.debug("expiry metric sink failed", exc_info=True) + + def count(self, name: str, value: int = 1, *, item_type: str = "", outcome: str = "") -> None: + self._emit(name, float(value), {"item_type": item_type, "outcome": outcome}) + + def observe(self, name: str, value: float, *, item_type: str = "", outcome: str = "") -> None: + self._emit(name, value, {"item_type": item_type, "outcome": outcome}) + + def gauge(self, name: str, value: float, *, item_type: str = "") -> None: + self._emit(name, value, {"item_type": item_type}) + + def log_invalid_source_items( + self, + *, + shard_id: int, + minute_ms: int, + reason: str, + count: int, + ) -> None: + logger.error( + "expiry discarded invalid source items shard_id=%d minute_ms=%d " + "reason=%s count=%d", + shard_id, + minute_ms, + reason, + count, + ) + + def log_exhausted( + self, + claim: Claim, + *, + max_attempts: int, + reason: str, + ) -> None: + identity = decode_item_id(claim.item_id) + logger.error( + "expiry callback exhausted item_type=%s account_id=%s user_id=%s " + "local_id=%s shard_id=%d minute_ms=%d claim_score=%d max_attempts=%d " + "reason=%s", + identity.item_type, + identity.account_id, + identity.user_id, + identity.local_id, + claim.bucket.shard_id, + claim.bucket.minute_ms, + claim.claim_score, + max_attempts, + reason, + ) diff --git a/openviking/scheduling/expiry/models.py b/openviking/scheduling/expiry/models.py new file mode 100644 index 0000000000..c81e56c344 --- /dev/null +++ b/openviking/scheduling/expiry/models.py @@ -0,0 +1,236 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Internal models and pure validation helpers for expiry scheduling.""" + +from __future__ import annotations + +import asyncio +import re +import time +from dataclasses import dataclass, field +from enum import Enum +from typing import Awaitable, Callable, Literal, Protocol, TypeVar + +from .api import ExpiryError, ExpiryMode, ExpiryRegistration, ItemIdentity + +T = TypeVar("T") +MINUTE_MS = 60_000 +MIN_METADATA_TTL_MS = 86_400_000 +SCORE_TIME_BASE = 1 << 40 +MAX_ATTEMPT = 255 +MAX_EXACT_INTEGER = (1 << 53) - 1 +_ITEM_TYPE_RE = re.compile(r"[a-z][a-z0-9_]{0,63}\Z") + + +class StoreOperationError(ExpiryError): + def __init__( + self, + code: str, + message: str, + *, + retryable: bool = False, + outcome_unknown: bool = False, + ) -> None: + super().__init__(code, message, retryable=retryable) + self.outcome_unknown = outcome_unknown + + +class OwnerRole(str, Enum): + SCHEDULER = "scheduler" + CALLBACK = "callback" + + +@dataclass(frozen=True) +class ItemMetadata: + identity: ItemIdentity + + +@dataclass(frozen=True) +class NamespaceSpec: + shard_count: int + hash_name: str + hash_seed: int + identity_encoding: str + + +@dataclass(frozen=True) +class OwnerToken: + role: OwnerRole + shard_id: int + token: str + expires_at_ms: int + item_type: str | None = None + + +@dataclass +class ShardLease: + owner: OwnerToken + lost: asyncio.Event + valid_until_monotonic: float + revoked: asyncio.Event = field(default_factory=asyncio.Event) + + def ensure_active(self) -> None: + if self.lost.is_set() or time.monotonic() >= self.valid_until_monotonic: + self.lost.set() + raise StoreOperationError("owner_lost", "Shard lease is no longer valid") + if self.revoked.is_set(): + raise StoreOperationError( + "owner_lost", "Shard lease is no longer valid (execution revoked)" + ) + + +@dataclass(frozen=True) +class ShardWindow: + progress_ms: int + upper_ms: int + redis_now_ms: int + + +@dataclass(frozen=True) +class SourceCursor: + due_at_ms: int + item_id: str + + +@dataclass(frozen=True) +class SourceItem: + item_id: str + due_at_ms: int + + +@dataclass(frozen=True) +class SourcePage: + items: tuple[SourceItem, ...] + next_cursor: SourceCursor | None + + +@dataclass(frozen=True) +class BucketRef: + item_type: str + shard_id: int + minute_ms: int + + +@dataclass(frozen=True) +class RegistryCursor: + score_ms: int + member: str + + +@dataclass(frozen=True) +class RegistryPage: + buckets: tuple[BucketRef, ...] + next_cursor: RegistryCursor | None + + +@dataclass(frozen=True) +class PublishResult: + status: Literal["published", "already_published"] + progress_ms: int + + +@dataclass(frozen=True) +class Claim: + bucket: BucketRef + item_id: str + claim_score: int + + +@dataclass(frozen=True) +class BucketItem: + item_id: str + score: int + + +@dataclass(frozen=True) +class BucketPage: + items: tuple[BucketItem, ...] + + +class AsyncExecutor(Protocol): + async def run(self, operation: Callable[[], Awaitable[T]]) -> T: ... + + +def task_is_cancelling(task: asyncio.Task | None) -> bool: + if task is None: + return False + cancelling = getattr(task, "cancelling", None) + return bool(cancelling()) if callable(cancelling) else task.cancelled() + + +def validate_item_type(item_type: str) -> None: + if not isinstance(item_type, str) or _ITEM_TYPE_RE.fullmatch(item_type) is None: + raise ExpiryError("invalid_identity", "Invalid resource item type") + + +def validate_identity(identity: ItemIdentity) -> None: + if not isinstance(identity, ItemIdentity): + raise ExpiryError("invalid_identity", "Invalid resource identity") + validate_item_type(identity.item_type) + for name in ("account_id", "user_id", "local_id"): + value = getattr(identity, name) + if not isinstance(value, str) or not value or "\x00" in value: + raise ExpiryError("invalid_identity", f"Invalid resource identity field: {name}") + + +def validate_report(task: ExpiryRegistration) -> None: + if not isinstance(task, ExpiryRegistration): + raise ExpiryError("invalid_report", "Invalid resource report") + validate_identity(task.identity) + if task.mode is not ExpiryMode.IDLE_ONESHOT: + raise ExpiryError("unsupported_mode", f"Unsupported timer mode: {task.mode}") + if type(task.due_at_ms) is not int or not 0 <= task.due_at_ms <= MAX_EXACT_INTEGER: + raise ExpiryError("invalid_report", "Due time is outside the supported range") + metadata_ttl_ms(task.timeout_seconds) + bucket_minute_ms(task.due_at_ms) + + +def minute_start_ms(timestamp_ms: int) -> int: + return timestamp_ms // MINUTE_MS * MINUTE_MS + + +def scan_cutoff_ms(now_ms: int) -> int: + return minute_start_ms(now_ms) + + +def bucket_minute_ms(due_at_ms: int) -> int: + if type(due_at_ms) is not int or not 0 <= due_at_ms <= MAX_EXACT_INTEGER: + raise ExpiryError("invalid_report", "Due time is outside the supported range") + minute = ((due_at_ms + MINUTE_MS - 1) // MINUTE_MS) * MINUTE_MS + if minute > MAX_EXACT_INTEGER: + raise ExpiryError("invalid_report", "Due time is outside the supported range") + return minute + + +def metadata_ttl_ms(timeout_seconds: int) -> int: + if type(timeout_seconds) is not int or not 0 < timeout_seconds <= MAX_EXACT_INTEGER // 2000: + raise ExpiryError("invalid_report", "Timeout is outside the supported range") + return max(timeout_seconds * 2000, MIN_METADATA_TTL_MS) + + +def encode_claim_score(attempt: int, started_at_seconds: int) -> int: + if type(attempt) is not int or not 0 <= attempt <= MAX_ATTEMPT: + raise ValueError("attempt out of range") + if type(started_at_seconds) is not int or not 0 < started_at_seconds < SCORE_TIME_BASE: + raise ValueError("claim timestamp out of range") + return attempt * SCORE_TIME_BASE + started_at_seconds + + +def decode_claim_score(score: int) -> tuple[int, int]: + if type(score) is not int or not 0 < score < (MAX_ATTEMPT + 1) * SCORE_TIME_BASE: + raise ValueError("invalid claim score") + attempt, started_at = divmod(score, SCORE_TIME_BASE) + if started_at == 0: + raise ValueError("invalid claim timestamp") + return attempt, started_at + + +def claim_layer_bounds(attempt: int) -> tuple[int, int]: + if type(attempt) is not int or not 0 <= attempt <= MAX_ATTEMPT: + raise ValueError("attempt out of range") + lower = 0 if attempt == 0 else attempt * SCORE_TIME_BASE + 1 + return lower, (attempt + 1) * SCORE_TIME_BASE - 1 + + +def encode_ready_claim_score(attempt: int) -> int: + return encode_claim_score(attempt, 1) diff --git a/openviking/scheduling/expiry/publisher.py b/openviking/scheduling/expiry/publisher.py new file mode 100644 index 0000000000..aa66fc3e3d --- /dev/null +++ b/openviking/scheduling/expiry/publisher.py @@ -0,0 +1,223 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Expiry publisher: source buckets to published callback buckets.""" + +from __future__ import annotations + +import asyncio +import time +from collections import defaultdict +from typing import Sequence + +from .api import SUPPORTED_ITEM_TYPES, ExpiryError +from .config import ExpiryConfig +from .keys import decode_item_id +from .metrics import ExpiryMetrics +from .models import ( + MINUTE_MS, + BucketRef, + OwnerRole, + PublishResult, + StoreOperationError, + bucket_minute_ms, +) +from .redis_store import RedisExpiryStore +from .shard_runner import ShardExecution, ShardRunner + + +async def _wait_until_next_round(stop: asyncio.Event, delay: float) -> None: + try: + await asyncio.wait_for(stop.wait(), delay) + except asyncio.TimeoutError: + pass + + +class ExpiryPublisher: + def __init__( + self, + store: RedisExpiryStore, + runner: ShardRunner, + config: ExpiryConfig, + metrics: ExpiryMetrics, + ) -> None: + self._store = store + self._runner = runner + self._config = config + self._metrics = metrics + + async def run(self, stop: asyncio.Event) -> None: + """Own shard traversal and scheduling cadence, not lease lifecycle.""" + while not stop.is_set(): + for shard_id in range(self._config.shard_count): + if stop.is_set(): + break + try: + await self._runner.run_if_acquired( + OwnerRole.SCHEDULER, shard_id, self.process_shard, stop + ) + except StoreOperationError as exc: + if exc.code == "shard_not_stopped": + raise + self._metrics.count("timer_shard_round", outcome="failed") + except Exception: + self._metrics.count("timer_shard_round", outcome="failed") + if not stop.is_set(): + delay = max(0.05, 60.0 - time.time() % 60.0) + await _wait_until_next_round(stop, delay) + + async def process_shard(self, lease: ShardExecution, stop: asyncio.Event) -> None: + while not stop.is_set(): + lease.ensure_active() + window = await self._store.read_scan_window(lease.owner.shard_id) + progress = window.progress_ms + lease.ensure_active() + await self._cleanup_source(lease, stop, progress) + if window.upper_ms <= progress: + return + after_ms = progress + published_any = False + while after_ms < window.upper_ms and not stop.is_set(): + lease.ensure_active() + page = await self._store.read_source_page( + lease.owner.shard_id, + lower_exclusive_ms=after_ms, + upper_inclusive_ms=window.upper_ms, + after=None, + limit=1, + ) + if not page.items: + break + minute = bucket_minute_ms(page.items[0].due_at_ms) + result = await self._process_minute( + lease, + stop, + progress_ms=progress, + minute_ms=minute, + ) + if result is not None: + progress = result + published_any = True + after_ms = minute + if not published_any: + return + + async def _process_minute( + self, + lease: ShardExecution, + stop: asyncio.Event, + *, + progress_ms: int, + minute_ms: int, + ) -> int | None: + lease.ensure_active() + await self._store.reset_unpublished_minute( + lease.owner, + minute_ms=minute_ms, + expected_progress_ms=progress_ms, + item_types=tuple(sorted(SUPPORTED_ITEM_TYPES)), + ) + cursor = None + types_with_data: set[str] = set() + while True: + lease.ensure_active() + page = await self._store.read_source_page( + lease.owner.shard_id, + lower_exclusive_ms=max(progress_ms, minute_ms - MINUTE_MS), + upper_inclusive_ms=minute_ms, + after=cursor, + limit=self._config.scan_batch_size, + ) + by_type: dict[str, list[str]] = defaultdict(list) + invalid_counts: dict[str, int] = defaultdict(int) + for item in page.items: + try: + identity = decode_item_id(item.item_id) + except ExpiryError: + invalid_counts["invalid_identity"] += 1 + continue + if identity.item_type not in SUPPORTED_ITEM_TYPES: + invalid_counts["unsupported_type"] += 1 + continue + by_type[identity.item_type].append(item.item_id) + for reason, count in sorted(invalid_counts.items()): + self._metrics.count( + "timer_source_invalid", + value=count, + outcome=reason, + ) + self._metrics.log_invalid_source_items( + shard_id=lease.owner.shard_id, + minute_ms=minute_ms, + reason=reason, + count=count, + ) + for item_type, item_ids in by_type.items(): + lease.ensure_active() + ref = BucketRef(item_type, lease.owner.shard_id, minute_ms) + await self._store.stage_items(lease.owner, ref, item_ids) + types_with_data.add(item_type) + cursor = page.next_cursor + if cursor is None: + break + if not types_with_data: + # Invalid members can be removed without advancing the published watermark. + # Their metadata lives in other cluster slots and expires by TTL. + await self._cleanup_source(lease, stop, minute_ms) + return None + refs = tuple( + BucketRef(item_type, lease.owner.shard_id, minute_ms) + for item_type in sorted(types_with_data) + ) + result = await self._publish_completed_minute( + lease, + expected_progress_ms=progress_ms, + minute_ms=minute_ms, + buckets=refs, + ) + lease.ensure_active() + # This also removes skipped invalid members; their metadata expires by TTL. + await self._cleanup_source(lease, stop, result.progress_ms) + return result.progress_ms + + async def _cleanup_source( + self, lease: ShardExecution, stop: asyncio.Event, through_ms: int + ) -> None: + limit = self._config.cleanup_batch_size + while not stop.is_set(): + lease.ensure_active() + removed = await self._store.cleanup_source_page( + lease.owner.shard_id, + through_ms, + limit=limit, + ) + if removed < limit: + return + + async def _publish_completed_minute( + self, + lease: ShardExecution, + *, + expected_progress_ms: int, + minute_ms: int, + buckets: Sequence[BucketRef], + ) -> PublishResult: + while True: + lease.ensure_active() + try: + return await self._store.publish_minute( + lease.owner, + expected_progress_ms=expected_progress_ms, + minute_ms=minute_ms, + buckets=buckets, + ) + except StoreOperationError as exc: + if not exc.outcome_unknown: + raise + lease.ensure_active() + progress = await self._store.read_progress(lease.owner.shard_id) + if progress >= minute_ms: + return PublishResult("already_published", progress) + if progress != expected_progress_ms: + raise StoreOperationError( + "progress_conflict", "Timer progress changed after uncertain publish" + ) from exc diff --git a/openviking/scheduling/expiry/redis_store.py b/openviking/scheduling/expiry/redis_store.py new file mode 100644 index 0000000000..cfade55136 --- /dev/null +++ b/openviking/scheduling/expiry/redis_store.py @@ -0,0 +1,784 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Redis persistence for distributed expiry scheduling.""" + +from __future__ import annotations + +import asyncio +import hashlib +import json +import time +from importlib import resources +from typing import Any, Literal, Sequence, cast +from urllib.parse import urlsplit + +from .api import ExpiryRegistration, ExpiryRegistrationResult, ItemIdentity +from .config import ExpiryRedisConfig +from .keys import ExpiryKeys +from .metrics import ExpiryMetrics +from .models import ( + MAX_EXACT_INTEGER, + SCORE_TIME_BASE, + AsyncExecutor, + BucketItem, + BucketPage, + BucketRef, + Claim, + ItemMetadata, + NamespaceSpec, + OwnerRole, + OwnerToken, + PublishResult, + RegistryCursor, + RegistryPage, + ShardWindow, + SourceCursor, + SourceItem, + SourcePage, + StoreOperationError, + claim_layer_bounds, + decode_claim_score, + encode_claim_score, + encode_ready_claim_score, + metadata_ttl_ms, +) + +RetryPolicy = Literal["read", "idempotent", "resolve_by_caller"] + + +class RedisExpiryStore: + def __init__( + self, + config: ExpiryRedisConfig, + keys: ExpiryKeys, + io: AsyncExecutor, + metrics: ExpiryMetrics, + ) -> None: + self._config = config + self._keys = keys + self._io = io + self._metrics = metrics + self._client: Any = None + self._scripts: dict[str, str] = {} + self._script_shas: dict[str, str] = {} + self._ready = False + + async def open(self) -> None: + await self._io.run(self._open_on_loop) + + async def close(self) -> None: + self._ready = False + if self._client is not None: + await self._io.run(self._close_on_loop) + + async def ensure_namespace(self, spec: NamespaceSpec) -> None: + key = self._keys.namespace() + expected = { + "shard_count": str(spec.shard_count), + "hash_name": spec.hash_name, + "hash_seed": str(spec.hash_seed), + "identity_encoding": spec.identity_encoding, + } + current = await self._command("hgetall", key) + initializing = True + if current: + normalized = {str(k): str(v) for k, v in current.items()} + for field, value in expected.items(): + if normalized.get(field) != value: + raise StoreOperationError( + "namespace_mismatch", f"Timer namespace mismatch: {field}" + ) + state = normalized.get("state") + if state not in {"INITIALIZING", "READY"}: + raise StoreOperationError("namespace_corrupt", "Invalid timer namespace state") + initializing = state == "INITIALIZING" + else: + created = dict(expected) + created["state"] = "INITIALIZING" + await self._command("hset", key, mapping=created, retry="idempotent") + for shard_id in range(spec.shard_count): + progress_key = self._keys.progress(shard_id) + if initializing: + await self._command("set", progress_key, "0", nx=True, retry="idempotent") + elif await self._command("get", progress_key) is None: + raise StoreOperationError( + "namespace_corrupt", f"Timer progress is missing for shard {shard_id}" + ) + if initializing: + await self._command("hset", key, "state", "READY", retry="idempotent") + self._ready = True + + async def ensure_metadata(self, item_id: str, task: ExpiryRegistration) -> None: + identity = task.identity + await self._eval( + "ensure_metadata", + [self._keys.metadata(item_id)], + [ + identity.item_type, + identity.local_id, + identity.account_id, + identity.user_id, + metadata_ttl_ms(task.timeout_seconds), + ], + retry="idempotent", + ) + + async def read_metadata(self, item_id: str) -> ItemMetadata | None: + raw = await self._command("hgetall", self._keys.metadata(item_id)) + if not raw: + return None + required = ("item_type", "account_id", "user_id", "local_id") + if any(name not in raw for name in required): + raise StoreOperationError("metadata_corrupt", "Timer metadata is incomplete") + return ItemMetadata(ItemIdentity(*(str(raw[name]) for name in required))) + + async def upsert_resource( + self, shard_id: int, item_id: str, task: ExpiryRegistration + ) -> ExpiryRegistrationResult: + raw = await self._eval( + "upsert_resource", + [self._keys.source(shard_id)], + [item_id, task.due_at_ms], + retry="resolve_by_caller", + ) + return ExpiryRegistrationResult(int(raw)) + + async def unregister_resource( + self, + shard_id: int, + item_id: str, + *, + delete_metadata: bool = True, + ) -> None: + await self._command("zrem", self._keys.source(shard_id), item_id, retry="idempotent") + if delete_metadata: + await self._single_key_command( + "DEL", + self._keys.metadata(item_id), + retry="idempotent", + ) + + async def read_progress(self, shard_id: int) -> int: + raw = await self._command("get", self._keys.progress(shard_id)) + if raw is None: + raise StoreOperationError("namespace_corrupt", "Timer progress is missing") + try: + return int(raw) + except (TypeError, ValueError) as exc: + raise StoreOperationError("namespace_corrupt", "Timer progress is invalid") from exc + + async def read_scan_window(self, shard_id: int) -> ShardWindow: + progress = await self.read_progress(shard_id) + source_now_ms = await self._read_time_ms_for_key(self._keys.source(shard_id)) + upper = source_now_ms // 60_000 * 60_000 + if upper < progress: + raise StoreOperationError( + "clock_regression", "Source time is behind published progress" + ) + return ShardWindow(progress, upper, source_now_ms) + + async def read_source_page( + self, + shard_id: int, + *, + lower_exclusive_ms: int, + upper_inclusive_ms: int, + after: SourceCursor | None, + limit: int, + ) -> SourcePage: + items, next_raw = await self._read_zset_page( + self._keys.source(shard_id), + lower_exclusive=lower_exclusive_ms, + upper_inclusive=upper_inclusive_ms, + after_score=None if after is None else after.due_at_ms, + after_member=None if after is None else after.item_id, + limit=limit, + ) + page_items = tuple(SourceItem(item, score) for item, score in items) + cursor = SourceCursor(*next_raw) if next_raw else None + return SourcePage(page_items, cursor) + + async def reset_unpublished_minute( + self, + owner: OwnerToken, + *, + minute_ms: int, + expected_progress_ms: int, + item_types: Sequence[str], + ) -> None: + await self.assert_owner(owner) + current = await self.read_progress(owner.shard_id) + if current != expected_progress_ms or minute_ms <= current: + raise StoreOperationError("progress_conflict", "Timer progress changed") + keys = [ + self._keys.bucket(BucketRef(item_type, owner.shard_id, minute_ms)) + for item_type in item_types + ] + # Bucket keys deliberately have no shared cluster hash tag. Delete them + # one at a time, bypassing redis-py's cross-slot pipeline helper. + for key in keys: + await self._single_key_command("UNLINK", key, retry="idempotent") + + async def stage_items( + self, owner: OwnerToken, bucket: BucketRef, item_ids: Sequence[str] + ) -> None: + if not item_ids: + return + await self.assert_owner(owner) + mapping = dict.fromkeys(item_ids, 0) + await self._command("zadd", self._keys.bucket(bucket), mapping, retry="resolve_by_caller") + + async def publish_minute( + self, + owner: OwnerToken, + *, + expected_progress_ms: int, + minute_ms: int, + buckets: Sequence[BucketRef], + ) -> PublishResult: + members = [self._keys.registry_member(ref) for ref in buckets] + raw = await self._eval( + "publish_minute", + [ + self._keys.registry(owner.shard_id), + self._keys.progress(owner.shard_id), + self._keys.owner(OwnerRole.SCHEDULER, owner.shard_id), + ], + [owner.token, expected_progress_ms, minute_ms, len(members), *members], + retry="resolve_by_caller", + ) + status = str(raw[0]) + if status not in {"published", "already_published"}: + raise StoreOperationError("namespace_corrupt", "Invalid timer publish result") + return PublishResult( + cast(Literal["published", "already_published"], status), + int(raw[1]), + ) + + async def cleanup_source_page(self, shard_id: int, through_ms: int, *, limit: int) -> int: + return int( + await self._eval( + "cleanup_source_page", + [self._keys.source(shard_id)], + [through_ms, limit], + retry="idempotent", + ) + ) + + async def read_registry_page( + self, shard_id: int, *, after: RegistryCursor | None, limit: int + ) -> RegistryPage: + items, next_raw = await self._read_zset_page( + self._keys.registry(shard_id), + lower_exclusive=-1, + upper_inclusive=None, + after_score=None if after is None else after.score_ms, + after_member=None if after is None else after.member, + limit=limit, + ) + buckets = tuple( + self._keys.parse_registry_member(str(member), expected_shard=shard_id) + for member, _score in items + ) + cursor = RegistryCursor(*next_raw) if next_raw else None + return RegistryPage(buckets, cursor) + + async def read_time_ms(self) -> int: + return self._decode_redis_time(await self._command("time")) + + @staticmethod + def _decode_redis_time(raw: Any) -> int: + try: + seconds, microseconds = (int(value) for value in raw) + except (TypeError, ValueError) as exc: + raise StoreOperationError("namespace_corrupt", "Invalid Redis time") from exc + if not 0 < seconds < SCORE_TIME_BASE or not 0 <= microseconds < 1_000_000: + raise StoreOperationError("namespace_corrupt", "Redis time is outside supported range") + return seconds * 1000 + microseconds // 1000 + + async def _read_time_ms_for_key(self, key: str) -> int: + """Read TIME from the primary owning key so source cutoffs share its clock.""" + + async def operation() -> Any: + if self._client is None: + raise StoreOperationError("not_started", "Redis timer store is closed") + kwargs: dict[str, object] = {} + if self._config.mode == "cluster": + node = self._client.get_node_from_key(key) + if node is None: + raise RuntimeError("Redis Cluster has no primary for the timer key") + kwargs["target_nodes"] = node + return await self._client.time(**kwargs) + + raw = await self._io.run(lambda: self._run_with_retry(operation, retry="read")) + return self._decode_redis_time(raw) + + async def _read_zset_page( + self, + key: str, + *, + lower_exclusive: int, + upper_inclusive: int | None, + after_score: int | None, + after_member: str | None, + limit: int, + ) -> tuple[tuple[tuple[str, int], ...], tuple[int, str] | None]: + if after_score is None: + upper = ( + await self._read_time_ms_for_key(key) + if upper_inclusive is None + else upper_inclusive + ) + raw_items = await self._command( + "zrangebyscore", + key, + f"({lower_exclusive}", + upper, + start=0, + num=limit + 1, + withscores=True, + ) + items = self._decode_zset_items( + raw_items, + lower_exclusive=lower_exclusive, + upper_inclusive=upper, + ) + has_more = len(items) > limit + items = items[:limit] + next_cursor = (items[-1][1], items[-1][0]) if has_more else None + return items, next_cursor + + raw = await self._eval( + "read_zset_page", + [key], + [ + lower_exclusive, + "" if upper_inclusive is None else upper_inclusive, + after_score, + after_member or "", + limit, + ], + retry="read", + ) + data = json.loads(raw) + items = tuple((str(member), int(score)) for member, score in data["items"]) + next_raw = data.get("next") + next_cursor = (int(next_raw[0]), str(next_raw[1])) if next_raw else None + return items, next_cursor + + @staticmethod + def _decode_zset_items( + raw_items: Any, + *, + lower_exclusive: int, + upper_inclusive: int, + ) -> tuple[tuple[str, int], ...]: + items: list[tuple[str, int]] = [] + try: + for member, raw_score in raw_items: + numeric_score = float(raw_score) + score = int(numeric_score) + if numeric_score != score or score <= lower_exclusive or score > upper_inclusive: + raise ValueError + items.append((str(member), score)) + except (TypeError, ValueError) as exc: + raise StoreOperationError( + "namespace_corrupt", "Timer sorted set contains an invalid score" + ) from exc + return tuple(items) + + async def read_bucket_page( + self, + bucket: BucketRef, + *, + attempt: int, + limit: int, + ) -> BucketPage: + lower, upper = claim_layer_bounds(attempt) + raw = await self._command( + "zrangebyscore", + self._keys.bucket(bucket), + lower, + upper, + start=0, + num=limit, + withscores=True, + ) + items: list[BucketItem] = [] + try: + for item_id, raw_score in raw: + numeric_score = float(raw_score) + score = int(numeric_score) + if numeric_score != score or not lower <= score <= upper: + raise ValueError + items.append(BucketItem(str(item_id), score)) + except (TypeError, ValueError) as exc: + raise StoreOperationError( + "namespace_corrupt", "Timer bucket contains an invalid score" + ) from exc + return BucketPage(tuple(items)) + + async def defer_registration( + self, + owner: OwnerToken, + bucket: BucketRef, + *, + next_poll_at_ms: int, + ) -> bool: + if ( + owner.role is not OwnerRole.CALLBACK + or owner.item_type != bucket.item_type + or owner.shard_id != bucket.shard_id + ): + raise StoreOperationError( + "invalid_owner", "Callback owner does not match the timer bucket" + ) + raw = await self._eval( + "defer_registration", + [ + self._keys.registry(bucket.shard_id), + self._keys.owner( + owner.role, + owner.shard_id, + item_type=owner.item_type, + ), + ], + [owner.token, self._keys.registry_member(bucket), next_poll_at_ms], + retry="idempotent", + ) + return bool(int(raw)) + + async def claim( + self, + bucket: BucketRef, + item: BucketItem, + *, + attempt: int, + started_at_seconds: int, + ) -> Claim | None: + if item.score == 0: + if attempt != 0: + raise StoreOperationError( + "namespace_corrupt", "Unclaimed expiry item is outside attempt zero" + ) + else: + item_attempt, _ = decode_claim_score(item.score) + if item_attempt != attempt: + raise StoreOperationError( + "namespace_corrupt", "Expiry item is outside its expected attempt layer" + ) + claim_score = encode_claim_score(attempt, started_at_seconds) + raw = await self._eval( + "cas_zset_member", + [self._keys.bucket(bucket)], + ["set", item.item_id, item.score, claim_score], + retry="resolve_by_caller", + ) + return Claim(bucket, item.item_id, claim_score) if bool(int(raw)) else None + + async def advance_retry(self, claim: Claim) -> bool: + attempt, _ = decode_claim_score(claim.claim_score) + retry_score = encode_ready_claim_score(attempt + 1) + raw = await self._eval( + "cas_zset_member", + [self._keys.bucket(claim.bucket)], + ["set", claim.item_id, claim.claim_score, retry_score], + retry="resolve_by_caller", + ) + return bool(int(raw)) + + async def ack(self, claim: Claim) -> bool: + raw = await self._eval( + "cas_zset_member", + [self._keys.bucket(claim.bucket)], + ["delete", claim.item_id, claim.claim_score, ""], + retry="idempotent", + ) + return bool(int(raw)) + + async def remove_registration_if_empty(self, bucket: BucketRef) -> bool: + if int(await self._command("zcard", self._keys.bucket(bucket))) != 0: + return False + await self._command( + "zrem", + self._keys.registry(bucket.shard_id), + self._keys.registry_member(bucket), + retry="idempotent", + ) + return True + + async def try_acquire_owner( + self, + role: OwnerRole, + shard_id: int, + token: str, + ttl_seconds: int, + *, + item_type: str | None = None, + ) -> OwnerToken | None: + ttl_ms = self._owner_ttl_ms(token, ttl_seconds) + # Wall-clock expiry is observational, so derive it from the conservative + # request start instead of adding TTL after the response arrives. + started_at_ms = time.time_ns() // 1_000_000 + owned = await self._command( + "set", + self._keys.owner(role, shard_id, item_type=item_type), + token, + nx=True, + px=ttl_ms, + retry="resolve_by_caller", + ) + if not owned: + return None + return OwnerToken(role, shard_id, token, started_at_ms + ttl_ms, item_type) + + async def renew_owner(self, owner: OwnerToken, ttl_seconds: int) -> OwnerToken | None: + ttl_ms = self._owner_ttl_ms(owner.token, ttl_seconds) + started_at_ms = time.time_ns() // 1_000_000 + raw = int( + await self._eval( + "compare_owner", + [ + self._keys.owner( + owner.role, + owner.shard_id, + item_type=owner.item_type, + ) + ], + ["renew", owner.token, ttl_ms], + retry="resolve_by_caller", + ) + ) + return ( + None + if raw == 0 + else OwnerToken( + owner.role, + owner.shard_id, + owner.token, + started_at_ms + ttl_ms, + owner.item_type, + ) + ) + + @staticmethod + def _owner_ttl_ms(token: str, ttl_seconds: int) -> int: + if not token: + raise StoreOperationError("invalid_owner", "Timer owner token is empty") + if type(ttl_seconds) is not int or not 0 < ttl_seconds <= MAX_EXACT_INTEGER // 1000: + raise StoreOperationError("invalid_integer", "Timer owner TTL is invalid") + return ttl_seconds * 1000 + + async def assert_owner(self, owner: OwnerToken) -> None: + raw = await self._command( + "get", + self._keys.owner( + owner.role, + owner.shard_id, + item_type=owner.item_type, + ), + ) + if raw != owner.token: + raise StoreOperationError("owner_lost", "Shard owner changed") + + async def release_owner(self, owner: OwnerToken) -> bool: + raw = await self._eval( + "compare_owner", + [ + self._keys.owner( + owner.role, + owner.shard_id, + item_type=owner.item_type, + ) + ], + ["release", owner.token, ""], + retry="idempotent", + ) + return bool(int(raw)) + + async def _open_on_loop(self) -> None: + try: + import redis.asyncio as redis_async + except ImportError as exc: + raise StoreOperationError( + "dependency_missing", + "Redis expiry scheduling requires the 'expiry' optional dependency", + ) from exc + endpoint = self._config.endpoints[0] + if "://" not in endpoint: + endpoint = ("rediss" if self._config.tls else "redis") + "://" + endpoint + password = self._config.resolve_password() + kwargs: dict[str, Any] = { + "username": self._config.username or None, + "password": password, + "decode_responses": True, + "socket_connect_timeout": self._config.connect_timeout_seconds, + "socket_timeout": self._config.socket_timeout_seconds, + "max_connections": self._config.max_connections, + } + if self._config.mode == "standalone": + self._client = redis_async.Redis.from_url(endpoint, **kwargs) + else: + from redis.asyncio.cluster import ClusterNode, RedisCluster + from redis.backoff import NoBackoff + from redis.retry import Retry + + nodes = [] + use_tls = self._config.tls + for raw in self._config.endpoints: + normalized = ( + raw + if "://" in raw + else (("rediss" if self._config.tls else "redis") + f"://{raw}") + ) + parsed = urlsplit(normalized) + use_tls = parsed.scheme == "rediss" + nodes.append(ClusterNode(parsed.hostname or "", parsed.port or 6379)) + # Store retry policies must be the only authority for replaying + # commands: a lost CAS/SET NX response cannot safely be replayed. + # These instance-local flags leave MOVED/ASK routing intact. + self._client = RedisCluster( + startup_nodes=nodes, + ssl=use_tls, + cluster_error_retry_attempts=0, + retry=Retry(NoBackoff(), 0), + **kwargs, + ) + try: + await self._client.ping() + except Exception as exc: + await self._close_on_loop() + raise self._translate_redis_error(exc, retry="read") from exc + + async def _close_on_loop(self) -> None: + client = self._client + if client is not None: + close = getattr(client, "aclose", None) or getattr(client, "close", None) + if close is not None: + result = close() + if result is not None: + await result + # Cancellation/failure is not proof of a closed client. Keep the + # reference until cleanup succeeds so partial startup can retry. + self._client = None + + async def _command( + self, name: str, *args: object, retry: RetryPolicy = "read", **kwargs: Any + ) -> Any: + async def operation() -> Any: + if self._client is None: + raise StoreOperationError("not_started", "Redis timer store is closed") + return await self._run_with_retry( + lambda: getattr(self._client, name)(*args, **kwargs), + retry=retry, + ) + + return await self._io.run(operation) + + async def _single_key_command( + self, + command: str, + key: str, + *, + retry: RetryPolicy, + ) -> Any: + # redis-py 5.3.1 routes DEL/UNLINK helpers through ClusterPipeline, + # which skips execution when cluster_error_retry_attempts is zero. + return await self._command("execute_command", command, key, retry=retry) + + async def _eval( + self, + script: str, + keys: Sequence[str], + args: Sequence[object], + *, + retry: RetryPolicy, + ) -> Any: + return await self._io.run(lambda: self._eval_on_loop(script, keys, args, retry=retry)) + + async def _eval_on_loop( + self, + script: str, + keys: Sequence[str], + args: Sequence[object], + *, + retry: RetryPolicy, + ) -> Any: + if self._client is None: + raise StoreOperationError("not_started", "Redis timer store is closed") + source = self._read_script(script) + sha = self._script_shas[script] + try: + return await self._run_with_retry( + lambda: self._client.evalsha(sha, len(keys), *keys, *args), + retry=retry, + ) + except StoreOperationError as exc: + if exc.code != "noscript": + raise + return await self._run_with_retry( + lambda: self._client.eval(source, len(keys), *keys, *args), + retry=retry, + ) + + async def _run_with_retry(self, operation: Any, *, retry: RetryPolicy) -> Any: + """Retry safe Redis operations inside the executor's total time budget.""" + delay = 0.01 + while True: + try: + return await operation() + except StoreOperationError: + raise + except Exception as exc: + translated = self._translate_redis_error(exc, retry=retry) + if retry == "resolve_by_caller" or not translated.retryable: + raise translated from exc + await asyncio.sleep(delay) + delay = min(delay * 2, 0.2) + + def _read_script(self, name: str) -> str: + if name not in self._scripts: + base = resources.files(__package__).joinpath("lua") + common = base.joinpath("_common.lua").read_text(encoding="utf-8") + body = base.joinpath(f"{name}.lua").read_text(encoding="utf-8") + self._scripts[name] = common + "\n" + body + self._script_shas[name] = hashlib.sha1( + self._scripts[name].encode("utf-8"), usedforsecurity=False + ).hexdigest() + return self._scripts[name] + + def _translate_redis_error(self, exc: Exception, *, retry: RetryPolicy) -> StoreOperationError: + message = str(exc) + normalized = message.lower() + known = ( + "noscript", + "invalid_identity", + "invalid_report", + "invalid_integer", + "invalid_owner", + "invalid_owner_action", + "invalid_claim_time", + "invalid_cas_action", + "invalid_publish", + "metadata_conflict", + "owner_lost", + "progress_conflict", + "namespace_corrupt", + "wrong_key_type", + ) + code = ( + "noscript" + if type(exc).__name__ == "NoScriptError" + else ( + "wrong_key_type" + if "wrongtype" in normalized + else next((item for item in known if item in normalized), "redis_error") + ) + ) + unknown = retry == "resolve_by_caller" and code == "redis_error" + return StoreOperationError( + code, + message or type(exc).__name__, + retryable=code == "redis_error", + outcome_unknown=unknown, + ) diff --git a/openviking/scheduling/expiry/registry.py b/openviking/scheduling/expiry/registry.py new file mode 100644 index 0000000000..2b70d71511 --- /dev/null +++ b/openviking/scheduling/expiry/registry.py @@ -0,0 +1,53 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Private resource handler registry.""" + +from __future__ import annotations + +import inspect +from dataclasses import dataclass +from types import MappingProxyType +from typing import Mapping + +from .api import SUPPORTED_ITEM_TYPES, ExpiryError, ExpiryHandler, HandlerOptions +from .models import validate_item_type + + +@dataclass(frozen=True) +class HandlerBinding: + handler: ExpiryHandler + options: HandlerOptions + + +class HandlerRegistry: + def __init__(self) -> None: + self._bindings: dict[str, HandlerBinding] = {} + self._frozen = False + + def register(self, item_type: str, handler: ExpiryHandler, options: HandlerOptions) -> None: + if self._frozen: + raise ExpiryError("registration_closed", "Timer handler registration is closed") + validate_item_type(item_type) + if item_type not in SUPPORTED_ITEM_TYPES: + raise ExpiryError("unregistered_type", f"Timer type is not supported: {item_type}") + handle = getattr(handler, "handle", None) + if not callable(handle) or not inspect.iscoroutinefunction(handle): + raise ExpiryError("invalid_handler", "Timer handler.handle must be async") + if not isinstance(options, HandlerOptions): + raise ExpiryError("invalid_handler", "Invalid timer handler options") + if type(options.concurrency) is not int or options.concurrency <= 0: + raise ExpiryError("invalid_handler", "Handler concurrency must be positive") + if type(options.max_attempts) is not int or not 1 <= options.max_attempts <= 256: + raise ExpiryError("invalid_handler", "Handler max_attempts must be in [1, 256]") + if type(options.metadata_enabled) is not bool: + raise ExpiryError("invalid_handler", "Handler metadata_enabled must be a boolean") + if item_type in self._bindings: + raise ExpiryError("duplicate_handler", f"Handler already registered: {item_type}") + self._bindings[item_type] = HandlerBinding(handler, options) + + def get(self, item_type: str) -> HandlerBinding | None: + return self._bindings.get(item_type) + + def freeze(self) -> Mapping[str, HandlerBinding]: + self._frozen = True + return MappingProxyType(dict(self._bindings)) diff --git a/openviking/scheduling/expiry/runtime.py b/openviking/scheduling/expiry/runtime.py new file mode 100644 index 0000000000..6e4e9a81c6 --- /dev/null +++ b/openviking/scheduling/expiry/runtime.py @@ -0,0 +1,459 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Expiry lifecycle and dedicated Redis event loop.""" + +from __future__ import annotations + +import asyncio +import concurrent.futures +import logging +import threading +import time +from dataclasses import dataclass, field +from typing import Any, Awaitable, Callable, Coroutine, TypeVar + +from .api import ExpiryScheduler +from .config import ExpiryConfig +from .consumer import ExpiryConsumer +from .keys import HASH_NAME, HASH_SEED, IDENTITY_ENCODING, ExpiryKeys, encode_component +from .leases import ShardLeaseCoordinator +from .metrics import ExpiryMetrics +from .models import NamespaceSpec, StoreOperationError +from .publisher import ExpiryPublisher +from .redis_store import RedisExpiryStore +from .registry import HandlerRegistry +from .service import ExpiryRuntimeState, _ExpirySchedulerFacade +from .shard_runner import ShardRunner + +T = TypeVar("T") +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class ExpiryRuntimeRoles: + report: bool = False + publish: bool = False + consume: bool = False + + def __post_init__(self) -> None: + if self.consume and not self.report: + raise ValueError("Expiry consumer requires reporting for deferred callbacks") + + +@dataclass(eq=False) +class _RedisRequest: + # Unlike run_coroutine_threadsafe's Future, this future is completed only + # when the Redis-loop task AND its operation have actually terminated. + finished: concurrent.futures.Future = field(default_factory=concurrent.futures.Future) + cancelled: threading.Event = field(default_factory=threading.Event) + task: asyncio.Task | None = None # Read and written only on the Redis loop. + + +class RedisLoopExecutor: + def __init__(self, *, max_inflight: int, request_budget_seconds: float) -> None: + self._max_inflight = max_inflight + self._request_budget_seconds = request_budget_seconds + self._loop: asyncio.AbstractEventLoop | None = None + self._thread: threading.Thread | None = None + self._gate: asyncio.Semaphore | None = None + self._accepting = False + self._ready = threading.Event() + self._requests: set[_RedisRequest] = set() + self._requests_lock = threading.Lock() + self._stop_requested = False + + async def start(self) -> None: + if self._accepting: + return + if self._thread is not None: + raise self._not_stopped() + self._ready.clear() + self._stop_requested = False + self._thread = threading.Thread( + target=self._thread_main, name="openviking-redis-expiry", daemon=True + ) + self._thread.start() + ready = await asyncio.to_thread(self._ready.wait, 5.0) + if not ready or self._loop is None: + raise RuntimeError("Expiry Redis loop failed to start") + with self._requests_lock: + self._accepting = True + + async def run(self, operation: Callable[[], Awaitable[T]]) -> T: + loop = self._loop + # Nested store.close calls during run_cleanup are already covered by + # their tracked parent request, and must be allowed after quiescing. + if loop is not None and asyncio.get_running_loop() is loop: + return await self._execute(operation) + request = self._submit(operation) + finished = asyncio.wrap_future(request.finished) + try: + return await asyncio.shield(finished) + except asyncio.CancelledError: + self._cancel(request) + # Keep the caller visible to ShardRunner until the actual operation + # exits. Repeated cancellation must not sever this lifetime bridge. + await self._join(finished) + raise + + async def quiesce(self, timeout_seconds: float) -> None: + """Reject new work and observe actual remote termination before closing Redis.""" + with self._requests_lock: + self._accepting = False + requests = list(self._requests) + if requests: + futures = {asyncio.wrap_future(request.finished) for request in requests} + done, pending = await asyncio.wait(futures, timeout=max(0.0, timeout_seconds)) + for future in done: + if not future.cancelled(): + future.exception() + # A caller might have already stopped awaiting a failed request. + for future in pending: + future.add_done_callback(self._consume_exception) + if pending: + raise self._not_stopped() + + async def run_cleanup(self, operation: Callable[[], Awaitable[T]], timeout_seconds: float) -> T: + """Run storage cleanup on its own loop without reopening public admission.""" + request = self._submit(operation, cleanup=True) + finished = asyncio.wrap_future(request.finished) + try: + done, _ = await asyncio.wait({finished}, timeout=max(0.0, timeout_seconds)) + except BaseException: + self._cancel(request) + finished.add_done_callback(self._consume_exception) + raise + if finished not in done: + self._cancel(request) + finished.add_done_callback(self._consume_exception) + raise self._not_stopped() + return finished.result() + + async def close(self, timeout_seconds: float) -> None: + deadline = time.monotonic() + max(0.0, timeout_seconds) + await self.quiesce(timeout_seconds) + loop, thread = self._loop, self._thread + if loop is None and thread is not None and thread.is_alive(): + # start() may have been cancelled before the new thread published + # its loop. Wait only within the close budget, then stop that loop. + await asyncio.to_thread(self._ready.wait, max(0.0, deadline - time.monotonic())) + loop = self._loop + if loop is not None and not self._stop_requested: + self._stop_requested = True + loop.call_soon_threadsafe(loop.stop) + if thread is not None: + await asyncio.to_thread(thread.join, max(0.0, deadline - time.monotonic())) + if thread.is_alive(): + # Retain both references so callers can retry, and cannot start + # another I/O thread while the old one still owns dependencies. + raise self._not_stopped() + self._loop = None + self._thread = None + self._gate = None + + def _submit( + self, operation: Callable[[], Awaitable[T]], *, cleanup: bool = False + ) -> _RedisRequest: + with self._requests_lock: + loop = self._loop + if (not cleanup and not self._accepting) or loop is None or self._stop_requested: + raise StoreOperationError("not_started", "Redis expiry I/O loop is not running") + request = _RedisRequest() + # Register before scheduling: shutdown must also see not-yet-started I/O. + self._requests.add(request) + + async def cleanup_operation() -> T: + return await operation() + + def start_request() -> None: + request.task = asyncio.create_task( + cleanup_operation() if cleanup else self._execute(operation) + ) + request.task.add_done_callback(lambda task: self._request_done(request, task)) + if request.cancelled.is_set(): + request.task.cancel() + + try: + loop.call_soon_threadsafe(start_request) + except BaseException: + self._requests.discard(request) + raise + return request + + def _cancel(self, request: _RedisRequest) -> None: + if request.cancelled.is_set(): + return + request.cancelled.set() + + def cancel_request() -> None: + if request.task is not None and not request.task.done(): + request.task.cancel() + + loop = self._loop + if loop is not None: + loop.call_soon_threadsafe(cancel_request) + + def _request_done(self, request: _RedisRequest, task: asyncio.Task) -> None: + if task.cancelled(): + request.finished.set_exception(asyncio.CancelledError()) + else: + error = task.exception() + if error is not None: + request.finished.set_exception(error) + else: + request.finished.set_result(task.result()) + with self._requests_lock: + self._requests.discard(request) + + @staticmethod + def _consume_exception(future: asyncio.Future) -> None: + if not future.cancelled(): + future.exception() + + @staticmethod + async def _join(future: asyncio.Future) -> None: + while not future.done(): + try: + await asyncio.shield(future) + except (asyncio.CancelledError, Exception): + pass + if not future.cancelled(): + future.exception() + + @staticmethod + def _not_stopped() -> StoreOperationError: + return StoreOperationError( + "shard_not_stopped", "Redis expiry work or its I/O thread has not terminated" + ) + + def _thread_main(self) -> None: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + self._loop = loop + self._gate = asyncio.Semaphore(self._max_inflight) + self._ready.set() + try: + loop.run_forever() + finally: + pending = asyncio.all_tasks(loop) + for task in pending: + task.cancel() + if pending: + loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.close() + + async def _execute(self, operation: Callable[[], Awaitable[T]]) -> T: + gate = self._gate + if gate is None: + raise StoreOperationError("not_started", "Redis expiry I/O loop is not ready") + + async def guarded() -> T: + async with gate: + return await operation() + + task = asyncio.create_task(guarded()) + try: + done, _ = await asyncio.wait({task}, timeout=self._request_budget_seconds) + if task not in done: + task.cancel() + await self._join(task) + raise StoreOperationError( + "redis_timeout", + "Redis expiry request exceeded its budget", + retryable=True, + outcome_unknown=True, + ) + return task.result() + finally: + if not task.done(): + task.cancel() + await self._join(task) + if not task.cancelled(): + task.exception() + + +class ExpiryRuntime: + def __init__( + self, + config: ExpiryConfig, + *, + roles: ExpiryRuntimeRoles, + metrics: ExpiryMetrics | None = None, + ) -> None: + self._config = config + self._roles = roles + self._metrics = metrics or ExpiryMetrics() + self._keys = ExpiryKeys(encode_component(config.service_name)) + self._io = RedisLoopExecutor( + max_inflight=config.redis.max_inflight_requests, + request_budget_seconds=config.redis.request_budget_seconds, + ) + self._store = RedisExpiryStore(config.redis, self._keys, self._io, self._metrics) + self._registry = HandlerRegistry() + self._state = ExpiryRuntimeState() + self._facade = _ExpirySchedulerFacade( + self._store, + self._registry, + self._state, + shard_count=config.shard_count, + reporting_enabled=roles.report, + metrics=self._metrics, + ) + self._leases = ShardLeaseCoordinator(self._store, config, self._metrics) + self._shard_runner = ShardRunner(self._leases, config, self._metrics) + self._publisher: ExpiryPublisher | None = None + self._consumers: dict[str, ExpiryConsumer] = {} + self._tasks: dict[str, asyncio.Task[None]] = {} + self._stop = asyncio.Event() + self._storage_started = False + self._storage_closed = False + # Readiness is not ownership: cancellation can leave resources behind + # before start_storage has completed. + self._storage_owned = False + self._store_open_attempted = False + self._storage_start_task: asyncio.Task[None] | None = None + + @property + def scheduler(self) -> ExpiryScheduler: + return self._facade + + async def start_storage(self) -> None: + if self._storage_started: + return + if self._storage_owned or self._storage_start_task is not None: + raise StoreOperationError( + "shard_not_stopped", "Previous expiry storage startup has not been cleaned up" + ) + self._storage_owned = True + self._storage_closed = False + # Keep the actual startup operation visible if it resists cancellation. + # The caller's cancellation must not make close assume it has stopped. + startup = asyncio.create_task(self._initialize_storage()) + startup.add_done_callback(RedisLoopExecutor._consume_exception) + self._storage_start_task = startup + try: + await asyncio.shield(startup) + except BaseException: + if not startup.done(): + startup.cancel() + try: + await self.close() + except BaseException: + # Preserve the original failure/cancellation. Ownership and the + # startup task remain available for a later close retry. + logger.warning("Expiry partial startup cleanup did not complete", exc_info=True) + raise + self._storage_start_task = None + self._storage_started = True + self._state.storage_ready = True + + async def _initialize_storage(self) -> None: + await self._io.start() + self._store_open_attempted = True + await self._store.open() + await self._store.ensure_namespace( + NamespaceSpec( + self._config.shard_count, + HASH_NAME, + HASH_SEED, + IDENTITY_ENCODING, + ) + ) + + async def _settle_storage_start(self) -> None: + startup = self._storage_start_task + if startup is None: + return + if not startup.done(): + startup.cancel() + done, _ = await asyncio.wait({startup}, timeout=self._config.shutdown_timeout_seconds) + if startup not in done: + raise StoreOperationError( + "shard_not_stopped", "Expiry storage startup is still running" + ) + if not startup.cancelled(): + startup.exception() + self._storage_start_task = None + + async def start_roles(self) -> None: + if not self._storage_started: + raise RuntimeError("Expiry storage is not started") + if self._state.accepting_calls: + return + bindings = self._registry.freeze() + self._stop.clear() + if self._roles.publish: + self._publisher = ExpiryPublisher( + self._store, self._shard_runner, self._config, self._metrics + ) + self._spawn_role("publisher", self._publisher.run(self._stop)) + if self._roles.consume: + for item_type, binding in bindings.items(): + consumer = ExpiryConsumer( + self._store, + self._shard_runner, + item_type, + binding, + self._config, + self._metrics, + ) + self._consumers[item_type] = consumer + self._spawn_role( + f"callback:{item_type}", + consumer.run(self._stop), + ) + self._state.accepting_calls = True + + async def wait(self) -> None: + if not self._tasks: + await self._stop.wait() + return + done, _ = await asyncio.wait(self._tasks.values(), return_when=asyncio.FIRST_COMPLETED) + for task in done: + await task + + async def stop_roles(self) -> None: + self._state.accepting_calls = False + self._stop.set() + tasks = set(self._tasks.values()) + if tasks: + _, pending = await asyncio.wait(tasks, timeout=self._config.shutdown_timeout_seconds) + for task in pending: + task.cancel() + if pending: + await asyncio.wait(pending, timeout=self._config.redis.request_budget_seconds) + for name, task in list(self._tasks.items()): + if task.done(): + if not task.cancelled(): + task.exception() + self._tasks.pop(name, None) + settled = await self._shard_runner.wait_pending(self._config.redis.request_budget_seconds) + if self._tasks or not settled: + raise StoreOperationError( + "shard_not_stopped", "Expiry work is still running; dependencies remain open" + ) + + async def close(self) -> None: + await self.stop_roles() + if not self._storage_owned: + return + await self._settle_storage_start() + # Stop admission and check real Redis-loop tasks, not just their + # application-loop wrappers, before closing any dependency. + await self._io.quiesce(self._config.shutdown_timeout_seconds) + if not self._storage_closed: + if self._store_open_attempted: + await self._io.run_cleanup(self._store.close, self._config.shutdown_timeout_seconds) + self._storage_closed = True + # Keep ownership if client cleanup or thread termination failed. Once + # client cleanup succeeds, retries only observe the I/O thread join. + await self._io.close(self._config.shutdown_timeout_seconds) + self._storage_owned = False + self._store_open_attempted = False + self._storage_started = False + self._state.storage_ready = False + + def _spawn_role(self, name: str, operation: Coroutine[Any, Any, None]) -> None: + task: asyncio.Task[None] = asyncio.create_task(operation, name=f"expiry-{name}") + self._tasks[name] = task diff --git a/openviking/scheduling/expiry/service.py b/openviking/scheduling/expiry/service.py new file mode 100644 index 0000000000..4bb264c5bf --- /dev/null +++ b/openviking/scheduling/expiry/service.py @@ -0,0 +1,104 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Private facade implementing the public ExpiryScheduler protocol.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from .api import ( + SUPPORTED_ITEM_TYPES, + ExpiryError, + ExpiryHandler, + ExpiryRegistration, + ExpiryRegistrationResult, + HandlerOptions, + ItemIdentity, +) +from .keys import encode_item_id, shard_for +from .metrics import ExpiryMetrics +from .models import StoreOperationError, validate_identity, validate_report +from .redis_store import RedisExpiryStore +from .registry import HandlerRegistry + + +@dataclass +class ExpiryRuntimeState: + storage_ready: bool = False + accepting_calls: bool = False + + +class _ExpirySchedulerFacade: + def __init__( + self, + store: RedisExpiryStore, + registry: HandlerRegistry, + state: ExpiryRuntimeState, + *, + shard_count: int, + reporting_enabled: bool, + metrics: ExpiryMetrics, + ) -> None: + self._store = store + self._registry = registry + self._state = state + self._metrics = metrics + self._shard_count = shard_count + self._reporting_enabled = reporting_enabled + + async def report(self, task: ExpiryRegistration) -> ExpiryRegistrationResult: + self._require_ready() + if not self._reporting_enabled: + raise ExpiryError("reporting_disabled", "Timer reporting is disabled") + validate_report(task) + self._ensure_supported_type(task.identity.item_type) + try: + item_id = encode_item_id(task.identity) + shard_id = shard_for(item_id, self._shard_count) + result = await self._store.upsert_resource(shard_id, item_id, task) + binding = self._registry.get(task.identity.item_type) + if binding is None or binding.options.metadata_enabled: + await self._store.ensure_metadata(item_id, task) + except StoreOperationError as exc: + raise self._business_error(exc) from exc + self._metrics.count("timer_report", item_type=task.identity.item_type, outcome="written") + return result + + async def unregister(self, identity: ItemIdentity) -> None: + self._require_ready() + validate_identity(identity) + self._ensure_supported_type(identity.item_type) + item_id = encode_item_id(identity) + try: + binding = self._registry.get(identity.item_type) + await self._store.unregister_resource( + shard_for(item_id, self._shard_count), + item_id, + delete_metadata=binding is None or binding.options.metadata_enabled, + ) + except StoreOperationError as exc: + raise self._business_error(exc) from exc + + def register_handler( + self, + item_type: str, + handler: ExpiryHandler, + *, + options: HandlerOptions | None = None, + ) -> None: + self._registry.register(item_type, handler, options or HandlerOptions()) + + @staticmethod + def _ensure_supported_type(item_type: str) -> None: + if item_type not in SUPPORTED_ITEM_TYPES: + raise ExpiryError("unregistered_type", f"Timer type is not supported: {item_type}") + + def _require_ready(self) -> None: + if not (self._state.storage_ready and self._state.accepting_calls): + raise ExpiryError("not_started", "Expiry scheduling is not ready") + + def _business_error(self, exc: StoreOperationError) -> ExpiryError: + public_codes = {"invalid_identity", "invalid_report", "unregistered_type"} + code = exc.code if exc.code in public_codes else "timer_unavailable" + self._metrics.count("timer_store_error", outcome=exc.code) + return ExpiryError(code, str(exc), retryable=exc.retryable or code == "timer_unavailable") diff --git a/openviking/scheduling/expiry/shard_runner.py b/openviking/scheduling/expiry/shard_runner.py new file mode 100644 index 0000000000..ae26e756b9 --- /dev/null +++ b/openviking/scheduling/expiry/shard_runner.py @@ -0,0 +1,379 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Bounded execution of business work under a temporary shard lease.""" + +from __future__ import annotations + +import asyncio +import logging +import time +from typing import Awaitable, Callable + +from .config import ExpiryConfig +from .leases import ShardLeaseCoordinator +from .metrics import ExpiryMetrics +from .models import ( + OwnerRole, + OwnerToken, + ShardLease, + StoreOperationError, + task_is_cancelling, +) + +ShardProcessor = Callable[["ShardExecution", asyncio.Event], Awaitable[None]] +RunnerScope = tuple[OwnerRole, str | None, int] + +logger = logging.getLogger(__name__) + + +class ShardExecution: + """A revocable checkpoint and the strong owner of a processor's children. + + Cancellation only stops local work. It does not prove that an external + request or callback has stopped, so cancelled work never authorizes release. + """ + + def __init__(self, lease: ShardLease) -> None: + self._lease = lease + self._children: set[asyncio.Task[None]] = set() + self._waiters: set[asyncio.Task[object]] = set() + self._error: BaseException | None = None + self._outcome_unknown = False + + @property + def owner(self) -> OwnerToken: + return self._lease.owner + + @property + def shard_id(self) -> int: + return self.owner.shard_id + + def ensure_active(self) -> None: + task = asyncio.current_task() + if task_is_cancelling(task): + # A handler may swallow child cancellation and return normally. Mark + # uncertainty before its next store operation, not in a late callback. + self._outcome_unknown = True + raise asyncio.CancelledError + self._lease.ensure_active() + + def spawn(self, operation: Callable[[], Awaitable[None]]) -> asyncio.Task[None]: + self.ensure_active() + + async def guarded() -> None: + self.ensure_active() + await operation() + + task = asyncio.create_task(guarded()) + self._children.add(task) + task.add_done_callback(self._child_done) + return task + + async def wait(self, delay_seconds: float, stop: asyncio.Event) -> None: + """Wait for a poll, stop, or invalidation, checking both sides of the wait.""" + self.ensure_active() + deadline = time.monotonic() + max(0.0, delay_seconds) + waiters = { + asyncio.create_task(stop.wait()), + asyncio.create_task(self._lease.lost.wait()), + asyncio.create_task(self._lease.revoked.wait()), + } + self._waiters.update(waiters) + for task in waiters: + task.add_done_callback(self._waiter_done) + try: + while True: + self.ensure_active() + now = time.monotonic() + if stop.is_set() or now >= deadline: + return + await asyncio.wait( + waiters, + timeout=max(0.0, min(deadline, self._lease.valid_until_monotonic) - now), + return_when=asyncio.FIRST_COMPLETED, + ) + self.ensure_active() + finally: + for task in waiters: + task.cancel() + + def _waiter_done(self, task: asyncio.Task[object]) -> None: + self._waiters.discard(task) + _consume_exception(task) + + def _child_done(self, task: asyncio.Task[None]) -> None: + self._children.discard(task) + if task.cancelled() or task_is_cancelling(task): + self._outcome_unknown = True + if task.cancelled(): + return + error = task.exception() + if error is not None: + self._error = self._error or error + if isinstance(error, StoreOperationError) and error.outcome_unknown: + self._outcome_unknown = True + + def _revoke(self) -> None: + self._lease.revoked.set() + + +def _consume_exception(task: asyncio.Task[object]) -> None: + if not task.cancelled(): + task.exception() + + +class ShardRunner: + """Run isolated shard executions while retaining work that cannot stop.""" + + def __init__( + self, + leases: ShardLeaseCoordinator, + config: ExpiryConfig, + metrics: ExpiryMetrics, + ) -> None: + self._leases = leases + self._config = config + self._metrics = metrics + self._active: set[RunnerScope] = set() + self._pending: dict[RunnerScope, set[asyncio.Task[object]]] = {} + + @property + def has_pending(self) -> bool: + return any(not task.done() for tasks in self._pending.values() for task in tasks) + + async def wait_pending(self, timeout: float) -> bool: + """Observe retained work for at most timeout seconds, without cancelling it.""" + tasks = {task for tasks in self._pending.values() for task in tasks if not task.done()} + if tasks: + await asyncio.wait(tasks, timeout=max(0.0, timeout)) + return not self.has_pending + + async def run_if_acquired( + self, + role: OwnerRole, + shard_id: int, + process: ShardProcessor, + stop: asyncio.Event, + *, + item_type: str | None = None, + ) -> bool: + scope = (role, item_type, shard_id) + if scope in self._active: + raise self._not_stopped(scope) + if stop.is_set(): + return False + self._active.add(scope) + try: + # Shield the request so caller cancellation cannot force an unbounded + # join of a Redis implementation which suppresses cancellation. + acquire = asyncio.create_task( + self._leases.try_acquire(role, shard_id, item_type=item_type) + ) + try: + lease = await asyncio.shield(acquire) + except BaseException: + acquire.cancel() + await self._settle(scope, {acquire}) + raise + if lease is None: + return False + await self._run(scope, ShardExecution(lease), process, stop) + return True + finally: + self._active.discard(scope) + + async def _run( + self, + scope: RunnerScope, + execution: ShardExecution, + process: ShardProcessor, + stop: asyncio.Event, + ) -> None: + lease = execution._lease + + async def guarded_process() -> None: + execution.ensure_active() + await process(execution, stop) + + processor = asyncio.create_task(guarded_process()) + renewing = asyncio.Event() + heartbeat = asyncio.create_task(self._heartbeat(lease, renewing)) + stopped = asyncio.create_task(stop.wait()) + lost = asyncio.create_task(lease.lost.wait()) + tasks = {processor, heartbeat, stopped, lost} + forced = False + error: BaseException | None = None + try: + grace_deadline: float | None = None + while True: + now = time.monotonic() + if lease.lost.is_set() or now >= lease.valid_until_monotonic: + lease.lost.set() + forced = True + break + if processor.done(): + break + if heartbeat.done(): + # A heartbeat is only supposed to finish when ownership is lost. + lease.lost.set() + forced = True + break + if stop.is_set() and grace_deadline is None: + grace_deadline = now + self._config.shutdown_timeout_seconds + if grace_deadline is not None and now >= grace_deadline: + forced = True + break + deadline = lease.valid_until_monotonic + watched = {processor, heartbeat, lost} + if grace_deadline is None: + watched.add(stopped) + else: + deadline = min(deadline, grace_deadline) + await asyncio.wait( + watched, + timeout=max(0.0, deadline - now), + return_when=asyncio.FIRST_COMPLETED, + ) + except BaseException as exc: + forced = True + error = exc + finally: + execution._revoke() + # A processor returning while it still owns detached operations is a + # lifecycle violation, even if those operations cooperate with cancel. + children = {task for task in execution._children if not task.done()} + if children: + if processor.done() and not forced: + error = error or StoreOperationError( + "shard_children_pending", + "Shard processor returned with unfinished children", + ) + forced = True + if not processor.done() or processor.cancelled() or task_is_cancelling(processor): + forced = True + # Mark synchronously before cancel: a Redis implementation may catch + # cancellation and return a success, which must not revive ownership. + if renewing.is_set() and not heartbeat.done(): + lease.lost.set() + local_tasks = tasks | children | execution._waiters + for task in local_tasks: + if not task.done(): + task.cancel() + await self._settle(scope, local_tasks, allow_orphans=True) + + # task.done() is terminal even if the task raised an ordinary exception. + # Read the exception before deciding release, then propagate it afterwards. + if processor.done() and not processor.cancelled(): + error = error or processor.exception() + error = error or execution._error + unknown = execution._outcome_unknown or ( + isinstance(error, StoreOperationError) and error.outcome_unknown + ) + if time.monotonic() >= lease.valid_until_monotonic: + lease.lost.set() + if not forced and not unknown and not lease.lost.is_set(): + release = asyncio.create_task(self._leases.release(lease)) + try: + done, _ = await asyncio.wait( + {release}, timeout=self._config.redis.request_budget_seconds + ) + except BaseException: + lease.lost.set() + release.cancel() + await self._settle(scope, {release}) + raise + if release not in done: + lease.lost.set() + release.cancel() + await self._settle(scope, {release}) + else: + release.result() + if error is not None: + raise error + + async def _heartbeat(self, lease: ShardLease, renewing: asyncio.Event) -> None: + while True: + await asyncio.sleep(self._config.owner_renew_interval_seconds) + renewing.set() + try: + if not await self._leases.renew(lease): + return + finally: + renewing.clear() + + async def _settle( + self, + scope: RunnerScope, + tasks: set[asyncio.Task], + *, + allow_orphans: bool = False, + ) -> None: + """Never couple the deadline to completion of a cancelled task.""" + if not tasks: + return + # Retain before awaiting: a second caller cancellation must not orphan + # work while cleanup itself is suspended. Each callback owns its exact + # set, not a later run's same-role state slot. + retained = self._pending.setdefault(scope, set()) + for task in tasks: + if not task.done() and task not in retained: + retained.add(task) + + def pending_done( + finished: asyncio.Task[object], + scope: RunnerScope = scope, + retained: set[asyncio.Task[object]] = retained, + ) -> None: + self._pending_done(scope, retained, finished) + + task.add_done_callback(pending_done) + try: + done, _ = await asyncio.wait(tasks, timeout=self._config.redis.request_budget_seconds) + except asyncio.CancelledError as exc: + if any(not task.done() for task in tasks): + raise self._not_stopped(scope) from exc + raise + for task in done: + _consume_exception(task) + retained.discard(task) + if not retained and self._pending.get(scope) is retained: + self._pending.pop(scope) + orphaned = sum(not task.done() for task in tasks) + if orphaned: + role, item_type, shard_id = scope + self._metrics.count( + "timer_shard_orphaned", + orphaned, + item_type=item_type or "", + outcome=role.value, + ) + logger.warning( + "expiry shard execution left orphaned tasks role=%s item_type=%s " + "shard_id=%d count=%d", + role.value, + item_type or "", + shard_id, + orphaned, + ) + if orphaned and not allow_orphans: + raise self._not_stopped(scope) + + def _pending_done( + self, + scope: RunnerScope, + retained: set[asyncio.Task[object]], + task: asyncio.Task[object], + ) -> None: + _consume_exception(task) + retained.discard(task) + if not retained and self._pending.get(scope) is retained: + self._pending.pop(scope) + + @staticmethod + def _not_stopped(scope: RunnerScope) -> StoreOperationError: + role, item_type, shard_id = scope + label = role.value if item_type is None else f"{role.value}:{item_type}" + return StoreOperationError( + "shard_not_stopped", + f"The {label} shard {shard_id} still has active local work", + ) diff --git a/openviking/server/bootstrap.py b/openviking/server/bootstrap.py index ff850605f2..b85deca763 100644 --- a/openviking/server/bootstrap.py +++ b/openviking/server/bootstrap.py @@ -7,14 +7,18 @@ import json import os import shutil +import signal import socket import subprocess import sys import tempfile +import threading import time -from dataclasses import dataclass +from contextlib import contextmanager +from dataclasses import dataclass, field from pathlib import Path -from typing import Optional +from types import FrameType +from typing import Iterator, NoReturn, Optional import uvicorn @@ -43,6 +47,30 @@ class BotProcess: log_file: Optional[object] = None +@dataclass +class ExpiryPublisherProcess: + process: subprocess.Popen + stopping: threading.Event = field(default_factory=threading.Event) + monitor_thread: Optional[threading.Thread] = None + + +def _raise_sigterm_as_system_exit(signum: int, _frame: FrameType | None) -> NoReturn: + raise SystemExit(128 + signum) + + +@contextmanager +def _sigterm_runs_finally() -> Iterator[None]: + """Make Uvicorn's replayed SIGTERM unwind through managed-process cleanup.""" + if os.name == "nt" or threading.current_thread() is not threading.main_thread(): + yield + return + previous = signal.signal(signal.SIGTERM, _raise_sigterm_as_system_exit) + try: + yield + finally: + signal.signal(signal.SIGTERM, previous) + + def _get_version() -> str: try: from openviking import __version__ @@ -131,6 +159,100 @@ def _resolve_cli_config_for_bot(config_path: Optional[str]) -> Optional[str]: return None +def _wait_for_expiry_publisher_ready( + process: subprocess.Popen, + status_path: Path, + *, + timeout_seconds: float = 60.0, +) -> None: + started = time.monotonic() + while time.monotonic() - started < timeout_seconds: + returncode = process.poll() + if returncode is not None: + raise RuntimeError(f"Expiry publisher exited before becoming ready (code {returncode})") + try: + status = json.loads(status_path.read_text(encoding="utf-8")) + except (FileNotFoundError, ValueError): + status = {} + if status.get("status") == "failed": + raise RuntimeError(status.get("error") or "Expiry publisher initialization failed") + if status.get("status") == "ready": + return + time.sleep(0.1) + raise TimeoutError("Expiry publisher readiness timed out") + + +def _monitor_expiry_publisher(managed: ExpiryPublisherProcess) -> None: + returncode = managed.process.wait() + if managed.stopping.is_set(): + return + print( + f"Error: managed expiry publisher exited unexpectedly (code {returncode}); " + "stopping OpenViking server.", + file=sys.stderr, + ) + os.kill(os.getpid(), signal.SIGTERM) + + +def _start_expiry_publisher(config_path: Path) -> ExpiryPublisherProcess: + startup_directory = tempfile.TemporaryDirectory(prefix="openviking-expiry-startup-") + status_path = Path(startup_directory.name) / "status.json" + process: subprocess.Popen | None = None + try: + process = subprocess.Popen( + [ + sys.executable, + "-m", + "openviking_cli.expiry_bootstrap", + "--config", + str(config_path), + "--startup-status", + str(status_path), + ], + env=os.environ.copy(), + start_new_session=os.name != "nt", + creationflags=( + getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0) if os.name == "nt" else 0 + ), + ) + _wait_for_expiry_publisher_ready(process, status_path) + managed = ExpiryPublisherProcess(process) + managed.monitor_thread = threading.Thread( + target=_monitor_expiry_publisher, + args=(managed,), + name="openviking-expiry-publisher-monitor", + daemon=True, + ) + managed.monitor_thread.start() + print(f"Expiry publisher started (PID: {process.pid})") + return managed + except BaseException: + if process is not None and process.poll() is None: + process.terminate() + try: + process.wait(timeout=60) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + raise + finally: + startup_directory.cleanup() + + +def _stop_expiry_publisher(managed: ExpiryPublisherProcess) -> None: + managed.stopping.set() + process = managed.process + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=60) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + if managed.monitor_thread is not None: + managed.monitor_thread.join(timeout=1) + + def main(): """Run the CLI with one recovery attempt for remote restart startup failures.""" from openviking.server.restart import RestartRecovery @@ -223,6 +345,8 @@ def _main(recovery): from openviking_cli.utils.config.open_viking_config import OpenVikingConfigSingleton # Load server config from ov.conf + resolved_config_path: Path | None = None + ov_config = None try: resolved_config_path = resolve_config_path( args.config, @@ -230,7 +354,7 @@ def _main(recovery): DEFAULT_OV_CONF, ) config = load_server_config(args.config) - OpenVikingConfigSingleton.initialize(config_path=args.config) + ov_config = OpenVikingConfigSingleton.initialize(config_path=args.config) except (FileNotFoundError, ValueError) as e: if isinstance(e, ValueError): print( @@ -264,7 +388,8 @@ def _main(recovery): try: from openviking_cli.utils.ollama import detect_ollama_in_config, ensure_ollama_for_server - ov_config = OpenVikingConfigSingleton.get_instance() + if ov_config is None: + ov_config = OpenVikingConfigSingleton.get_instance() uses_ollama, ollama_host, ollama_port = detect_ollama_in_config(ov_config) if uses_ollama: result = ensure_ollama_for_server(ollama_host, ollama_port) @@ -329,53 +454,71 @@ def _main(recovery): ) sys.exit(1) + expiry_publisher_process: Optional[ExpiryPublisherProcess] = None restart_requested = False try: - # Create and run server app - app = create_app( - config, - server_overrides={ - key: getattr(config, key) - for key in ("host", "port", "workers", "with_bot") - if getattr(args, key) is not None and (key != "with_bot" or args.with_bot) - }, - config_path=( - str(resolved_config_path) if resolved_config_path is not None else args.config - ), - ) - workers_info = f" (workers: {config.workers})" if config.workers > 1 else "" - print(f"OpenViking HTTP Server is running on {config.host}:{config.port}{workers_info}") - - workers = config.workers - if workers > 1: - # Multi-worker mode requires an import string so each worker - # can independently import the application. We stash the - # resolved config path in an env-var so that the factory can - # pick it up (ServerConfig already reads OPENVIKING_CONFIG_FILE). - os.environ[WORKER_HOST_ENV] = config.host or "" - os.environ[WORKER_PORT_ENV] = str(config.port) - os.environ[WORKER_WITH_BOT_ENV] = "1" if config.with_bot else "0" - os.environ[WORKER_BOT_API_URL_ENV] = config.bot_api_url - uvicorn.run( - "openviking.server.app:create_worker_app", - factory=True, - host=config.host, - port=config.port, - workers=workers, - timeout_keep_alive=config.timeout_keep_alive, - log_config=None, - ) - else: - restart_requested = _run_restartable_server( - app, - recovery=recovery, - host=config.host, - port=config.port, - timeout_keep_alive=config.timeout_keep_alive, - log_config=None, + with _sigterm_runs_finally(): + idle = getattr(getattr(ov_config, "memory", None), "session_auto_commit", None) + expiry = getattr(ov_config, "expiry", None) + if ( + idle is not None + and idle.enabled + and idle.idle_backend == "redis" + and expiry is not None + and expiry.publisher_separate_process + ): + if resolved_config_path is None: + raise RuntimeError( + "Cannot start expiry publisher without a resolved ov.conf path" + ) + expiry_publisher_process = _start_expiry_publisher(resolved_config_path) + + # Create and run server app only after managed dependencies are ready. + app = create_app( + config, + server_overrides={ + key: getattr(config, key) + for key in ("host", "port", "workers", "with_bot") + if getattr(args, key) is not None and (key != "with_bot" or args.with_bot) + }, + config_path=( + str(resolved_config_path) if resolved_config_path is not None else args.config + ), ) + workers_info = f" (workers: {config.workers})" if config.workers > 1 else "" + print(f"OpenViking HTTP Server is running on {config.host}:{config.port}{workers_info}") + + workers = config.workers + if workers > 1: + # Multi-worker mode requires an import string so each worker + # can independently import the application. We stash the + # resolved config path in an env-var so that the factory can + # pick it up (ServerConfig already reads OPENVIKING_CONFIG_FILE). + os.environ[WORKER_HOST_ENV] = config.host or "" + os.environ[WORKER_PORT_ENV] = str(config.port) + os.environ[WORKER_WITH_BOT_ENV] = "1" if config.with_bot else "0" + os.environ[WORKER_BOT_API_URL_ENV] = config.bot_api_url + uvicorn.run( + "openviking.server.app:create_worker_app", + factory=True, + host=config.host, + port=config.port, + workers=workers, + timeout_keep_alive=config.timeout_keep_alive, + log_config=None, + ) + else: + restart_requested = _run_restartable_server( + app, + recovery=recovery, + host=config.host, + port=config.port, + timeout_keep_alive=config.timeout_keep_alive, + log_config=None, + ) finally: - # Cleanup vikingbot process on shutdown + if expiry_publisher_process is not None: + _stop_expiry_publisher(expiry_publisher_process) if bot_process is not None: _stop_vikingbot_gateway(bot_process) if restart_requested: diff --git a/openviking/server/config.py b/openviking/server/config.py index 7759acd0fd..3f8161b1d1 100644 --- a/openviking/server/config.py +++ b/openviking/server/config.py @@ -164,9 +164,16 @@ class UserConfigDefaults(UserConfig): @field_validator("auto_commit_policy", mode="before") @classmethod def validate_auto_commit_policy(cls, value: Any) -> Optional[Dict[str, Any]]: - from openviking.session.auto_commit_policy import AutoCommitPolicy + from openviking.session.auto_commit_policy import ( + AutoCommitPolicy, + validate_idle_timeout_input, + ) - return None if value is None else AutoCommitPolicy.from_dict(value).to_dict() + if value is None: + return None + if isinstance(value, dict) and "idle_timeout_seconds" in value: + validate_idle_timeout_input(value["idle_timeout_seconds"]) + return AutoCommitPolicy.from_dict(value).to_dict() class MetricsAccountDimensionConfig(BaseModel): diff --git a/openviking/server/routers/sessions.py b/openviking/server/routers/sessions.py index 3423bdfc7d..a17471c6f8 100644 --- a/openviking/server/routers/sessions.py +++ b/openviking/server/routers/sessions.py @@ -17,6 +17,7 @@ from openviking.server.models import Response from openviking.server.responses import error_response from openviking.server.telemetry import run_operation +from openviking.session.auto_commit_policy import validate_idle_timeout_input from openviking.telemetry import TelemetryRequest from openviking.utils.image_search import is_viking_uri from openviking_cli.utils import get_logger @@ -78,9 +79,9 @@ class ToolPartRequest(BaseModel): class AutoCommitPolicyRequest(BaseModel): """Session-level auto-commit policy overrides. - All fields are optional. Bounds are not enforced here: ``AutoCommitPolicy`` - clamps every value into ``[0, max]`` so clamping lives in one place across - the HTTP, SDK, and CLI entrypoints. + All fields are optional. ``idle_timeout_seconds`` validates its write-time + minimum here so legacy stored policies remain readable. Other bounds are + clamped by ``AutoCommitPolicy``. """ pending_token_threshold: Optional[int] = None @@ -91,6 +92,11 @@ class AutoCommitPolicyRequest(BaseModel): model_config = {"extra": "forbid"} + @field_validator("idle_timeout_seconds") + @classmethod + def validate_idle_timeout_seconds(cls, value: Optional[int]) -> Optional[int]: + return None if value is None else validate_idle_timeout_input(value) + @model_validator(mode="after") def reject_explicit_null_fields(self) -> "AutoCommitPolicyRequest": null_fields = [field for field in self.model_fields_set if getattr(self, field) is None] diff --git a/openviking/service/core.py b/openviking/service/core.py index 434ed1b64b..593c4c2308 100644 --- a/openviking/service/core.py +++ b/openviking/service/core.py @@ -60,6 +60,7 @@ if TYPE_CHECKING: from openviking.config.vlm import VLMResolver + from openviking.scheduling.expiry.runtime import ExpiryRuntime from openviking.session.compressor_v3 import SessionCompressorV3 @@ -103,6 +104,7 @@ def __init__( self._uri_mutation_coordinator = UriMutationCoordinator() self._watch_scheduler: Optional[WatchScheduler] = None self._session_auto_commit_scheduler: Optional[SessionAutoCommitScheduler] = None + self._expiry_runtime: Optional["ExpiryRuntime"] = None self._encryptor: Optional[Any] = None self._privacy_config_service: Optional[UserPrivacyConfigService] = None self._runtime_config_manager: Optional[Any] = None @@ -577,7 +579,15 @@ async def initialize(self) -> None: except Exception: session_auto_commit_config = SessionAutoCommitConfig() self._session_service.set_session_auto_commit_config(session_auto_commit_config) - if session_auto_commit_config.enabled: + expiry_config = get_openviking_config().expiry + session_expiry_enabled = ( + session_auto_commit_config.enabled + and getattr(session_auto_commit_config, "idle_backend", "local") == "redis" + ) + if ( + session_auto_commit_config.enabled + and getattr(session_auto_commit_config, "idle_backend", "local") == "local" + ): self._session_auto_commit_scheduler = SessionAutoCommitScheduler( self._session_service, session_auto_commit_config, @@ -649,6 +659,29 @@ async def initialize(self) -> None: self._queue_manager.start() logger.info("QueueManager workers started") + if session_expiry_enabled: + from openviking.scheduling.expiry.runtime import ExpiryRuntime, ExpiryRuntimeRoles + + self._expiry_runtime = ExpiryRuntime( + expiry_config, + roles=ExpiryRuntimeRoles( + report=True, + publish=not expiry_config.publisher_separate_process, + consume=True, + ), + ) + try: + self._session_service.configure_expiry( + self._expiry_runtime.scheduler, + consumer_concurrency=expiry_config.session_consumer_concurrency, + ) + await self._expiry_runtime.start_storage() + await self._expiry_runtime.start_roles() + except Exception: + await self._stop_expiry() + raise + logger.info("Expiry roles started") + # Preflight the MinerU endpoint when it will be used, so endpoint # misconfiguration or a stopped service surfaces now instead of on the # first PDF import. Required for strategy="mineru"; advisory for "auto". @@ -663,6 +696,7 @@ async def initialize(self) -> None: logger.info("MinerU preflight passed: %s", pdf_config.mineru_endpoint) except RuntimeError as exc: if pdf_config.strategy == "mineru": + await self._stop_expiry() raise logger.warning( "MinerU preflight failed (fallback will retry on first parse): %s", exc @@ -673,11 +707,32 @@ async def initialize(self) -> None: self._initialized = True logger.info("OpenVikingService initialized") + async def _stop_expiry(self) -> None: + """Stop expiry work before releasing dependencies used by callbacks.""" + runtime = self._expiry_runtime + if runtime is None: + self._session_service.set_session_expiry_adapter(None) + return + try: + await runtime.close() + except Exception as exc: + logger.exception("Expiry shutdown failed") + if getattr(exc, "code", None) == "shard_not_stopped": + # Retain the runtime and leave QueueFS/AGFS open for still-running + # callbacks. The caller can retry close after they terminate. + raise + else: + self._session_service.set_session_expiry_adapter(None) + self._expiry_runtime = None + logger.info("Expiry stopped") + async def close(self) -> None: """Close OpenViking and release resources.""" await self._resource_service.close_background_tasks() vlm_resolver = getattr(self, "_vlm_resolver", None) + await self._stop_expiry() + if self._runtime_config_manager: await self._runtime_config_manager.stop_refresh_loop() self._runtime_config_manager = None diff --git a/openviking/service/session_expiry.py b/openviking/service/session_expiry.py new file mode 100644 index 0000000000..cd5faa82a9 --- /dev/null +++ b/openviking/service/session_expiry.py @@ -0,0 +1,156 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Session business adapter for distributed expiry scheduling.""" + +from __future__ import annotations + +import asyncio +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Any + +from openviking.scheduling.expiry import ( + SESSION_ITEM_TYPE, + CallbackRequest, + CallbackResult, + CallbackStatus, + ExpiryMode, + ExpiryRegistration, + ExpiryScheduler, + HandlerOptions, + ItemIdentity, +) +from openviking.server.identity import RequestContext, Role +from openviking.session.auto_commit_policy import AutoCommitPolicy +from openviking.session.session import SessionExpirySnapshot +from openviking_cli.session.user_id import UserIdentifier +from openviking_cli.utils import get_logger + +if TYPE_CHECKING: + from openviking.service.session_service import SessionService + +logger = get_logger(__name__) + + +def install_session_expiry( + scheduler: ExpiryScheduler, + sessions: "SessionService", + *, + concurrency: int = 8, +) -> None: + scheduler.register_handler( + SESSION_ITEM_TYPE, + SessionTimeoutHandler(sessions, scheduler), + options=HandlerOptions( + concurrency=concurrency, + max_attempts=3, + metadata_enabled=False, + ), + ) + sessions.set_session_expiry_adapter(SessionExpiryAdapter(scheduler)) + + +def effective_idle_timeout(policy: dict | None) -> int | None: + if policy is None: + return None + timeout = AutoCommitPolicy.from_dict(policy).idle_timeout_seconds + return timeout if timeout > 0 else None + + +def timestamp_ms(value: str) -> int: + normalized = value[:-1] + "+00:00" if value.endswith("Z") else value + parsed = datetime.fromisoformat(normalized) + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + return int(parsed.astimezone(timezone.utc).timestamp() * 1000) + + +class SessionExpiryAdapter: + def __init__(self, scheduler: ExpiryScheduler) -> None: + self._scheduler = scheduler + + def build_report(self, snapshot: SessionExpirySnapshot) -> ExpiryRegistration | None: + timeout = effective_idle_timeout(snapshot.auto_commit_policy) + if timeout is None: + return None + return ExpiryRegistration( + identity=ItemIdentity( + SESSION_ITEM_TYPE, + snapshot.account_id, + snapshot.user_id, + snapshot.session_id, + ), + due_at_ms=timestamp_ms(snapshot.last_message_at) + timeout * 1000, + timeout_seconds=timeout, + mode=ExpiryMode.IDLE_ONESHOT, + ) + + async def report_after_append(self, snapshot: SessionExpirySnapshot) -> None: + try: + report = self.build_report(snapshot) + if report is not None: + await self._scheduler.report(report) + except asyncio.CancelledError: + raise + except Exception as exc: + self.record_report_failure("report_failed", exc) + + async def unregister_after_delete(self, identity: ItemIdentity) -> None: + try: + await self._scheduler.unregister(identity) + except asyncio.CancelledError: + raise + except Exception as exc: + self.record_report_failure("unregister_failed", exc) + + def record_report_failure(self, reason: str, exc: Exception | None = None) -> None: + logger.warning( + "Session expiry operation failed reason=%s error=%s", + reason, + type(exc).__name__ if exc is not None else "unknown", + exc_info=exc is not None, + ) + + +class SessionTimeoutHandler: + def __init__(self, sessions: "SessionService", scheduler: ExpiryScheduler) -> None: + self._sessions = sessions + self._scheduler = scheduler + + async def handle(self, request: CallbackRequest) -> CallbackResult: + result = await self._sessions.commit_on_timeout( + request.identity.local_id, self._make_context(request.identity) + ) + if result.get("status") == "deferred": + await self._report_deferred(request.identity, result) + return CallbackResult(CallbackStatus.COMPLETED, "idle_rescheduled") + return self._to_callback_result(result) + + async def _report_deferred(self, identity: ItemIdentity, result: dict[str, Any]) -> None: + if result.get("reason") != "idle_cooldown": + raise ValueError("Unexpected deferred session result") + await self._scheduler.report( + ExpiryRegistration( + identity=identity, + due_at_ms=result["due_at_ms"], + timeout_seconds=result["timeout_seconds"], + mode=ExpiryMode.IDLE_ONESHOT, + ) + ) + + def _make_context(self, identity: ItemIdentity) -> RequestContext: + return RequestContext( + user=UserIdentifier(account_id=identity.account_id, user_id=identity.user_id), + role=Role.USER, + ) + + def _to_callback_result(self, result: dict[str, Any]) -> CallbackResult: + status = result.get("status") + if status == "accepted": + return CallbackResult(CallbackStatus.COMPLETED, "accepted") + if status == "skipped" and result.get("reason") in { + "not_found", + "idle_disabled", + "no_messages", + }: + return CallbackResult(CallbackStatus.COMPLETED, str(result["reason"])) + return CallbackResult(CallbackStatus.RETRY, "unexpected_commit_result") diff --git a/openviking/service/session_service.py b/openviking/service/session_service.py index c33f071290..ee32755555 100644 --- a/openviking/service/session_service.py +++ b/openviking/service/session_service.py @@ -8,7 +8,7 @@ import asyncio from dataclasses import replace -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Dict, List, Optional from openviking.core.namespace import canonical_session_uri @@ -28,6 +28,7 @@ ) from openviking.service.task_tracker import get_task_tracker from openviking.session import Session +from openviking.session.archive_store import is_storage_not_found from openviking.session.auto_commit_policy import AutoCommitPolicy from openviking.session.memory.memory_type_registry import get_default_registry from openviking.session.memory_policy import MemoryPolicy @@ -46,8 +47,18 @@ logger = get_logger(__name__) + +def _parse_utc_timestamp(value: str) -> datetime: + parsed = parse_iso_datetime(value) + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + return parsed.astimezone(timezone.utc) + + if TYPE_CHECKING: from openviking.config.vlm import VLMResolver + from openviking.scheduling.expiry import ExpiryScheduler + from openviking.service.session_expiry import SessionExpiryAdapter from openviking.session.compressor_v3 import SessionCompressorV3 from openviking.usage_reporter import UsageReporter @@ -79,6 +90,24 @@ def __init__( self._auto_commit_inflight_lock = asyncio.Lock() # Strong refs so spawned tasks aren't GC'd mid-await. self._auto_commit_tasks: set[asyncio.Task] = set() + self._session_expiry_adapter: Optional["SessionExpiryAdapter"] = None + + def configure_expiry( + self, + scheduler: "ExpiryScheduler", + *, + consumer_concurrency: int = 8, + ) -> None: + from openviking.service.session_expiry import install_session_expiry + + install_session_expiry( + scheduler, + self, + concurrency=consumer_concurrency, + ) + + def set_session_expiry_adapter(self, adapter: Optional["SessionExpiryAdapter"]) -> None: + self._session_expiry_adapter = adapter def set_dependencies( self, @@ -218,6 +247,7 @@ def session( memory_policy_provider=lambda: self._get_user_memory_policy(ctx), usage_reporter=self._usage_reporter, vlm_resolver=self._vlm_resolver, + expiry_reporter=self._session_expiry_adapter, ) async def _get_user_memory_policy(self, ctx: RequestContext) -> Optional[Dict[str, Any]]: @@ -362,6 +392,12 @@ async def delete(self, session_id: str, ctx: RequestContext) -> bool: raise NotFoundError(session_id, "session") await self._viking_fs.rm(session_uri, recursive=True, ctx=ctx) + if self._session_expiry_adapter is not None: + from openviking.scheduling.expiry import SESSION_ITEM_TYPE, ItemIdentity + + await self._session_expiry_adapter.unregister_after_delete( + ItemIdentity(SESSION_ITEM_TYPE, ctx.account_id, ctx.user.user_id, session_id) + ) logger.info(f"Deleted session: {session_id}") self._record_lifecycle_metric("delete", "ok") return True @@ -454,6 +490,58 @@ async def commit_async( self._record_archive_metric("ok" if result.get("archived") else "skip") return result + async def commit_on_timeout(self, session_id: str, ctx: RequestContext) -> Dict[str, Any]: + """Treat an idle-timeout callback as a signal to commit current messages. + + The callback does not recheck ``last_message_at``. Messages appended + after the timeout was scheduled are included in the commit snapshot; + later appends schedule their own timeout callback. + """ + try: + session = await self.get(session_id, ctx, auto_create=False) + except Exception as exc: + if not is_storage_not_found(exc): + raise + return {"status": "skipped", "reason": "not_found"} + decision = self._check_timeout_commit(session, datetime.now(timezone.utc)) + if decision is not None: + return decision + result = await session.commit_async( + keep_recent_count=0, + persist_keep_recent_count=False, + record_auto_commit_success=True, + record_idle_commit_success=True, + ) + self._record_lifecycle_metric("commit", "ok" if result.get("status") else "error") + self._record_archive_metric("ok" if result.get("archived") else "skip") + return result + + @staticmethod + def _check_timeout_commit(session: Session, now: datetime) -> Optional[Dict[str, Any]]: + """Check terminal skips and cooldown without rechecking idle activity.""" + policy_data = session.meta.auto_commit_policy + if policy_data is None: + return {"status": "skipped", "reason": "idle_disabled"} + timeout_seconds = AutoCommitPolicy.from_dict(policy_data).idle_timeout_seconds + if timeout_seconds <= 0: + return {"status": "skipped", "reason": "idle_disabled"} + if not session.messages: + return {"status": "skipped", "reason": "no_messages"} + if now.tzinfo is None: + now = now.replace(tzinfo=timezone.utc) + now_ms = int(now.timestamp() * 1000) + accepted_at = session.meta.last_idle_commit_accepted_at + if accepted_at: + accepted = _parse_utc_timestamp(accepted_at) + if (now - accepted).total_seconds() < timeout_seconds: + return { + "status": "deferred", + "reason": "idle_cooldown", + "timeout_seconds": timeout_seconds, + "due_at_ms": now_ms + timeout_seconds * 1000, + } + return None + async def get_commit_task(self, task_id: str, ctx: RequestContext) -> Optional[Dict[str, Any]]: """Query background commit task status by task_id for the calling owner.""" task = await get_task_tracker().get( diff --git a/openviking/session/auto_commit_policy.py b/openviking/session/auto_commit_policy.py index e01eb34ece..050fdcad39 100644 --- a/openviking/session/auto_commit_policy.py +++ b/openviking/session/auto_commit_policy.py @@ -22,7 +22,8 @@ DEFAULT_KEEP_RECENT_COUNT = 0 DEFAULT_MIN_COMMIT_INTERVAL_SECONDS = 0 -# PRD upper bounds. +# PRD bounds. +MIN_IDLE_TIMEOUT_SECONDS = 300 # 5 minutes MAX_PENDING_TOKEN_THRESHOLD = 1000000 MAX_MESSAGE_COUNT_THRESHOLD = 1000 MAX_IDLE_TIMEOUT_SECONDS = 604800 # 7 days @@ -48,6 +49,20 @@ def _coerce_int(value: Any, *, field: str, minimum: int, maximum: int) -> int: return parsed +def validate_idle_timeout_input(value: Any) -> int: + """Validate a newly supplied idle timeout without changing stored policies.""" + try: + parsed = int(value) + except (TypeError, ValueError): + raise ValueError("auto_commit_policy.idle_timeout_seconds must be an integer") + if parsed != 0 and parsed < MIN_IDLE_TIMEOUT_SECONDS: + raise ValueError( + "auto_commit_policy.idle_timeout_seconds must be 0 or at least " + f"{MIN_IDLE_TIMEOUT_SECONDS}" + ) + return parsed + + @dataclass class AutoCommitPolicy: """Effective automatic-commit policy for one session.""" @@ -103,9 +118,7 @@ def from_dict(cls, data: Any) -> "AutoCommitPolicy": maximum=MAX_MESSAGE_COUNT_THRESHOLD, ), min_commit_interval_seconds=_coerce_int( - data.get( - "min_commit_interval_seconds", default.min_commit_interval_seconds - ), + data.get("min_commit_interval_seconds", default.min_commit_interval_seconds), field="min_commit_interval_seconds", minimum=0, maximum=MAX_IDLE_TIMEOUT_SECONDS, diff --git a/openviking/session/session.py b/openviking/session/session.py index 084115951a..c66cce97b3 100644 --- a/openviking/session/session.py +++ b/openviking/session/session.py @@ -6,12 +6,13 @@ """ import asyncio +import copy import inspect import json import re from dataclasses import dataclass, field, replace from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, List, Optional, Protocol from uuid import uuid4 from openviking.core.context import ContextLevel @@ -320,6 +321,9 @@ class SessionMeta: # Timestamp of the most recent successful auto-commit, surfaced via session # GET and used to throttle auto-commit frequency. last_auto_commit_at: str = "" + # Timeout commits stamp their accepted time after the Phase 1 ready marker; + # this remains separate from the legacy auto-commit statistic. + last_idle_commit_accepted_at: str = "" # Default custom scalar tags applied to event memories extracted from this # session. Maps to config.memory_extraction_config.events.tags in the API. # None means no session default; a commit may still override per-call. @@ -350,6 +354,7 @@ def to_dict(self) -> Dict[str, Any]: ), "last_message_at": self.last_message_at, "last_auto_commit_at": self.last_auto_commit_at, + "last_idle_commit_accepted_at": self.last_idle_commit_accepted_at, } if self.total_message_count is not None: data["total_message_count"] = self.total_message_count @@ -404,10 +409,26 @@ def from_dict(cls, data: Dict[str, Any]) -> "SessionMeta": auto_commit_policy=data.get("auto_commit_policy"), last_message_at=data.get("last_message_at", ""), last_auto_commit_at=data.get("last_auto_commit_at", ""), + last_idle_commit_accepted_at=data.get("last_idle_commit_accepted_at", ""), event_search_tags=data.get("event_search_tags"), ) +@dataclass(frozen=True) +class SessionExpirySnapshot: + account_id: str + user_id: str + session_id: str + last_message_at: str + auto_commit_policy: Optional[Dict[str, Any]] + + +class SessionExpiryReporter(Protocol): + async def report_after_append(self, snapshot: SessionExpirySnapshot) -> None: ... + + def record_report_failure(self, reason: str, exc: Optional[Exception] = None) -> None: ... + + class Session: """Session management class - Message = role + parts.""" @@ -427,6 +448,7 @@ def __init__( agent_evolution_enabled_provider: Optional[Callable[[], bool | Awaitable[bool]]] = None, memory_policy_provider: Optional[MemoryPolicyProvider] = None, vlm_resolver: Optional["VLMResolver"] = None, + expiry_reporter: Optional[SessionExpiryReporter] = None, ): self._viking_fs = viking_fs self._vikingdb_manager = vikingdb_manager @@ -466,6 +488,7 @@ def __init__( self._memory_policy_provider = memory_policy_provider self._usage_reporter = usage_reporter self._vlm_resolver = vlm_resolver + self._expiry_reporter = expiry_reporter async def _get_vlm_config(self): if self._vlm_resolver is None: @@ -719,6 +742,9 @@ async def _append_messages_authoritatively(self, messages: List[Message]) -> Non lease = await self._viking_fs._async_agfs.pathlock_acquire_exact( session_path, timeout_secs=_SESSION_PHASE1_LOCK_TIMEOUT_SECONDS ) + snapshot: Optional[SessionExpirySnapshot] = None + snapshot_error: Optional[Exception] = None + metadata_reload_error: Optional[Exception] = None try: live_messages_missing = False try: @@ -735,10 +761,11 @@ async def _append_messages_authoritatively(self, messages: List[Message]) -> Non ctx=self.ctx, ) self._meta = SessionMeta.from_dict(json.loads(meta_content)) - except Exception: + except Exception as exc: # Legacy/malformed metadata must not prevent an otherwise safe # append. Message correctness remains rooted in messages.jsonl. self._meta = in_memory_meta + metadata_reload_error = exc await self._apply_appended_messages_to_state(messages) batch_content = "".join(message.to_jsonl() + "\n" for message in messages) @@ -755,8 +782,42 @@ async def _append_messages_authoritatively(self, messages: List[Message]) -> Non ctx=self.ctx, ) await self._save_meta() + if self._expiry_reporter is not None: + try: + # _save_meta() made this exact in-memory state authoritative + # under the Session path lock, even when the preceding reload + # failed. Reporting it prevents the previous due time from + # firing early and repairs a missing expiry registration. + snapshot = self._build_expiry_snapshot_locked() + except Exception as exc: + snapshot_error = exc finally: await self._viking_fs._async_agfs.pathlock_release(lease) + if self._expiry_reporter is not None: + if metadata_reload_error is not None: + logger.warning( + "Session metadata reload failed before append; reporting repaired state", + exc_info=metadata_reload_error, + ) + if snapshot is not None: + await self._expiry_reporter.report_after_append(snapshot) + else: + self._expiry_reporter.record_report_failure("snapshot_build_failed", snapshot_error) + + def _build_expiry_snapshot_locked(self) -> SessionExpirySnapshot: + if not self._meta.last_message_at: + raise ValueError("Session last_message_at is missing") + account_id = self._meta.created_by_account_id or self.ctx.account_id + user_id = self._meta.created_by_user_id or self.ctx.user.user_id + if not account_id or not user_id: + raise ValueError("Session owner identity is incomplete") + return SessionExpirySnapshot( + account_id=account_id, + user_id=user_id, + session_id=self.session_id, + last_message_at=self._meta.last_message_at, + auto_commit_policy=copy.deepcopy(self._meta.auto_commit_policy), + ) async def _apply_appended_messages_to_state(self, messages: List[Message]) -> None: """Update in-memory counters after an authoritative root reload.""" @@ -1255,6 +1316,7 @@ async def commit_async( min_raw_tail_steps: Optional[int] = None, persist_keep_recent_count: bool = True, record_auto_commit_success: bool = False, + record_idle_commit_success: bool = False, event_tags: Optional[List[str]] = None, reset_context: bool = False, ) -> Dict[str, Any]: @@ -1283,6 +1345,9 @@ async def commit_async( record_auto_commit_success: When ``True``, clear the auto-commit error fields and stamp ``last_auto_commit_at`` in the same lock-protected meta update as the commit boundary. + record_idle_commit_success: When ``True``, best-effort stamp + ``last_idle_commit_accepted_at`` after the Phase 1 ready marker + is durable and before releasing the Session lock. event_tags: Per-commit override for the custom scalar tags applied to event memories. ``None`` uses the session default (``meta.event_search_tags``); an empty list disables default-tag @@ -1621,6 +1686,20 @@ async def commit_async( self._meta.last_auto_commit_at = get_current_timestamp() await self._save_meta() await self._write_phase1_ready_marker(archive_uri) + if record_idle_commit_success: + previous_accepted_at = self._meta.last_idle_commit_accepted_at + previous_updated_at = self._meta.updated_at + try: + self._meta.last_idle_commit_accepted_at = get_current_timestamp() + await self._save_meta() + except Exception: + self._meta.last_idle_commit_accepted_at = previous_accepted_at + self._meta.updated_at = previous_updated_at + logger.warning( + "Failed to record idle commit time for %s", + self.session_id, + exc_info=True, + ) except Exception as e: logger.error(f"[commit] Failed during {phase1_stage}: {e}") # Whether the queue write failed or a queued Phase 1 stopped diff --git a/openviking_cli/expiry_bootstrap.py b/openviking_cli/expiry_bootstrap.py new file mode 100644 index 0000000000..34038c9cc5 --- /dev/null +++ b/openviking_cli/expiry_bootstrap.py @@ -0,0 +1,127 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Internal distributed expiry publisher process.""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import logging +import os +import signal +from pathlib import Path + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Internal OpenViking expiry publisher") + parser.add_argument("--config", required=True, type=Path) + parser.add_argument("--startup-status", required=True, type=Path) + return parser.parse_args(argv) + + +def load_expiry_config(path: Path): + from openviking_cli.utils.config.open_viking_config import OpenVikingConfig + + raw = json.loads(os.path.expandvars(path.expanduser().read_text(encoding="utf-8-sig"))) + config = OpenVikingConfig.from_dict(raw) + idle = config.memory.session_auto_commit + if not idle.enabled or idle.idle_backend != "redis": + raise ValueError( + "expiry publisher requires memory.session_auto_commit.enabled=true " + "and idle_backend='redis'" + ) + if not config.expiry.publisher_separate_process: + raise ValueError("expiry publisher process is disabled by configuration") + return config.expiry + + +def write_startup_status(path: Path, status: str, error: str | None = None) -> None: + payload: dict[str, object] = {"pid": os.getpid(), "status": status} + if error: + payload["error"] = error + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") + temporary.write_text(json.dumps(payload), encoding="utf-8") + os.replace(temporary, path) + + +def install_signal_handlers(stop: asyncio.Event) -> None: + loop = asyncio.get_running_loop() + for signum in (signal.SIGINT, signal.SIGTERM): + try: + loop.add_signal_handler(signum, stop.set) + except NotImplementedError: + signal.signal(signum, lambda *_args: loop.call_soon_threadsafe(stop.set)) + + +async def run_publisher(config, stop: asyncio.Event, startup_status: Path) -> None: + from openviking.scheduling.expiry.runtime import ExpiryRuntime, ExpiryRuntimeRoles + + runtime = ExpiryRuntime(config, roles=ExpiryRuntimeRoles(publish=True)) + wait_task: asyncio.Task[None] | None = None + stop_task: asyncio.Task[bool] | None = None + try: + await runtime.start_storage() + await runtime.start_roles() + write_startup_status(startup_status, "ready") + wait_task = asyncio.create_task(runtime.wait()) + stop_task = asyncio.create_task(stop.wait()) + try: + done, _ = await asyncio.wait( + {wait_task, stop_task}, return_when=asyncio.FIRST_COMPLETED + ) + if wait_task in done: + await wait_task + finally: + stop.set() + await runtime.close() + finally: + if wait_task is not None: + if not wait_task.done(): + wait_task.cancel() + await asyncio.gather(wait_task, return_exceptions=True) + if stop_task is not None: + if not stop_task.done(): + stop_task.cancel() + await asyncio.gather(stop_task, return_exceptions=True) + if wait_task is None: + await runtime.close() + + +def main(argv: list[str] | None = None) -> int: + startup_status: Path | None = None + try: + args = parse_args(argv) + current_startup_status: Path = args.startup_status + startup_status = current_startup_status + try: + config = load_expiry_config(args.config) + except ValueError as exc: + # Validation errors are already safely formatted. Avoid emitting + # chained validation tracebacks containing raw configuration input. + write_startup_status(current_startup_status, "failed", str(exc)) + logging.error("Expiry publisher configuration failed: %s", exc) + return 1 + write_startup_status(current_startup_status, "starting") + + async def runner() -> None: + stop = asyncio.Event() + install_signal_handlers(stop) + await run_publisher(config, stop, current_startup_status) + + asyncio.run(runner()) + return 0 + except KeyboardInterrupt: + return 0 + except Exception as exc: + if startup_status is not None: + try: + write_startup_status(startup_status, "failed", str(exc)) + except Exception: + pass + logging.exception("Expiry publisher failed") + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/openviking_cli/utils/config/__init__.py b/openviking_cli/utils/config/__init__.py index 75b31354a8..93e82f17d5 100644 --- a/openviking_cli/utils/config/__init__.py +++ b/openviking_cli/utils/config/__init__.py @@ -53,6 +53,7 @@ SYSTEM_CONFIG_DIR, ) from .embedding_config import EmbeddingConfig +from .expiry_config import ExpiryConfig, ExpiryRedisConfig from .git_config import GitConfig, GitLocalConfig, GitS3Config from .glob_config import GlobConfig, GlobEngine from .grep_config import GrepConfig, GrepEngine @@ -146,6 +147,8 @@ "GrepEngine", "RerankConfig", "RetrievalConfig", + "ExpiryConfig", + "ExpiryRedisConfig", "StorageConfig", "OpenGaussConfig", "VectorDBBackendConfig", diff --git a/openviking_cli/utils/config/expiry_config.py b/openviking_cli/utils/config/expiry_config.py new file mode 100644 index 0000000000..fd7ca560a6 --- /dev/null +++ b/openviking_cli/utils/config/expiry_config.py @@ -0,0 +1,5 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +from openviking.scheduling.expiry.config import ExpiryConfig, ExpiryRedisConfig + +__all__ = ["ExpiryConfig", "ExpiryRedisConfig"] diff --git a/openviking_cli/utils/config/memory_config.py b/openviking_cli/utils/config/memory_config.py index 92e358f8fd..130b761bf6 100644 --- a/openviking_cli/utils/config/memory_config.py +++ b/openviking_cli/utils/config/memory_config.py @@ -21,6 +21,13 @@ class SessionAutoCommitConfig(BaseModel): "started. When disabled, neither happens." ), ) + idle_backend: Literal["local", "redis"] = Field( + default="local", + description=( + "Idle timeout discovery backend. 'local' scans session metadata; " + "'redis' uses distributed expiry scheduling." + ), + ) check_interval_seconds: float = Field(default=600.0, gt=0) scan_rate_limit_files_per_second: float = Field( default=2.0, diff --git a/openviking_cli/utils/config/open_viking_config.py b/openviking_cli/utils/config/open_viking_config.py index 3c2b6d0abc..60185bf45a 100644 --- a/openviking_cli/utils/config/open_viking_config.py +++ b/openviking_cli/utils/config/open_viking_config.py @@ -31,6 +31,7 @@ ) from .embedding_config import EmbeddingConfig from .encryption_config import EncryptionConfig +from .expiry_config import ExpiryConfig from .git_config import GitConfig from .glob_config import GlobConfig from .grep_config import GrepConfig @@ -464,6 +465,21 @@ def _inherit_git_defaults_from_agfs(cls, data: Any) -> Any: memory: MemoryConfig = Field(default_factory=MemoryConfig, description="Memory configuration") + expiry: ExpiryConfig = Field( + default_factory=ExpiryConfig, + description="Distributed expiry scheduling configuration", + ) + + @model_validator(mode="after") + def _validate_expiry_backend(self) -> "OpenVikingConfig": + self._check_expiry_backend() + return self + + def _check_expiry_backend(self) -> None: + idle = self.memory.session_auto_commit + if idle.enabled and idle.idle_backend == "redis": + self.expiry.validate_for_runtime() + agent_evolution: AgentEvolutionConfig = RuntimeField( default_factory=AgentEvolutionConfig, description="Dynamic cluster default for Agent Evolution.", @@ -495,7 +511,7 @@ def _inherit_git_defaults_from_agfs(cls, data: Any) -> Any: ), ) - model_config = {"arbitrary_types_allowed": True} + model_config = {"arbitrary_types_allowed": True, "hide_input_in_errors": True} @classmethod def from_dict(cls, config: Dict[str, Any]) -> "OpenVikingConfig": @@ -582,7 +598,8 @@ def from_dict(cls, config: Dict[str, Any]) -> "OpenVikingConfig": error=e, path_prefix="memory", ) - ) from e + ) from None + instance._check_expiry_backend() # Apply parser configurations for parser_type in parser_types: @@ -615,7 +632,9 @@ def from_dict(cls, config: Dict[str, Any]) -> "OpenVikingConfig": ) return instance except ValidationError as e: - raise ValueError(format_validation_error(root_model=cls, error=e)) from e + # The formatted error omits input values; do not reattach a raw + # validation traceback which could expose nested credentials. + raise ValueError(format_validation_error(root_model=cls, error=e)) from None def to_dict(self) -> Dict[str, Any]: """Convert configuration to dictionary.""" diff --git a/pyproject.toml b/pyproject.toml index 57f65f8181..94f5172e48 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -94,6 +94,9 @@ dependencies = [ [project.optional-dependencies] +expiry = [ + "redis>=5,<7", +] gateway-fast = [ "uvloop>=0.21; sys_platform != 'win32'", "httptools>=0.6", @@ -239,6 +242,7 @@ exclude = ["tests*", "docs*", "examples*"] [tool.setuptools.package-data] openviking_gateway = ["client-rules.json"] +"openviking.scheduling.expiry" = ["lua/*.lua"] "vikingbot.studio.providers.feishu" = ["icon.png"] openviking = [ "prompts/templates/**/*.yaml", diff --git a/setup.py b/setup.py index d6a5260817..74ba4393f6 100644 --- a/setup.py +++ b/setup.py @@ -579,6 +579,7 @@ def has_ext_modules(self): cmdclass=cmdclass, package_data={ "openviking_gateway": ["client-rules.json"], + "openviking.scheduling.expiry": ["lua/*.lua"], "openviking": [ "lib/ragfs_python*.so", "lib/ragfs_python*.pyd", diff --git a/tests/cli/conftest.py b/tests/cli/conftest.py index 35ff6a0b78..1a5480c949 100644 --- a/tests/cli/conftest.py +++ b/tests/cli/conftest.py @@ -493,8 +493,6 @@ def ov_add_resource(path, to_uri, *extra_args, attempts=15, interval=10): to_uri, *extra_args, "--wait", - "--timeout", - ADD_RESOURCE_WAIT_TIMEOUT, "-o", "json", ] diff --git a/tests/server/test_agent_evolution_global_setting.py b/tests/server/test_agent_evolution_global_setting.py index b7c036fc3c..7536f2a58e 100644 --- a/tests/server/test_agent_evolution_global_setting.py +++ b/tests/server/test_agent_evolution_global_setting.py @@ -7,6 +7,7 @@ import httpx import pytest import pytest_asyncio +from pydantic import ValidationError from openviking.pyagfs import AGFSNotFoundError from openviking.server.app import create_app @@ -18,6 +19,7 @@ from openviking.service.session_service import SessionService from openviking.session import Session from openviking_cli.session.user_id import UserIdentifier +from openviking_cli.utils.config.open_viking_config import OpenVikingConfig class _FakeAGFS: @@ -195,13 +197,15 @@ async def test_server_agent_evolution_updates_initialized_runtime_baseline(): OpenVikingConfigSingleton.reset_instance() -def test_server_default_memory_policy_is_configured_on_session_service( - fake_viking_fs, monkeypatch -): +def test_server_default_memory_policy_is_configured_on_session_service(fake_viking_fs, monkeypatch): monkeypatch.setattr( "openviking.service.session_service.get_default_registry", lambda: SimpleNamespace(list_names=lambda **_: ["profile"]), ) + monkeypatch.setattr( + "openviking.server.app.get_openviking_config", + lambda: OpenVikingConfig(), + ) sessions = SessionService(viking_fs=fake_viking_fs) service = SimpleNamespace(sessions=sessions) config = ServerConfig( @@ -214,13 +218,18 @@ def test_server_default_memory_policy_is_configured_on_session_service( "self": {"enabled": True}, "peer": {"enabled": True}, "memory_types": ["profile"], + "working_memory": {"enabled": False}, } async def test_auto_commit_defaults_apply_to_new_sessions(fake_viking_fs, monkeypatch): + ov_config = OpenVikingConfig() + monkeypatch.setattr("openviking.server.app.get_openviking_config", lambda: ov_config) + with pytest.raises(ValidationError, match="must be 0 or at least 300"): + ServerConfig(user_config_defaults={"auto_commit_policy": {"idle_timeout_seconds": 299}}) sessions = SessionService(viking_fs=fake_viking_fs) config = ServerConfig( - user_config_defaults={"auto_commit_policy": {"idle_timeout_seconds": 7200}} + user_config_defaults={"auto_commit_policy": {"idle_timeout_seconds": 300}} ) create_app( config=config, @@ -311,9 +320,7 @@ async def test_account_settings_admin_api_reads_and_updates_effective_value( "agent_evolution": {"enabled": True}, "acl": {"enabled": True}, } - assert ( - await service.runtime_config_manager.get_account("default", "agent_evolution") - ).enabled + assert (await service.runtime_config_manager.get_account("default", "agent_evolution")).enabled assert (await service.runtime_config_manager.get_account("default", "acl")).enabled assert await service.viking_fs.acl_manager.is_enabled("default") @@ -375,9 +382,7 @@ async def test_account_configuration_exposes_three_state_layer(settings_http): json={"settings": {"acl": None}}, ) assert removed.status_code == 200, removed.text - assert removed.json()["result"]["settings"] == { - "github": {"token": "account-token"} - } + assert removed.json()["result"]["settings"] == {"github": {"token": "account-token"}} from openviking.server.auth import get_request_context @@ -418,9 +423,7 @@ async def test_cluster_agent_evolution_override_is_account_fallback(settings_htt json={"settings": {"agent_evolution": {"enabled": True}}}, ) assert updated.status_code == 200, updated.text - assert updated.json()["result"]["settings"] == { - "agent_evolution": {"enabled": True} - } + assert updated.json()["result"]["settings"] == {"agent_evolution": {"enabled": True}} assert await service.sessions.get_agent_evolution_enabled("default") await service.runtime_config_manager.patch_account( @@ -428,9 +431,7 @@ async def test_cluster_agent_evolution_override_is_account_fallback(settings_htt ) assert not await service.sessions.get_agent_evolution_enabled("default") - await service.runtime_config_manager.patch_account( - "default", {"agent_evolution": None} - ) + await service.runtime_config_manager.patch_account("default", {"agent_evolution": None}) assert await service.sessions.get_agent_evolution_enabled("default") @@ -443,9 +444,9 @@ async def test_legacy_admin_configuration_routes_are_deprecated(settings_http): legacy = schema["paths"]["/api/v1/admin/accounts/{account_id}/settings"] assert legacy["get"]["deprecated"] assert legacy["patch"]["deprecated"] - assert not schema["paths"][ - "/api/v1/admin/accounts/{account_id}/configuration" - ]["patch"].get("deprecated", False) + assert not schema["paths"]["/api/v1/admin/accounts/{account_id}/configuration"]["patch"].get( + "deprecated", False + ) assert "/api/v1/admin/configuration" in schema["paths"] assert "/api/v1/admin/settings" not in schema["paths"] @@ -480,9 +481,7 @@ async def test_agent_evolution_endpoint_keeps_name_and_uses_account_settings( assert response.status_code == 200, response.text assert response.json()["result"]["enabled"] is True - assert ( - await service.runtime_config_manager.get_account("default", "agent_evolution") - ).enabled + assert (await service.runtime_config_manager.get_account("default", "agent_evolution")).enabled def test_deprecated_user_agent_evolution_config_is_not_persisted(): diff --git a/tests/server/test_api_sessions.py b/tests/server/test_api_sessions.py index 1240aef32b..8f80c504d7 100644 --- a/tests/server/test_api_sessions.py +++ b/tests/server/test_api_sessions.py @@ -21,6 +21,7 @@ from openviking.server.dependencies import set_service from openviking.server.identity import RequestContext, Role from openviking.server.routers import sessions as sessions_router +from openviking.session.auto_commit_policy import MAX_PENDING_TOKEN_THRESHOLD, AutoCommitPolicy from openviking.session.session import Session from openviking_cli.session.user_id import UserIdentifier from openviking_cli.utils.config import OPENVIKING_CONFIG_ENV @@ -673,13 +674,7 @@ async def test_create_session_uses_default_policy_when_server_default_enabled( resp = await client.post("/api/v1/sessions", json={}) assert resp.status_code == 200 - assert resp.json()["result"]["auto_commit_policy"] == { - "pending_token_threshold": 10000, - "message_count_threshold": 50, - "idle_timeout_seconds": 86400, - "keep_recent_count": 2, - "min_commit_interval_seconds": 0, - } + assert resp.json()["result"]["auto_commit_policy"] == AutoCommitPolicy.default().to_dict() async def test_create_session_can_disable_server_default_auto_commit( @@ -722,13 +717,10 @@ async def test_auto_created_session_uses_default_policy_when_server_default_enab session_resp = await client.get(f"/api/v1/sessions/{session_id}") assert session_resp.status_code == 200 - assert session_resp.json()["result"]["auto_commit_policy"] == { - "pending_token_threshold": 10000, - "message_count_threshold": 50, - "idle_timeout_seconds": 86400, - "keep_recent_count": 2, - "min_commit_interval_seconds": 0, - } + assert ( + session_resp.json()["result"]["auto_commit_policy"] + == AutoCommitPolicy.default().to_dict() + ) async def test_create_session_applies_config_and_fills_defaults(client: httpx.AsyncClient): @@ -743,13 +735,9 @@ async def test_create_session_applies_config_and_fills_defaults(client: httpx.As ) assert resp.status_code == 200 result = resp.json()["result"] - assert result["auto_commit_policy"] == { - "pending_token_threshold": 8000, - "message_count_threshold": 50, - "idle_timeout_seconds": 86400, - "keep_recent_count": 10, - "min_commit_interval_seconds": 0, - } + expected_policy = AutoCommitPolicy.default().to_dict() + expected_policy.update({"pending_token_threshold": 8000, "keep_recent_count": 10}) + assert result["auto_commit_policy"] == expected_policy session_resp = await client.get(f"/api/v1/sessions/{result['session_id']}") session_result = session_resp.json()["result"] @@ -764,10 +752,19 @@ async def test_create_session_clamps_config_above_bounds(client: httpx.AsyncClie ) assert resp.status_code == 200 policy = resp.json()["result"]["auto_commit_policy"] - assert policy["pending_token_threshold"] == 50000 + assert policy["pending_token_threshold"] == MAX_PENDING_TOKEN_THRESHOLD + +async def test_session_idle_policy_validates_writes_and_preserves_stored_values( + client: httpx.AsyncClient, service +): + rejected = await client.post( + "/api/v1/sessions", + json={"auto_commit_policy": {"idle_timeout_seconds": 299}}, + ) + assert rejected.status_code == 400 + assert "must be 0 or at least 300" in rejected.json()["error"]["message"] -async def test_get_session_returns_effective_config(client: httpx.AsyncClient): create_resp = await client.post("/api/v1/sessions", json={}) session_id = create_resp.json()["result"]["session_id"] @@ -775,6 +772,40 @@ async def test_get_session_returns_effective_config(client: httpx.AsyncClient): assert resp.status_code == 200 assert resp.json()["result"]["auto_commit_policy"] is None + # Emulate a stored policy from before the HTTP write-time minimum existed. + ctx = RequestContext(user=DEFAULT_USER, role=Role.ROOT) + session = await service.sessions.get(session_id, ctx) + await session.update_config( + auto_commit_policy={"idle_timeout_seconds": 60}, + update_auto_commit_policy=True, + ) + config_url = f"/api/v1/sessions/{session_id}/config" + updated = await client.patch( + config_url, json={"auto_commit_policy": {"message_count_threshold": 25}} + ) + assert updated.status_code == 200 + policy = updated.json()["result"]["auto_commit_policy"] + assert policy["idle_timeout_seconds"] == 60 + assert policy["message_count_threshold"] == 25 + + rejected = await client.patch( + config_url, json={"auto_commit_policy": {"idle_timeout_seconds": 299}} + ) + assert rejected.status_code == 400 + persisted = await client.get(f"/api/v1/sessions/{session_id}") + assert persisted.json()["result"]["auto_commit_policy"] == policy + + updated = await client.patch( + config_url, json={"auto_commit_policy": {"idle_timeout_seconds": 300}} + ) + assert updated.status_code == 200 + assert updated.json()["result"]["auto_commit_policy"]["idle_timeout_seconds"] == 300 + disabled = await client.patch( + config_url, json={"auto_commit_policy": {"idle_timeout_seconds": 0}} + ) + assert disabled.status_code == 200 + assert disabled.json()["result"]["auto_commit_policy"]["idle_timeout_seconds"] == 0 + async def test_patch_session_config_is_not_supported(client: httpx.AsyncClient): create_resp = await client.post( diff --git a/tests/server/test_bootstrap.py b/tests/server/test_bootstrap.py index 98cd1bdf4a..a7c666a0a3 100644 --- a/tests/server/test_bootstrap.py +++ b/tests/server/test_bootstrap.py @@ -11,8 +11,10 @@ import openviking.server.app as app_module import openviking.server.bootstrap as bootstrap +from openviking.scheduling.expiry.config import ExpiryConfig from openviking.server.config import ServerConfig from openviking.utils.agfs_utils import resolve_queuefs_mount_point +from openviking_cli.utils.config.memory_config import SessionAutoCommitConfig from openviking_cli.utils.config.open_viking_config import OpenVikingConfigSingleton from openviking_cli.utils.config.storage_config import StorageConfig @@ -37,10 +39,9 @@ def test_main_keeps_config_host_when_cli_host_is_omitted(monkeypatch): lambda: None, ) monkeypatch.setattr( - bootstrap, - "OpenVikingConfigSingleton", - SimpleNamespace(initialize=lambda config_path: None), - raising=False, + OpenVikingConfigSingleton, + "initialize", + classmethod(lambda cls, config_path: None), ) monkeypatch.setattr( bootstrap.argparse.ArgumentParser, @@ -91,10 +92,9 @@ def test_main_coerces_cli_host_all_to_none(monkeypatch): lambda: None, ) monkeypatch.setattr( - bootstrap, - "OpenVikingConfigSingleton", - SimpleNamespace(initialize=lambda config_path: None), - raising=False, + OpenVikingConfigSingleton, + "initialize", + classmethod(lambda cls, config_path: None), ) monkeypatch.setattr( bootstrap.argparse.ArgumentParser, @@ -387,7 +387,9 @@ def run(): @pytest.mark.parametrize("app_startup_fails", [False, True]) -def test_main_restart_cleans_up_bot_and_preserves_launch_command(monkeypatch, app_startup_fails): +def test_main_restart_cleans_up_managed_children_and_preserves_launch_command( + monkeypatch, app_startup_fails +): import sys config = ServerConfig(host="127.0.0.1", port=1933, with_bot=True) @@ -398,34 +400,60 @@ def test_main_restart_cleans_up_bot_and_preserves_launch_command(monkeypatch, ap monkeypatch.setattr(bootstrap, "load_server_config", lambda path: config) monkeypatch.setattr(bootstrap, "resolve_config_path", lambda *a: Path("/tmp/explicit.conf")) monkeypatch.setattr( - OpenVikingConfigSingleton, "initialize", classmethod(lambda cls, **kw: None) + OpenVikingConfigSingleton, + "initialize", + classmethod( + lambda cls, **kw: SimpleNamespace( + memory=SimpleNamespace( + session_auto_commit=SessionAutoCommitConfig(enabled=True, idle_backend="redis") + ), + expiry=ExpiryConfig(), + ) + ), ) monkeypatch.setattr( "openviking_cli.utils.ollama.detect_ollama_in_config", lambda c: (False, "", 0) ) monkeypatch.setattr(bootstrap, "configure_uvicorn_logging", lambda: None) - monkeypatch.setattr(bootstrap, "create_app", lambda *a, **kw: "app") + + def create_app(*args, **kwargs): + assert steps == [("start", Path("/tmp/explicit.conf"))] + if app_startup_fails: + raise RuntimeError("application initialization failed") + return "app" + + monkeypatch.setattr(bootstrap, "create_app", create_app) monkeypatch.setattr(bootstrap, "_abort_if_port_in_use", lambda *a: None) monkeypatch.setattr(bootstrap, "_start_vikingbot_gateway", lambda *a, **kw: "bot") + monkeypatch.setattr(bootstrap, "_start_expiry_publisher", lambda *a, **kw: None) monkeypatch.setattr( bootstrap, "_stop_vikingbot_gateway", lambda bot: steps.append(("stop", bot)) ) + monkeypatch.setattr( + bootstrap, + "_start_expiry_publisher", + lambda path: steps.append(("start", path)) or "publisher", + ) + monkeypatch.setattr( + bootstrap, "_stop_expiry_publisher", lambda child: steps.append(("stop", child)) + ) monkeypatch.setattr(bootstrap, "_run_restartable_server", lambda *a, **kw: True) monkeypatch.setattr( os, "execv", lambda executable, args: steps.append(("exec", executable, args)) ) monkeypatch.setenv("OPENVIKING_CONFIG_FILE", "/tmp/explicit.conf") + lifecycle = [ + ("start", Path("/tmp/explicit.conf")), + ("stop", "publisher"), + ("stop", "bot"), + ] if app_startup_fails: - def fail_app(*args, **kwargs): - raise RuntimeError("application initialization failed") - - monkeypatch.setattr(bootstrap, "create_app", fail_app) with pytest.raises(RuntimeError, match="application initialization failed"): bootstrap.main() - assert steps == [("stop", "bot")] + assert steps == lifecycle return bootstrap.main() - assert steps == [("stop", "bot"), ("exec", sys.executable, original_argv)] + assert steps == [*lifecycle, ("exec", sys.executable, original_argv)] assert os.environ["OPENVIKING_CONFIG_FILE"] == "/tmp/explicit.conf" diff --git a/tests/server/test_store_assembly.py b/tests/server/test_store_assembly.py index 136707a2c2..35dc991688 100644 --- a/tests/server/test_store_assembly.py +++ b/tests/server/test_store_assembly.py @@ -3,9 +3,14 @@ import ssl from types import SimpleNamespace +from unittest.mock import AsyncMock import pytest +from openviking.scheduling.expiry.config import ExpiryRedisConfig +from openviking.scheduling.expiry.keys import ExpiryKeys +from openviking.scheduling.expiry.metrics import ExpiryMetrics +from openviking.scheduling.expiry.redis_store import RedisExpiryStore from openviking.server.account_stores.redis_client import create_redis_client from openviking.server.api_keys.legacy import FileStore from openviking.server.config import AccountStoreConfig @@ -17,6 +22,11 @@ from tests.server.account_store_fakes import InMemoryAGFS +class _InlineExecutor: + async def run(self, operation): + return await operation() + + async def test_account_store_connection_configuration(): with pytest.raises(ValueError, match="account_store.provider must be one of"): AccountStoreConfig(provider="missing") @@ -41,6 +51,27 @@ async def test_account_store_connection_configuration(): await handle.close() +async def test_expiry_cluster_disables_client_level_command_retries(monkeypatch): + cluster = pytest.importorskip("redis.asyncio.cluster") + ping = AsyncMock(return_value=True) + monkeypatch.setattr(cluster.RedisCluster, "ping", ping) + store = RedisExpiryStore( + ExpiryRedisConfig(mode="cluster", endpoints=["127.0.0.1:6379"]), + ExpiryKeys("resource"), + _InlineExecutor(), + ExpiryMetrics(), + ) + await store.open() + other_cluster = cluster.RedisCluster(host="127.0.0.1", port=6379) + try: + assert store._client.retry._retries == 0 + assert other_cluster.retry._retries == 3 + assert store._client.retry is not other_cluster.retry + finally: + await store.close() + await other_cluster.aclose() + + @pytest.mark.parametrize("injected", [False, True]) async def test_assembled_manager_supports_account_authentication(injected): account_store = FileStore(SimpleNamespace(agfs=InMemoryAGFS())) diff --git a/tests/session/test_session_commit.py b/tests/session/test_session_commit.py index 4f5c1454d7..a9b90cc717 100644 --- a/tests/session/test_session_commit.py +++ b/tests/session/test_session_commit.py @@ -9,10 +9,17 @@ from unittest.mock import AsyncMock, MagicMock from openviking.message import TextPart +from openviking.scheduling.expiry import ( + CallbackRequest, + CallbackStatus, + ExpiryScheduler, + ItemIdentity, +) from openviking.server.identity import RequestContext from openviking.service.core import OpenVikingService from openviking.service.task_tracker import get_task_tracker from openviking.session import Session +from openviking.utils.time_utils import parse_iso_datetime async def _wait_for_task(task_id: str, timeout: float = 30.0) -> dict: @@ -67,53 +74,61 @@ async def test_commit_preserves_unicode_separators_and_accepts_later_messages( if part["type"] == "text" } - async def test_commit_extracts_memories( + async def test_idle_callback_commits_current_backlog_without_activity_recheck( self, wm_session_with_messages: Session, service: OpenVikingService, + monkeypatch, ): - """Test commit kicks off background memory extraction""" - - async def extract_long_term_memories(**kwargs): - archive_uri = kwargs["archive_uri"] - await wm_session_with_messages._viking_fs.write_file( - uri=f"{archive_uri}/memory_diff.json", - content=json.dumps({"archive_uri": archive_uri}), - ctx=wm_session_with_messages.ctx, - ) - return [] - - wm_session_with_messages._session_compressor.extract_long_term_memories = AsyncMock( - side_effect=extract_long_term_memories + """An expiry callback archives the latest lock-protected snapshot.""" + session = wm_session_with_messages + await session.update_config( + auto_commit_policy={"idle_timeout_seconds": 300, "keep_recent_count": 2}, + update_auto_commit_policy=True, ) - if hasattr(wm_session_with_messages._session_compressor, "extract_execution_memories"): - wm_session_with_messages._session_compressor.extract_execution_memories = AsyncMock( - return_value={"contexts": [], "session_skills": []} - ) - - result = await wm_session_with_messages.commit_async() - task_id = result["task_id"] - - # Wait for background memory extraction to complete - task_result = await _wait_for_task(task_id) - assert task_result["status"] == "completed" - assert ( - task_result["result"]["memory_diff_uri"] - == f"{task_result['result']['archive_uri']}/memory_diff.json" + session.meta.keep_recent_count = 2 + await session._save_meta() + scheduler = MagicMock(spec=ExpiryScheduler) + service.sessions.configure_expiry(scheduler) + handler = scheduler.register_handler.call_args.args[1] + session = await service.sessions.get(session.session_id, session.ctx) + await session.add_message_async("user", [TextPart("latest activity")]) + report = scheduler.report.await_args.args[0] + last_message_at = parse_iso_datetime(session.meta.last_message_at) + assert report.identity == ItemIdentity( + "session", session.ctx.account_id, session.ctx.user.user_id, session.session_id ) - memory_diff = json.loads( - await wm_session_with_messages._viking_fs.read_file( - task_result["result"]["memory_diff_uri"], - ctx=wm_session_with_messages.ctx, - ) + assert report.due_at_ms == int(last_message_at.timestamp() * 1000) + 300_000 + + original_commit = Session.commit_async + commit_results = [] + + async def append_before_commit(loaded_session, **kwargs): + scheduler.report.side_effect = RuntimeError("expiry unavailable") + await loaded_session.add_message_async("user", [TextPart("append after idle check")]) + result = await original_commit(loaded_session, **kwargs) + commit_results.append(result) + return result + + monkeypatch.setattr(Session, "commit_async", append_before_commit) + callback = await handler.handle(CallbackRequest(report.identity)) + assert callback.status is CallbackStatus.COMPLETED + assert callback.reason == "accepted" + assert scheduler.report.await_count == 2 + result = commit_results[0] + assert result["status"] == "accepted" + archive_content = await session._viking_fs.read_file( + f"{result['archive_uri']}/messages.jsonl", ctx=session.ctx ) - assert memory_diff["archive_uri"] == task_result["result"]["archive_uri"] - assert "memories_extracted" in task_result["result"] - memory_counts = task_result["result"]["memories_extracted"] - assert isinstance(memory_counts, dict) - - # The extraction stub emits no semantic or embedding work. Its task - # completion above is the relevant processing boundary for this test. + archived_messages = [json.loads(line) for line in archive_content.split("\n") if line] + assert len(archived_messages) == 6 + assert archived_messages[-1]["parts"][0]["text"] == "append after idle check" + persisted = await service.sessions.get(session.session_id, session.ctx) + assert persisted.messages == [] + assert persisted.meta.keep_recent_count == 2 + assert persisted.meta.last_idle_commit_accepted_at + phase1 = await persisted._read_phase1_meta(result["archive_uri"]) + assert phase1["status"] == "ready" async def test_phase2_splits_with_the_committed_auto_commit_policy( self, diff --git a/tests/session/test_session_retention_integration.py b/tests/session/test_session_retention_integration.py index cf8c0703c5..ec0af852ba 100644 --- a/tests/session/test_session_retention_integration.py +++ b/tests/session/test_session_retention_integration.py @@ -1743,7 +1743,7 @@ async def enqueue(self, *_args, **_kwargs): ) with pytest.raises(RuntimeError, match="queue unavailable"): - await session.commit_async() + await session.commit_async(record_idle_commit_success=True) states = await session._archives.scan_states() context = await session.get_session_context() @@ -1761,6 +1761,7 @@ async def enqueue(self, *_args, **_kwargs): assert [message.content for message in raw] == ["do not lose me"] assert [message["parts"][0]["text"] for message in context["messages"]] == ["do not lose me"] assert context["stats"]["failedArchives"] == 1 + assert session.meta.last_idle_commit_accepted_at == "" async def test_phase1_root_rewrite_failure_marks_orphan_archive_failed( diff --git a/uv.lock b/uv.lock index d88b650c02..2b554c17e5 100644 --- a/uv.lock +++ b/uv.lock @@ -3655,6 +3655,9 @@ eval = [ { name = "pandas", version = "3.0.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "ragas" }, ] +expiry = [ + { name = "redis" }, +] gateway-fast = [ { name = "httptools" }, { name = "uvloop", marker = "sys_platform != 'win32'" }, @@ -3804,6 +3807,7 @@ requires-dist = [ { name = "ragas", marker = "extra == 'test'", specifier = ">=0.1.0" }, { name = "rapidfuzz", specifier = ">=3.14,<4" }, { name = "readability-lxml", marker = "extra == 'bot'", specifier = ">=0.8.0" }, + { name = "redis", marker = "extra == 'expiry'", specifier = ">=5,<7" }, { name = "redis", marker = "extra == 'redis-mysql'", specifier = ">=5.0,<7" }, { name = "redis", marker = "extra == 'test'", specifier = ">=5.0,<7" }, { name = "requests", specifier = ">=2.33.0" }, @@ -3848,7 +3852,7 @@ requires-dist = [ { name = "wheel", marker = "extra == 'build'" }, { name = "xxhash", specifier = ">=3.0.0" }, ] -provides-extras = ["gateway-fast", "mysql", "redis-mysql", "test", "auth", "dev", "doc", "eval", "gemini", "gemini-async", "ocr", "build", "bot", "benchmark", "local-embed", "opengauss"] +provides-extras = ["expiry", "gateway-fast", "mysql", "redis-mysql", "test", "auth", "dev", "doc", "eval", "gemini", "gemini-async", "ocr", "build", "bot", "benchmark", "local-embed", "opengauss"] [package.metadata.requires-dev] dev = [{ name = "pytest", specifier = ">=9.0.2" }]