i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-03 01:40:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: 引入 aether-runtime/cache/data/http/testkit 基础 crate，完善并发门控与审计系统

Browse Source 
新增 crate：
- aether-runtime: 服务运行时基础设施（并发门控、分布式并发、指标、队列、优雅关闭、tracing）
- aether-cache: 通用 TTL 缓存与命名空间抽象
- aether-data: 数据访问层（PostgreSQL/Redis 后端、repository 模式）
- aether-http: HTTP 客户端封装（重试、配置）
- aether-testkit: 集成测试工具集（gateway/executor/hub/proxy fixture、等待、负载测试）

gateway 扩展：
- 引入 audit 模块（shadow 执行审计、决策链路追踪、请求审计 bundle）
- 引入 cache 模块（AuthContext 缓存、direct-plan bypass 缓存）
- 引入 data 模块（auth/candidates/config/usage/video_tasks 数据访问）
- 集成 ConcurrencyGate/DistributedConcurrencyGate 请求门控
- 新增本地 auth 拒绝、过载响应构建器
- 补充 control/auth_cache/video/concurrency 集成测试

aether-proxy 扩展：
- AppState 集成 stream_gate / distributed_stream_gate 并发门控
- 新增 ProxyAdmissionError 及准入拒绝流程
- stream_handler 补充门控饱和/不可用场景测试
- 配置与注册客户端逻辑完善

aether-hub 扩展：
- main.rs 引入运行时初始化、指标端点、健康检查
- local_relay 重构为 lib.rs 暴露公共接口
This commit is contained in:
fawney19
2026-03-24 15:12:56 +08:00
parent eaf8475f9e
commit b5a0070023
157 changed files with 22097 additions and 448 deletions
Download Patch File Download Diff File Expand all files Collapse all files

549
Cargo.lock generated
View File

@@ -8,6 +8,10 @@ version = "2.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa"
[[package]]
name = "aether-cache"
version = "0.1.0"
[[package]] [[package]]
name = "aether-contracts" name = "aether-contracts"
version = "0.1.0" version = "0.1.0"
@@ -17,11 +21,30 @@ dependencies = [
"thiserror 2.0.18", "thiserror 2.0.18",
] ]
[[package]]
name = "aether-data"
version = "0.1.0"
dependencies = [
"aether-cache",
"async-trait",
"futures-util",
"redis",
"serde",
"serde_json",
"sqlx",
"thiserror 2.0.18",
"tokio",
"url",
"uuid",
]
[[package]] [[package]]
name = "aether-executor" name = "aether-executor"
version = "0.1.0" version = "0.1.0"
dependencies = [ dependencies = [
"aether-contracts", "aether-contracts",
"aether-http",
"aether-runtime",
"async-stream", "async-stream",
"axum", "axum",
"base64", "base64",
@@ -37,7 +60,6 @@ dependencies = [
"thiserror 2.0.18", "thiserror 2.0.18",
"tokio", "tokio",
"tracing", "tracing",
"tracing-subscriber",
"webpki-roots 0.26.11", "webpki-roots 0.26.11",
] ]
@@ -45,7 +67,11 @@ dependencies = [
name = "aether-gateway" name = "aether-gateway"
version = "0.1.0" version = "0.1.0"
dependencies = [ dependencies = [
"aether-cache",
"aether-contracts", "aether-contracts",
"aether-data",
"aether-http",
"aether-runtime",
"async-stream", "async-stream",
"axum", "axum",
"base64", "base64",
@@ -60,15 +86,24 @@ dependencies = [
"tokio", "tokio",
"tokio-util", "tokio-util",
"tracing", "tracing",
"tracing-subscriber",
"url", "url",
"uuid", "uuid",
] ]
[[package]]
name = "aether-http"
version = "0.1.0"
dependencies = [
"reqwest",
"serde",
]
[[package]] [[package]]
name = "aether-hub" name = "aether-hub"
version = "0.2.0" version = "0.2.0"
dependencies = [ dependencies = [
"aether-http",
"aether-runtime",
"async-stream", "async-stream",
"axum", "axum",
"bytes", "bytes",
@@ -82,14 +117,16 @@ dependencies = [
"serde", "serde",
"serde_json", "serde_json",
"tokio", "tokio",
"tokio-tungstenite 0.28.0",
"tracing", "tracing",
"tracing-subscriber",
] ]
[[package]] [[package]]
name = "aether-proxy" name = "aether-proxy"
version = "0.2.5" version = "0.2.5"
dependencies = [ dependencies = [
"aether-http",
"aether-runtime",
"anyhow", "anyhow",
"arc-swap", "arc-swap",
"base64", "base64",
@@ -119,11 +156,50 @@ dependencies = [
"toml", "toml",
"tower-service", "tower-service",
"tracing", "tracing",
"tracing-subscriber",
"url", "url",
"webpki-roots 0.26.11", "webpki-roots 0.26.11",
] ]
[[package]]
name = "aether-runtime"
version = "0.1.0"
dependencies = [
"async-stream",
"axum",
"futures-util",
"redis",
"thiserror 2.0.18",
"tokio",
"tracing",
"tracing-subscriber",
"url",
"uuid",
]
[[package]]
name = "aether-testkit"
version = "0.1.0"
dependencies = [
"aether-contracts",
"aether-data",
"aether-executor",
"aether-gateway",
"aether-http",
"aether-hub",
"aether-runtime",
"async-stream",
"axum",
"bytes",
"futures-util",
"http",
"reqwest",
"serde",
"serde_json",
"sqlx",
"tokio",
"tokio-tungstenite 0.28.0",
]
[[package]] [[package]]
name = "aho-corasick" name = "aho-corasick"
version = "1.1.4" version = "1.1.4"
@@ -226,6 +302,26 @@ dependencies = [
"syn 2.0.117", "syn 2.0.117",
] ]
[[package]]
name = "async-trait"
version = "0.1.89"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.117",
]
[[package]]
name = "atoi"
version = "2.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f28d99ec8bfea296261ca1af174f24225171fea9664ba9003cbebee704810528"
dependencies = [
"num-traits",
]
[[package]] [[package]]
name = "atomic" name = "atomic"
version = "0.6.1" version = "0.6.1"
@@ -478,6 +574,20 @@ version = "1.0.5"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570"
[[package]]
name = "combine"
version = "4.6.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd"
dependencies = [
"bytes",
"futures-core",
"memchr",
"pin-project-lite",
"tokio",
"tokio-util",
]
[[package]] [[package]]
name = "compact_str" name = "compact_str"
version = "0.9.0" version = "0.9.0"
@@ -492,6 +602,15 @@ dependencies = [
"static_assertions", "static_assertions",
] ]
[[package]]
name = "concurrent-queue"
version = "2.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973"
dependencies = [
"crossbeam-utils",
]
[[package]] [[package]]
name = "convert_case" name = "convert_case"
version = "0.10.0" version = "0.10.0"
@@ -516,6 +635,21 @@ dependencies = [
"libc", "libc",
] ]
[[package]]
name = "crc"
version = "3.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5eb8a2a1cd12ab0d987a5d5e825195d372001a4094a0376319d5a0ad71c1ba0d"
dependencies = [
"crc-catalog",
]
[[package]]
name = "crc-catalog"
version = "2.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "19d374276b40fb8bbdee95aef7c7fa6b5316ec764510eb64b8dd0e2ed0d7e7f5"
[[package]] [[package]]
name = "crc32fast" name = "crc32fast"
version = "1.5.0" version = "1.5.0"
@@ -544,6 +678,15 @@ dependencies = [
"crossbeam-utils", "crossbeam-utils",
] ]
[[package]]
name = "crossbeam-queue"
version = "0.3.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0f58bbc28f91df819d0aa2a2c00cd19754769c2fad90579b3592b1c9ba7a3115"
dependencies = [
"crossbeam-utils",
]
[[package]] [[package]]
name = "crossbeam-utils" name = "crossbeam-utils"
version = "0.8.21" version = "0.8.21"
@@ -712,6 +855,7 @@ checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292"
dependencies = [ dependencies = [
"block-buffer", "block-buffer",
"crypto-common", "crypto-common",
"subtle",
] ]
[[package]] [[package]]
@@ -734,6 +878,12 @@ dependencies = [
"litrs", "litrs",
] ]
[[package]]
name = "dotenvy"
version = "0.15.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b"
[[package]] [[package]]
name = "dunce" name = "dunce"
version = "1.0.5" version = "1.0.5"
@@ -745,6 +895,9 @@ name = "either"
version = "1.15.0" version = "1.15.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
dependencies = [
"serde",
]
[[package]] [[package]]
name = "equivalent" name = "equivalent"
@@ -762,6 +915,17 @@ dependencies = [
"windows-sys 0.61.2", "windows-sys 0.61.2",
] ]
[[package]]
name = "etcetera"
version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "136d1b5283a1ab77bd9257427ffd09d8667ced0570b6f938942bc7568ed5b943"
dependencies = [
"cfg-if",
"home",
"windows-sys 0.48.0",
]
[[package]] [[package]]
name = "euclid" name = "euclid"
version = "0.22.14" version = "0.22.14"
@@ -771,6 +935,17 @@ dependencies = [
"num-traits", "num-traits",
] ]
[[package]]
name = "event-listener"
version = "5.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e13b66accf52311f30a0db42147dadea9850cb48cd070028831ae5f5d4b856ab"
dependencies = [
"concurrent-queue",
"parking",
"pin-project-lite",
]
[[package]] [[package]]
name = "fancy-regex" name = "fancy-regex"
version = "0.11.0" version = "0.11.0"
@@ -871,6 +1046,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d"
dependencies = [ dependencies = [
"futures-core", "futures-core",
"futures-sink",
] ]
[[package]] [[package]]
@@ -879,6 +1055,17 @@ version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d"
[[package]]
name = "futures-intrusive"
version = "0.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1d930c203dd0b6ff06e0201a4a2fe9149b43c684fd4420555b26d21b1a02956f"
dependencies = [
"futures-core",
"lock_api",
"parking_lot",
]
[[package]] [[package]]
name = "futures-io" name = "futures-io"
version = "0.3.32" version = "0.3.32"
@@ -1005,6 +1192,8 @@ version = "0.15.5"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1"
dependencies = [ dependencies = [
"allocator-api2",
"equivalent",
"foldhash 0.1.5", "foldhash 0.1.5",
] ]
@@ -1019,6 +1208,15 @@ dependencies = [
"foldhash 0.2.0", "foldhash 0.2.0",
] ]
[[package]]
name = "hashlink"
version = "0.10.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7382cf6263419f2d8df38c55d7da83da5c18aef87fc7a7fc1fb1e344edfe14c1"
dependencies = [
"hashbrown 0.15.5",
]
[[package]] [[package]]
name = "heck" name = "heck"
version = "0.5.0" version = "0.5.0"
@@ -1031,6 +1229,33 @@ version = "0.4.3"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70"
[[package]]
name = "hkdf"
version = "0.12.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7b5f8eb2ad728638ea2c7d47a21db23b7b58a72ed6a38256b8a1849f15fbbdf7"
dependencies = [
"hmac",
]
[[package]]
name = "hmac"
version = "0.12.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e"
dependencies = [
"digest",
]
[[package]]
name = "home"
version = "0.5.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d"
dependencies = [
"windows-sys 0.61.2",
]
[[package]] [[package]]
name = "http" name = "http"
version = "1.4.0" version = "1.4.0"
@@ -1479,6 +1704,16 @@ version = "0.8.4"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3"
[[package]]
name = "md-5"
version = "0.10.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d89e7ee0cfbedfc4da3340218492196241d89eefb6dab27de5df917a6d2e78cf"
dependencies = [
"cfg-if",
"digest",
]
[[package]] [[package]]
name = "memchr" name = "memchr"
version = "2.8.0" version = "2.8.0"
@@ -1575,6 +1810,16 @@ dependencies = [
"windows-sys 0.61.2", "windows-sys 0.61.2",
] ]
[[package]]
name = "num-bigint"
version = "0.4.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9"
dependencies = [
"num-integer",
"num-traits",
]
[[package]] [[package]]
name = "num-conv" name = "num-conv"
version = "0.2.0" version = "0.2.0"
@@ -1592,6 +1837,15 @@ dependencies = [
"syn 2.0.117", "syn 2.0.117",
] ]
[[package]]
name = "num-integer"
version = "0.1.46"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f"
dependencies = [
"num-traits",
]
[[package]] [[package]]
name = "num-traits" name = "num-traits"
version = "0.2.19" version = "0.2.19"
@@ -1631,6 +1885,12 @@ dependencies = [
"num-traits", "num-traits",
] ]
[[package]]
name = "parking"
version = "2.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba"
[[package]] [[package]]
name = "parking_lot" name = "parking_lot"
version = "0.12.5" version = "0.12.5"
@@ -2062,6 +2322,27 @@ dependencies = [
"crossbeam-utils", "crossbeam-utils",
] ]
[[package]]
name = "redis"
version = "0.28.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e37ec3fd44bea2ec947ba6cc7634d7999a6590aca7c35827c250bc0de502bda6"
dependencies = [
"arc-swap",
"bytes",
"combine",
"futures-util",
"itoa",
"num-bigint",
"percent-encoding",
"pin-project-lite",
"ryu",
"sha1_smol",
"tokio",
"tokio-util",
"url",
]
[[package]] [[package]]
name = "redox_syscall" name = "redox_syscall"
version = "0.5.18" version = "0.5.18"
@@ -2354,6 +2635,12 @@ dependencies = [
"digest", "digest",
] ]
[[package]]
name = "sha1_smol"
version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d"
[[package]] [[package]]
name = "sha2" name = "sha2"
version = "0.10.9" version = "0.10.9"
@@ -2434,6 +2721,9 @@ name = "smallvec"
version = "1.15.1" version = "1.15.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03"
dependencies = [
"serde",
]
[[package]] [[package]]
name = "socket2" name = "socket2"
@@ -2455,6 +2745,126 @@ dependencies = [
"windows-sys 0.61.2", "windows-sys 0.61.2",
] ]
[[package]]
name = "sqlx"
version = "0.8.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1fefb893899429669dcdd979aff487bd78f4064e5e7907e4269081e0ef7d97dc"
dependencies = [
"sqlx-core",
"sqlx-macros",
"sqlx-postgres",
]
[[package]]
name = "sqlx-core"
version = "0.8.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ee6798b1838b6a0f69c007c133b8df5866302197e404e8b6ee8ed3e3a5e68dc6"
dependencies = [
"base64",
"bytes",
"crc",
"crossbeam-queue",
"either",
"event-listener",
"futures-core",
"futures-intrusive",
"futures-io",
"futures-util",
"hashbrown 0.15.5",
"hashlink",
"indexmap",
"log",
"memchr",
"once_cell",
"percent-encoding",
"rustls",
"serde",
"serde_json",
"sha2",
"smallvec",
"thiserror 2.0.18",
"tokio",
"tokio-stream",
"tracing",
"url",
"webpki-roots 0.26.11",
]
[[package]]
name = "sqlx-macros"
version = "0.8.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a2d452988ccaacfbf5e0bdbc348fb91d7c8af5bee192173ac3636b5fb6e6715d"
dependencies = [
"proc-macro2",
"quote",
"sqlx-core",
"sqlx-macros-core",
"syn 2.0.117",
]
[[package]]
name = "sqlx-macros-core"
version = "0.8.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "19a9c1841124ac5a61741f96e1d9e2ec77424bf323962dd894bdb93f37d5219b"
dependencies = [
"dotenvy",
"either",
"heck",
"hex",
"once_cell",
"proc-macro2",
"quote",
"serde",
"serde_json",
"sha2",
"sqlx-core",
"sqlx-postgres",
"syn 2.0.117",
"tokio",
"url",
]
[[package]]
name = "sqlx-postgres"
version = "0.8.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "db58fcd5a53cf07c184b154801ff91347e4c30d17a3562a635ff028ad5deda46"
dependencies = [
"atoi",
"base64",
"bitflags 2.11.0",
"byteorder",
"crc",
"dotenvy",
"etcetera",
"futures-channel",
"futures-core",
"futures-util",
"hex",
"hkdf",
"hmac",
"home",
"itoa",
"log",
"md-5",
"memchr",
"once_cell",
"rand 0.8.5",
"serde",
"serde_json",
"sha2",
"smallvec",
"sqlx-core",
"stringprep",
"thiserror 2.0.18",
"tracing",
"whoami",
]
[[package]] [[package]]
name = "stable_deref_trait" name = "stable_deref_trait"
version = "1.2.1" version = "1.2.1"
@@ -2467,6 +2877,17 @@ version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f"
[[package]]
name = "stringprep"
version = "0.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7b4df3d392d81bd458a8a621b8bffbd2302a12ffe288a9d931670948749463b1"
dependencies = [
"unicode-bidi",
"unicode-normalization",
"unicode-properties",
]
[[package]] [[package]]
name = "strsim" name = "strsim"
version = "0.11.1" version = "0.11.1"
@@ -2763,6 +3184,17 @@ dependencies = [
"tokio", "tokio",
] ]
[[package]]
name = "tokio-stream"
version = "0.1.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70"
dependencies = [
"futures-core",
"pin-project-lite",
"tokio",
]
[[package]] [[package]]
name = "tokio-tungstenite" name = "tokio-tungstenite"
version = "0.24.0" version = "0.24.0"
@@ -2787,8 +3219,12 @@ checksum = "d25a406cddcc431a75d3d9afc6a7c0f7428d4891dd973e4d54c56b46127bf857"
dependencies = [ dependencies = [
"futures-util", "futures-util",
"log", "log",
"rustls",
"rustls-pki-types",
"tokio", "tokio",
"tokio-rustls",
"tungstenite 0.28.0", "tungstenite 0.28.0",
"webpki-roots 0.26.11",
] ]
[[package]] [[package]]
@@ -3004,6 +3440,8 @@ dependencies = [
"httparse", "httparse",
"log", "log",
"rand 0.9.2", "rand 0.9.2",
"rustls",
"rustls-pki-types",
"sha1", "sha1",
"thiserror 2.0.18", "thiserror 2.0.18",
"utf-8", "utf-8",
@@ -3021,12 +3459,33 @@ version = "0.1.7"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971" checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971"
[[package]]
name = "unicode-bidi"
version = "0.3.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c1cb5db39152898a79168971543b1cb5020dff7fe43c8dc468b0885f5e29df5"
[[package]] [[package]]
name = "unicode-ident" name = "unicode-ident"
version = "1.0.24" version = "1.0.24"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
[[package]]
name = "unicode-normalization"
version = "0.1.25"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5fd4f6878c9cb28d874b009da9e8d183b5abc80117c40bbd187a1fde336be6e8"
dependencies = [
"tinyvec",
]
[[package]]
name = "unicode-properties"
version = "0.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7df058c713841ad818f1dc5d3fd88063241cc61f49f5fbea4b951e8cf5a8d71d"
[[package]] [[package]]
name = "unicode-segmentation" name = "unicode-segmentation"
version = "1.12.0" version = "1.12.0"
@@ -3159,6 +3618,12 @@ dependencies = [
"wit-bindgen", "wit-bindgen",
] ]
[[package]]
name = "wasite"
version = "0.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b8dad83b4f25e74f184f64c43b150b91efe7647395b42289f38e50566d82855b"
[[package]] [[package]]
name = "wasm-bindgen" name = "wasm-bindgen"
version = "0.2.114" version = "0.2.114"
@@ -3375,6 +3840,16 @@ dependencies = [
"wezterm-dynamic", "wezterm-dynamic",
] ]
[[package]]
name = "whoami"
version = "1.6.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5d4a4db5077702ca3015d3d02d74974948aba2ad9e12ab7df718ee64ccd7e97d"
dependencies = [
"libredox",
"wasite",
]
[[package]] [[package]]
name = "winapi" name = "winapi"
version = "0.3.9" version = "0.3.9"
@@ -3456,6 +3931,15 @@ dependencies = [
"windows-targets 0.52.6", "windows-targets 0.52.6",
] ]
[[package]]
name = "windows-sys"
version = "0.48.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9"
dependencies = [
"windows-targets 0.48.5",
]
[[package]] [[package]]
name = "windows-sys" name = "windows-sys"
version = "0.52.0" version = "0.52.0"
@@ -3492,6 +3976,21 @@ dependencies = [
"windows-link", "windows-link",
] ]
[[package]]
name = "windows-targets"
version = "0.48.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c"
dependencies = [
"windows_aarch64_gnullvm 0.48.5",
"windows_aarch64_msvc 0.48.5",
"windows_i686_gnu 0.48.5",
"windows_i686_msvc 0.48.5",
"windows_x86_64_gnu 0.48.5",
"windows_x86_64_gnullvm 0.48.5",
"windows_x86_64_msvc 0.48.5",
]
[[package]] [[package]]
name = "windows-targets" name = "windows-targets"
version = "0.52.6" version = "0.52.6"
@@ -3525,6 +4024,12 @@ dependencies = [
"windows_x86_64_msvc 0.53.1", "windows_x86_64_msvc 0.53.1",
] ]
[[package]]
name = "windows_aarch64_gnullvm"
version = "0.48.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8"
[[package]] [[package]]
name = "windows_aarch64_gnullvm" name = "windows_aarch64_gnullvm"
version = "0.52.6" version = "0.52.6"
@@ -3537,6 +4042,12 @@ version = "0.53.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53"
[[package]]
name = "windows_aarch64_msvc"
version = "0.48.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc"
[[package]] [[package]]
name = "windows_aarch64_msvc" name = "windows_aarch64_msvc"
version = "0.52.6" version = "0.52.6"
@@ -3549,6 +4060,12 @@ version = "0.53.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006"
[[package]]
name = "windows_i686_gnu"
version = "0.48.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e"
[[package]] [[package]]
name = "windows_i686_gnu" name = "windows_i686_gnu"
version = "0.52.6" version = "0.52.6"
@@ -3573,6 +4090,12 @@ version = "0.53.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c"
[[package]]
name = "windows_i686_msvc"
version = "0.48.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406"
[[package]] [[package]]
name = "windows_i686_msvc" name = "windows_i686_msvc"
version = "0.52.6" version = "0.52.6"
@@ -3585,6 +4108,12 @@ version = "0.53.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2"
[[package]]
name = "windows_x86_64_gnu"
version = "0.48.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e"
[[package]] [[package]]
name = "windows_x86_64_gnu" name = "windows_x86_64_gnu"
version = "0.52.6" version = "0.52.6"
@@ -3597,6 +4126,12 @@ version = "0.53.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499"
[[package]]
name = "windows_x86_64_gnullvm"
version = "0.48.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc"
[[package]] [[package]]
name = "windows_x86_64_gnullvm" name = "windows_x86_64_gnullvm"
version = "0.52.6" version = "0.52.6"
@@ -3609,6 +4144,12 @@ version = "0.53.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1"
[[package]]
name = "windows_x86_64_msvc"
version = "0.48.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538"
[[package]] [[package]]
name = "windows_x86_64_msvc" name = "windows_x86_64_msvc"
version = "0.52.6" version = "0.52.6"

21
Cargo.toml
View File

@@ -2,9 +2,14 @@
members = [ members = [
"aether-hub", "aether-hub",
"aether-proxy", "aether-proxy",
"crates/aether-cache",
"crates/aether-contracts", "crates/aether-contracts",
"crates/aether-data",
"crates/aether-executor", "crates/aether-executor",
"crates/aether-gateway", "crates/aether-gateway",
"crates/aether-http",
"crates/aether-runtime",
"crates/aether-testkit",
] ]
resolver = "2" resolver = "2"
@@ -14,20 +19,34 @@ license = "LicenseRef-Aether-NonCommercial"
repository = "https://github.com/fawney19/Aether.git" repository = "https://github.com/fawney19/Aether.git"
[workspace.dependencies] [workspace.dependencies]
aether-hub = { path = "aether-hub" }
aether-cache = { path = "crates/aether-cache" }
aether-contracts = { path = "crates/aether-contracts" } aether-contracts = { path = "crates/aether-contracts" }
aether-data = { path = "crates/aether-data" }
aether-executor = { path = "crates/aether-executor" }
aether-gateway = { path = "crates/aether-gateway" }
aether-http = { path = "crates/aether-http" }
aether-runtime = { path = "crates/aether-runtime" }
aether-testkit = { path = "crates/aether-testkit" }
async-stream = "0.3" async-stream = "0.3"
async-trait = "0.1"
axum = "0.8"
base64 = "0.22" base64 = "0.22"
bytes = "1" bytes = "1"
flate2 = "1" flate2 = "1"
futures-util = "0.3" futures-util = "0.3"
http = "1" http = "1"
reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "rustls-tls", "http2"] } reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "rustls-tls", "http2"] }
redis = { version = "0.28", default-features = false, features = ["tokio-comp", "script", "streams"] }
rustls = { version = "0.23", features = ["ring"] } rustls = { version = "0.23", features = ["ring"] }
serde = { version = "1", features = ["derive"] } serde = { version = "1", features = ["derive"] }
serde_json = "1" serde_json = "1"
sqlx = { version = "0.8", default-features = false, features = ["postgres", "runtime-tokio-rustls"] }
thiserror = "2" thiserror = "2"
tokio = { version = "1", features = ["macros", "net", "rt-multi-thread", "sync", "time"] } tokio = { version = "1", features = ["macros", "net", "rt-multi-thread", "signal", "sync", "time"] }
tokio-util = { version = "0.7", features = ["codec", "io-util"] } tokio-util = { version = "0.7", features = ["codec", "io-util"] }
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
uuid = { version = "1", features = ["serde", "v4"] } uuid = { version = "1", features = ["serde", "v4"] }
webpki-roots = "0.26" webpki-roots = "0.26"
url = "2" url = "2"

8
aether-hub/Cargo.toml
View File

@@ -5,12 +5,13 @@ edition = "2021"
description = "Tunnel Hub for Aether - frame router between workers and proxies" description = "Tunnel Hub for Aether - frame router between workers and proxies"
[dependencies] [dependencies]
aether-http.workspace = true
aether-runtime.workspace = true
tokio = { version = "1", features = ["full"] } tokio = { version = "1", features = ["full"] }
axum = { version = "0.8", features = ["ws"] } axum = { version = "0.8", features = ["ws"] }
serde = { version = "1", features = ["derive"] } serde = { version = "1", features = ["derive"] }
serde_json = "1" serde_json = "1"
tracing = "0.1" tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
clap = { version = "4", features = ["derive", "env"] } clap = { version = "4", features = ["derive", "env"] }
dashmap = "6" dashmap = "6"
parking_lot = "0.12" parking_lot = "0.12"
@@ -19,4 +20,7 @@ futures-util = "0.3"
bytes = "1" bytes = "1"
async-stream = "0.3" async-stream = "0.3"
http-body-util = "0.1" http-body-util = "0.1"
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] } reqwest.workspace = true
[dev-dependencies]
tokio-tungstenite = { version = "0.28", features = ["rustls-tls-webpki-roots"] }

9
aether-hub/src/control_plane.rs
View File

@@ -1,3 +1,4 @@
use aether_http::{build_http_client, HttpClientConfig};
use reqwest::Client; use reqwest::Client;
#[derive(Clone)] #[derive(Clone)]
@@ -8,9 +9,11 @@ pub struct ControlPlaneClient {
impl ControlPlaneClient { impl ControlPlaneClient {
pub fn new(base_url: String) -> Self { pub fn new(base_url: String) -> Self {
let client = Client::builder() let client = build_http_client(&HttpClientConfig {
.timeout(std::time::Duration::from_secs(10)) request_timeout_ms: Some(10_000),
.build() user_agent: Some("aether-hub/control-plane".to_string()),
..HttpClientConfig::default()
})
.ok(); .ok();
Self { client, base_url } Self { client, base_url }
} }

43
aether-hub/src/hub.rs
View File

@@ -3,12 +3,12 @@ use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, AtomicUsize, Ordering}
use std::sync::Arc; use std::sync::Arc;
use std::time::Duration; use std::time::Duration;
use aether_runtime::{BoundedQueueSender, MetricKind, MetricSample, QueueSendError};
use axum::extract::ws::Message; use axum::extract::ws::Message;
use bytes::Bytes; use bytes::Bytes;
use dashmap::DashMap; use dashmap::DashMap;
use parking_lot::{Mutex, RwLock}; use parking_lot::{Mutex, RwLock};
use tokio::sync::mpsc; use tokio::sync::mpsc;
use tokio::sync::mpsc::error::TrySendError;
use tokio::sync::{watch, Notify}; use tokio::sync::{watch, Notify};
use tracing::{debug, info, warn}; use tracing::{debug, info, warn};
@@ -32,13 +32,13 @@ pub struct ConnConfig {
} }
pub struct BoundedOutbound { pub struct BoundedOutbound {
tx: mpsc::Sender<Message>, tx: BoundedQueueSender<Message>,
close_tx: watch::Sender<bool>, close_tx: watch::Sender<bool>,
closing: AtomicBool, closing: AtomicBool,
} }
impl BoundedOutbound { impl BoundedOutbound {
pub fn new(tx: mpsc::Sender<Message>, close_tx: watch::Sender<bool>) -> Self { pub fn new(tx: BoundedQueueSender<Message>, close_tx: watch::Sender<bool>) -> Self {
Self { Self {
tx, tx,
close_tx, close_tx,
@@ -53,11 +53,11 @@ impl BoundedOutbound {
match self.tx.try_send(msg) { match self.tx.try_send(msg) {
Ok(()) => SendStatus::Queued, Ok(()) => SendStatus::Queued,
Err(TrySendError::Closed(_)) => { Err(QueueSendError::Closed(_)) => {
self.mark_closing(); self.mark_closing();
SendStatus::Closed SendStatus::Closed
} }
Err(TrySendError::Full(_)) => { Err(QueueSendError::Full(_)) => {
self.mark_closing(); self.mark_closing();
SendStatus::Congested SendStatus::Congested
} }
@@ -92,7 +92,7 @@ impl ProxyConn {
id: u64, id: u64,
node_id: String, node_id: String,
node_name: String, node_name: String,
tx: mpsc::Sender<Message>, tx: BoundedQueueSender<Message>,
close_tx: watch::Sender<bool>, close_tx: watch::Sender<bool>,
max_streams: usize, max_streams: usize,
) -> Self { ) -> Self {
@@ -786,8 +786,35 @@ pub struct HubStats {
pub active_streams: usize, pub active_streams: usize,
} }
impl HubStats {
pub fn to_metric_samples(&self) -> Vec<MetricSample> {
vec![
MetricSample::new(
"hub_proxy_connections",
"Current number of connected proxy sockets.",
MetricKind::Gauge,
self.proxy_connections as u64,
),
MetricSample::new(
"hub_nodes",
"Current number of connected logical nodes.",
MetricKind::Gauge,
self.nodes as u64,
),
MetricSample::new(
"hub_active_streams",
"Current number of active local relay streams.",
MetricKind::Gauge,
self.active_streams as u64,
),
]
}
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use aether_runtime::bounded_queue;
use super::*; use super::*;
fn build_meta() -> protocol::RequestMeta { fn build_meta() -> protocol::RequestMeta {
@@ -803,7 +830,7 @@ mod tests {
async fn cancel_local_stream_notifies_proxy() { async fn cancel_local_stream_notifies_proxy() {
let hub = HubRouter::new(ControlPlaneClient::disabled()); let hub = HubRouter::new(ControlPlaneClient::disabled());
let (proxy_tx, mut proxy_rx) = mpsc::channel(8); let (proxy_tx, mut proxy_rx) = bounded_queue(8);
let (proxy_close_tx, _) = watch::channel(false); let (proxy_close_tx, _) = watch::channel(false);
let proxy = Arc::new(ProxyConn::new( let proxy = Arc::new(ProxyConn::new(
100, 100,
@@ -838,7 +865,7 @@ mod tests {
async fn push_local_request_body_splits_large_payload_and_marks_end() { async fn push_local_request_body_splits_large_payload_and_marks_end() {
let hub = HubRouter::new(ControlPlaneClient::disabled()); let hub = HubRouter::new(ControlPlaneClient::disabled());
let (proxy_tx, mut proxy_rx) = mpsc::channel(8); let (proxy_tx, mut proxy_rx) = bounded_queue(8);
let (proxy_close_tx, _) = watch::channel(false); let (proxy_close_tx, _) = watch::channel(false);
let proxy = Arc::new(ProxyConn::new( let proxy = Arc::new(ProxyConn::new(
200, 200,

242
aether-hub/src/lib.rs Normal file
View File

@@ -0,0 +1,242 @@
mod control_plane;
mod hub;
mod local_relay;
pub mod protocol;
mod proxy_conn;
use std::sync::Arc;
use aether_runtime::{
hold_admission_permit_until, prometheus_response, service_up_sample, AdmissionPermit,
ConcurrencyError, ConcurrencyGate, ConcurrencySnapshot, DistributedConcurrencyError,
DistributedConcurrencyGate, DistributedConcurrencySnapshot, MetricKind, MetricLabel,
MetricSample,
};
use axum::extract::ws::WebSocketUpgrade;
use axum::extract::State;
use axum::response::{IntoResponse, Json};
use axum::routing::{get, post};
use axum::Router;
use tracing::warn;
pub use control_plane::ControlPlaneClient;
pub use hub::{ConnConfig, HubRouter};
#[derive(Clone)]
pub struct AppState {
pub hub: Arc<HubRouter>,
pub proxy_conn_cfg: ConnConfig,
pub max_streams: usize,
request_gate: Option<Arc<ConcurrencyGate>>,
distributed_request_gate: Option<Arc<DistributedConcurrencyGate>>,
}
#[derive(Debug)]
enum RequestAdmissionError {
Local(ConcurrencyError),
Distributed(DistributedConcurrencyError),
}
impl AppState {
pub fn new(
control_plane: ControlPlaneClient,
proxy_conn_cfg: ConnConfig,
max_streams: usize,
) -> Self {
Self {
hub: HubRouter::new(control_plane),
proxy_conn_cfg,
max_streams,
request_gate: None,
distributed_request_gate: None,
}
}
pub fn with_request_concurrency_limit(mut self, limit: Option<usize>) -> Self {
self.request_gate = limit
.filter(|limit| *limit > 0)
.map(|limit| Arc::new(ConcurrencyGate::new("hub_requests", limit)));
self
}
pub fn with_distributed_request_gate(mut self, gate: DistributedConcurrencyGate) -> Self {
self.distributed_request_gate = Some(Arc::new(gate));
self
}
fn request_concurrency_snapshot(&self) -> Option<ConcurrencySnapshot> {
self.request_gate.as_ref().map(|gate| gate.snapshot())
}
async fn distributed_request_concurrency_snapshot(
&self,
) -> Result<Option<DistributedConcurrencySnapshot>, DistributedConcurrencyError> {
match self.distributed_request_gate.as_ref() {
Some(gate) => gate.snapshot().await.map(Some),
None => Ok(None),
}
}
async fn metric_samples(&self) -> Vec<MetricSample> {
let mut samples = vec![service_up_sample("aether-hub")];
if let Some(snapshot) = self.request_concurrency_snapshot() {
samples.extend(snapshot.to_metric_samples("hub_requests"));
}
if let Some(gate) = self.distributed_request_gate.as_ref() {
match gate.snapshot().await {
Ok(snapshot) => {
samples.extend(snapshot.to_metric_samples("hub_requests_distributed"));
}
Err(_) => samples.push(
MetricSample::new(
"concurrency_unavailable",
"Whether the distributed concurrency gate is currently unavailable.",
MetricKind::Gauge,
1,
)
.with_labels(vec![MetricLabel::new("gate", "hub_requests_distributed")]),
),
}
}
samples.extend(self.hub.stats().to_metric_samples());
samples
}
async fn try_acquire_request_permit(
&self,
) -> Result<Option<AdmissionPermit>, RequestAdmissionError> {
let local = self
.request_gate
.as_ref()
.map(|gate| gate.try_acquire())
.transpose()
.map_err(RequestAdmissionError::Local)?;
let distributed = match self.distributed_request_gate.as_ref() {
Some(gate) => Some(
gate.try_acquire()
.await
.map_err(RequestAdmissionError::Distributed)?,
),
None => None,
};
Ok(AdmissionPermit::from_parts(local, distributed))
}
}
pub fn build_router_with_state(state: AppState) -> Router {
Router::new()
.route("/health", get(health))
.route("/metrics", get(metrics))
.route("/stats", get(stats))
.route("/proxy", get(ws_proxy))
.route("/local/relay/{node_id}", post(local_relay::relay_request))
.with_state(state)
}
async fn health(State(state): State<AppState>) -> impl IntoResponse {
let request_concurrency = state.request_concurrency_snapshot().map(|snapshot| {
serde_json::json!({
"limit": snapshot.limit,
"in_flight": snapshot.in_flight,
"available_permits": snapshot.available_permits,
"high_watermark": snapshot.high_watermark,
"rejected": snapshot.rejected,
})
});
let distributed_request_concurrency = state
.distributed_request_concurrency_snapshot()
.await
.ok()
.flatten()
.map(|snapshot| {
serde_json::json!({
"limit": snapshot.limit,
"in_flight": snapshot.in_flight,
"available_permits": snapshot.available_permits,
"high_watermark": snapshot.high_watermark,
"rejected": snapshot.rejected,
})
});
Json(serde_json::json!({
"status": "ok",
"request_concurrency": request_concurrency,
"distributed_request_concurrency": distributed_request_concurrency,
}))
}
async fn stats(State(state): State<AppState>) -> impl IntoResponse {
Json(state.hub.stats())
}
async fn metrics(State(state): State<AppState>) -> impl IntoResponse {
prometheus_response(&state.metric_samples().await)
}
async fn ws_proxy(
ws: WebSocketUpgrade,
State(state): State<AppState>,
headers: axum::http::HeaderMap,
) -> impl IntoResponse {
let node_id = headers
.get("x-node-id")
.and_then(|v| v.to_str().ok())
.unwrap_or("")
.trim()
.to_string();
let node_name = headers
.get("x-node-name")
.and_then(|v| v.to_str().ok())
.unwrap_or(&node_id)
.trim()
.to_string();
let max_streams: usize = headers
.get("x-tunnel-max-streams")
.and_then(|v| v.to_str().ok())
.and_then(|v| v.parse().ok())
.unwrap_or(state.max_streams)
.clamp(64, 2048);
if node_id.is_empty() {
warn!("proxy connection rejected: missing X-Node-ID header");
return axum::http::StatusCode::BAD_REQUEST.into_response();
}
let request_permit = match state.try_acquire_request_permit().await {
Ok(permit) => permit,
Err(RequestAdmissionError::Local(ConcurrencyError::Saturated { .. }))
| Err(RequestAdmissionError::Distributed(DistributedConcurrencyError::Saturated {
..
}))
| Err(RequestAdmissionError::Distributed(DistributedConcurrencyError::Unavailable {
..
})) => return axum::http::StatusCode::SERVICE_UNAVAILABLE.into_response(),
Err(RequestAdmissionError::Local(ConcurrencyError::Closed { gate })) => {
warn!(gate = gate, "hub request concurrency gate is closed");
return axum::http::StatusCode::SERVICE_UNAVAILABLE.into_response();
}
Err(RequestAdmissionError::Distributed(
DistributedConcurrencyError::InvalidConfiguration(message),
)) => {
warn!(error = %message, "hub distributed request gate is invalid");
return axum::http::StatusCode::SERVICE_UNAVAILABLE.into_response();
}
};
ws.max_frame_size(64 * 1024 * 1024)
.on_upgrade(move |socket| {
hold_admission_permit_until(request_permit, async move {
proxy_conn::handle_proxy_connection(
socket,
state.hub,
node_id,
node_name,
max_streams,
state.proxy_conn_cfg,
)
.await
})
})
.into_response()
}

99
aether-hub/src/local_relay.rs
View File

@@ -2,6 +2,7 @@ use std::io;
use std::net::SocketAddr; use std::net::SocketAddr;
use std::time::Duration; use std::time::Duration;
use aether_runtime::{maybe_hold_axum_response_permit, AdmissionPermit};
use async_stream::stream; use async_stream::stream;
use axum::body::{Body, Bytes}; use axum::body::{Body, Bytes};
use axum::extract::{ConnectInfo, Path, Request, State}; use axum::extract::{ConnectInfo, Path, Request, State};
@@ -47,6 +48,43 @@ pub async fn relay_request(
); );
} }
let request_permit = match state.try_acquire_request_permit().await {
Ok(permit) => permit,
Err(crate::RequestAdmissionError::Local(aether_runtime::ConcurrencyError::Saturated {
..
}))
| Err(crate::RequestAdmissionError::Distributed(
aether_runtime::DistributedConcurrencyError::Saturated { .. },
))
| Err(crate::RequestAdmissionError::Distributed(
aether_runtime::DistributedConcurrencyError::Unavailable { .. },
)) => {
return tunnel_error_response(
StatusCode::SERVICE_UNAVAILABLE,
"overloaded",
"hub relay overloaded",
);
}
Err(crate::RequestAdmissionError::Local(aether_runtime::ConcurrencyError::Closed {
..
})) => {
return tunnel_error_response(
StatusCode::SERVICE_UNAVAILABLE,
"overloaded",
"hub relay gate closed",
);
}
Err(crate::RequestAdmissionError::Distributed(
aether_runtime::DistributedConcurrencyError::InvalidConfiguration(_),
)) => {
return tunnel_error_response(
StatusCode::SERVICE_UNAVAILABLE,
"overloaded",
"hub relay distributed gate invalid",
);
}
};
let mut body_stream = request.into_body().into_data_stream(); let mut body_stream = request.into_body().into_data_stream();
let mut envelope_buf = BytesMut::new(); let mut envelope_buf = BytesMut::new();
let mut meta: Option<protocol::RequestMeta> = None; let mut meta: Option<protocol::RequestMeta> = None;
@@ -62,10 +100,13 @@ pub async fn relay_request(
.cancel_local_stream(active_stream.id, "failed to read relay request body"); .cancel_local_stream(active_stream.id, "failed to read relay request body");
} }
warn!(error = %error, "failed to read local relay request body"); warn!(error = %error, "failed to read local relay request body");
return tunnel_error_response( return release_permit_response(
tunnel_error_response(
StatusCode::BAD_GATEWAY, StatusCode::BAD_GATEWAY,
"relay", "relay",
"failed to read relay request body", "failed to read relay request body",
),
request_permit,
); );
} }
}; };
@@ -75,7 +116,10 @@ pub async fn relay_request(
let Some((parsed_meta, body_offset)) = (match try_decode_envelope_meta(&envelope_buf) { let Some((parsed_meta, body_offset)) = (match try_decode_envelope_meta(&envelope_buf) {
Ok(result) => result, Ok(result) => result,
Err(error) => { Err(error) => {
return tunnel_error_response(StatusCode::BAD_REQUEST, "bad_request", &error); return release_permit_response(
tunnel_error_response(StatusCode::BAD_REQUEST, "bad_request", &error),
request_permit,
);
} }
}) else { }) else {
continue; continue;
@@ -84,10 +128,9 @@ pub async fn relay_request(
let opened_stream = match state.hub.open_local_stream(&node_id, &parsed_meta) { let opened_stream = match state.hub.open_local_stream(&node_id, &parsed_meta) {
Ok(stream) => stream, Ok(stream) => stream,
Err(error) => { Err(error) => {
return tunnel_error_response( return release_permit_response(
StatusCode::SERVICE_UNAVAILABLE, tunnel_error_response(StatusCode::SERVICE_UNAVAILABLE, "connect", &error),
"connect", request_permit,
&error,
); );
} }
}; };
@@ -100,10 +143,9 @@ pub async fn relay_request(
.push_local_request_body(opened_stream.id, first_body_chunk, false) .push_local_request_body(opened_stream.id, first_body_chunk, false)
{ {
state.hub.cancel_local_stream(opened_stream.id, &error); state.hub.cancel_local_stream(opened_stream.id, &error);
return tunnel_error_response( return release_permit_response(
StatusCode::SERVICE_UNAVAILABLE, tunnel_error_response(StatusCode::SERVICE_UNAVAILABLE, "connect", &error),
"connect", request_permit,
&error,
); );
} }
} }
@@ -122,17 +164,23 @@ pub async fn relay_request(
.push_local_request_body(active_stream.id, chunk, false) .push_local_request_body(active_stream.id, chunk, false)
{ {
state.hub.cancel_local_stream(active_stream.id, &error); state.hub.cancel_local_stream(active_stream.id, &error);
return tunnel_error_response(StatusCode::SERVICE_UNAVAILABLE, "connect", &error); return release_permit_response(
tunnel_error_response(StatusCode::SERVICE_UNAVAILABLE, "connect", &error),
request_permit,
);
} }
} }
let (meta, stream) = match (meta, stream) { let (meta, stream) = match (meta, stream) {
(Some(meta), Some(stream)) => (meta, stream), (Some(meta), Some(stream)) => (meta, stream),
_ => { _ => {
return tunnel_error_response( return release_permit_response(
tunnel_error_response(
StatusCode::BAD_REQUEST, StatusCode::BAD_REQUEST,
"bad_request", "bad_request",
"relay envelope metadata truncated", "relay envelope metadata truncated",
),
request_permit,
); );
} }
}; };
@@ -142,7 +190,10 @@ pub async fn relay_request(
.push_local_request_body(stream.id, Bytes::new(), true) .push_local_request_body(stream.id, Bytes::new(), true)
{ {
state.hub.cancel_local_stream(stream.id, &error); state.hub.cancel_local_stream(stream.id, &error);
return tunnel_error_response(StatusCode::SERVICE_UNAVAILABLE, "connect", &error); return release_permit_response(
tunnel_error_response(StatusCode::SERVICE_UNAVAILABLE, "connect", &error),
request_permit,
);
} }
let request_guard = StreamGuard { let request_guard = StreamGuard {
@@ -156,7 +207,10 @@ pub async fn relay_request(
Ok(response) => response, Ok(response) => response,
Err(error) => { Err(error) => {
state.hub.cancel_local_stream(stream.id, &error); state.hub.cancel_local_stream(stream.id, &error);
return tunnel_error_response(StatusCode::GATEWAY_TIMEOUT, "timeout", &error); return release_permit_response(
tunnel_error_response(StatusCode::GATEWAY_TIMEOUT, "timeout", &error),
request_permit,
);
} }
}; };
@@ -164,10 +218,13 @@ pub async fn relay_request(
state state
.hub .hub
.cancel_local_stream(stream.id, "missing relay response body receiver"); .cancel_local_stream(stream.id, "missing relay response body receiver");
return tunnel_error_response( return release_permit_response(
tunnel_error_response(
StatusCode::BAD_GATEWAY, StatusCode::BAD_GATEWAY,
"relay", "relay",
"missing relay response body receiver", "missing relay response body receiver",
),
request_permit,
); );
}; };
@@ -199,18 +256,28 @@ pub async fn relay_request(
append_headers(headers, &response_head.headers); append_headers(headers, &response_head.headers);
} }
match builder.body(Body::from_stream(body_stream)) { match builder.body(Body::from_stream(body_stream)) {
Ok(response) => response, Ok(response) => maybe_hold_axum_response_permit(response, request_permit),
Err(error) => { Err(error) => {
warn!(error = %error, "failed to build relay response"); warn!(error = %error, "failed to build relay response");
release_permit_response(
tunnel_error_response( tunnel_error_response(
StatusCode::BAD_GATEWAY, StatusCode::BAD_GATEWAY,
"relay", "relay",
"failed to build relay response", "failed to build relay response",
),
request_permit,
) )
} }
} }
} }
fn release_permit_response(
response: Response<Body>,
_request_permit: Option<AdmissionPermit>,
) -> Response<Body> {
response
}
fn try_decode_envelope_meta( fn try_decode_envelope_meta(
buffer: &BytesMut, buffer: &BytesMut,
) -> Result<Option<(protocol::RequestMeta, usize)>, String> { ) -> Result<Option<(protocol::RequestMeta, usize)>, String> {

289
aether-hub/src/main.rs
View File

@@ -1,44 +1,29 @@
mod control_plane;
mod hub;
mod local_relay;
mod protocol;
mod proxy_conn;
use std::net::SocketAddr; use std::net::SocketAddr;
use std::time::Duration; use std::time::Duration;
use axum::extract::ws::WebSocketUpgrade; use aether_hub::{build_router_with_state, AppState, ConnConfig, ControlPlaneClient};
use axum::extract::State; use aether_runtime::{
use axum::response::{IntoResponse, Json}; init_service_runtime, DistributedConcurrencyGate, RedisDistributedConcurrencyConfig,
use axum::routing::{get, post}; ServiceRuntimeConfig,
use axum::Router; };
use clap::Parser; use clap::Parser;
use tracing::{info, warn}; use tracing::info;
use crate::control_plane::ControlPlaneClient;
use crate::hub::{ConnConfig, HubRouter};
use crate::local_relay::relay_request;
#[derive(Parser, Debug)] #[derive(Parser, Debug)]
#[command(name = "aether-hub", about = "Tunnel Hub for Aether")] #[command(name = "aether-hub", about = "Tunnel Hub for Aether")]
struct Args { struct Args {
/// Bind address
#[arg(long, default_value = "0.0.0.0:8085", env = "TUNNEL_HUB_BIND")] #[arg(long, default_value = "0.0.0.0:8085", env = "TUNNEL_HUB_BIND")]
bind: String, bind: String,
/// Proxy-side idle timeout in seconds (0 to disable)
#[arg(long, default_value_t = 0, env = "TUNNEL_HUB_PROXY_IDLE_TIMEOUT")] #[arg(long, default_value_t = 0, env = "TUNNEL_HUB_PROXY_IDLE_TIMEOUT")]
proxy_idle_timeout: u64, proxy_idle_timeout: u64,
/// Ping interval in seconds (for both sides)
#[arg(long, default_value_t = 15, env = "TUNNEL_HUB_PING_INTERVAL")] #[arg(long, default_value_t = 15, env = "TUNNEL_HUB_PING_INTERVAL")]
ping_interval: u64, ping_interval: u64,
/// Max concurrent streams per proxy connection
#[arg(long, default_value_t = 2048, env = "TUNNEL_HUB_MAX_STREAMS")] #[arg(long, default_value_t = 2048, env = "TUNNEL_HUB_MAX_STREAMS")]
max_streams: usize, max_streams: usize,
/// Per-connection outbound queue capacity before treating the socket as congested
#[arg( #[arg(
long, long,
default_value_t = 128, default_value_t = 128,
@@ -46,54 +31,96 @@ struct Args {
)] )]
outbound_queue_capacity: usize, outbound_queue_capacity: usize,
/// Local Aether app base URL for control-plane callbacks
#[arg( #[arg(
long, long,
default_value = "http://127.0.0.1:8084", default_value = "http://127.0.0.1:8084",
env = "TUNNEL_HUB_APP_BASE_URL" env = "TUNNEL_HUB_APP_BASE_URL"
)] )]
app_base_url: String, app_base_url: String,
}
#[derive(Clone)] #[arg(long, env = "TUNNEL_HUB_MAX_IN_FLIGHT_REQUESTS")]
pub struct AppState { max_in_flight_requests: Option<usize>,
pub hub: std::sync::Arc<HubRouter>,
pub proxy_conn_cfg: ConnConfig, #[arg(long, env = "TUNNEL_HUB_DISTRIBUTED_REQUEST_LIMIT")]
pub max_streams: usize, distributed_request_limit: Option<usize>,
#[arg(long, env = "TUNNEL_HUB_DISTRIBUTED_REQUEST_REDIS_URL")]
distributed_request_redis_url: Option<String>,
#[arg(long, env = "TUNNEL_HUB_DISTRIBUTED_REQUEST_REDIS_KEY_PREFIX")]
distributed_request_redis_key_prefix: Option<String>,
#[arg(
long,
env = "TUNNEL_HUB_DISTRIBUTED_REQUEST_LEASE_TTL_MS",
default_value_t = 30_000
)]
distributed_request_lease_ttl_ms: u64,
#[arg(
long,
env = "TUNNEL_HUB_DISTRIBUTED_REQUEST_RENEW_INTERVAL_MS",
default_value_t = 10_000
)]
distributed_request_renew_interval_ms: u64,
#[arg(
long,
env = "TUNNEL_HUB_DISTRIBUTED_REQUEST_COMMAND_TIMEOUT_MS",
default_value_t = 1_000
)]
distributed_request_command_timeout_ms: u64,
} }
#[tokio::main] #[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> { async fn main() -> Result<(), Box<dyn std::error::Error>> {
// Initialize tracing init_service_runtime(ServiceRuntimeConfig::new("aether-hub", "aether_hub=info"))?;
tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| "aether_hub=info".into()),
)
.init();
let args = Args::parse(); let args = Args::parse();
let hub = HubRouter::new(ControlPlaneClient::new(args.app_base_url));
let outbound_queue_capacity = args.outbound_queue_capacity.clamp(8, 4096); let outbound_queue_capacity = args.outbound_queue_capacity.clamp(8, 4096);
let ping_interval = Duration::from_secs(args.ping_interval); let ping_interval = Duration::from_secs(args.ping_interval);
let state = AppState { let mut state = AppState::new(
hub, ControlPlaneClient::new(args.app_base_url),
proxy_conn_cfg: ConnConfig { ConnConfig {
ping_interval, ping_interval,
idle_timeout: Duration::from_secs(args.proxy_idle_timeout), idle_timeout: Duration::from_secs(args.proxy_idle_timeout),
outbound_queue_capacity, outbound_queue_capacity,
}, },
max_streams: args.max_streams, args.max_streams,
}; )
.with_request_concurrency_limit(args.max_in_flight_requests);
let app = Router::new() if let Some(limit) = args.distributed_request_limit.filter(|limit| *limit > 0) {
.route("/health", get(health)) let redis_url = args
.route("/stats", get(stats)) .distributed_request_redis_url
.route("/proxy", get(ws_proxy)) .as_deref()
.route("/local/relay/{node_id}", post(relay_request)) .map(str::trim)
.with_state(state); .filter(|value| !value.is_empty())
.ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"TUNNEL_HUB_DISTRIBUTED_REQUEST_REDIS_URL is required when distributed request limit is enabled",
)
})?;
state = state.with_distributed_request_gate(DistributedConcurrencyGate::new_redis(
"hub_requests_distributed",
limit,
RedisDistributedConcurrencyConfig {
url: redis_url.to_string(),
key_prefix: args
.distributed_request_redis_key_prefix
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
lease_ttl_ms: args.distributed_request_lease_ttl_ms.max(1),
renew_interval_ms: args.distributed_request_renew_interval_ms.max(1),
command_timeout_ms: Some(args.distributed_request_command_timeout_ms.max(1)),
},
)?);
}
let app = build_router_with_state(state);
let listener = tokio::net::TcpListener::bind(&args.bind).await?; let listener = tokio::net::TcpListener::bind(&args.bind).await?;
info!(bind = %args.bind, "aether-hub started"); info!(bind = %args.bind, "aether-hub started");
@@ -105,63 +132,119 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
Ok(()) Ok(())
} }
// --------------------------------------------------------------------------- #[cfg(test)]
// HTTP endpoints mod tests {
// --------------------------------------------------------------------------- use super::*;
use axum::http::{HeaderValue, StatusCode};
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
async fn health() -> impl IntoResponse { async fn start_server(app: axum::Router) -> (String, tokio::task::JoinHandle<()>) {
Json(serde_json::json!({"status": "ok"})) let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
} .await
.expect("listener should bind");
async fn stats(State(state): State<AppState>) -> impl IntoResponse { let addr = listener.local_addr().expect("local addr should resolve");
Json(state.hub.stats()) let handle = tokio::spawn(async move {
} axum::serve(listener, app).await.expect("server should run");
});
// --------------------------------------------------------------------------- (format!("http://{addr}"), handle)
// WebSocket endpoints
// ---------------------------------------------------------------------------
async fn ws_proxy(
ws: WebSocketUpgrade,
State(state): State<AppState>,
headers: axum::http::HeaderMap,
) -> impl IntoResponse {
let node_id = headers
.get("x-node-id")
.and_then(|v| v.to_str().ok())
.unwrap_or("")
.trim()
.to_string();
let node_name = headers
.get("x-node-name")
.and_then(|v| v.to_str().ok())
.unwrap_or(&node_id)
.trim()
.to_string();
let max_streams: usize = headers
.get("x-tunnel-max-streams")
.and_then(|v| v.to_str().ok())
.and_then(|v| v.parse().ok())
.unwrap_or(state.max_streams)
.clamp(64, 2048);
if node_id.is_empty() {
warn!("proxy connection rejected: missing X-Node-ID header");
return axum::http::StatusCode::BAD_REQUEST.into_response();
} }
ws.max_frame_size(64 * 1024 * 1024) #[tokio::test]
.on_upgrade(move |socket| { async fn hub_exposes_metrics_endpoint() {
proxy_conn::handle_proxy_connection( let app = build_router_with_state(AppState::new(
socket, ControlPlaneClient::disabled(),
state.hub, ConnConfig {
node_id, ping_interval: Duration::from_secs(15),
node_name, idle_timeout: Duration::from_secs(0),
max_streams, outbound_queue_capacity: 128,
state.proxy_conn_cfg, },
128,
));
let (base_url, handle) = start_server(app).await;
let response = reqwest::Client::new()
.get(format!("{base_url}/metrics"))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), axum::http::StatusCode::OK);
assert_eq!(
response
.headers()
.get(axum::http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok()),
Some("text/plain; version=0.0.4; charset=utf-8")
);
let body = response.text().await.expect("body should read");
assert!(body.contains("service_up{service=\"aether-hub\"} 1"));
assert!(body.contains("hub_proxy_connections 0"));
assert!(body.contains("hub_active_streams 0"));
handle.abort();
}
#[tokio::test]
async fn hub_rejects_second_proxy_connection_with_distributed_overload() {
let distributed_gate =
DistributedConcurrencyGate::new_in_memory("hub_requests_distributed", 1);
let app_a = build_router_with_state(
AppState::new(
ControlPlaneClient::disabled(),
ConnConfig {
ping_interval: Duration::from_secs(15),
idle_timeout: Duration::from_secs(0),
outbound_queue_capacity: 128,
},
128,
) )
}) .with_distributed_request_gate(distributed_gate.clone()),
.into_response() );
let app_b = build_router_with_state(
AppState::new(
ControlPlaneClient::disabled(),
ConnConfig {
ping_interval: Duration::from_secs(15),
idle_timeout: Duration::from_secs(0),
outbound_queue_capacity: 128,
},
128,
)
.with_distributed_request_gate(distributed_gate),
);
let (base_a, handle_a) = start_server(app_a).await;
let (base_b, handle_b) = start_server(app_b).await;
let request_a = format!("{}/proxy", base_a.replace("http://", "ws://"))
.into_client_request()
.expect("request should build");
let mut request_a = request_a;
request_a
.headers_mut()
.insert("x-node-id", HeaderValue::from_static("node-a"));
let (socket, _) = tokio_tungstenite::connect_async(request_a)
.await
.expect("first websocket should connect");
let request_b = format!("{}/proxy", base_b.replace("http://", "ws://"))
.into_client_request()
.expect("request should build");
let mut request_b = request_b;
request_b
.headers_mut()
.insert("x-node-id", HeaderValue::from_static("node-b"));
let error = tokio_tungstenite::connect_async(request_b)
.await
.expect_err("second websocket should be rejected");
match error {
tokio_tungstenite::tungstenite::Error::Http(response) => {
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
}
other => panic!("unexpected websocket error: {other}"),
}
drop(socket);
handle_a.abort();
handle_b.abort();
}
} }

5
aether-hub/src/proxy_conn.rs
View File

@@ -5,9 +5,10 @@
use std::sync::Arc; use std::sync::Arc;
use std::time::Duration; use std::time::Duration;
use aether_runtime::bounded_queue;
use axum::extract::ws::{Message, WebSocket}; use axum::extract::ws::{Message, WebSocket};
use futures_util::{SinkExt, StreamExt}; use futures_util::{SinkExt, StreamExt};
use tokio::sync::{mpsc, watch}; use tokio::sync::watch;
use tracing::{debug, info, warn}; use tracing::{debug, info, warn};
use crate::hub::{ConnConfig, HubRouter, ProxyConn, SendStatus}; use crate::hub::{ConnConfig, HubRouter, ProxyConn, SendStatus};
@@ -27,7 +28,7 @@ pub async fn handle_proxy_connection(
let conn_id = hub.alloc_conn_id(); let conn_id = hub.alloc_conn_id();
let (mut ws_tx, ws_rx) = ws.split(); let (mut ws_tx, ws_rx) = ws.split();
let (tx, mut rx) = mpsc::channel::<Message>(cfg.outbound_queue_capacity); let (tx, mut rx) = bounded_queue::<Message>(cfg.outbound_queue_capacity);
let (close_tx, mut close_rx) = watch::channel(false); let (close_tx, mut close_rx) = watch::channel(false);
let conn = Arc::new(ProxyConn::new( let conn = Arc::new(ProxyConn::new(

5
aether-proxy/Cargo.toml
View File

@@ -5,8 +5,10 @@ edition = "2021"
description = "Tunnel proxy for Aether" description = "Tunnel proxy for Aether"
[dependencies] [dependencies]
aether-http.workspace = true
aether-runtime.workspace = true
tokio = { version = "1", features = ["full"] } tokio = { version = "1", features = ["full"] }
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls", "stream", "http2"] } reqwest.workspace = true
hyper = { version = "1", features = ["client", "http1", "http2"] } hyper = { version = "1", features = ["client", "http1", "http2"] }
hyper-util = { version = "0.1", features = ["client", "client-legacy", "http1", "http2", "tokio"] } hyper-util = { version = "0.1", features = ["client", "client-legacy", "http1", "http2", "tokio"] }
http-body-util = "0.1" http-body-util = "0.1"
@@ -16,7 +18,6 @@ futures-util = "0.3"
base64 = "0.22" base64 = "0.22"
clap = { version = "4", features = ["derive", "env"] } clap = { version = "4", features = ["derive", "env"] }
tracing = "0.1" tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
serde = { version = "1", features = ["derive"] } serde = { version = "1", features = ["derive"] }
serde_json = "1" serde_json = "1"
thiserror = "2" thiserror = "2"

88
aether-proxy/src/app.rs
View File

@@ -4,8 +4,11 @@ use std::sync::atomic::AtomicU64;
use std::sync::{Arc, RwLock}; use std::sync::{Arc, RwLock};
use std::time::Duration; use std::time::Duration;
use aether_runtime::{
init_reloadable_tracing, wait_for_shutdown_signal, ConcurrencyGate, DistributedConcurrencyGate,
LogFormat, RedisDistributedConcurrencyConfig,
};
use arc_swap::ArcSwap; use arc_swap::ArcSwap;
use tokio::signal;
use tokio::sync::{watch, Mutex}; use tokio::sync::{watch, Mutex};
use tracing::{error, info, warn}; use tracing::{error, info, warn};
@@ -140,12 +143,38 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
// Build shared application state // Build shared application state
let tunnel_tls_config = Arc::new(crate::tunnel::client::build_tls_config()); let tunnel_tls_config = Arc::new(crate::tunnel::client::build_tls_config());
let state = Arc::new(AppState { let mut state = AppState {
config: Arc::new(config), config: Arc::new(config),
dns_cache, dns_cache,
upstream_client, upstream_client,
tunnel_tls_config, tunnel_tls_config,
}); stream_gate: None,
distributed_stream_gate: None,
};
if let Some(limit) = state.config.max_in_flight_streams {
state = state
.with_stream_concurrency_gate(Arc::new(ConcurrencyGate::new("proxy_streams", limit)));
}
if let Some(limit) = state.config.distributed_stream_limit {
let redis_url = state
.config
.distributed_stream_redis_url
.clone()
.expect("distributed stream redis url should be validated");
let distributed_gate = DistributedConcurrencyGate::new_redis(
"proxy_streams_distributed",
limit,
RedisDistributedConcurrencyConfig {
url: redis_url,
key_prefix: state.config.distributed_stream_redis_key_prefix.clone(),
lease_ttl_ms: state.config.distributed_stream_lease_ttl_ms,
renew_interval_ms: state.config.distributed_stream_renew_interval_ms,
command_timeout_ms: Some(state.config.distributed_stream_command_timeout_ms),
},
)?;
state = state.with_distributed_stream_concurrency_gate(Arc::new(distributed_gate));
}
let state = Arc::new(state);
// Shutdown signal channel // Shutdown signal channel
let (shutdown_tx, shutdown_rx) = watch::channel(false); let (shutdown_tx, shutdown_rx) = watch::channel(false);
@@ -309,52 +338,19 @@ async fn retry_failed_registrations(
} }
fn init_tracing(config: &Config) { fn init_tracing(config: &Config) {
use tracing_subscriber::prelude::*; let format = if config.log_json {
use tracing_subscriber::{reload, EnvFilter}; LogFormat::Json
let filter = EnvFilter::try_new(&config.log_level).unwrap_or_else(|_| EnvFilter::new("info"));
let (filter_layer, reload_handle) = reload::Layer::new(filter);
runtime::set_log_reloader(Box::new(move |level: &str| {
if let Ok(new_filter) = EnvFilter::try_new(level) {
let _ = reload_handle.modify(|f| *f = new_filter);
}
}));
if config.log_json {
tracing_subscriber::registry()
.with(filter_layer)
.with(tracing_subscriber::fmt::layer().json())
.init();
} else { } else {
tracing_subscriber::registry() LogFormat::Pretty
.with(filter_layer) };
.with(tracing_subscriber::fmt::layer())
.init(); let reloader = init_reloadable_tracing(&config.log_level, format)
} .expect("proxy tracing should initialize");
runtime::set_log_reloader(reloader);
} }
async fn wait_for_shutdown() { async fn wait_for_shutdown() {
let ctrl_c = async { wait_for_shutdown_signal()
signal::ctrl_c()
.await .await
.expect("failed to install Ctrl+C handler"); .expect("failed to install shutdown signal handler");
};
#[cfg(unix)]
let terminate = async {
signal::unix::signal(signal::unix::SignalKind::terminate())
.expect("failed to install SIGTERM handler")
.recv()
.await;
};
#[cfg(not(unix))]
let terminate = std::future::pending::<()>();
tokio::select! {
_ = ctrl_c => {},
_ = terminate => {},
}
} }

65
aether-proxy/src/config.rs
View File

@@ -146,6 +146,46 @@ pub struct Config {
#[arg(long, env = "AETHER_PROXY_MAX_CONCURRENT_CONNECTIONS")] #[arg(long, env = "AETHER_PROXY_MAX_CONCURRENT_CONNECTIONS")]
pub max_concurrent_connections: Option<u64>, pub max_concurrent_connections: Option<u64>,
/// Maximum in-flight tunneled streams accepted by this proxy instance.
#[arg(long, env = "AETHER_PROXY_MAX_IN_FLIGHT_STREAMS")]
pub max_in_flight_streams: Option<usize>,
/// Maximum in-flight tunneled streams admitted across all proxy instances.
#[arg(long, env = "AETHER_PROXY_DISTRIBUTED_STREAM_LIMIT")]
pub distributed_stream_limit: Option<usize>,
/// Redis URL used for cross-instance stream admission.
#[arg(long, env = "AETHER_PROXY_DISTRIBUTED_STREAM_REDIS_URL")]
pub distributed_stream_redis_url: Option<String>,
/// Optional key prefix for cross-instance stream admission state.
#[arg(long, env = "AETHER_PROXY_DISTRIBUTED_STREAM_REDIS_KEY_PREFIX")]
pub distributed_stream_redis_key_prefix: Option<String>,
/// Lease TTL in milliseconds for distributed stream admission permits.
#[arg(
long,
env = "AETHER_PROXY_DISTRIBUTED_STREAM_LEASE_TTL_MS",
default_value_t = 30_000
)]
pub distributed_stream_lease_ttl_ms: u64,
/// Renew interval in milliseconds for distributed stream admission permits.
#[arg(
long,
env = "AETHER_PROXY_DISTRIBUTED_STREAM_RENEW_INTERVAL_MS",
default_value_t = 10_000
)]
pub distributed_stream_renew_interval_ms: u64,
/// Command timeout in milliseconds for distributed stream admission Redis calls.
#[arg(
long,
env = "AETHER_PROXY_DISTRIBUTED_STREAM_COMMAND_TIMEOUT_MS",
default_value_t = 1_000
)]
pub distributed_stream_command_timeout_ms: u64,
/// DNS cache TTL in seconds /// DNS cache TTL in seconds
#[arg(long, env = "AETHER_PROXY_DNS_CACHE_TTL", default_value_t = 60)] #[arg(long, env = "AETHER_PROXY_DNS_CACHE_TTL", default_value_t = 60)]
pub dns_cache_ttl_secs: u64, pub dns_cache_ttl_secs: u64,
@@ -291,6 +331,31 @@ impl Config {
if self.upstream_connect_timeout_secs == 0 { if self.upstream_connect_timeout_secs == 0 {
anyhow::bail!("upstream_connect_timeout_secs must be > 0"); anyhow::bail!("upstream_connect_timeout_secs must be > 0");
} }
if matches!(self.max_in_flight_streams, Some(0)) {
anyhow::bail!("max_in_flight_streams must be > 0");
}
if matches!(self.distributed_stream_limit, Some(0)) {
anyhow::bail!("distributed_stream_limit must be > 0");
}
if self.distributed_stream_limit.is_some() && self.distributed_stream_redis_url.is_none() {
anyhow::bail!(
"distributed_stream_redis_url must be set when distributed_stream_limit is enabled"
);
}
if self.distributed_stream_lease_ttl_ms == 0 {
anyhow::bail!("distributed_stream_lease_ttl_ms must be > 0");
}
if self.distributed_stream_renew_interval_ms == 0 {
anyhow::bail!("distributed_stream_renew_interval_ms must be > 0");
}
if self.distributed_stream_renew_interval_ms >= self.distributed_stream_lease_ttl_ms {
anyhow::bail!(
"distributed_stream_renew_interval_ms must be < distributed_stream_lease_ttl_ms"
);
}
if self.distributed_stream_command_timeout_ms == 0 {
anyhow::bail!("distributed_stream_command_timeout_ms must be > 0");
}
Ok(()) Ok(())
} }
} }

18
aether-proxy/src/net.rs
View File

@@ -2,7 +2,7 @@
//! //!
//! These are standalone helpers not tied to any specific client or service. //! These are standalone helpers not tied to any specific client or service.
use reqwest::Client; use aether_http::{build_http_client, HttpClientConfig};
use tracing::{debug, info}; use tracing::{debug, info};
/// Auto-detect public IP by querying external services. /// Auto-detect public IP by querying external services.
@@ -13,9 +13,11 @@ pub async fn detect_public_ip() -> anyhow::Result<String> {
"https://icanhazip.com", "https://icanhazip.com",
]; ];
let client = Client::builder() let client = build_http_client(&HttpClientConfig {
.timeout(std::time::Duration::from_secs(5)) request_timeout_ms: Some(5_000),
.build()?; user_agent: Some("aether-proxy/net".to_string()),
..HttpClientConfig::default()
})?;
for endpoint in &endpoints { for endpoint in &endpoints {
match client.get(*endpoint).send().await { match client.get(*endpoint).send().await {
@@ -48,9 +50,11 @@ pub async fn detect_region(ip: &str) -> Option<String> {
// Try HTTPS provider first // Try HTTPS provider first
let https_url = format!("https://ipinfo.io/{}/country", ip); let https_url = format!("https://ipinfo.io/{}/country", ip);
let client = Client::builder() let client = build_http_client(&HttpClientConfig {
.timeout(std::time::Duration::from_secs(5)) request_timeout_ms: Some(5_000),
.build() user_agent: Some("aether-proxy/net".to_string()),
..HttpClientConfig::default()
})
.ok()?; .ok()?;
// Try ipinfo.io (HTTPS, returns plain text country code) // Try ipinfo.io (HTTPS, returns plain text country code)

78
aether-proxy/src/registration/client.rs
View File

@@ -1,5 +1,4 @@
use std::time::{Duration, SystemTime, UNIX_EPOCH}; use aether_http::{build_http_client, jittered_delay_for_retry, HttpClientConfig, HttpRetryConfig};
use reqwest::{Client, StatusCode}; use reqwest::{Client, StatusCode};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use tokio::time::sleep; use tokio::time::sleep;
@@ -49,44 +48,40 @@ pub struct AetherClient {
http: Client, http: Client,
base_url: String, base_url: String,
token: String, token: String,
retry_max_attempts: u32, retry: HttpRetryConfig,
retry_base_delay: Duration,
retry_max_delay: Duration,
} }
impl AetherClient { impl AetherClient {
pub fn new(config: &Config, aether_url: &str, management_token: &str) -> Self { pub fn new(config: &Config, aether_url: &str, management_token: &str) -> Self {
let mut builder = Client::builder() let http = build_http_client(&HttpClientConfig {
.timeout(Duration::from_secs(config.aether_request_timeout_secs)) connect_timeout_ms: Some(config.aether_connect_timeout_secs.saturating_mul(1_000)),
.connect_timeout(Duration::from_secs(config.aether_connect_timeout_secs)) request_timeout_ms: Some(config.aether_request_timeout_secs.saturating_mul(1_000)),
.pool_max_idle_per_host(config.aether_pool_max_idle_per_host) pool_idle_timeout_ms: Some(config.aether_pool_idle_timeout_secs.saturating_mul(1_000)),
.pool_idle_timeout(Duration::from_secs(config.aether_pool_idle_timeout_secs)) pool_max_idle_per_host: Some(config.aether_pool_max_idle_per_host),
.tcp_nodelay(config.aether_tcp_nodelay); tcp_keepalive_ms: if config.aether_tcp_keepalive_secs > 0 {
Some(config.aether_tcp_keepalive_secs.saturating_mul(1_000))
if config.aether_tcp_keepalive_secs > 0 {
builder =
builder.tcp_keepalive(Some(Duration::from_secs(config.aether_tcp_keepalive_secs)));
} else { } else {
builder = builder.tcp_keepalive(None); None
},
tcp_nodelay: config.aether_tcp_nodelay,
http2_adaptive_window: config.aether_http2,
user_agent: Some(format!("aether-proxy/{}", env!("CARGO_PKG_VERSION"))),
..HttpClientConfig::default()
})
.expect("failed to create HTTP client");
let retry = HttpRetryConfig {
max_attempts: config.aether_retry_max_attempts,
base_delay_ms: config.aether_retry_base_delay_ms,
max_delay_ms: config.aether_retry_max_delay_ms,
} }
.normalized();
if config.aether_http2 {
builder = builder.http2_adaptive_window(true);
}
let http = builder.build().expect("failed to create HTTP client");
let retry_base_delay = Duration::from_millis(config.aether_retry_base_delay_ms);
let retry_max_delay =
Duration::from_millis(config.aether_retry_max_delay_ms).max(retry_base_delay);
Self { Self {
http, http,
base_url: aether_url.trim_end_matches('/').to_string(), base_url: aether_url.trim_end_matches('/').to_string(),
token: management_token.to_string(), token: management_token.to_string(),
retry_max_attempts: config.aether_retry_max_attempts.max(1), retry,
retry_base_delay,
retry_max_delay,
} }
} }
@@ -193,15 +188,14 @@ impl AetherClient {
F: FnMut() -> reqwest::RequestBuilder, F: FnMut() -> reqwest::RequestBuilder,
{ {
let mut attempt: u32 = 0; let mut attempt: u32 = 0;
let mut delay = self.retry_base_delay;
loop { loop {
attempt = attempt.saturating_add(1); attempt = attempt.saturating_add(1);
let resp = make_req().send().await; let resp = make_req().send().await;
match resp { match resp {
Ok(resp) => { Ok(resp) => {
if should_retry_status(resp.status()) && attempt < self.retry_max_attempts { if should_retry_status(resp.status()) && attempt < self.retry.max_attempts {
let sleep_for = jitter_delay(delay); let sleep_for = jittered_delay_for_retry(self.retry, attempt - 1);
debug!( debug!(
attempt, attempt,
status = %resp.status(), status = %resp.status(),
@@ -210,15 +204,13 @@ impl AetherClient {
"Aether request retrying" "Aether request retrying"
); );
sleep(sleep_for).await; sleep(sleep_for).await;
let next_delay = delay.checked_mul(2).unwrap_or(self.retry_max_delay);
delay = std::cmp::min(next_delay, self.retry_max_delay);
continue; continue;
} }
return Ok(resp); return Ok(resp);
} }
Err(e) => { Err(e) => {
if attempt < self.retry_max_attempts { if attempt < self.retry.max_attempts {
let sleep_for = jitter_delay(delay); let sleep_for = jittered_delay_for_retry(self.retry, attempt - 1);
debug!( debug!(
attempt, attempt,
error = %e, error = %e,
@@ -227,8 +219,6 @@ impl AetherClient {
"Aether request retrying" "Aether request retrying"
); );
sleep(sleep_for).await; sleep(sleep_for).await;
let next_delay = delay.checked_mul(2).unwrap_or(self.retry_max_delay);
delay = std::cmp::min(next_delay, self.retry_max_delay);
continue; continue;
} }
return Err(e); return Err(e);
@@ -243,15 +233,3 @@ fn should_retry_status(status: StatusCode) -> bool {
|| status == StatusCode::TOO_MANY_REQUESTS || status == StatusCode::TOO_MANY_REQUESTS
|| status == StatusCode::REQUEST_TIMEOUT || status == StatusCode::REQUEST_TIMEOUT
} }
fn jitter_delay(base: Duration) -> Duration {
if base.is_zero() {
return base;
}
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.subsec_nanos() as u64)
.unwrap_or(0);
let jitter_ms = nanos % 100;
base + Duration::from_millis(jitter_ms)
}

13
aether-proxy/src/setup/upgrade.rs
View File

@@ -5,6 +5,7 @@
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use aether_http::{apply_http_client_config, HttpClientConfig};
use sha2::{Digest, Sha256}; use sha2::{Digest, Sha256};
const GITHUB_API_BASE: &str = "https://api.github.com"; const GITHUB_API_BASE: &str = "https://api.github.com";
@@ -56,10 +57,14 @@ fn build_github_client() -> anyhow::Result<reqwest::Client> {
reqwest::header::HeaderValue::from_static("application/vnd.github+json"), reqwest::header::HeaderValue::from_static("application/vnd.github+json"),
); );
Ok(reqwest::Client::builder() Ok(apply_http_client_config(
.timeout(std::time::Duration::from_secs(300)) reqwest::Client::builder().default_headers(headers),
.user_agent(format!("aether-proxy/{}", CURRENT_VERSION)) &HttpClientConfig {
.default_headers(headers) request_timeout_ms: Some(300_000),
user_agent: Some(format!("aether-proxy/{}", CURRENT_VERSION)),
..HttpClientConfig::default()
},
)
.build()?) .build()?)
} }

106
aether-proxy/src/state.rs
View File

@@ -4,6 +4,11 @@ use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, RwLock}; use std::sync::{Arc, RwLock};
use std::time::Duration; use std::time::Duration;
use aether_runtime::{
AdmissionPermit, ConcurrencyError, ConcurrencyGate, ConcurrencySnapshot,
DistributedConcurrencyError, DistributedConcurrencyGate, DistributedConcurrencySnapshot,
};
use crate::config::Config; use crate::config::Config;
use crate::registration::client::AetherClient; use crate::registration::client::AetherClient;
use crate::runtime::SharedDynamicConfig; use crate::runtime::SharedDynamicConfig;
@@ -19,6 +24,10 @@ pub struct AppState {
pub upstream_client: UpstreamClient, pub upstream_client: UpstreamClient,
/// Shared TLS config for tunnel WebSocket connections (avoids re-parsing root CAs on each reconnect). /// Shared TLS config for tunnel WebSocket connections (avoids re-parsing root CAs on each reconnect).
pub tunnel_tls_config: Arc<rustls::ClientConfig>, pub tunnel_tls_config: Arc<rustls::ClientConfig>,
/// Optional per-process stream admission gate.
pub stream_gate: Option<Arc<ConcurrencyGate>>,
/// Optional cross-instance stream admission gate.
pub distributed_stream_gate: Option<Arc<DistributedConcurrencyGate>>,
} }
/// Per-server state: one instance per Aether server connection. /// Per-server state: one instance per Aether server connection.
@@ -75,3 +84,100 @@ impl ProxyMetrics {
self.total_latency_ns.fetch_add(nanos, Ordering::Release); self.total_latency_ns.fetch_add(nanos, Ordering::Release);
} }
} }
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum ProxyAdmissionError {
#[error("proxy stream admission saturated at {limit} for gate {gate}")]
Saturated { gate: &'static str, limit: usize },
#[error("proxy stream admission unavailable for gate {gate}: {message}")]
Unavailable {
gate: &'static str,
limit: usize,
message: String,
},
}
impl AppState {
pub fn with_stream_concurrency_gate(mut self, gate: Arc<ConcurrencyGate>) -> Self {
self.stream_gate = Some(gate);
self
}
pub fn with_distributed_stream_concurrency_gate(
mut self,
gate: Arc<DistributedConcurrencyGate>,
) -> Self {
self.distributed_stream_gate = Some(gate);
self
}
pub fn stream_concurrency_snapshot(&self) -> Option<ConcurrencySnapshot> {
self.stream_gate.as_ref().map(|gate| gate.snapshot())
}
pub async fn distributed_stream_concurrency_snapshot(
&self,
) -> Result<Option<DistributedConcurrencySnapshot>, DistributedConcurrencyError> {
match &self.distributed_stream_gate {
Some(gate) => gate.snapshot().await.map(Some),
None => Ok(None),
}
}
pub async fn try_acquire_stream_permit(
&self,
) -> Result<Option<AdmissionPermit>, ProxyAdmissionError> {
let local = match &self.stream_gate {
Some(gate) => Some(gate.try_acquire().map_err(|err| {
match err {
ConcurrencyError::Saturated { gate, limit } => {
ProxyAdmissionError::Saturated { gate, limit }
}
ConcurrencyError::Closed { gate } => ProxyAdmissionError::Unavailable {
gate,
limit: self
.stream_gate
.as_ref()
.map(|inner| inner.snapshot().limit)
.unwrap_or(0),
message: "local stream gate is closed".to_string(),
},
}
})?),
None => None,
};
let distributed = match &self.distributed_stream_gate {
Some(gate) => Some(gate.try_acquire().await.map_err(|err| {
match err {
DistributedConcurrencyError::Saturated { gate, limit } => {
ProxyAdmissionError::Saturated { gate, limit }
}
DistributedConcurrencyError::Unavailable {
gate,
limit,
message,
} => ProxyAdmissionError::Unavailable {
gate,
limit,
message,
},
DistributedConcurrencyError::InvalidConfiguration(message) => {
ProxyAdmissionError::Unavailable {
gate: "proxy_streams_distributed",
limit: self
.distributed_stream_gate
.as_ref()
.map(|inner| inner.limit())
.unwrap_or(0),
message,
}
}
}
})?),
None => None,
};
Ok(AdmissionPermit::from_parts(local, distributed))
}
}

2
aether-proxy/src/tunnel/client.rs
View File

@@ -136,7 +136,7 @@ pub async fn connect_and_run(
// resetting shared atomic metrics via swap(0)) // resetting shared atomic metrics via swap(0))
let hb_handle = if conn_idx == 0 { let hb_handle = if conn_idx == 0 {
heartbeat::spawn( heartbeat::spawn(
Arc::clone(&state.config), Arc::clone(state),
Arc::clone(server), Arc::clone(server),
frame_tx.clone(), frame_tx.clone(),
shutdown.clone(), shutdown.clone(),

41
aether-proxy/src/tunnel/heartbeat.rs
View File

@@ -10,9 +10,9 @@ use bytes::Bytes;
use tokio::sync::watch; use tokio::sync::watch;
use tracing::{debug, info, warn}; use tracing::{debug, info, warn};
use crate::config::Config;
use crate::registration::client::RemoteConfig; use crate::registration::client::RemoteConfig;
use crate::runtime; use crate::runtime;
use crate::state::AppState;
use crate::state::ServerContext; use crate::state::ServerContext;
use super::protocol::{Frame, MsgType}; use super::protocol::{Frame, MsgType};
@@ -62,7 +62,7 @@ struct HeartbeatSnapshot {
/// Spawn the heartbeat task. Returns a handle for forwarding ACKs. /// Spawn the heartbeat task. Returns a handle for forwarding ACKs.
pub fn spawn( pub fn spawn(
_config: Arc<Config>, state: Arc<AppState>,
server: Arc<ServerContext>, server: Arc<ServerContext>,
frame_tx: FrameSender, frame_tx: FrameSender,
mut shutdown: watch::Receiver<bool>, mut shutdown: watch::Receiver<bool>,
@@ -107,11 +107,12 @@ pub fn spawn(
}; };
let payload = build_heartbeat_payload( let payload = build_heartbeat_payload(
&state,
&server, &server,
&heartbeat_session_id, &heartbeat_session_id,
heartbeat_id, heartbeat_id,
snapshot snapshot
); ).await;
let frame = Frame::control(MsgType::HeartbeatData, payload); let frame = Frame::control(MsgType::HeartbeatData, payload);
if frame_tx.send(frame).await.is_err() { if frame_tx.send(frame).await.is_err() {
if let Some((_, snap)) = pending.take() { if let Some((_, snap)) = pending.take() {
@@ -216,7 +217,8 @@ fn restore_snapshot(server: &ServerContext, snap: HeartbeatSnapshot) {
} }
} }
fn build_heartbeat_payload( async fn build_heartbeat_payload(
state: &AppState,
server: &ServerContext, server: &ServerContext,
heartbeat_session_id: &str, heartbeat_session_id: &str,
heartbeat_id: u64, heartbeat_id: u64,
@@ -230,6 +232,36 @@ fn build_heartbeat_payload(
None None
}; };
let local_admission = state.stream_concurrency_snapshot().map(|snapshot| {
serde_json::json!({
"limit": snapshot.limit,
"in_flight": snapshot.in_flight,
"available_permits": snapshot.available_permits,
"high_watermark": snapshot.high_watermark,
"rejected_total": snapshot.rejected,
})
});
let distributed_admission = match state.distributed_stream_concurrency_snapshot().await {
Ok(Some(snapshot)) => Some(serde_json::json!({
"limit": snapshot.limit,
"in_flight": snapshot.in_flight,
"available_permits": snapshot.available_permits,
"high_watermark": snapshot.high_watermark,
"rejected_total": snapshot.rejected,
})),
Ok(None) => None,
Err(err) => Some(serde_json::json!({
"error": err.to_string(),
})),
};
let admission = match (local_admission, distributed_admission) {
(None, None) => None,
(local, distributed) => Some(serde_json::json!({
"local_streams": local,
"distributed_streams": distributed,
})),
};
let payload = serde_json::json!({ let payload = serde_json::json!({
"node_id": node_id, "node_id": node_id,
"heartbeat_session_id": heartbeat_session_id, "heartbeat_session_id": heartbeat_session_id,
@@ -242,6 +274,7 @@ fn build_heartbeat_payload(
"stream_errors": snapshot.stream_errors, "stream_errors": snapshot.stream_errors,
"proxy_metadata": { "proxy_metadata": {
"version": CURRENT_VERSION, "version": CURRENT_VERSION,
"admission": admission,
}, },
}); });

209
aether-proxy/src/tunnel/stream_handler.rs
View File

@@ -9,6 +9,7 @@ use std::sync::atomic::Ordering;
use std::sync::Arc; use std::sync::Arc;
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
use aether_runtime::hold_admission_permit_until;
use bytes::Bytes; use bytes::Bytes;
use futures_util::stream; use futures_util::stream;
use futures_util::StreamExt; use futures_util::StreamExt;
@@ -72,10 +73,26 @@ pub async fn handle_stream(
body_rx: mpsc::Receiver<TunnelFrame>, body_rx: mpsc::Receiver<TunnelFrame>,
frame_tx: FrameSender, frame_tx: FrameSender,
) { ) {
let permit = match state.try_acquire_stream_permit().await {
Ok(permit) => permit,
Err(err) => {
let message = match err {
crate::state::ProxyAdmissionError::Saturated { .. } => "proxy overloaded",
crate::state::ProxyAdmissionError::Unavailable { .. } => {
"proxy admission unavailable"
}
};
send_error(&frame_tx, stream_id, message).await;
return;
}
};
server.active_connections.fetch_add(1, Ordering::Release); server.active_connections.fetch_add(1, Ordering::Release);
let connect_elapsed = let connect_elapsed = hold_admission_permit_until(permit, async {
handle_stream_inner(&state, &server, stream_id, meta, body_rx, &frame_tx).await; handle_stream_inner(&state, &server, stream_id, meta, body_rx, &frame_tx).await
})
.await;
server.active_connections.fetch_sub(1, Ordering::Release); server.active_connections.fetch_sub(1, Ordering::Release);
if let Some(d) = connect_elapsed { if let Some(d) = connect_elapsed {
@@ -431,7 +448,20 @@ fn build_streaming_request_body(
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use std::collections::HashMap;
use std::sync::atomic::AtomicU64;
use std::sync::Once;
use aether_runtime::{bounded_queue, ConcurrencyGate, DistributedConcurrencyGate};
use arc_swap::ArcSwap;
use super::*; use super::*;
use crate::config::Config;
use crate::registration::client::AetherClient;
use crate::runtime::DynamicConfig;
use crate::state::ProxyMetrics;
use crate::target_filter::DnsCache;
use crate::tunnel::client::build_tls_config;
#[tokio::test] #[tokio::test]
async fn streaming_request_body_yields_chunks_and_tracks_size() { async fn streaming_request_body_yields_chunks_and_tracks_size() {
@@ -503,4 +533,179 @@ mod tests {
assert!(body.frame().await.is_none()); assert!(body.frame().await.is_none());
assert_eq!(body_size.load(Ordering::Relaxed), 0); assert_eq!(body_size.load(Ordering::Relaxed), 0);
} }
#[tokio::test]
async fn rejects_stream_when_local_admission_gate_is_saturated() {
let gate = Arc::new(ConcurrencyGate::new("proxy_streams", 1));
let _permit = gate.try_acquire().expect("first permit");
let state = sample_state(Some(gate), None);
let server = sample_server(&state);
let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(4);
let (_body_tx, body_rx) = mpsc::channel(1);
handle_stream(
Arc::clone(&state),
server,
7,
sample_request_meta(),
body_rx,
frame_tx,
)
.await;
let frame = frame_rx.recv().await.expect("overload frame");
assert_eq!(frame.stream_id, 7);
assert_eq!(frame.msg_type, MsgType::StreamError);
assert_eq!(frame.payload, Bytes::from_static(b"proxy overloaded"));
assert_eq!(
state
.stream_gate
.as_ref()
.expect("stream gate")
.snapshot()
.rejected,
1
);
}
#[tokio::test]
async fn rejects_stream_when_distributed_admission_gate_is_saturated() {
let gate = Arc::new(DistributedConcurrencyGate::new_in_memory(
"proxy_streams_distributed",
1,
));
let _permit = gate.try_acquire().await.expect("first permit");
let state = sample_state(None, Some(gate));
let server = sample_server(&state);
let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(4);
let (_body_tx, body_rx) = mpsc::channel(1);
handle_stream(
Arc::clone(&state),
server,
9,
sample_request_meta(),
body_rx,
frame_tx,
)
.await;
let frame = frame_rx.recv().await.expect("overload frame");
assert_eq!(frame.stream_id, 9);
assert_eq!(frame.msg_type, MsgType::StreamError);
assert_eq!(frame.payload, Bytes::from_static(b"proxy overloaded"));
assert_eq!(
state
.distributed_stream_gate
.as_ref()
.expect("distributed gate")
.snapshot()
.await
.expect("distributed snapshot")
.rejected,
1
);
}
fn sample_request_meta() -> RequestMeta {
RequestMeta {
method: "GET".to_string(),
url: "https://example.com/ok".to_string(),
headers: HashMap::new(),
timeout: 30,
}
}
fn sample_state(
stream_gate: Option<Arc<ConcurrencyGate>>,
distributed_stream_gate: Option<Arc<DistributedConcurrencyGate>>,
) -> Arc<AppState> {
ensure_rustls_provider();
let config = Arc::new(sample_config());
let dns_cache = Arc::new(DnsCache::new(Duration::from_secs(60), 128));
let upstream_client =
upstream_client::build_upstream_client(&config, Arc::clone(&dns_cache));
Arc::new(AppState {
config,
dns_cache,
upstream_client,
tunnel_tls_config: Arc::new(build_tls_config()),
stream_gate,
distributed_stream_gate,
})
}
fn sample_server(state: &Arc<AppState>) -> Arc<ServerContext> {
let config = Arc::clone(&state.config);
Arc::new(ServerContext {
server_label: "server".to_string(),
aether_url: config.aether_url.clone(),
management_token: config.management_token.clone(),
node_name: config.node_name.clone(),
node_id: Arc::new(std::sync::RwLock::new("node-1".to_string())),
aether_client: Arc::new(AetherClient::new(
&config,
&config.aether_url,
&config.management_token,
)),
dynamic: Arc::new(ArcSwap::from_pointee(DynamicConfig::from_config(&config))),
active_connections: Arc::new(AtomicU64::new(0)),
metrics: Arc::new(ProxyMetrics::new()),
})
}
fn sample_config() -> Config {
Config {
aether_url: "https://aether.example.com".to_string(),
management_token: "token".to_string(),
public_ip: None,
node_name: "proxy-test".to_string(),
node_region: None,
heartbeat_interval: 30,
allowed_ports: vec![80, 443],
aether_request_timeout_secs: 10,
aether_connect_timeout_secs: 10,
aether_pool_max_idle_per_host: 8,
aether_pool_idle_timeout_secs: 90,
aether_tcp_keepalive_secs: 60,
aether_tcp_nodelay: true,
aether_http2: true,
aether_retry_max_attempts: 3,
aether_retry_base_delay_ms: 200,
aether_retry_max_delay_ms: 2_000,
max_concurrent_connections: None,
max_in_flight_streams: None,
distributed_stream_limit: None,
distributed_stream_redis_url: None,
distributed_stream_redis_key_prefix: None,
distributed_stream_lease_ttl_ms: 30_000,
distributed_stream_renew_interval_ms: 10_000,
distributed_stream_command_timeout_ms: 1_000,
dns_cache_ttl_secs: 60,
dns_cache_capacity: 128,
upstream_connect_timeout_secs: 30,
upstream_pool_max_idle_per_host: 4,
upstream_pool_idle_timeout_secs: 60,
upstream_tcp_keepalive_secs: 60,
upstream_tcp_nodelay: true,
log_level: "info".to_string(),
log_json: false,
tunnel_reconnect_base_ms: 500,
tunnel_reconnect_max_ms: 30_000,
tunnel_ping_interval_secs: 15,
tunnel_max_streams: Some(8),
tunnel_connect_timeout_secs: 15,
tunnel_tcp_keepalive_secs: 30,
tunnel_tcp_nodelay: true,
tunnel_stale_timeout_secs: 45,
tunnel_connections: 1,
}
}
fn ensure_rustls_provider() {
static INIT: Once = Once::new();
INIT.call_once(|| {
let _ = rustls::crypto::ring::default_provider().install_default();
});
}
} }

6
aether-proxy/src/tunnel/writer.rs
View File

@@ -7,8 +7,8 @@
use std::time::Duration; use std::time::Duration;
use aether_runtime::{bounded_queue, BoundedQueueSender};
use futures_util::SinkExt; use futures_util::SinkExt;
use tokio::sync::mpsc;
use tokio::task::JoinHandle; use tokio::task::JoinHandle;
use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::Message;
use tracing::{debug, error, trace}; use tracing::{debug, error, trace};
@@ -16,7 +16,7 @@ use tracing::{debug, error, trace};
use super::protocol::Frame; use super::protocol::Frame;
/// Sender half — cloned by stream handlers and heartbeat. /// Sender half — cloned by stream handlers and heartbeat.
pub type FrameSender = mpsc::Sender<Frame>; pub type FrameSender = BoundedQueueSender<Frame>;
/// Spawn the writer task. Returns the sender and a JoinHandle for cleanup. /// Spawn the writer task. Returns the sender and a JoinHandle for cleanup.
/// ///
@@ -26,7 +26,7 @@ pub fn spawn_writer<S>(mut sink: S, ping_interval: Duration) -> (FrameSender, Jo
where where
S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static, S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
{ {
let (tx, mut rx) = mpsc::channel::<Frame>(256); let (tx, mut rx) = bounded_queue::<Frame>(256);
let handle = tokio::spawn(async move { let handle = tokio::spawn(async move {
let mut ping_ticker = tokio::time::interval(ping_interval); let mut ping_ticker = tokio::time::interval(ping_interval);

9
crates/aether-cache/Cargo.toml Normal file
View File

@@ -0,0 +1,9 @@
[package]
name = "aether-cache"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
description = "Shared in-memory cache primitives for Aether Rust services"
[dependencies]

5
crates/aether-cache/src/lib.rs Normal file
View File

@@ -0,0 +1,5 @@
mod namespace;
mod ttl_map;
pub use namespace::CacheKeyNamespace;
pub use ttl_map::ExpiringMap;

51
crates/aether-cache/src/namespace.rs Normal file
View File

@@ -0,0 +1,51 @@
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct CacheKeyNamespace {
prefix: String,
}
impl CacheKeyNamespace {
pub fn new(prefix: impl Into<String>) -> Self {
Self {
prefix: prefix.into(),
}
}
pub fn child(&self, suffix: &str) -> Self {
if self.prefix.is_empty() {
return Self::new(suffix);
}
if suffix.is_empty() {
return self.clone();
}
Self::new(format!("{}:{}", self.prefix, suffix))
}
pub fn key(&self, raw_key: &str) -> String {
if self.prefix.is_empty() {
return raw_key.to_string();
}
if raw_key.is_empty() {
return self.prefix.clone();
}
format!("{}:{}", self.prefix, raw_key)
}
pub fn prefix(&self) -> &str {
&self.prefix
}
}
#[cfg(test)]
mod tests {
use super::CacheKeyNamespace;
#[test]
fn composes_scoped_keys() {
let root = CacheKeyNamespace::new("aether");
let child = root.child("auth");
assert_eq!(root.key("user-1"), "aether:user-1");
assert_eq!(child.key("user-1"), "aether:auth:user-1");
assert_eq!(child.prefix(), "aether:auth");
}
}

175
crates/aether-cache/src/ttl_map.rs Normal file
View File

@@ -0,0 +1,175 @@
use std::collections::HashMap;
use std::hash::Hash;
use std::sync::Mutex;
use std::time::{Duration, Instant};
#[derive(Debug, Clone)]
struct TimedEntry<V> {
value: V,
inserted_at: Instant,
}
#[derive(Debug)]
pub struct ExpiringMap<K, V> {
entries: Mutex<HashMap<K, TimedEntry<V>>>,
}
impl<K, V> Default for ExpiringMap<K, V> {
fn default() -> Self {
Self {
entries: Mutex::new(HashMap::new()),
}
}
}
impl<K, V> ExpiringMap<K, V>
where
K: Eq + Hash + Clone,
{
pub fn new() -> Self {
Self::default()
}
pub fn insert(&self, key: K, value: V, ttl: Duration, max_entries: usize) {
let Ok(mut entries) = self.entries.lock() else {
return;
};
prune_expired(&mut entries, ttl);
while max_entries > 0 && entries.len() >= max_entries {
let Some(oldest_key) = entries
.iter()
.min_by_key(|(_, entry)| entry.inserted_at)
.map(|(key, _)| key.clone())
else {
break;
};
entries.remove(&oldest_key);
}
entries.insert(
key,
TimedEntry {
value,
inserted_at: Instant::now(),
},
);
}
pub fn remove(&self, key: &K) -> Option<V> {
let Ok(mut entries) = self.entries.lock() else {
return None;
};
entries.remove(key).map(|entry| entry.value)
}
pub fn len(&self) -> usize {
self.entries
.lock()
.map(|entries| entries.len())
.unwrap_or(0)
}
}
impl<K, V> ExpiringMap<K, V>
where
K: Eq + Hash + Clone,
V: Clone,
{
pub fn get_fresh(&self, key: &K, ttl: Duration) -> Option<V> {
let Ok(mut entries) = self.entries.lock() else {
return None;
};
let Some(entry) = entries.get(key).cloned() else {
return None;
};
if entry.inserted_at.elapsed() > ttl {
entries.remove(key);
return None;
}
Some(entry.value)
}
pub fn contains_fresh(&self, key: &K, ttl: Duration) -> bool {
self.get_fresh(key, ttl).is_some()
}
}
fn prune_expired<K, V>(entries: &mut HashMap<K, TimedEntry<V>>, ttl: Duration)
where
K: Eq + Hash,
{
if ttl.is_zero() {
entries.clear();
return;
}
entries.retain(|_, entry| entry.inserted_at.elapsed() <= ttl);
}
#[cfg(test)]
mod tests {
use std::thread::sleep;
use super::ExpiringMap;
#[test]
fn evicts_expired_entries_on_read() {
let cache = ExpiringMap::new();
cache.insert(
"hello".to_string(),
42_u32,
std::time::Duration::from_millis(10),
16,
);
sleep(std::time::Duration::from_millis(20));
assert_eq!(
cache.get_fresh(&"hello".to_string(), std::time::Duration::from_millis(10)),
None
);
assert_eq!(cache.len(), 0);
}
#[test]
fn evicts_oldest_entry_when_capacity_is_hit() {
let cache = ExpiringMap::new();
cache.insert(
"one".to_string(),
1_u32,
std::time::Duration::from_secs(60),
2,
);
sleep(std::time::Duration::from_millis(2));
cache.insert(
"two".to_string(),
2_u32,
std::time::Duration::from_secs(60),
2,
);
sleep(std::time::Duration::from_millis(2));
cache.insert(
"three".to_string(),
3_u32,
std::time::Duration::from_secs(60),
2,
);
assert_eq!(
cache.get_fresh(&"one".to_string(), std::time::Duration::from_secs(60)),
None
);
assert_eq!(
cache.get_fresh(&"two".to_string(), std::time::Duration::from_secs(60)),
Some(2)
);
assert_eq!(
cache.get_fresh(&"three".to_string(), std::time::Duration::from_secs(60)),
Some(3)
);
}
}

20
crates/aether-data/Cargo.toml Normal file
View File

@@ -0,0 +1,20 @@
[package]
name = "aether-data"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
description = "Shared data access contracts and config for Aether Rust services"
[dependencies]
aether-cache.workspace = true
async-trait.workspace = true
futures-util.workspace = true
redis.workspace = true
serde.workspace = true
serde_json.workspace = true
sqlx.workspace = true
thiserror.workspace = true
tokio.workspace = true
url.workspace = true
uuid.workspace = true

66
crates/aether-data/src/backends/leases.rs Normal file
View File

@@ -0,0 +1,66 @@
use std::fmt;
use super::PostgresBackend;
use crate::postgres::{PostgresLeaseRunner, PostgresLeaseRunnerConfig};
use crate::DataLayerError;
#[derive(Clone, Default)]
pub struct DataLeaseBackends {
postgres: Option<PostgresLeaseRunner>,
}
impl fmt::Debug for DataLeaseBackends {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("DataLeaseBackends")
.field("has_postgres", &self.postgres.is_some())
.finish()
}
}
impl DataLeaseBackends {
pub(crate) fn from_postgres(
postgres: Option<&PostgresBackend>,
) -> Result<Self, DataLayerError> {
Ok(Self {
postgres: postgres
.map(|backend| backend.lease_runner(PostgresLeaseRunnerConfig::default()))
.transpose()?,
})
}
pub fn postgres(&self) -> Option<PostgresLeaseRunner> {
self.postgres.clone()
}
pub fn has_any(&self) -> bool {
self.postgres.is_some()
}
}
#[cfg(test)]
mod tests {
use super::DataLeaseBackends;
use crate::backends::PostgresBackend;
use crate::postgres::PostgresPoolConfig;
#[tokio::test]
async fn builds_postgres_lease_runner_from_backend() {
let backend = PostgresBackend::from_config(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("postgres backend should build");
let leases =
DataLeaseBackends::from_postgres(Some(&backend)).expect("lease backends should build");
assert!(leases.has_any());
assert!(leases.postgres().is_some());
}
}

58
crates/aether-data/src/backends/locks.rs Normal file
View File

@@ -0,0 +1,58 @@
use std::fmt;
use super::RedisBackend;
use crate::redis::{RedisLockRunner, RedisLockRunnerConfig};
use crate::DataLayerError;
#[derive(Clone, Default)]
pub struct DataLockBackends {
redis: Option<RedisLockRunner>,
}
impl fmt::Debug for DataLockBackends {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("DataLockBackends")
.field("has_redis", &self.redis.is_some())
.finish()
}
}
impl DataLockBackends {
pub(crate) fn from_redis(redis: Option<&RedisBackend>) -> Result<Self, DataLayerError> {
Ok(Self {
redis: redis
.map(|backend| backend.lock_runner(RedisLockRunnerConfig::default()))
.transpose()?,
})
}
pub fn redis(&self) -> Option<RedisLockRunner> {
self.redis.clone()
}
pub fn has_any(&self) -> bool {
self.redis.is_some()
}
}
#[cfg(test)]
mod tests {
use super::DataLockBackends;
use crate::backends::RedisBackend;
use crate::redis::RedisClientConfig;
#[test]
fn builds_redis_lock_runner_from_backend() {
let backend = RedisBackend::from_config(RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
})
.expect("redis backend should build");
let locks =
DataLockBackends::from_redis(Some(&backend)).expect("lock backends should build");
assert!(locks.has_any());
assert!(locks.redis().is_some());
}
}

195
crates/aether-data/src/backends/mod.rs Normal file
View File

@@ -0,0 +1,195 @@
mod leases;
mod locks;
mod postgres;
mod read;
mod redis;
mod transactions;
mod workers;
mod write;
pub use leases::DataLeaseBackends;
pub use locks::DataLockBackends;
pub use postgres::PostgresBackend;
pub use read::DataReadRepositories;
pub use redis::RedisBackend;
pub use transactions::DataTransactionBackends;
pub use workers::DataWorkerBackends;
pub use write::DataWriteRepositories;
use crate::{DataLayerConfig, DataLayerError};
#[derive(Debug, Clone, Default)]
pub struct DataBackends {
config: DataLayerConfig,
postgres: Option<PostgresBackend>,
redis: Option<RedisBackend>,
leases: DataLeaseBackends,
locks: DataLockBackends,
read: DataReadRepositories,
transactions: DataTransactionBackends,
workers: DataWorkerBackends,
write: DataWriteRepositories,
}
impl DataBackends {
pub fn from_config(config: DataLayerConfig) -> Result<Self, DataLayerError> {
config.validate()?;
let postgres = config
.postgres
.clone()
.map(PostgresBackend::from_config)
.transpose()?;
let redis = config
.redis
.clone()
.map(RedisBackend::from_config)
.transpose()?;
let leases = DataLeaseBackends::from_postgres(postgres.as_ref())?;
let locks = DataLockBackends::from_redis(redis.as_ref())?;
let read = DataReadRepositories::from_postgres(postgres.as_ref());
let transactions = DataTransactionBackends::from_postgres(postgres.as_ref());
let workers = DataWorkerBackends::from_redis(redis.as_ref())?;
let write = DataWriteRepositories::from_postgres(postgres.as_ref());
Ok(Self {
config,
postgres,
redis,
leases,
locks,
read,
transactions,
workers,
write,
})
}
pub fn config(&self) -> &DataLayerConfig {
&self.config
}
pub fn postgres(&self) -> Option<&PostgresBackend> {
self.postgres.as_ref()
}
pub fn redis(&self) -> Option<&RedisBackend> {
self.redis.as_ref()
}
pub fn read(&self) -> &DataReadRepositories {
&self.read
}
pub fn leases(&self) -> &DataLeaseBackends {
&self.leases
}
pub fn locks(&self) -> &DataLockBackends {
&self.locks
}
pub fn transactions(&self) -> &DataTransactionBackends {
&self.transactions
}
pub fn workers(&self) -> &DataWorkerBackends {
&self.workers
}
pub fn write(&self) -> &DataWriteRepositories {
&self.write
}
pub fn has_runtime_backends(&self) -> bool {
self.postgres.is_some()
|| self.redis.is_some()
|| self.leases.has_any()
|| self.locks.has_any()
|| self.read.has_any()
|| self.transactions.has_any()
|| self.workers.has_any()
|| self.write.has_any()
}
}
#[cfg(test)]
mod tests {
use super::DataBackends;
use crate::{postgres::PostgresPoolConfig, DataLayerConfig};
#[test]
fn builds_empty_backends_from_default_config() {
let backends = DataBackends::from_config(DataLayerConfig::default())
.expect("empty config should be accepted");
assert!(!backends.has_runtime_backends());
assert!(backends.postgres().is_none());
assert!(backends.redis().is_none());
assert!(backends.leases().postgres().is_none());
assert!(backends.locks().redis().is_none());
assert!(backends.read().auth_api_keys().is_none());
assert!(backends.read().request_candidates().is_none());
assert!(backends.read().provider_catalog().is_none());
assert!(backends.read().usage().is_none());
assert!(backends.read().video_tasks().is_none());
assert!(backends.read().shadow_results().is_none());
assert!(backends.transactions().postgres().is_none());
assert!(backends.workers().redis().is_none());
assert!(backends.write().shadow_results().is_none());
}
#[tokio::test]
async fn builds_postgres_backend_from_config() {
let backends = DataBackends::from_config(DataLayerConfig {
postgres: Some(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
}),
redis: None,
})
.expect("postgres backend should build");
assert!(backends.has_runtime_backends());
assert!(backends.postgres().is_some());
assert!(backends.leases().postgres().is_some());
assert!(backends.read().auth_api_keys().is_some());
assert!(backends.read().request_candidates().is_some());
assert!(backends.read().provider_catalog().is_some());
assert!(backends.read().usage().is_some());
assert!(backends.read().video_tasks().is_some());
assert!(backends.read().shadow_results().is_some());
assert!(backends.transactions().postgres().is_some());
assert!(backends.write().shadow_results().is_some());
assert!(backends.config().postgres.is_some());
}
#[test]
fn builds_redis_backend_from_config() {
let backends = DataBackends::from_config(DataLayerConfig {
postgres: None,
redis: Some(crate::redis::RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
}),
})
.expect("redis backend should build");
assert!(backends.has_runtime_backends());
assert!(backends.postgres().is_none());
assert!(backends.redis().is_some());
assert!(backends.leases().postgres().is_none());
assert!(backends.locks().redis().is_some());
assert!(backends.workers().redis().is_some());
assert!(backends.read().auth_api_keys().is_none());
assert!(backends.transactions().postgres().is_none());
assert!(backends.write().shadow_results().is_none());
assert!(backends.config().redis.is_some());
}
}

123
crates/aether-data/src/backends/postgres.rs Normal file
View File

@@ -0,0 +1,123 @@
use std::sync::Arc;
use crate::postgres::{
PostgresLeaseRunner, PostgresLeaseRunnerConfig, PostgresPool, PostgresPoolConfig,
PostgresPoolFactory, PostgresTransactionRunner,
};
use crate::repository::auth::{AuthApiKeyReadRepository, SqlxAuthApiKeySnapshotReadRepository};
use crate::repository::candidates::{
RequestCandidateReadRepository, SqlxRequestCandidateReadRepository,
};
use crate::repository::provider_catalog::{
ProviderCatalogReadRepository, SqlxProviderCatalogReadRepository,
};
use crate::repository::shadow_results::{
ShadowResultReadRepository, ShadowResultWriteRepository, SqlxShadowResultRepository,
};
use crate::repository::usage::{SqlxUsageReadRepository, UsageReadRepository};
use crate::repository::video_tasks::{SqlxVideoTaskReadRepository, VideoTaskReadRepository};
use crate::DataLayerError;
#[derive(Debug, Clone)]
pub struct PostgresBackend {
config: PostgresPoolConfig,
pool: PostgresPool,
}
impl PostgresBackend {
pub fn from_config(config: PostgresPoolConfig) -> Result<Self, DataLayerError> {
let factory = PostgresPoolFactory::new(config.clone())?;
let pool = factory.connect_lazy()?;
Ok(Self { config, pool })
}
pub fn config(&self) -> &PostgresPoolConfig {
&self.config
}
pub fn pool(&self) -> &PostgresPool {
&self.pool
}
pub fn pool_clone(&self) -> PostgresPool {
self.pool.clone()
}
pub fn auth_api_key_read_repository(&self) -> Arc<dyn AuthApiKeyReadRepository> {
Arc::new(SqlxAuthApiKeySnapshotReadRepository::new(self.pool_clone()))
}
pub fn request_candidate_read_repository(&self) -> Arc<dyn RequestCandidateReadRepository> {
Arc::new(SqlxRequestCandidateReadRepository::new(self.pool_clone()))
}
pub fn provider_catalog_read_repository(&self) -> Arc<dyn ProviderCatalogReadRepository> {
Arc::new(SqlxProviderCatalogReadRepository::new(self.pool_clone()))
}
pub fn usage_read_repository(&self) -> Arc<dyn UsageReadRepository> {
Arc::new(SqlxUsageReadRepository::new(self.pool_clone()))
}
pub fn video_task_read_repository(&self) -> Arc<dyn VideoTaskReadRepository> {
Arc::new(SqlxVideoTaskReadRepository::new(self.pool_clone()))
}
pub fn transaction_runner(&self) -> PostgresTransactionRunner {
PostgresTransactionRunner::new(self.pool_clone())
}
pub fn lease_runner(
&self,
config: PostgresLeaseRunnerConfig,
) -> Result<PostgresLeaseRunner, DataLayerError> {
PostgresLeaseRunner::new(self.transaction_runner(), config)
}
pub fn shadow_result_write_repository(&self) -> Arc<dyn ShadowResultWriteRepository> {
Arc::new(SqlxShadowResultRepository::new(self.pool_clone()))
}
pub fn shadow_result_read_repository(&self) -> Arc<dyn ShadowResultReadRepository> {
Arc::new(SqlxShadowResultRepository::new(self.pool_clone()))
}
}
#[cfg(test)]
mod tests {
use super::PostgresBackend;
use crate::postgres::{PostgresLeaseRunnerConfig, PostgresPoolConfig};
#[tokio::test]
async fn backend_retains_config_and_pool() {
let config = PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
};
let backend =
PostgresBackend::from_config(config.clone()).expect("backend should build lazily");
assert_eq!(backend.config(), &config);
let _pool = backend.pool();
let _pool_clone = backend.pool_clone();
let _auth_api_key_reader = backend.auth_api_key_read_repository();
let _request_candidate_reader = backend.request_candidate_read_repository();
let _provider_catalog_reader = backend.provider_catalog_read_repository();
let _usage_reader = backend.usage_read_repository();
let _video_task_reader = backend.video_task_read_repository();
let _transaction_runner = backend.transaction_runner();
let _lease_runner = backend
.lease_runner(PostgresLeaseRunnerConfig::default())
.expect("lease runner should build");
let _shadow_result_reader = backend.shadow_result_read_repository();
let _shadow_result_writer = backend.shadow_result_write_repository();
}
}

111
crates/aether-data/src/backends/read.rs Normal file
View File

@@ -0,0 +1,111 @@
use std::fmt;
use std::sync::Arc;
use super::PostgresBackend;
use crate::repository::auth::AuthApiKeyReadRepository;
use crate::repository::candidates::RequestCandidateReadRepository;
use crate::repository::provider_catalog::ProviderCatalogReadRepository;
use crate::repository::shadow_results::ShadowResultReadRepository;
use crate::repository::usage::UsageReadRepository;
use crate::repository::video_tasks::VideoTaskReadRepository;
#[derive(Clone, Default)]
pub struct DataReadRepositories {
auth_api_keys: Option<Arc<dyn AuthApiKeyReadRepository>>,
request_candidates: Option<Arc<dyn RequestCandidateReadRepository>>,
provider_catalog: Option<Arc<dyn ProviderCatalogReadRepository>>,
usage: Option<Arc<dyn UsageReadRepository>>,
video_tasks: Option<Arc<dyn VideoTaskReadRepository>>,
shadow_results: Option<Arc<dyn ShadowResultReadRepository>>,
}
impl fmt::Debug for DataReadRepositories {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("DataReadRepositories")
.field("has_auth_api_keys", &self.auth_api_keys.is_some())
.field("has_request_candidates", &self.request_candidates.is_some())
.field("has_provider_catalog", &self.provider_catalog.is_some())
.field("has_usage", &self.usage.is_some())
.field("has_video_tasks", &self.video_tasks.is_some())
.field("has_shadow_results", &self.shadow_results.is_some())
.finish()
}
}
impl DataReadRepositories {
pub(crate) fn from_postgres(postgres: Option<&PostgresBackend>) -> Self {
Self {
auth_api_keys: postgres.map(PostgresBackend::auth_api_key_read_repository),
request_candidates: postgres.map(PostgresBackend::request_candidate_read_repository),
provider_catalog: postgres.map(PostgresBackend::provider_catalog_read_repository),
usage: postgres.map(PostgresBackend::usage_read_repository),
video_tasks: postgres.map(PostgresBackend::video_task_read_repository),
shadow_results: postgres.map(PostgresBackend::shadow_result_read_repository),
}
}
pub fn auth_api_keys(&self) -> Option<Arc<dyn AuthApiKeyReadRepository>> {
self.auth_api_keys.clone()
}
pub fn request_candidates(&self) -> Option<Arc<dyn RequestCandidateReadRepository>> {
self.request_candidates.clone()
}
pub fn provider_catalog(&self) -> Option<Arc<dyn ProviderCatalogReadRepository>> {
self.provider_catalog.clone()
}
pub fn usage(&self) -> Option<Arc<dyn UsageReadRepository>> {
self.usage.clone()
}
pub fn video_tasks(&self) -> Option<Arc<dyn VideoTaskReadRepository>> {
self.video_tasks.clone()
}
pub fn shadow_results(&self) -> Option<Arc<dyn ShadowResultReadRepository>> {
self.shadow_results.clone()
}
pub fn has_any(&self) -> bool {
self.auth_api_keys.is_some()
|| self.request_candidates.is_some()
|| self.provider_catalog.is_some()
|| self.usage.is_some()
|| self.video_tasks.is_some()
|| self.shadow_results.is_some()
}
}
#[cfg(test)]
mod tests {
use super::DataReadRepositories;
use crate::backends::PostgresBackend;
use crate::postgres::PostgresPoolConfig;
#[tokio::test]
async fn builds_read_repositories_from_postgres_backend() {
let backend = PostgresBackend::from_config(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("postgres backend should build");
let read = DataReadRepositories::from_postgres(Some(&backend));
assert!(read.has_any());
assert!(read.auth_api_keys().is_some());
assert!(read.request_candidates().is_some());
assert!(read.provider_catalog().is_some());
assert!(read.usage().is_some());
assert!(read.video_tasks().is_some());
assert!(read.shadow_results().is_some());
}
}

76
crates/aether-data/src/backends/redis.rs Normal file
View File

@@ -0,0 +1,76 @@
use crate::redis::{
RedisClient, RedisClientConfig, RedisClientFactory, RedisKeyspace, RedisLockRunner,
RedisLockRunnerConfig, RedisStreamRunner, RedisStreamRunnerConfig,
};
use crate::DataLayerError;
#[derive(Debug, Clone)]
pub struct RedisBackend {
config: RedisClientConfig,
client: RedisClient,
}
impl RedisBackend {
pub fn from_config(config: RedisClientConfig) -> Result<Self, DataLayerError> {
let factory = RedisClientFactory::new(config.clone())?;
let client = factory.connect_lazy()?;
Ok(Self { config, client })
}
pub fn config(&self) -> &RedisClientConfig {
&self.config
}
pub fn client(&self) -> &RedisClient {
&self.client
}
pub fn client_clone(&self) -> RedisClient {
self.client.clone()
}
pub fn keyspace(&self) -> RedisKeyspace {
self.config.keyspace()
}
pub fn lock_runner(
&self,
config: RedisLockRunnerConfig,
) -> Result<RedisLockRunner, DataLayerError> {
RedisLockRunner::new(self.client_clone(), self.keyspace(), config)
}
pub fn stream_runner(
&self,
config: RedisStreamRunnerConfig,
) -> Result<RedisStreamRunner, DataLayerError> {
RedisStreamRunner::new(self.client_clone(), self.keyspace(), config)
}
}
#[cfg(test)]
mod tests {
use super::RedisBackend;
use crate::redis::{RedisClientConfig, RedisLockRunnerConfig, RedisStreamRunnerConfig};
#[test]
fn backend_retains_config_client_and_shared_runners() {
let config = RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
};
let backend = RedisBackend::from_config(config.clone()).expect("backend should build");
assert_eq!(backend.config(), &config);
assert_eq!(backend.keyspace().key("audit"), "aether:audit");
let _client_ref = backend.client();
let _client_clone = backend.client_clone();
let _lock_runner = backend
.lock_runner(RedisLockRunnerConfig::default())
.expect("lock runner should build");
let _stream_runner = backend
.stream_runner(RedisStreamRunnerConfig::default())
.expect("stream runner should build");
}
}

60
crates/aether-data/src/backends/transactions.rs Normal file
View File

@@ -0,0 +1,60 @@
use std::fmt;
use super::PostgresBackend;
use crate::postgres::PostgresTransactionRunner;
#[derive(Clone, Default)]
pub struct DataTransactionBackends {
postgres: Option<PostgresTransactionRunner>,
}
impl fmt::Debug for DataTransactionBackends {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("DataTransactionBackends")
.field("has_postgres", &self.postgres.is_some())
.finish()
}
}
impl DataTransactionBackends {
pub(crate) fn from_postgres(postgres: Option<&PostgresBackend>) -> Self {
Self {
postgres: postgres.map(PostgresBackend::transaction_runner),
}
}
pub fn postgres(&self) -> Option<PostgresTransactionRunner> {
self.postgres.clone()
}
pub fn has_any(&self) -> bool {
self.postgres.is_some()
}
}
#[cfg(test)]
mod tests {
use super::DataTransactionBackends;
use crate::backends::PostgresBackend;
use crate::postgres::PostgresPoolConfig;
#[tokio::test]
async fn builds_postgres_transaction_runner_from_backend() {
let backend = PostgresBackend::from_config(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("postgres backend should build");
let transactions = DataTransactionBackends::from_postgres(Some(&backend));
assert!(transactions.has_any());
assert!(transactions.postgres().is_some());
}
}

58
crates/aether-data/src/backends/workers.rs Normal file
View File

@@ -0,0 +1,58 @@
use std::fmt;
use super::RedisBackend;
use crate::redis::{RedisStreamRunner, RedisStreamRunnerConfig};
use crate::DataLayerError;
#[derive(Clone, Default)]
pub struct DataWorkerBackends {
redis: Option<RedisStreamRunner>,
}
impl fmt::Debug for DataWorkerBackends {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("DataWorkerBackends")
.field("has_redis", &self.redis.is_some())
.finish()
}
}
impl DataWorkerBackends {
pub(crate) fn from_redis(redis: Option<&RedisBackend>) -> Result<Self, DataLayerError> {
Ok(Self {
redis: redis
.map(|backend| backend.stream_runner(RedisStreamRunnerConfig::default()))
.transpose()?,
})
}
pub fn redis(&self) -> Option<RedisStreamRunner> {
self.redis.clone()
}
pub fn has_any(&self) -> bool {
self.redis.is_some()
}
}
#[cfg(test)]
mod tests {
use super::DataWorkerBackends;
use crate::backends::RedisBackend;
use crate::redis::RedisClientConfig;
#[test]
fn builds_redis_stream_runner_from_backend() {
let backend = RedisBackend::from_config(RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
})
.expect("redis backend should build");
let workers =
DataWorkerBackends::from_redis(Some(&backend)).expect("worker backends should build");
assert!(workers.has_any());
assert!(workers.redis().is_some());
}
}

61
crates/aether-data/src/backends/write.rs Normal file
View File

@@ -0,0 +1,61 @@
use std::fmt;
use std::sync::Arc;
use super::PostgresBackend;
use crate::repository::shadow_results::ShadowResultWriteRepository;
#[derive(Clone, Default)]
pub struct DataWriteRepositories {
shadow_results: Option<Arc<dyn ShadowResultWriteRepository>>,
}
impl fmt::Debug for DataWriteRepositories {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("DataWriteRepositories")
.field("has_shadow_results", &self.shadow_results.is_some())
.finish()
}
}
impl DataWriteRepositories {
pub(crate) fn from_postgres(postgres: Option<&PostgresBackend>) -> Self {
Self {
shadow_results: postgres.map(PostgresBackend::shadow_result_write_repository),
}
}
pub fn shadow_results(&self) -> Option<Arc<dyn ShadowResultWriteRepository>> {
self.shadow_results.clone()
}
pub fn has_any(&self) -> bool {
self.shadow_results.is_some()
}
}
#[cfg(test)]
mod tests {
use super::DataWriteRepositories;
use crate::backends::PostgresBackend;
use crate::postgres::PostgresPoolConfig;
#[tokio::test]
async fn builds_shadow_result_writer_from_postgres_backend() {
let backend = PostgresBackend::from_config(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("postgres backend should build");
let write = DataWriteRepositories::from_postgres(Some(&backend));
assert!(write.has_any());
assert!(write.shadow_results().is_some());
}
}

74
crates/aether-data/src/config.rs Normal file
View File

@@ -0,0 +1,74 @@
use crate::postgres::PostgresPoolConfig;
use crate::redis::RedisClientConfig;
use crate::DataLayerError;
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
pub struct DataLayerConfig {
pub postgres: Option<PostgresPoolConfig>,
pub redis: Option<RedisClientConfig>,
}
impl DataLayerConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
if let Some(postgres) = &self.postgres {
postgres.validate()?;
}
if let Some(redis) = &self.redis {
redis.validate()?;
}
Ok(())
}
pub fn has_persistent_backends(&self) -> bool {
self.postgres.is_some() || self.redis.is_some()
}
}
#[cfg(test)]
mod tests {
use super::DataLayerConfig;
use crate::postgres::PostgresPoolConfig;
use crate::redis::RedisClientConfig;
#[test]
fn validates_nested_backend_configs() {
let config = DataLayerConfig {
postgres: Some(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 2,
max_connections: 8,
acquire_timeout_ms: 1_500,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
}),
redis: Some(RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
}),
};
assert!(config.validate().is_ok());
assert!(config.has_persistent_backends());
}
#[test]
fn rejects_invalid_nested_backend_configs() {
let config = DataLayerConfig {
postgres: Some(PostgresPoolConfig {
database_url: String::new(),
min_connections: 4,
max_connections: 2,
acquire_timeout_ms: 1_500,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
}),
redis: None,
};
assert!(config.validate().is_err());
}
}

20
crates/aether-data/src/error.rs Normal file
View File

@@ -0,0 +1,20 @@
#[derive(Debug, thiserror::Error)]
pub enum DataLayerError {
#[error("invalid configuration: {0}")]
InvalidConfiguration(String),
#[error("invalid input: {0}")]
InvalidInput(String),
#[error("postgres error: {0}")]
Postgres(#[from] sqlx::Error),
#[error("redis error: {0}")]
Redis(#[from] redis::RedisError),
#[error("operation timed out: {0}")]
TimedOut(String),
#[error("unexpected database value: {0}")]
UnexpectedValue(String),
}

14
crates/aether-data/src/lib.rs Normal file
View File

@@ -0,0 +1,14 @@
pub mod backends;
mod config;
mod error;
pub mod postgres;
pub mod redis;
pub mod repository;
pub use backends::{
DataBackends, DataLeaseBackends, DataLockBackends, DataReadRepositories,
DataTransactionBackends, DataWorkerBackends, DataWriteRepositories, PostgresBackend,
RedisBackend,
};
pub use config::DataLayerConfig;
pub use error::DataLayerError;

417
crates/aether-data/src/postgres/lease.rs Normal file
View File

@@ -0,0 +1,417 @@
use crate::postgres::{DatabaseRecordId, PostgresTransactionOptions, PostgresTransactionRunner};
use crate::DataLayerError;
use futures_util::FutureExt;
use sqlx::query_scalar;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PostgresLeaseClaimOptions {
pub batch_size: usize,
pub lease_ms: u64,
}
impl PostgresLeaseClaimOptions {
pub fn validate(&self) -> Result<(), DataLayerError> {
if self.batch_size == 0 {
return Err(DataLayerError::InvalidConfiguration(
"postgres lease batch_size must be positive".to_string(),
));
}
if self.lease_ms == 0 {
return Err(DataLayerError::InvalidConfiguration(
"postgres lease lease_ms must be positive".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PostgresLeaseClaimSpec {
pub table: &'static str,
pub id_column: &'static str,
pub lease_owner_column: &'static str,
pub lease_expires_at_column: &'static str,
pub eligibility_predicate_sql: &'static str,
pub order_by_sql: &'static str,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct PostgresLeaseRunnerConfig {
pub statement_timeout_ms: Option<u64>,
pub lock_timeout_ms: Option<u64>,
}
impl PostgresLeaseRunnerConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
if matches!(self.statement_timeout_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"postgres lease statement_timeout_ms must be positive".to_string(),
));
}
if matches!(self.lock_timeout_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"postgres lease lock_timeout_ms must be positive".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct PostgresLeaseRunner {
transaction_runner: PostgresTransactionRunner,
config: PostgresLeaseRunnerConfig,
}
impl PostgresLeaseRunner {
pub fn new(
transaction_runner: PostgresTransactionRunner,
config: PostgresLeaseRunnerConfig,
) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self {
transaction_runner,
config,
})
}
pub fn config(&self) -> PostgresLeaseRunnerConfig {
self.config
}
pub fn transaction_runner(&self) -> &PostgresTransactionRunner {
&self.transaction_runner
}
pub async fn claim_ids(
&self,
spec: &PostgresLeaseClaimSpec,
options: PostgresLeaseClaimOptions,
owner: &str,
) -> Result<Vec<DatabaseRecordId>, DataLayerError> {
validate_lease_owner(owner)?;
let sql = build_postgres_lease_claim_sql(spec, options)?;
let owner = owner.trim().to_string();
let lease_ms = i64::try_from(options.lease_ms).map_err(|_| {
DataLayerError::InvalidInput("postgres lease lease_ms exceeds i64 range".to_string())
})?;
let tx_options = PostgresTransactionOptions {
statement_timeout_ms: self.config.statement_timeout_ms,
lock_timeout_ms: self.config.lock_timeout_ms,
..PostgresTransactionOptions::read_write()
};
self.transaction_runner
.run(tx_options, |tx| {
async move {
let rows = query_scalar::<_, String>(&sql)
.bind(owner)
.bind(lease_ms)
.fetch_all(&mut **tx)
.await?;
Ok(rows.into_iter().map(DatabaseRecordId).collect())
}
.boxed()
})
.await
}
pub async fn release_ids(
&self,
spec: &PostgresLeaseClaimSpec,
ids: &[DatabaseRecordId],
owner: &str,
) -> Result<Vec<DatabaseRecordId>, DataLayerError> {
validate_lease_owner(owner)?;
if ids.is_empty() {
return Ok(Vec::new());
}
let sql = build_postgres_lease_release_sql(spec)?;
let owner = owner.trim().to_string();
let ids = ids.iter().map(|id| id.0.clone()).collect::<Vec<_>>();
let tx_options = PostgresTransactionOptions {
statement_timeout_ms: self.config.statement_timeout_ms,
lock_timeout_ms: self.config.lock_timeout_ms,
..PostgresTransactionOptions::read_write()
};
self.transaction_runner
.run(tx_options, |tx| {
async move {
let rows = query_scalar::<_, String>(&sql)
.bind(ids)
.bind(owner)
.fetch_all(&mut **tx)
.await?;
Ok(rows.into_iter().map(DatabaseRecordId).collect())
}
.boxed()
})
.await
}
pub async fn renew_ids(
&self,
spec: &PostgresLeaseClaimSpec,
ids: &[DatabaseRecordId],
owner: &str,
lease_ms: u64,
) -> Result<Vec<DatabaseRecordId>, DataLayerError> {
validate_lease_owner(owner)?;
if lease_ms == 0 {
return Err(DataLayerError::InvalidInput(
"postgres lease lease_ms must be positive".to_string(),
));
}
if ids.is_empty() {
return Ok(Vec::new());
}
let sql = build_postgres_lease_renew_sql(spec)?;
let owner = owner.trim().to_string();
let lease_ms = i64::try_from(lease_ms).map_err(|_| {
DataLayerError::InvalidInput("postgres lease lease_ms exceeds i64 range".to_string())
})?;
let ids = ids.iter().map(|id| id.0.clone()).collect::<Vec<_>>();
let tx_options = PostgresTransactionOptions {
statement_timeout_ms: self.config.statement_timeout_ms,
lock_timeout_ms: self.config.lock_timeout_ms,
..PostgresTransactionOptions::read_write()
};
self.transaction_runner
.run(tx_options, |tx| {
async move {
let rows = query_scalar::<_, String>(&sql)
.bind(ids)
.bind(owner)
.bind(lease_ms)
.fetch_all(&mut **tx)
.await?;
Ok(rows.into_iter().map(DatabaseRecordId).collect())
}
.boxed()
})
.await
}
}
pub fn build_postgres_lease_claim_sql(
spec: &PostgresLeaseClaimSpec,
options: PostgresLeaseClaimOptions,
) -> Result<String, DataLayerError> {
options.validate()?;
validate_lease_spec(spec)?;
Ok(format!(
"WITH claimable AS (\
SELECT {id_column} \
FROM {table} \
WHERE ({eligibility_predicate_sql}) \
AND ({lease_expires_at_column} IS NULL OR {lease_expires_at_column} <= NOW()) \
ORDER BY {order_by_sql} \
FOR UPDATE SKIP LOCKED \
LIMIT {batch_size}\
) \
UPDATE {table} AS target \
SET {lease_owner_column} = $1, \
{lease_expires_at_column} = NOW() + ($2::bigint * INTERVAL '1 millisecond') \
FROM claimable \
WHERE target.{id_column} = claimable.{id_column} \
RETURNING target.{id_column}",
id_column = spec.id_column,
table = spec.table,
eligibility_predicate_sql = spec.eligibility_predicate_sql,
lease_expires_at_column = spec.lease_expires_at_column,
order_by_sql = spec.order_by_sql,
batch_size = options.batch_size,
lease_owner_column = spec.lease_owner_column,
))
}
pub fn build_postgres_lease_release_sql(
spec: &PostgresLeaseClaimSpec,
) -> Result<String, DataLayerError> {
validate_lease_spec(spec)?;
Ok(format!(
"UPDATE {table} \
SET {lease_owner_column} = NULL, \
{lease_expires_at_column} = NULL \
WHERE {id_column} = ANY($1) \
AND {lease_owner_column} = $2 \
RETURNING {id_column}",
table = spec.table,
id_column = spec.id_column,
lease_owner_column = spec.lease_owner_column,
lease_expires_at_column = spec.lease_expires_at_column,
))
}
pub fn build_postgres_lease_renew_sql(
spec: &PostgresLeaseClaimSpec,
) -> Result<String, DataLayerError> {
validate_lease_spec(spec)?;
Ok(format!(
"UPDATE {table} \
SET {lease_expires_at_column} = NOW() + ($3::bigint * INTERVAL '1 millisecond') \
WHERE {id_column} = ANY($1) \
AND {lease_owner_column} = $2 \
RETURNING {id_column}",
table = spec.table,
id_column = spec.id_column,
lease_owner_column = spec.lease_owner_column,
lease_expires_at_column = spec.lease_expires_at_column,
))
}
fn validate_lease_spec(spec: &PostgresLeaseClaimSpec) -> Result<(), DataLayerError> {
for (field, value) in [
("table", spec.table),
("id_column", spec.id_column),
("lease_owner_column", spec.lease_owner_column),
("lease_expires_at_column", spec.lease_expires_at_column),
("eligibility_predicate_sql", spec.eligibility_predicate_sql),
("order_by_sql", spec.order_by_sql),
] {
if value.trim().is_empty() {
return Err(DataLayerError::InvalidConfiguration(format!(
"postgres lease {field} cannot be empty"
)));
}
}
Ok(())
}
fn validate_lease_owner(owner: &str) -> Result<(), DataLayerError> {
if owner.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"postgres lease owner cannot be empty".to_string(),
));
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::{
build_postgres_lease_claim_sql, build_postgres_lease_release_sql,
build_postgres_lease_renew_sql, PostgresLeaseClaimOptions, PostgresLeaseClaimSpec,
PostgresLeaseRunner, PostgresLeaseRunnerConfig,
};
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory, PostgresTransactionRunner};
fn sample_spec() -> PostgresLeaseClaimSpec {
PostgresLeaseClaimSpec {
table: "video_tasks",
id_column: "id",
lease_owner_column: "lease_owner",
lease_expires_at_column: "lease_expires_at",
eligibility_predicate_sql: "status IN ('submitted', 'processing')",
order_by_sql: "updated_at ASC",
}
}
#[test]
fn builds_skip_locked_claim_sql() {
let sql = build_postgres_lease_claim_sql(
&sample_spec(),
PostgresLeaseClaimOptions {
batch_size: 16,
lease_ms: 15_000,
},
)
.expect("claim SQL should build");
assert!(sql.contains("FOR UPDATE SKIP LOCKED"));
assert!(sql.contains("LIMIT 16"));
assert!(sql.contains("lease_owner = $1"));
assert!(sql.contains("NOW() + ($2::bigint * INTERVAL '1 millisecond')"));
}
#[test]
fn builds_release_sql() {
let sql =
build_postgres_lease_release_sql(&sample_spec()).expect("release SQL should build");
assert!(sql.contains("id = ANY($1)"));
assert!(sql.contains("lease_owner = $2"));
assert!(sql.contains("lease_expires_at = NULL"));
}
#[test]
fn builds_renew_sql() {
let sql = build_postgres_lease_renew_sql(&sample_spec()).expect("renew SQL should build");
assert!(sql.contains("id = ANY($1)"));
assert!(sql.contains("lease_owner = $2"));
assert!(sql.contains("NOW() + ($3::bigint * INTERVAL '1 millisecond')"));
}
#[tokio::test]
async fn lease_runner_reuses_transaction_runner() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let pool = factory.connect_lazy().expect("lazy pool should build");
let transaction_runner = PostgresTransactionRunner::new(pool);
let lease_runner = PostgresLeaseRunner::new(
transaction_runner.clone(),
PostgresLeaseRunnerConfig {
statement_timeout_ms: Some(2_000),
lock_timeout_ms: Some(500),
},
)
.expect("lease runner should build");
assert_eq!(
lease_runner.config(),
PostgresLeaseRunnerConfig {
statement_timeout_ms: Some(2_000),
lock_timeout_ms: Some(500),
}
);
let _runner_ref = lease_runner.transaction_runner();
}
#[tokio::test]
async fn release_and_renew_empty_ids_short_circuit() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let pool = factory.connect_lazy().expect("lazy pool should build");
let runner = PostgresLeaseRunner::new(
PostgresTransactionRunner::new(pool),
PostgresLeaseRunnerConfig::default(),
)
.expect("lease runner should build");
assert!(runner
.release_ids(&sample_spec(), &[], "worker-1")
.await
.expect("empty release should succeed")
.is_empty());
assert!(runner
.renew_ids(&sample_spec(), &[], "worker-1", 5_000)
.await
.expect("empty renew should succeed")
.is_empty());
}
}

15
crates/aether-data/src/postgres/mod.rs Normal file
View File

@@ -0,0 +1,15 @@
mod lease;
mod pool;
mod tx;
mod types;
pub use lease::{
build_postgres_lease_claim_sql, build_postgres_lease_release_sql,
build_postgres_lease_renew_sql, PostgresLeaseClaimOptions, PostgresLeaseClaimSpec,
PostgresLeaseRunner, PostgresLeaseRunnerConfig,
};
pub use pool::{PostgresPool, PostgresPoolConfig, PostgresPoolFactory};
pub use tx::{
PostgresTransaction, PostgresTransactionOptions, PostgresTransactionRunner, TransactionMode,
};
pub use types::DatabaseRecordId;

122
crates/aether-data/src/postgres/pool.rs Normal file
View File

@@ -0,0 +1,122 @@
use crate::DataLayerError;
use sqlx::postgres::{PgConnectOptions, PgPoolOptions, PgSslMode};
use sqlx::PgPool;
use std::str::FromStr;
use std::time::Duration;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
pub struct PostgresPoolConfig {
pub database_url: String,
pub min_connections: u32,
pub max_connections: u32,
pub acquire_timeout_ms: u64,
pub idle_timeout_ms: u64,
pub max_lifetime_ms: u64,
pub statement_cache_capacity: usize,
pub require_ssl: bool,
}
impl Default for PostgresPoolConfig {
fn default() -> Self {
Self {
database_url: String::new(),
min_connections: 1,
max_connections: 20,
acquire_timeout_ms: 5_000,
idle_timeout_ms: 60_000,
max_lifetime_ms: 30 * 60_000,
statement_cache_capacity: 100,
require_ssl: false,
}
}
}
impl PostgresPoolConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
if self.database_url.trim().is_empty() {
return Err(DataLayerError::InvalidConfiguration(
"postgres database_url cannot be empty".to_string(),
));
}
if self.min_connections > self.max_connections {
return Err(DataLayerError::InvalidConfiguration(
"postgres min_connections cannot exceed max_connections".to_string(),
));
}
if self.statement_cache_capacity == 0 {
return Err(DataLayerError::InvalidConfiguration(
"postgres statement_cache_capacity must be positive".to_string(),
));
}
Ok(())
}
pub fn connect_options(&self) -> Result<PgConnectOptions, DataLayerError> {
self.validate()?;
let ssl_mode = if self.require_ssl {
PgSslMode::Require
} else {
PgSslMode::Prefer
};
let options = PgConnectOptions::from_str(self.database_url.trim()).map_err(|err| {
DataLayerError::InvalidConfiguration(format!("invalid postgres database_url: {err}"))
})?;
Ok(options
.ssl_mode(ssl_mode)
.statement_cache_capacity(self.statement_cache_capacity))
}
}
pub type PostgresPool = PgPool;
#[derive(Debug, Clone)]
pub struct PostgresPoolFactory {
config: PostgresPoolConfig,
}
impl PostgresPoolFactory {
pub fn new(config: PostgresPoolConfig) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self { config })
}
pub fn config(&self) -> &PostgresPoolConfig {
&self.config
}
pub fn connect_lazy(&self) -> Result<PostgresPool, DataLayerError> {
let options = self.config.connect_options()?;
Ok(PgPoolOptions::new()
.min_connections(self.config.min_connections)
.max_connections(self.config.max_connections)
.acquire_timeout(Duration::from_millis(self.config.acquire_timeout_ms))
.idle_timeout(Duration::from_millis(self.config.idle_timeout_ms))
.max_lifetime(Duration::from_millis(self.config.max_lifetime_ms))
.connect_lazy_with(options))
}
}
#[cfg(test)]
mod tests {
use super::{PostgresPoolConfig, PostgresPoolFactory};
#[tokio::test]
async fn factory_builds_lazy_pool_from_valid_config() {
let config = PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
};
let factory = PostgresPoolFactory::new(config).expect("factory should build");
let _pool = factory.connect_lazy().expect("lazy pool should build");
}
}

202
crates/aether-data/src/postgres/tx.rs Normal file
View File

@@ -0,0 +1,202 @@
use futures_util::future::BoxFuture;
use sqlx::{Postgres, Transaction};
use crate::postgres::PostgresPool;
use crate::DataLayerError;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum TransactionMode {
ReadOnly,
#[default]
ReadWrite,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct PostgresTransactionOptions {
pub mode: TransactionMode,
pub statement_timeout_ms: Option<u64>,
pub lock_timeout_ms: Option<u64>,
}
impl PostgresTransactionOptions {
pub fn read_only() -> Self {
Self {
mode: TransactionMode::ReadOnly,
..Self::default()
}
}
pub fn read_write() -> Self {
Self {
mode: TransactionMode::ReadWrite,
..Self::default()
}
}
pub fn validate(&self) -> Result<(), DataLayerError> {
if matches!(self.statement_timeout_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"postgres statement_timeout_ms must be positive".to_string(),
));
}
if matches!(self.lock_timeout_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"postgres lock_timeout_ms must be positive".to_string(),
));
}
Ok(())
}
}
pub type PostgresTransaction = Transaction<'static, Postgres>;
#[derive(Debug, Clone)]
pub struct PostgresTransactionRunner {
pool: PostgresPool,
}
impl PostgresTransactionRunner {
pub fn new(pool: PostgresPool) -> Self {
Self { pool }
}
pub fn pool(&self) -> &PostgresPool {
&self.pool
}
pub async fn begin(
&self,
options: PostgresTransactionOptions,
) -> Result<PostgresTransaction, DataLayerError> {
options.validate()?;
let mut tx = self.pool.begin().await?;
for statement in build_transaction_setup_statements(options) {
sqlx::query(statement.as_str()).execute(&mut *tx).await?;
}
Ok(tx)
}
pub async fn run<T, F>(
&self,
options: PostgresTransactionOptions,
f: F,
) -> Result<T, DataLayerError>
where
F: for<'tx> FnOnce(
&'tx mut PostgresTransaction,
) -> BoxFuture<'tx, Result<T, DataLayerError>>,
{
let mut tx = self.begin(options).await?;
match f(&mut tx).await {
Ok(value) => {
tx.commit().await?;
Ok(value)
}
Err(err) => {
let _ = tx.rollback().await;
Err(err)
}
}
}
pub async fn run_read_only<T, F>(&self, f: F) -> Result<T, DataLayerError>
where
F: for<'tx> FnOnce(
&'tx mut PostgresTransaction,
) -> BoxFuture<'tx, Result<T, DataLayerError>>,
{
self.run(PostgresTransactionOptions::read_only(), f).await
}
pub async fn run_read_write<T, F>(&self, f: F) -> Result<T, DataLayerError>
where
F: for<'tx> FnOnce(
&'tx mut PostgresTransaction,
) -> BoxFuture<'tx, Result<T, DataLayerError>>,
{
self.run(PostgresTransactionOptions::read_write(), f).await
}
}
pub(crate) fn build_transaction_setup_statements(
options: PostgresTransactionOptions,
) -> Vec<String> {
let mut statements = Vec::new();
if options.mode == TransactionMode::ReadOnly {
statements.push("SET TRANSACTION READ ONLY".to_string());
}
if let Some(statement_timeout_ms) = options.statement_timeout_ms {
statements.push(format!(
"SET LOCAL statement_timeout = {}",
statement_timeout_ms
));
}
if let Some(lock_timeout_ms) = options.lock_timeout_ms {
statements.push(format!("SET LOCAL lock_timeout = {}", lock_timeout_ms));
}
statements
}
#[cfg(test)]
mod tests {
use super::{
build_transaction_setup_statements, PostgresTransactionOptions, PostgresTransactionRunner,
TransactionMode,
};
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
#[test]
fn validates_transaction_options() {
assert!(PostgresTransactionOptions {
statement_timeout_ms: Some(0),
..PostgresTransactionOptions::default()
}
.validate()
.is_err());
assert!(PostgresTransactionOptions {
lock_timeout_ms: Some(0),
..PostgresTransactionOptions::default()
}
.validate()
.is_err());
}
#[test]
fn builds_expected_setup_statements() {
let statements = build_transaction_setup_statements(PostgresTransactionOptions {
mode: TransactionMode::ReadOnly,
statement_timeout_ms: Some(1_500),
lock_timeout_ms: Some(750),
});
assert_eq!(
statements,
vec![
"SET TRANSACTION READ ONLY".to_string(),
"SET LOCAL statement_timeout = 1500".to_string(),
"SET LOCAL lock_timeout = 750".to_string(),
]
);
}
#[tokio::test]
async fn runner_reuses_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let pool = factory.connect_lazy().expect("lazy pool should build");
let runner = PostgresTransactionRunner::new(pool.clone());
let _pool_ref = runner.pool();
}
}

2
crates/aether-data/src/postgres/types.rs Normal file
View File

@@ -0,0 +1,2 @@
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct DatabaseRecordId(pub String);

68
crates/aether-data/src/redis/client.rs Normal file
View File

@@ -0,0 +1,68 @@
use crate::redis::RedisKeyspace;
use crate::DataLayerError;
pub type RedisClient = redis::Client;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
pub struct RedisClientConfig {
pub url: String,
pub key_prefix: Option<String>,
}
impl RedisClientConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
let raw = self.url.trim();
if raw.is_empty() {
return Err(DataLayerError::InvalidConfiguration(
"redis url cannot be empty".to_string(),
));
}
url::Url::parse(raw).map_err(|err| {
DataLayerError::InvalidConfiguration(format!("invalid redis url: {err}"))
})?;
Ok(())
}
pub fn keyspace(&self) -> RedisKeyspace {
RedisKeyspace::new(self.key_prefix.as_deref())
}
}
#[derive(Debug, Clone)]
pub struct RedisClientFactory {
config: RedisClientConfig,
}
impl RedisClientFactory {
pub fn new(config: RedisClientConfig) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self { config })
}
pub fn config(&self) -> &RedisClientConfig {
&self.config
}
pub fn connect_lazy(&self) -> Result<RedisClient, DataLayerError> {
Ok(RedisClient::open(self.config.url.clone())?)
}
}
#[cfg(test)]
mod tests {
use super::{RedisClientConfig, RedisClientFactory};
#[test]
fn factory_builds_lazy_client_from_valid_config() {
let config = RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
};
let factory = RedisClientFactory::new(config.clone()).expect("factory should build");
assert_eq!(factory.config(), &config);
let _client = factory
.connect_lazy()
.expect("lazy redis client should build");
}
}

307
crates/aether-data/src/redis/lock.rs Normal file
View File

@@ -0,0 +1,307 @@
use std::future::Future;
use std::time::Duration;
use crate::redis::{RedisClient, RedisKeyspace};
use crate::DataLayerError;
use uuid::Uuid;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RedisLockKey(pub String);
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RedisLockLease {
pub key: RedisLockKey,
pub owner: String,
pub token: String,
pub ttl_ms: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RedisLockRunnerConfig {
pub command_timeout_ms: Option<u64>,
pub default_ttl_ms: u64,
}
impl Default for RedisLockRunnerConfig {
fn default() -> Self {
Self {
command_timeout_ms: Some(1_000),
default_ttl_ms: 15_000,
}
}
}
impl RedisLockRunnerConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
if matches!(self.command_timeout_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"redis lock command_timeout_ms must be positive".to_string(),
));
}
if self.default_ttl_ms == 0 {
return Err(DataLayerError::InvalidConfiguration(
"redis lock default_ttl_ms must be positive".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct RedisLockRunner {
client: RedisClient,
keyspace: RedisKeyspace,
config: RedisLockRunnerConfig,
}
impl RedisLockRunner {
pub fn new(
client: RedisClient,
keyspace: RedisKeyspace,
config: RedisLockRunnerConfig,
) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self {
client,
keyspace,
config,
})
}
pub fn client(&self) -> &RedisClient {
&self.client
}
pub fn keyspace(&self) -> &RedisKeyspace {
&self.keyspace
}
pub fn config(&self) -> RedisLockRunnerConfig {
self.config
}
pub async fn try_acquire(
&self,
key: &RedisLockKey,
owner: &str,
ttl_ms: Option<u64>,
) -> Result<Option<RedisLockLease>, DataLayerError> {
validate_owner(owner)?;
validate_key(key)?;
let ttl_ms = self.resolve_ttl_ms(ttl_ms)?;
let token = format!("{owner}:{}", Uuid::new_v4());
self.run_with_timeout("redis lock acquire", async {
let mut connection = self.client.get_multiplexed_async_connection().await?;
let status = redis::cmd("SET")
.arg(&key.0)
.arg(&token)
.arg("NX")
.arg("PX")
.arg(ttl_ms)
.query_async::<Option<String>>(&mut connection)
.await?;
Ok(status.map(|_| RedisLockLease {
key: key.clone(),
owner: owner.to_string(),
token,
ttl_ms,
}))
})
.await
}
pub async fn release(&self, lease: &RedisLockLease) -> Result<bool, DataLayerError> {
validate_lease(lease)?;
self.run_with_timeout("redis lock release", async {
let mut connection = self.client.get_multiplexed_async_connection().await?;
let deleted = redis::Script::new(
"if redis.call('get', KEYS[1]) == ARGV[1] then \
return redis.call('del', KEYS[1]) \
else \
return 0 \
end",
)
.key(&lease.key.0)
.arg(&lease.token)
.invoke_async::<i32>(&mut connection)
.await?;
Ok(deleted > 0)
})
.await
}
pub async fn renew(
&self,
lease: &RedisLockLease,
ttl_ms: Option<u64>,
) -> Result<bool, DataLayerError> {
validate_lease(lease)?;
let ttl_ms = self.resolve_ttl_ms(ttl_ms)?;
self.run_with_timeout("redis lock renew", async {
let mut connection = self.client.get_multiplexed_async_connection().await?;
let renewed = redis::Script::new(
"if redis.call('get', KEYS[1]) == ARGV[1] then \
return redis.call('pexpire', KEYS[1], ARGV[2]) \
else \
return 0 \
end",
)
.key(&lease.key.0)
.arg(&lease.token)
.arg(ttl_ms)
.invoke_async::<i32>(&mut connection)
.await?;
Ok(renewed > 0)
})
.await
}
async fn run_with_timeout<T, F>(
&self,
operation: &'static str,
future: F,
) -> Result<T, DataLayerError>
where
F: Future<Output = Result<T, DataLayerError>>,
{
if let Some(timeout_ms) = self.config.command_timeout_ms {
tokio::time::timeout(Duration::from_millis(timeout_ms), future)
.await
.map_err(|_| {
DataLayerError::TimedOut(format!("{operation} exceeded {timeout_ms}ms timeout"))
})?
} else {
future.await
}
}
fn resolve_ttl_ms(&self, ttl_ms: Option<u64>) -> Result<u64, DataLayerError> {
let ttl_ms = ttl_ms.unwrap_or(self.config.default_ttl_ms);
if ttl_ms == 0 {
return Err(DataLayerError::InvalidInput(
"redis lock ttl_ms must be positive".to_string(),
));
}
Ok(ttl_ms)
}
}
fn validate_owner(owner: &str) -> Result<(), DataLayerError> {
if owner.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis lock owner cannot be empty".to_string(),
));
}
Ok(())
}
fn validate_key(key: &RedisLockKey) -> Result<(), DataLayerError> {
if key.0.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis lock key cannot be empty".to_string(),
));
}
Ok(())
}
fn validate_lease(lease: &RedisLockLease) -> Result<(), DataLayerError> {
validate_key(&lease.key)?;
validate_owner(&lease.owner)?;
if lease.token.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis lock token cannot be empty".to_string(),
));
}
if lease.ttl_ms == 0 {
return Err(DataLayerError::InvalidInput(
"redis lock ttl_ms must be positive".to_string(),
));
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::{RedisLockKey, RedisLockLease, RedisLockRunner, RedisLockRunnerConfig};
use crate::redis::{RedisClientConfig, RedisClientFactory};
fn sample_runner() -> RedisLockRunner {
let client = RedisClientFactory::new(RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
})
.expect("factory should build")
.connect_lazy()
.expect("client should build");
RedisLockRunner::new(
client,
RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
}
.keyspace(),
RedisLockRunnerConfig::default(),
)
.expect("runner should build")
}
#[test]
fn validates_runner_config() {
assert!(RedisLockRunnerConfig {
command_timeout_ms: Some(0),
..RedisLockRunnerConfig::default()
}
.validate()
.is_err());
assert!(RedisLockRunnerConfig {
default_ttl_ms: 0,
..RedisLockRunnerConfig::default()
}
.validate()
.is_err());
}
#[test]
fn runner_reuses_client_and_keyspace() {
let runner = sample_runner();
assert_eq!(runner.config(), RedisLockRunnerConfig::default());
assert_eq!(runner.keyspace().lock_key("poller").0, "aether:lock:poller");
let _client_ref = runner.client();
}
#[tokio::test]
async fn rejects_invalid_owner_or_lease_before_network() {
let runner = sample_runner();
assert!(runner
.try_acquire(&RedisLockKey("aether:lock:poller".to_string()), "", None)
.await
.is_err());
assert!(runner
.release(&RedisLockLease {
key: RedisLockKey("aether:lock:poller".to_string()),
owner: "worker-1".to_string(),
token: String::new(),
ttl_ms: 1_000,
})
.await
.is_err());
assert!(runner
.renew(
&RedisLockLease {
key: RedisLockKey("aether:lock:poller".to_string()),
owner: "worker-1".to_string(),
token: "token-1".to_string(),
ttl_ms: 1_000,
},
Some(0),
)
.await
.is_err());
}
}

12
crates/aether-data/src/redis/mod.rs Normal file
View File

@@ -0,0 +1,12 @@
mod client;
mod lock;
mod namespace;
mod stream;
pub use client::{RedisClient, RedisClientConfig, RedisClientFactory};
pub use lock::{RedisLockKey, RedisLockLease, RedisLockRunner, RedisLockRunnerConfig};
pub use namespace::RedisKeyspace;
pub use stream::{
RedisConsumerGroup, RedisConsumerName, RedisStreamEntry, RedisStreamName,
RedisStreamReclaimConfig, RedisStreamReclaimResult, RedisStreamRunner, RedisStreamRunnerConfig,
};

43
crates/aether-data/src/redis/namespace.rs Normal file
View File

@@ -0,0 +1,43 @@
use aether_cache::CacheKeyNamespace;
use crate::redis::{RedisLockKey, RedisStreamName};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RedisKeyspace {
namespace: CacheKeyNamespace,
}
impl RedisKeyspace {
pub fn new(prefix: Option<&str>) -> Self {
let normalized = prefix.unwrap_or_default().trim().trim_matches(':');
Self {
namespace: CacheKeyNamespace::new(normalized),
}
}
pub fn key(&self, raw_key: &str) -> String {
self.namespace.key(raw_key)
}
pub fn lock_key(&self, raw_key: &str) -> RedisLockKey {
RedisLockKey(self.namespace.child("lock").key(raw_key))
}
pub fn stream_name(&self, raw_name: &str) -> RedisStreamName {
RedisStreamName(self.namespace.child("stream").key(raw_name))
}
}
#[cfg(test)]
mod tests {
use super::RedisKeyspace;
#[test]
fn composes_prefixed_lock_and_stream_names() {
let keyspace = RedisKeyspace::new(Some("aether"));
assert_eq!(keyspace.key("auth:user"), "aether:auth:user");
assert_eq!(keyspace.lock_key("poller").0, "aether:lock:poller");
assert_eq!(keyspace.stream_name("audit").0, "aether:stream:audit");
}
}

658
crates/aether-data/src/redis/stream.rs Normal file
View File

@@ -0,0 +1,658 @@
use std::collections::BTreeMap;
use std::future::Future;
use std::time::Duration;
use redis::from_redis_value;
use redis::streams::StreamReadReply;
use redis::Value as RedisValue;
use crate::redis::{RedisClient, RedisKeyspace};
use crate::DataLayerError;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RedisStreamName(pub String);
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RedisConsumerGroup(pub String);
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RedisConsumerName(pub String);
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RedisStreamEntry {
pub id: String,
pub fields: BTreeMap<String, String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RedisStreamReclaimResult {
pub next_start_id: String,
pub entries: Vec<RedisStreamEntry>,
pub deleted_ids: Vec<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RedisStreamReclaimConfig {
pub min_idle_ms: u64,
pub count: usize,
}
impl Default for RedisStreamReclaimConfig {
fn default() -> Self {
Self {
min_idle_ms: 60_000,
count: 32,
}
}
}
impl RedisStreamReclaimConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
if self.min_idle_ms == 0 {
return Err(DataLayerError::InvalidConfiguration(
"redis stream reclaim min_idle_ms must be positive".to_string(),
));
}
if self.count == 0 {
return Err(DataLayerError::InvalidConfiguration(
"redis stream reclaim count must be positive".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RedisStreamRunnerConfig {
pub command_timeout_ms: Option<u64>,
pub read_block_ms: Option<u64>,
pub read_count: usize,
}
impl Default for RedisStreamRunnerConfig {
fn default() -> Self {
Self {
command_timeout_ms: Some(1_000),
read_block_ms: Some(1_000),
read_count: 32,
}
}
}
impl RedisStreamRunnerConfig {
pub fn validate(&self) -> Result<(), DataLayerError> {
if matches!(self.command_timeout_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"redis stream command_timeout_ms must be positive".to_string(),
));
}
if matches!(self.read_block_ms, Some(0)) {
return Err(DataLayerError::InvalidConfiguration(
"redis stream read_block_ms must be positive".to_string(),
));
}
if self.read_count == 0 {
return Err(DataLayerError::InvalidConfiguration(
"redis stream read_count must be positive".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct RedisStreamRunner {
client: RedisClient,
keyspace: RedisKeyspace,
config: RedisStreamRunnerConfig,
}
impl RedisStreamRunner {
pub fn new(
client: RedisClient,
keyspace: RedisKeyspace,
config: RedisStreamRunnerConfig,
) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self {
client,
keyspace,
config,
})
}
pub fn client(&self) -> &RedisClient {
&self.client
}
pub fn keyspace(&self) -> &RedisKeyspace {
&self.keyspace
}
pub fn config(&self) -> RedisStreamRunnerConfig {
self.config
}
pub async fn ensure_consumer_group(
&self,
stream: &RedisStreamName,
group: &RedisConsumerGroup,
start_id: &str,
) -> Result<(), DataLayerError> {
validate_stream_name(stream)?;
validate_group(group)?;
validate_stream_position(start_id)?;
self.run_with_timeout("redis stream ensure consumer group", async {
let mut connection = self.client.get_multiplexed_async_connection().await?;
let result = redis::cmd("XGROUP")
.arg("CREATE")
.arg(&stream.0)
.arg(&group.0)
.arg(start_id)
.arg("MKSTREAM")
.query_async::<String>(&mut connection)
.await;
match result {
Ok(_) => Ok(()),
Err(err) if err.code() == Some("BUSYGROUP") => Ok(()),
Err(err) => Err(DataLayerError::Redis(err)),
}
})
.await
}
pub async fn append_fields(
&self,
stream: &RedisStreamName,
fields: &BTreeMap<String, String>,
) -> Result<String, DataLayerError> {
validate_stream_name(stream)?;
if fields.is_empty() {
return Err(DataLayerError::InvalidInput(
"redis stream fields cannot be empty".to_string(),
));
}
self.run_with_timeout("redis stream append", async {
let mut connection = self.client.get_multiplexed_async_connection().await?;
let mut command = redis::cmd("XADD");
command.arg(&stream.0).arg("*");
for (key, value) in fields {
command.arg(key).arg(value);
}
Ok(command.query_async::<String>(&mut connection).await?)
})
.await
}
pub async fn append_json(
&self,
stream: &RedisStreamName,
field: &str,
payload: &serde_json::Value,
) -> Result<String, DataLayerError> {
if field.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis stream json field cannot be empty".to_string(),
));
}
let mut fields = BTreeMap::new();
fields.insert(
field.to_string(),
serde_json::to_string(payload).map_err(|err| {
DataLayerError::UnexpectedValue(format!(
"failed to serialize redis stream payload: {err}"
))
})?,
);
self.append_fields(stream, &fields).await
}
pub async fn read_group(
&self,
stream: &RedisStreamName,
group: &RedisConsumerGroup,
consumer: &RedisConsumerName,
) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
validate_stream_name(stream)?;
validate_group(group)?;
validate_consumer(consumer)?;
self.run_with_timeout("redis stream read group", async {
let mut connection = self.client.get_multiplexed_async_connection().await?;
let mut command = redis::cmd("XREADGROUP");
command
.arg("GROUP")
.arg(&group.0)
.arg(&consumer.0)
.arg("COUNT")
.arg(self.config.read_count);
if let Some(block_ms) = self.config.read_block_ms {
command.arg("BLOCK").arg(block_ms);
}
command.arg("STREAMS").arg(&stream.0).arg(">");
let reply = command
.query_async::<StreamReadReply>(&mut connection)
.await?;
Ok(reply
.keys
.into_iter()
.flat_map(|key| key.ids.into_iter())
.map(|id| RedisStreamEntry {
id: id.id,
fields: id
.map
.into_iter()
.filter_map(|(field, value)| {
redis::from_redis_value::<String>(&value)
.ok()
.map(|text| (field, text))
})
.collect(),
})
.collect())
})
.await
}
pub async fn ack(
&self,
stream: &RedisStreamName,
group: &RedisConsumerGroup,
ids: &[String],
) -> Result<usize, DataLayerError> {
validate_stream_name(stream)?;
validate_group(group)?;
if ids.is_empty() {
return Ok(0);
}
self.run_with_timeout("redis stream ack", async {
let mut connection = self.client.get_multiplexed_async_connection().await?;
let mut command = redis::cmd("XACK");
command.arg(&stream.0).arg(&group.0);
for id in ids {
command.arg(id);
}
Ok(command.query_async::<usize>(&mut connection).await?)
})
.await
}
pub async fn claim_stale(
&self,
stream: &RedisStreamName,
group: &RedisConsumerGroup,
consumer: &RedisConsumerName,
start_id: &str,
config: RedisStreamReclaimConfig,
) -> Result<RedisStreamReclaimResult, DataLayerError> {
validate_stream_name(stream)?;
validate_group(group)?;
validate_consumer(consumer)?;
validate_stream_position(start_id)?;
config.validate()?;
self.run_with_timeout("redis stream reclaim", async {
let mut connection = self.client.get_multiplexed_async_connection().await?;
let reply = redis::cmd("XAUTOCLAIM")
.arg(&stream.0)
.arg(&group.0)
.arg(&consumer.0)
.arg(config.min_idle_ms)
.arg(start_id)
.arg("COUNT")
.arg(config.count)
.query_async::<RedisValue>(&mut connection)
.await?;
parse_reclaim_result(reply)
})
.await
}
async fn run_with_timeout<T, F>(
&self,
operation: &'static str,
future: F,
) -> Result<T, DataLayerError>
where
F: Future<Output = Result<T, DataLayerError>>,
{
if let Some(timeout_ms) = self.config.command_timeout_ms {
tokio::time::timeout(Duration::from_millis(timeout_ms), future)
.await
.map_err(|_| {
DataLayerError::TimedOut(format!("{operation} exceeded {timeout_ms}ms timeout"))
})?
} else {
future.await
}
}
}
fn validate_stream_name(stream: &RedisStreamName) -> Result<(), DataLayerError> {
if stream.0.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis stream name cannot be empty".to_string(),
));
}
Ok(())
}
fn validate_group(group: &RedisConsumerGroup) -> Result<(), DataLayerError> {
if group.0.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis consumer group cannot be empty".to_string(),
));
}
Ok(())
}
fn validate_consumer(consumer: &RedisConsumerName) -> Result<(), DataLayerError> {
if consumer.0.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis consumer name cannot be empty".to_string(),
));
}
Ok(())
}
fn validate_stream_position(position: &str) -> Result<(), DataLayerError> {
if position.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
"redis stream position cannot be empty".to_string(),
));
}
Ok(())
}
fn parse_reclaim_result(value: RedisValue) -> Result<RedisStreamReclaimResult, DataLayerError> {
let RedisValue::Array(parts) = value else {
return Err(DataLayerError::UnexpectedValue(
"redis xautoclaim returned non-array payload".to_string(),
));
};
if parts.len() < 2 || parts.len() > 3 {
return Err(DataLayerError::UnexpectedValue(format!(
"redis xautoclaim returned {} top-level fields, expected 2 or 3",
parts.len()
)));
}
let next_start_id = parse_string_value(&parts[0], "redis xautoclaim next_start_id")?;
let entries = parse_reclaim_entries(&parts[1])?;
let deleted_ids = match parts.get(2) {
Some(value) => parse_string_array(value, "redis xautoclaim deleted_ids")?,
None => Vec::new(),
};
Ok(RedisStreamReclaimResult {
next_start_id,
entries,
deleted_ids,
})
}
fn parse_reclaim_entries(value: &RedisValue) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
match value {
RedisValue::Array(entries) => entries.iter().map(parse_reclaim_entry).collect(),
RedisValue::Nil => Ok(Vec::new()),
_ => Err(DataLayerError::UnexpectedValue(
"redis xautoclaim entries payload was not an array".to_string(),
)),
}
}
fn parse_reclaim_entry(value: &RedisValue) -> Result<RedisStreamEntry, DataLayerError> {
let RedisValue::Array(parts) = value else {
return Err(DataLayerError::UnexpectedValue(
"redis xautoclaim entry was not an array".to_string(),
));
};
if parts.len() != 2 {
return Err(DataLayerError::UnexpectedValue(format!(
"redis xautoclaim entry had {} fields, expected 2",
parts.len()
)));
}
let id = parse_string_value(&parts[0], "redis xautoclaim entry id")?;
let fields = parse_string_map(&parts[1], "redis xautoclaim entry fields")?;
Ok(RedisStreamEntry { id, fields })
}
fn parse_string_map(
value: &RedisValue,
context: &str,
) -> Result<BTreeMap<String, String>, DataLayerError> {
match value {
RedisValue::Array(values) => {
if values.len() % 2 != 0 {
return Err(DataLayerError::UnexpectedValue(format!(
"{context} expected an even number of field elements, got {}",
values.len()
)));
}
let mut fields = BTreeMap::new();
for pair in values.chunks(2) {
let key = parse_string_value(&pair[0], context)?;
let value = parse_string_value(&pair[1], context)?;
fields.insert(key, value);
}
Ok(fields)
}
RedisValue::Map(entries) => entries
.iter()
.map(|(key, value)| {
Ok((
parse_string_value(key, context)?,
parse_string_value(value, context)?,
))
})
.collect(),
RedisValue::Nil => Ok(BTreeMap::new()),
_ => Err(DataLayerError::UnexpectedValue(format!(
"{context} expected a redis array/map payload"
))),
}
}
fn parse_string_array(value: &RedisValue, context: &str) -> Result<Vec<String>, DataLayerError> {
match value {
RedisValue::Array(values) => values
.iter()
.map(|value| parse_string_value(value, context))
.collect(),
RedisValue::Nil => Ok(Vec::new()),
_ => Err(DataLayerError::UnexpectedValue(format!(
"{context} expected a redis array payload"
))),
}
}
fn parse_string_value(value: &RedisValue, context: &str) -> Result<String, DataLayerError> {
from_redis_value::<String>(value).map_err(|err| {
DataLayerError::UnexpectedValue(format!(
"{context} was not a string-compatible redis value: {err}"
))
})
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use super::{
parse_reclaim_result, RedisConsumerGroup, RedisConsumerName, RedisStreamName,
RedisStreamReclaimConfig, RedisStreamReclaimResult, RedisStreamRunner,
RedisStreamRunnerConfig,
};
use crate::redis::{RedisClientConfig, RedisClientFactory};
use redis::Value as RedisValue;
fn sample_runner() -> RedisStreamRunner {
let config = RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
};
let client = RedisClientFactory::new(config.clone())
.expect("factory should build")
.connect_lazy()
.expect("client should build");
RedisStreamRunner::new(
client,
config.keyspace(),
RedisStreamRunnerConfig::default(),
)
.expect("runner should build")
}
#[test]
fn validates_stream_runner_config() {
assert!(RedisStreamRunnerConfig {
command_timeout_ms: Some(0),
..RedisStreamRunnerConfig::default()
}
.validate()
.is_err());
assert!(RedisStreamRunnerConfig {
read_block_ms: Some(0),
..RedisStreamRunnerConfig::default()
}
.validate()
.is_err());
assert!(RedisStreamRunnerConfig {
read_count: 0,
..RedisStreamRunnerConfig::default()
}
.validate()
.is_err());
}
#[test]
fn validates_reclaim_config() {
assert!(RedisStreamReclaimConfig {
min_idle_ms: 0,
..RedisStreamReclaimConfig::default()
}
.validate()
.is_err());
assert!(RedisStreamReclaimConfig {
count: 0,
..RedisStreamReclaimConfig::default()
}
.validate()
.is_err());
}
#[test]
fn runner_reuses_client_and_keyspace() {
let runner = sample_runner();
assert_eq!(runner.config(), RedisStreamRunnerConfig::default());
assert_eq!(
runner.keyspace().stream_name("audit").0,
"aether:stream:audit"
);
let _client_ref = runner.client();
}
#[tokio::test]
async fn rejects_invalid_inputs_before_network() {
let runner = sample_runner();
let stream = RedisStreamName("aether:stream:audit".to_string());
let group = RedisConsumerGroup("audit-workers".to_string());
let consumer = RedisConsumerName("worker-1".to_string());
assert!(runner
.ensure_consumer_group(&stream, &group, "")
.await
.is_err());
assert!(runner
.append_fields(&stream, &BTreeMap::new())
.await
.is_err());
assert!(runner
.append_json(&stream, "", &serde_json::json!({"ok": true}))
.await
.is_err());
assert!(runner
.read_group(&stream, &group, &RedisConsumerName(String::new()))
.await
.is_err());
assert_eq!(
runner.ack(&stream, &group, &[]).await.expect("empty ack"),
0
);
assert!(runner
.claim_stale(
&stream,
&group,
&consumer,
"",
RedisStreamReclaimConfig::default()
)
.await
.is_err());
let _ = consumer;
}
#[test]
fn parses_reclaim_result_with_deleted_ids() {
let parsed = parse_reclaim_result(RedisValue::Array(vec![
RedisValue::BulkString(b"0-0".to_vec()),
RedisValue::Array(vec![RedisValue::Array(vec![
RedisValue::BulkString(b"1710000000000-0".to_vec()),
RedisValue::Array(vec![
RedisValue::BulkString(b"payload".to_vec()),
RedisValue::BulkString(br#"{"ok":true}"#.to_vec()),
RedisValue::BulkString(b"kind".to_vec()),
RedisValue::BulkString(b"shadow".to_vec()),
]),
])]),
RedisValue::Array(vec![RedisValue::BulkString(b"1709999999999-0".to_vec())]),
]))
.expect("reclaim result should parse");
assert_eq!(
parsed,
RedisStreamReclaimResult {
next_start_id: "0-0".to_string(),
entries: vec![super::RedisStreamEntry {
id: "1710000000000-0".to_string(),
fields: BTreeMap::from([
("kind".to_string(), "shadow".to_string()),
("payload".to_string(), r#"{"ok":true}"#.to_string()),
]),
}],
deleted_ids: vec!["1709999999999-0".to_string()],
}
);
}
#[test]
fn parses_reclaim_result_without_deleted_ids() {
let parsed = parse_reclaim_result(RedisValue::Array(vec![
RedisValue::BulkString(b"0-0".to_vec()),
RedisValue::Array(vec![]),
]))
.expect("reclaim result should parse");
assert_eq!(
parsed,
RedisStreamReclaimResult {
next_start_id: "0-0".to_string(),
entries: Vec::new(),
deleted_ids: Vec::new(),
}
);
}
}

133
crates/aether-data/src/repository/auth/memory.rs Normal file
View File

@@ -0,0 +1,133 @@
use std::collections::BTreeMap;
use std::sync::RwLock;
use async_trait::async_trait;
use super::types::{AuthApiKeyLookupKey, AuthApiKeyReadRepository, StoredAuthApiKeySnapshot};
use crate::DataLayerError;
#[derive(Debug, Default)]
struct MemoryAuthApiKeyIndex {
by_api_key_id: BTreeMap<String, StoredAuthApiKeySnapshot>,
by_key_hash: BTreeMap<String, String>,
}
#[derive(Debug, Default)]
pub struct InMemoryAuthApiKeySnapshotRepository {
index: RwLock<MemoryAuthApiKeyIndex>,
}
impl InMemoryAuthApiKeySnapshotRepository {
pub fn seed<I>(items: I) -> Self
where
I: IntoIterator<Item = (Option<String>, StoredAuthApiKeySnapshot)>,
{
let mut by_api_key_id = BTreeMap::new();
let mut by_key_hash = BTreeMap::new();
for (key_hash, snapshot) in items {
if let Some(key_hash) = key_hash {
by_key_hash.insert(key_hash, snapshot.api_key_id.clone());
}
by_api_key_id.insert(snapshot.api_key_id.clone(), snapshot);
}
Self {
index: RwLock::new(MemoryAuthApiKeyIndex {
by_api_key_id,
by_key_hash,
}),
}
}
}
#[async_trait]
impl AuthApiKeyReadRepository for InMemoryAuthApiKeySnapshotRepository {
async fn find_api_key_snapshot(
&self,
key: AuthApiKeyLookupKey<'_>,
) -> Result<Option<StoredAuthApiKeySnapshot>, DataLayerError> {
let index = self
.index
.read()
.expect("auth api key snapshot repository lock");
Ok(match key {
AuthApiKeyLookupKey::KeyHash(key_hash) => index
.by_key_hash
.get(key_hash)
.and_then(|api_key_id| index.by_api_key_id.get(api_key_id))
.cloned(),
AuthApiKeyLookupKey::ApiKeyId(api_key_id) => {
index.by_api_key_id.get(api_key_id).cloned()
}
AuthApiKeyLookupKey::UserApiKeyIds {
user_id,
api_key_id,
} => index
.by_api_key_id
.get(api_key_id)
.filter(|snapshot| snapshot.user_id == user_id)
.cloned(),
})
}
}
#[cfg(test)]
mod tests {
use super::InMemoryAuthApiKeySnapshotRepository;
use crate::repository::auth::{
AuthApiKeyLookupKey, AuthApiKeyReadRepository, StoredAuthApiKeySnapshot,
};
fn sample_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-4.1"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(200),
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-4.1"])),
)
.expect("snapshot should build")
}
#[tokio::test]
async fn reads_auth_snapshot_by_all_supported_keys() {
let repository = InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-1".to_string()),
sample_snapshot("key-1", "user-1"),
)]);
assert!(repository
.find_api_key_snapshot(AuthApiKeyLookupKey::KeyHash("hash-1"))
.await
.expect("find by hash should succeed")
.is_some());
assert!(repository
.find_api_key_snapshot(AuthApiKeyLookupKey::ApiKeyId("key-1"))
.await
.expect("find by api key id should succeed")
.is_some());
assert!(repository
.find_api_key_snapshot(AuthApiKeyLookupKey::UserApiKeyIds {
user_id: "user-1",
api_key_id: "key-1",
})
.await
.expect("find by user/api key ids should succeed")
.is_some());
}
}

9
crates/aether-data/src/repository/auth/mod.rs Normal file
View File

@@ -0,0 +1,9 @@
mod memory;
mod sql;
mod types;
pub use memory::InMemoryAuthApiKeySnapshotRepository;
pub use sql::SqlxAuthApiKeySnapshotReadRepository;
pub use types::{
AuthApiKeyLookupKey, AuthApiKeyReadRepository, AuthRepository, StoredAuthApiKeySnapshot,
};

202
crates/aether-data/src/repository/auth/sql.rs Normal file
View File

@@ -0,0 +1,202 @@
use async_trait::async_trait;
use sqlx::{PgPool, Row};
use super::types::{AuthApiKeyLookupKey, AuthApiKeyReadRepository, StoredAuthApiKeySnapshot};
use crate::DataLayerError;
const FIND_BY_KEY_HASH_SQL: &str = r#"
SELECT
users.id AS user_id,
users.username,
users.email,
users.role::text AS user_role,
users.auth_source::text AS user_auth_source,
users.is_active AS user_is_active,
users.is_deleted AS user_is_deleted,
users.allowed_providers AS user_allowed_providers,
users.allowed_api_formats AS user_allowed_api_formats,
users.allowed_models AS user_allowed_models,
api_keys.id AS api_key_id,
api_keys.name AS api_key_name,
api_keys.is_active AS api_key_is_active,
api_keys.is_locked AS api_key_is_locked,
api_keys.is_standalone AS api_key_is_standalone,
api_keys.rate_limit AS api_key_rate_limit,
api_keys.concurrent_limit AS api_key_concurrent_limit,
CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs,
api_keys.allowed_providers AS api_key_allowed_providers,
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
api_keys.allowed_models AS api_key_allowed_models
FROM api_keys
JOIN users ON users.id = api_keys.user_id
WHERE api_keys.key_hash = $1
LIMIT 1
"#;
const FIND_BY_API_KEY_ID_SQL: &str = r#"
SELECT
users.id AS user_id,
users.username,
users.email,
users.role::text AS user_role,
users.auth_source::text AS user_auth_source,
users.is_active AS user_is_active,
users.is_deleted AS user_is_deleted,
users.allowed_providers AS user_allowed_providers,
users.allowed_api_formats AS user_allowed_api_formats,
users.allowed_models AS user_allowed_models,
api_keys.id AS api_key_id,
api_keys.name AS api_key_name,
api_keys.is_active AS api_key_is_active,
api_keys.is_locked AS api_key_is_locked,
api_keys.is_standalone AS api_key_is_standalone,
api_keys.rate_limit AS api_key_rate_limit,
api_keys.concurrent_limit AS api_key_concurrent_limit,
CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs,
api_keys.allowed_providers AS api_key_allowed_providers,
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
api_keys.allowed_models AS api_key_allowed_models
FROM api_keys
JOIN users ON users.id = api_keys.user_id
WHERE api_keys.id = $1
LIMIT 1
"#;
const FIND_BY_USER_API_KEY_IDS_SQL: &str = r#"
SELECT
users.id AS user_id,
users.username,
users.email,
users.role::text AS user_role,
users.auth_source::text AS user_auth_source,
users.is_active AS user_is_active,
users.is_deleted AS user_is_deleted,
users.allowed_providers AS user_allowed_providers,
users.allowed_api_formats AS user_allowed_api_formats,
users.allowed_models AS user_allowed_models,
api_keys.id AS api_key_id,
api_keys.name AS api_key_name,
api_keys.is_active AS api_key_is_active,
api_keys.is_locked AS api_key_is_locked,
api_keys.is_standalone AS api_key_is_standalone,
api_keys.rate_limit AS api_key_rate_limit,
api_keys.concurrent_limit AS api_key_concurrent_limit,
CAST(EXTRACT(EPOCH FROM api_keys.expires_at) AS BIGINT) AS api_key_expires_at_unix_secs,
api_keys.allowed_providers AS api_key_allowed_providers,
api_keys.allowed_api_formats AS api_key_allowed_api_formats,
api_keys.allowed_models AS api_key_allowed_models
FROM api_keys
JOIN users ON users.id = api_keys.user_id
WHERE api_keys.id = $1 AND users.id = $2
LIMIT 1
"#;
#[derive(Debug, Clone)]
pub struct SqlxAuthApiKeySnapshotReadRepository {
pool: PgPool,
}
impl SqlxAuthApiKeySnapshotReadRepository {
pub fn new(pool: PgPool) -> Self {
Self { pool }
}
pub fn pool(&self) -> &PgPool {
&self.pool
}
pub async fn find_api_key_snapshot(
&self,
key: AuthApiKeyLookupKey<'_>,
) -> Result<Option<StoredAuthApiKeySnapshot>, DataLayerError> {
let row = match key {
AuthApiKeyLookupKey::KeyHash(key_hash) => {
sqlx::query(FIND_BY_KEY_HASH_SQL)
.bind(key_hash)
.fetch_optional(&self.pool)
.await?
}
AuthApiKeyLookupKey::ApiKeyId(api_key_id) => {
sqlx::query(FIND_BY_API_KEY_ID_SQL)
.bind(api_key_id)
.fetch_optional(&self.pool)
.await?
}
AuthApiKeyLookupKey::UserApiKeyIds {
user_id,
api_key_id,
} => {
sqlx::query(FIND_BY_USER_API_KEY_IDS_SQL)
.bind(api_key_id)
.bind(user_id)
.fetch_optional(&self.pool)
.await?
}
};
row.as_ref().map(map_auth_api_key_snapshot_row).transpose()
}
}
#[async_trait]
impl AuthApiKeyReadRepository for SqlxAuthApiKeySnapshotReadRepository {
async fn find_api_key_snapshot(
&self,
key: AuthApiKeyLookupKey<'_>,
) -> Result<Option<StoredAuthApiKeySnapshot>, DataLayerError> {
Self::find_api_key_snapshot(self, key).await
}
}
fn map_auth_api_key_snapshot_row(
row: &sqlx::postgres::PgRow,
) -> Result<StoredAuthApiKeySnapshot, DataLayerError> {
StoredAuthApiKeySnapshot::new(
row.try_get("user_id")?,
row.try_get("username")?,
row.try_get("email")?,
row.try_get("user_role")?,
row.try_get("user_auth_source")?,
row.try_get("user_is_active")?,
row.try_get("user_is_deleted")?,
row.try_get("user_allowed_providers")?,
row.try_get("user_allowed_api_formats")?,
row.try_get("user_allowed_models")?,
row.try_get("api_key_id")?,
row.try_get("api_key_name")?,
row.try_get("api_key_is_active")?,
row.try_get("api_key_is_locked")?,
row.try_get("api_key_is_standalone")?,
row.try_get("api_key_rate_limit")?,
row.try_get("api_key_concurrent_limit")?,
row.try_get("api_key_expires_at_unix_secs")?,
row.try_get("api_key_allowed_providers")?,
row.try_get("api_key_allowed_api_formats")?,
row.try_get("api_key_allowed_models")?,
)
}
#[cfg(test)]
mod tests {
use super::SqlxAuthApiKeySnapshotReadRepository;
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let pool = factory.connect_lazy().expect("pool should build");
let repository = SqlxAuthApiKeySnapshotReadRepository::new(pool);
let _ = repository.pool();
}
}

225
crates/aether-data/src/repository/auth/types.rs Normal file
View File

@@ -0,0 +1,225 @@
use async_trait::async_trait;
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredAuthApiKeySnapshot {
pub user_id: String,
pub username: String,
pub email: Option<String>,
pub user_role: String,
pub user_auth_source: String,
pub user_is_active: bool,
pub user_is_deleted: bool,
pub user_allowed_providers: Option<Vec<String>>,
pub user_allowed_api_formats: Option<Vec<String>>,
pub user_allowed_models: Option<Vec<String>>,
pub api_key_id: String,
pub api_key_name: Option<String>,
pub api_key_is_active: bool,
pub api_key_is_locked: bool,
pub api_key_is_standalone: bool,
pub api_key_rate_limit: Option<i32>,
pub api_key_concurrent_limit: Option<i32>,
pub api_key_expires_at_unix_secs: Option<u64>,
pub api_key_allowed_providers: Option<Vec<String>>,
pub api_key_allowed_api_formats: Option<Vec<String>>,
pub api_key_allowed_models: Option<Vec<String>>,
}
impl StoredAuthApiKeySnapshot {
#[allow(clippy::too_many_arguments)]
pub fn new(
user_id: String,
username: String,
email: Option<String>,
user_role: String,
user_auth_source: String,
user_is_active: bool,
user_is_deleted: bool,
user_allowed_providers: Option<serde_json::Value>,
user_allowed_api_formats: Option<serde_json::Value>,
user_allowed_models: Option<serde_json::Value>,
api_key_id: String,
api_key_name: Option<String>,
api_key_is_active: bool,
api_key_is_locked: bool,
api_key_is_standalone: bool,
api_key_rate_limit: Option<i32>,
api_key_concurrent_limit: Option<i32>,
api_key_expires_at_unix_secs: Option<i64>,
api_key_allowed_providers: Option<serde_json::Value>,
api_key_allowed_api_formats: Option<serde_json::Value>,
api_key_allowed_models: Option<serde_json::Value>,
) -> Result<Self, crate::DataLayerError> {
Ok(Self {
user_id,
username,
email,
user_role,
user_auth_source,
user_is_active,
user_is_deleted,
user_allowed_providers: parse_string_list(
user_allowed_providers,
"users.allowed_providers",
)?,
user_allowed_api_formats: parse_string_list(
user_allowed_api_formats,
"users.allowed_api_formats",
)?,
user_allowed_models: parse_string_list(user_allowed_models, "users.allowed_models")?,
api_key_id,
api_key_name,
api_key_is_active,
api_key_is_locked,
api_key_is_standalone,
api_key_rate_limit,
api_key_concurrent_limit,
api_key_expires_at_unix_secs: api_key_expires_at_unix_secs
.map(|value| {
u64::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid api_keys.expires_at_unix_secs: {value}"
))
})
})
.transpose()?,
api_key_allowed_providers: parse_string_list(
api_key_allowed_providers,
"api_keys.allowed_providers",
)?,
api_key_allowed_api_formats: parse_string_list(
api_key_allowed_api_formats,
"api_keys.allowed_api_formats",
)?,
api_key_allowed_models: parse_string_list(
api_key_allowed_models,
"api_keys.allowed_models",
)?,
})
}
pub fn is_currently_usable(&self, now_unix_secs: u64) -> bool {
if !self.user_is_active || self.user_is_deleted {
return false;
}
if !self.api_key_is_active {
return false;
}
if self.api_key_is_locked && !self.api_key_is_standalone {
return false;
}
if let Some(expires_at_unix_secs) = self.api_key_expires_at_unix_secs {
if expires_at_unix_secs < now_unix_secs {
return false;
}
}
true
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AuthApiKeyLookupKey<'a> {
KeyHash(&'a str),
ApiKeyId(&'a str),
UserApiKeyIds {
user_id: &'a str,
api_key_id: &'a str,
},
}
#[async_trait]
pub trait AuthApiKeyReadRepository: Send + Sync {
async fn find_api_key_snapshot(
&self,
key: AuthApiKeyLookupKey<'_>,
) -> Result<Option<StoredAuthApiKeySnapshot>, crate::DataLayerError>;
}
pub trait AuthRepository: AuthApiKeyReadRepository + Send + Sync {}
impl<T> AuthRepository for T where T: AuthApiKeyReadRepository + Send + Sync {}
fn parse_string_list(
value: Option<serde_json::Value>,
field_name: &str,
) -> Result<Option<Vec<String>>, crate::DataLayerError> {
let Some(value) = value else {
return Ok(None);
};
let array = value.as_array().ok_or_else(|| {
crate::DataLayerError::UnexpectedValue(format!("{field_name} is not a JSON array"))
})?;
let mut items = Vec::with_capacity(array.len());
for item in array {
let Some(item) = item.as_str() else {
return Err(crate::DataLayerError::UnexpectedValue(format!(
"{field_name} contains a non-string item"
)));
};
items.push(item.to_string());
}
Ok(Some(items))
}
#[cfg(test)]
mod tests {
use super::StoredAuthApiKeySnapshot;
#[test]
fn rejects_non_array_allowed_providers() {
assert!(StoredAuthApiKeySnapshot::new(
"user-1".to_string(),
"alice".to_string(),
None,
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!({"bad": true})),
None,
None,
"key-1".to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
None,
None,
None,
None,
)
.is_err());
}
#[test]
fn expired_non_standalone_key_is_not_usable() {
let snapshot = StoredAuthApiKeySnapshot::new(
"user-1".to_string(),
"alice".to_string(),
None,
"user".to_string(),
"local".to_string(),
true,
false,
None,
None,
None,
"key-1".to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(100),
None,
None,
None,
)
.expect("snapshot should build");
assert!(!snapshot.is_currently_usable(101));
}
}

148
crates/aether-data/src/repository/candidates/memory.rs Normal file
View File

@@ -0,0 +1,148 @@
use std::collections::BTreeMap;
use std::sync::RwLock;
use async_trait::async_trait;
use super::types::{RequestCandidateReadRepository, StoredRequestCandidate};
use crate::DataLayerError;
#[derive(Debug, Default)]
pub struct InMemoryRequestCandidateRepository {
by_id: RwLock<BTreeMap<String, StoredRequestCandidate>>,
}
impl InMemoryRequestCandidateRepository {
pub fn seed<I>(items: I) -> Self
where
I: IntoIterator<Item = StoredRequestCandidate>,
{
let mut by_id = BTreeMap::new();
for item in items {
by_id.insert(item.id.clone(), item);
}
Self {
by_id: RwLock::new(by_id),
}
}
}
#[async_trait]
impl RequestCandidateReadRepository for InMemoryRequestCandidateRepository {
async fn list_by_request_id(
&self,
request_id: &str,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
let mut rows = self
.by_id
.read()
.expect("request candidate repository lock")
.values()
.filter(|row| row.request_id == request_id)
.cloned()
.collect::<Vec<_>>();
rows.sort_by(|left, right| {
left.candidate_index
.cmp(&right.candidate_index)
.then(left.retry_index.cmp(&right.retry_index))
.then(left.created_at_unix_secs.cmp(&right.created_at_unix_secs))
});
Ok(rows)
}
async fn list_recent(
&self,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
if limit == 0 {
return Ok(Vec::new());
}
let mut rows = self
.by_id
.read()
.expect("request candidate repository lock")
.values()
.cloned()
.collect::<Vec<_>>();
rows.sort_by(|left, right| right.created_at_unix_secs.cmp(&left.created_at_unix_secs));
rows.truncate(limit);
Ok(rows)
}
}
#[cfg(test)]
mod tests {
use super::InMemoryRequestCandidateRepository;
use crate::repository::candidates::{
RequestCandidateReadRepository, RequestCandidateStatus, StoredRequestCandidate,
};
fn sample_candidate(
id: &str,
request_id: &str,
created_at_unix_secs: i64,
) -> StoredRequestCandidate {
StoredRequestCandidate::new(
id.to_string(),
request_id.to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("key-1".to_string()),
RequestCandidateStatus::Success,
None,
false,
Some(200),
None,
None,
Some(10),
Some(1),
None,
None,
created_at_unix_secs,
Some(created_at_unix_secs),
Some(created_at_unix_secs + 1),
)
.expect("candidate should build")
}
#[tokio::test]
async fn lists_request_candidates_by_request_id_in_candidate_order() {
let repository = InMemoryRequestCandidateRepository::seed(vec![
sample_candidate("cand-2", "req-1", 200),
sample_candidate("cand-1", "req-1", 100),
sample_candidate("cand-3", "req-2", 300),
]);
let rows = repository
.list_by_request_id("req-1")
.await
.expect("list should succeed");
assert_eq!(rows.len(), 2);
assert_eq!(rows[0].request_id, "req-1");
assert_eq!(rows[1].request_id, "req-1");
}
#[tokio::test]
async fn lists_recent_request_candidates_in_descending_created_order() {
let repository = InMemoryRequestCandidateRepository::seed(vec![
sample_candidate("cand-1", "req-1", 100),
sample_candidate("cand-2", "req-2", 200),
]);
let rows = repository
.list_recent(10)
.await
.expect("list recent should succeed");
assert_eq!(rows.len(), 2);
assert_eq!(rows[0].id, "cand-2");
assert_eq!(rows[1].id, "cand-1");
}
}

10
crates/aether-data/src/repository/candidates/mod.rs Normal file
View File

@@ -0,0 +1,10 @@
mod memory;
mod sql;
mod types;
pub use memory::InMemoryRequestCandidateRepository;
pub use sql::SqlxRequestCandidateReadRepository;
pub use types::{
RequestCandidateReadRepository, RequestCandidateRepository, RequestCandidateStatus,
StoredRequestCandidate,
};

189
crates/aether-data/src/repository/candidates/sql.rs Normal file
View File

@@ -0,0 +1,189 @@
use async_trait::async_trait;
use sqlx::{PgPool, Row};
use super::types::{
RequestCandidateReadRepository, RequestCandidateStatus, StoredRequestCandidate,
};
use crate::DataLayerError;
const LIST_BY_REQUEST_ID_SQL: &str = r#"
SELECT
id,
request_id,
user_id,
api_key_id,
username,
api_key_name,
candidate_index,
retry_index,
provider_id,
endpoint_id,
key_id,
status,
skip_reason,
is_cached,
status_code,
error_type,
error_message,
latency_ms,
concurrent_requests,
extra_data,
required_capabilities,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM started_at) AS BIGINT) AS started_at_unix_secs,
CAST(EXTRACT(EPOCH FROM finished_at) AS BIGINT) AS finished_at_unix_secs
FROM request_candidates
WHERE request_id = $1
ORDER BY candidate_index ASC, retry_index ASC, created_at ASC
"#;
const LIST_RECENT_SQL: &str = r#"
SELECT
id,
request_id,
user_id,
api_key_id,
username,
api_key_name,
candidate_index,
retry_index,
provider_id,
endpoint_id,
key_id,
status,
skip_reason,
is_cached,
status_code,
error_type,
error_message,
latency_ms,
concurrent_requests,
extra_data,
required_capabilities,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM started_at) AS BIGINT) AS started_at_unix_secs,
CAST(EXTRACT(EPOCH FROM finished_at) AS BIGINT) AS finished_at_unix_secs
FROM request_candidates
ORDER BY created_at DESC
LIMIT $1
"#;
#[derive(Debug, Clone)]
pub struct SqlxRequestCandidateReadRepository {
pool: PgPool,
}
impl SqlxRequestCandidateReadRepository {
pub fn new(pool: PgPool) -> Self {
Self { pool }
}
pub fn pool(&self) -> &PgPool {
&self.pool
}
pub async fn list_by_request_id(
&self,
request_id: &str,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
let rows = sqlx::query(LIST_BY_REQUEST_ID_SQL)
.bind(request_id)
.fetch_all(&self.pool)
.await?;
rows.iter().map(map_request_candidate_row).collect()
}
pub async fn list_recent(
&self,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
if limit == 0 {
return Ok(Vec::new());
}
let rows = sqlx::query(LIST_RECENT_SQL)
.bind(i64::try_from(limit).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid recent request candidate limit: {limit}"
))
})?)
.fetch_all(&self.pool)
.await?;
rows.iter().map(map_request_candidate_row).collect()
}
}
#[async_trait]
impl RequestCandidateReadRepository for SqlxRequestCandidateReadRepository {
async fn list_by_request_id(
&self,
request_id: &str,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
Self::list_by_request_id(self, request_id).await
}
async fn list_recent(
&self,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
Self::list_recent(self, limit).await
}
}
fn map_request_candidate_row(
row: &sqlx::postgres::PgRow,
) -> Result<StoredRequestCandidate, DataLayerError> {
let status =
RequestCandidateStatus::from_database(row.try_get::<String, _>("status")?.as_str())?;
StoredRequestCandidate::new(
row.try_get("id")?,
row.try_get("request_id")?,
row.try_get("user_id")?,
row.try_get("api_key_id")?,
row.try_get("username")?,
row.try_get("api_key_name")?,
row.try_get("candidate_index")?,
row.try_get("retry_index")?,
row.try_get("provider_id")?,
row.try_get("endpoint_id")?,
row.try_get("key_id")?,
status,
row.try_get("skip_reason")?,
row.try_get("is_cached")?,
row.try_get("status_code")?,
row.try_get("error_type")?,
row.try_get("error_message")?,
row.try_get("latency_ms")?,
row.try_get("concurrent_requests")?,
row.try_get("extra_data")?,
row.try_get("required_capabilities")?,
row.try_get("created_at_unix_secs")?,
row.try_get("started_at_unix_secs")?,
row.try_get("finished_at_unix_secs")?,
)
}
#[cfg(test)]
mod tests {
use super::SqlxRequestCandidateReadRepository;
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let pool = factory.connect_lazy().expect("pool should build");
let repository = SqlxRequestCandidateReadRepository::new(pool);
let _ = repository.pool();
}
}

289
crates/aether-data/src/repository/candidates/types.rs Normal file
View File

@@ -0,0 +1,289 @@
use async_trait::async_trait;
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RequestCandidateStatus {
Available,
Unused,
Pending,
Streaming,
Success,
Failed,
Cancelled,
Skipped,
}
impl RequestCandidateStatus {
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
match value.trim().to_ascii_lowercase().as_str() {
"available" => Ok(Self::Available),
"unused" => Ok(Self::Unused),
"pending" => Ok(Self::Pending),
"streaming" => Ok(Self::Streaming),
"success" => Ok(Self::Success),
"failed" => Ok(Self::Failed),
"cancelled" => Ok(Self::Cancelled),
"skipped" => Ok(Self::Skipped),
other => Err(crate::DataLayerError::UnexpectedValue(format!(
"unsupported request_candidates.status: {other}"
))),
}
}
pub fn is_attempted(self, started_at_unix_secs: Option<u64>) -> bool {
match self {
Self::Available | Self::Unused | Self::Skipped => false,
Self::Pending => started_at_unix_secs.is_some(),
Self::Streaming | Self::Success | Self::Failed | Self::Cancelled => true,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredRequestCandidate {
pub id: String,
pub request_id: String,
pub user_id: Option<String>,
pub api_key_id: Option<String>,
pub username: Option<String>,
pub api_key_name: Option<String>,
pub candidate_index: u32,
pub retry_index: u32,
pub provider_id: Option<String>,
pub endpoint_id: Option<String>,
pub key_id: Option<String>,
pub status: RequestCandidateStatus,
pub skip_reason: Option<String>,
pub is_cached: bool,
pub status_code: Option<u16>,
pub error_type: Option<String>,
pub error_message: Option<String>,
pub latency_ms: Option<u64>,
pub concurrent_requests: Option<u32>,
pub extra_data: Option<serde_json::Value>,
pub required_capabilities: Option<serde_json::Value>,
pub created_at_unix_secs: u64,
pub started_at_unix_secs: Option<u64>,
pub finished_at_unix_secs: Option<u64>,
}
impl StoredRequestCandidate {
#[allow(clippy::too_many_arguments)]
pub fn new(
id: String,
request_id: String,
user_id: Option<String>,
api_key_id: Option<String>,
username: Option<String>,
api_key_name: Option<String>,
candidate_index: i32,
retry_index: i32,
provider_id: Option<String>,
endpoint_id: Option<String>,
key_id: Option<String>,
status: RequestCandidateStatus,
skip_reason: Option<String>,
is_cached: bool,
status_code: Option<i32>,
error_type: Option<String>,
error_message: Option<String>,
latency_ms: Option<i32>,
concurrent_requests: Option<i32>,
extra_data: Option<serde_json::Value>,
required_capabilities: Option<serde_json::Value>,
created_at_unix_secs: i64,
started_at_unix_secs: Option<i64>,
finished_at_unix_secs: Option<i64>,
) -> Result<Self, crate::DataLayerError> {
let candidate_index = u32::try_from(candidate_index).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid request_candidates.candidate_index: {candidate_index}"
))
})?;
let retry_index = u32::try_from(retry_index).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid request_candidates.retry_index: {retry_index}"
))
})?;
let status_code = status_code
.map(|value| {
u16::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid request_candidates.status_code: {value}"
))
})
})
.transpose()?;
let latency_ms = latency_ms
.map(|value| {
u64::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid request_candidates.latency_ms: {value}"
))
})
})
.transpose()?;
let concurrent_requests = concurrent_requests
.map(|value| {
u32::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid request_candidates.concurrent_requests: {value}"
))
})
})
.transpose()?;
let created_at_unix_secs = u64::try_from(created_at_unix_secs).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid request_candidates.created_at_unix_secs: {created_at_unix_secs}"
))
})?;
let started_at_unix_secs = started_at_unix_secs
.map(|value| {
u64::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid request_candidates.started_at_unix_secs: {value}"
))
})
})
.transpose()?;
let finished_at_unix_secs = finished_at_unix_secs
.map(|value| {
u64::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid request_candidates.finished_at_unix_secs: {value}"
))
})
})
.transpose()?;
Ok(Self {
id,
request_id,
user_id,
api_key_id,
username,
api_key_name,
candidate_index,
retry_index,
provider_id,
endpoint_id,
key_id,
status,
skip_reason,
is_cached,
status_code,
error_type,
error_message,
latency_ms,
concurrent_requests,
extra_data,
required_capabilities,
created_at_unix_secs,
started_at_unix_secs,
finished_at_unix_secs,
})
}
}
#[async_trait]
pub trait RequestCandidateReadRepository: Send + Sync {
async fn list_by_request_id(
&self,
request_id: &str,
) -> Result<Vec<StoredRequestCandidate>, crate::DataLayerError>;
async fn list_recent(
&self,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, crate::DataLayerError>;
}
pub trait RequestCandidateRepository: RequestCandidateReadRepository + Send + Sync {}
impl<T> RequestCandidateRepository for T where T: RequestCandidateReadRepository + Send + Sync {}
#[cfg(test)]
mod tests {
use super::{RequestCandidateStatus, StoredRequestCandidate};
#[test]
fn parses_status_from_database_text() {
assert_eq!(
RequestCandidateStatus::from_database("streaming").expect("status should parse"),
RequestCandidateStatus::Streaming
);
}
#[test]
fn rejects_invalid_database_status() {
assert!(RequestCandidateStatus::from_database("mystery").is_err());
}
#[test]
fn rejects_negative_candidate_index() {
assert!(StoredRequestCandidate::new(
"cand-1".to_string(),
"req-1".to_string(),
None,
None,
None,
None,
-1,
0,
None,
None,
None,
RequestCandidateStatus::Pending,
None,
false,
Some(200),
None,
None,
Some(10),
Some(1),
None,
None,
100,
None,
None,
)
.is_err());
}
#[test]
fn rejects_negative_created_at() {
assert!(StoredRequestCandidate::new(
"cand-1".to_string(),
"req-1".to_string(),
None,
None,
None,
None,
0,
0,
None,
None,
None,
RequestCandidateStatus::Pending,
None,
false,
Some(200),
None,
None,
Some(10),
Some(1),
None,
None,
-1,
None,
None,
)
.is_err());
}
#[test]
fn pending_without_started_at_is_not_attempted() {
assert!(!RequestCandidateStatus::Pending.is_attempted(None));
assert!(RequestCandidateStatus::Pending.is_attempted(Some(1)));
}
}

6
crates/aether-data/src/repository/mod.rs Normal file
View File

@@ -0,0 +1,6 @@
pub mod auth;
pub mod candidates;
pub mod provider_catalog;
pub mod shadow_results;
pub mod usage;
pub mod video_tasks;

157
crates/aether-data/src/repository/provider_catalog/memory.rs Normal file
View File

@@ -0,0 +1,157 @@
use std::collections::BTreeMap;
use std::sync::RwLock;
use async_trait::async_trait;
use super::types::{
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogProvider,
};
use crate::DataLayerError;
#[derive(Debug, Default)]
struct MemoryProviderCatalogIndex {
providers: BTreeMap<String, StoredProviderCatalogProvider>,
endpoints: BTreeMap<String, StoredProviderCatalogEndpoint>,
keys: BTreeMap<String, StoredProviderCatalogKey>,
}
#[derive(Debug, Default)]
pub struct InMemoryProviderCatalogReadRepository {
index: RwLock<MemoryProviderCatalogIndex>,
}
impl InMemoryProviderCatalogReadRepository {
pub fn seed(
providers: Vec<StoredProviderCatalogProvider>,
endpoints: Vec<StoredProviderCatalogEndpoint>,
keys: Vec<StoredProviderCatalogKey>,
) -> Self {
Self {
index: RwLock::new(MemoryProviderCatalogIndex {
providers: providers
.into_iter()
.map(|provider| (provider.id.clone(), provider))
.collect(),
endpoints: endpoints
.into_iter()
.map(|endpoint| (endpoint.id.clone(), endpoint))
.collect(),
keys: keys.into_iter().map(|key| (key.id.clone(), key)).collect(),
}),
}
}
}
#[async_trait]
impl ProviderCatalogReadRepository for InMemoryProviderCatalogReadRepository {
async fn list_providers_by_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
let index = self.index.read().expect("provider catalog repository lock");
Ok(provider_ids
.iter()
.filter_map(|id| index.providers.get(id).cloned())
.collect())
}
async fn list_endpoints_by_ids(
&self,
endpoint_ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
let index = self.index.read().expect("provider catalog repository lock");
Ok(endpoint_ids
.iter()
.filter_map(|id| index.endpoints.get(id).cloned())
.collect())
}
async fn list_keys_by_ids(
&self,
key_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
let index = self.index.read().expect("provider catalog repository lock");
Ok(key_ids
.iter()
.filter_map(|id| index.keys.get(id).cloned())
.collect())
}
}
#[cfg(test)]
mod tests {
use super::InMemoryProviderCatalogReadRepository;
use crate::repository::provider_catalog::{
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogProvider,
};
fn sample_provider(id: &str) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
id.to_string(),
format!("provider-{id}"),
Some("https://example.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
}
fn sample_endpoint(id: &str, provider_id: &str) -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
id.to_string(),
provider_id.to_string(),
"openai:chat".to_string(),
Some("openai".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
}
fn sample_key(id: &str, provider_id: &str) -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
id.to_string(),
provider_id.to_string(),
"default".to_string(),
"api_key".to_string(),
Some(serde_json::json!({"cache_1h": true})),
true,
)
.expect("key should build")
}
#[tokio::test]
async fn reads_provider_catalog_items_by_id() {
let repository = InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-1")],
vec![sample_endpoint("endpoint-1", "provider-1")],
vec![sample_key("key-1", "provider-1")],
);
assert_eq!(
repository
.list_providers_by_ids(&["provider-1".to_string()])
.await
.expect("providers should read")
.len(),
1
);
assert_eq!(
repository
.list_endpoints_by_ids(&["endpoint-1".to_string()])
.await
.expect("endpoints should read")
.len(),
1
);
assert_eq!(
repository
.list_keys_by_ids(&["key-1".to_string()])
.await
.expect("keys should read")
.len(),
1
);
}
}

10
crates/aether-data/src/repository/provider_catalog/mod.rs Normal file
View File

@@ -0,0 +1,10 @@
mod memory;
mod sql;
mod types;
pub use memory::InMemoryProviderCatalogReadRepository;
pub use sql::SqlxProviderCatalogReadRepository;
pub use types::{
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogProvider,
};

209
crates/aether-data/src/repository/provider_catalog/sql.rs Normal file
View File

@@ -0,0 +1,209 @@
use async_trait::async_trait;
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
use super::types::{
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogProvider,
};
use crate::DataLayerError;
const LIST_PROVIDERS_BY_IDS_PREFIX: &str = r#"
SELECT
id,
name,
website,
provider_type
FROM providers
WHERE id IN (
"#;
const LIST_ENDPOINTS_BY_IDS_PREFIX: &str = r#"
SELECT
id,
provider_id,
api_format,
api_family,
endpoint_kind,
is_active
FROM provider_endpoints
WHERE id IN (
"#;
const LIST_KEYS_BY_IDS_PREFIX: &str = r#"
SELECT
id,
provider_id,
name,
auth_type,
capabilities,
is_active
FROM provider_api_keys
WHERE id IN (
"#;
#[derive(Debug, Clone)]
pub struct SqlxProviderCatalogReadRepository {
pool: PgPool,
}
impl SqlxProviderCatalogReadRepository {
pub fn new(pool: PgPool) -> Self {
Self { pool }
}
pub fn pool(&self) -> &PgPool {
&self.pool
}
pub async fn list_providers_by_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
if provider_ids.is_empty() {
return Ok(Vec::new());
}
let rows = build_list_query(
LIST_PROVIDERS_BY_IDS_PREFIX,
provider_ids,
" ORDER BY name ASC",
)
.build()
.fetch_all(&self.pool)
.await?;
rows.iter().map(map_provider_row).collect()
}
pub async fn list_endpoints_by_ids(
&self,
endpoint_ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
if endpoint_ids.is_empty() {
return Ok(Vec::new());
}
let rows = build_list_query(
LIST_ENDPOINTS_BY_IDS_PREFIX,
endpoint_ids,
" ORDER BY api_format ASC, id ASC",
)
.build()
.fetch_all(&self.pool)
.await?;
rows.iter().map(map_endpoint_row).collect()
}
pub async fn list_keys_by_ids(
&self,
key_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
if key_ids.is_empty() {
return Ok(Vec::new());
}
let rows = build_list_query(
LIST_KEYS_BY_IDS_PREFIX,
key_ids,
" ORDER BY name ASC, id ASC",
)
.build()
.fetch_all(&self.pool)
.await?;
rows.iter().map(map_key_row).collect()
}
}
#[async_trait]
impl ProviderCatalogReadRepository for SqlxProviderCatalogReadRepository {
async fn list_providers_by_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
Self::list_providers_by_ids(self, provider_ids).await
}
async fn list_endpoints_by_ids(
&self,
endpoint_ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
Self::list_endpoints_by_ids(self, endpoint_ids).await
}
async fn list_keys_by_ids(
&self,
key_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
Self::list_keys_by_ids(self, key_ids).await
}
}
fn build_list_query<'a>(
prefix: &'static str,
ids: &'a [String],
suffix: &'static str,
) -> QueryBuilder<'a, Postgres> {
let mut builder = QueryBuilder::<Postgres>::new(prefix);
let mut separated = builder.separated(", ");
for id in ids {
separated.push_bind(id);
}
separated.push_unseparated(")");
builder.push(suffix);
builder
}
fn map_provider_row(row: &PgRow) -> Result<StoredProviderCatalogProvider, DataLayerError> {
StoredProviderCatalogProvider::new(
row.try_get("id")?,
row.try_get("name")?,
row.try_get("website")?,
row.try_get("provider_type")?,
)
}
fn map_endpoint_row(row: &PgRow) -> Result<StoredProviderCatalogEndpoint, DataLayerError> {
StoredProviderCatalogEndpoint::new(
row.try_get("id")?,
row.try_get("provider_id")?,
row.try_get("api_format")?,
row.try_get("api_family")?,
row.try_get("endpoint_kind")?,
row.try_get("is_active")?,
)
}
fn map_key_row(row: &PgRow) -> Result<StoredProviderCatalogKey, DataLayerError> {
StoredProviderCatalogKey::new(
row.try_get("id")?,
row.try_get("provider_id")?,
row.try_get("name")?,
row.try_get("auth_type")?,
row.try_get("capabilities")?,
row.try_get("is_active")?,
)
}
#[cfg(test)]
mod tests {
use super::SqlxProviderCatalogReadRepository;
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let pool = factory.connect_lazy().expect("pool should build");
let repository = SqlxProviderCatalogReadRepository::new(pool);
let _ = repository.pool();
}
}

175
crates/aether-data/src/repository/provider_catalog/types.rs Normal file
View File

@@ -0,0 +1,175 @@
use async_trait::async_trait;
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredProviderCatalogProvider {
pub id: String,
pub name: String,
pub website: Option<String>,
pub provider_type: String,
}
impl StoredProviderCatalogProvider {
pub fn new(
id: String,
name: String,
website: Option<String>,
provider_type: String,
) -> Result<Self, crate::DataLayerError> {
if name.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"providers.name is empty".to_string(),
));
}
if provider_type.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"providers.provider_type is empty".to_string(),
));
}
Ok(Self {
id,
name,
website,
provider_type,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredProviderCatalogEndpoint {
pub id: String,
pub provider_id: String,
pub api_format: String,
pub api_family: Option<String>,
pub endpoint_kind: Option<String>,
pub is_active: bool,
}
impl StoredProviderCatalogEndpoint {
pub fn new(
id: String,
provider_id: String,
api_format: String,
api_family: Option<String>,
endpoint_kind: Option<String>,
is_active: bool,
) -> Result<Self, crate::DataLayerError> {
if api_format.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"provider_endpoints.api_format is empty".to_string(),
));
}
Ok(Self {
id,
provider_id,
api_format,
api_family,
endpoint_kind,
is_active,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredProviderCatalogKey {
pub id: String,
pub provider_id: String,
pub name: String,
pub auth_type: String,
pub capabilities: Option<serde_json::Value>,
pub is_active: bool,
}
impl StoredProviderCatalogKey {
pub fn new(
id: String,
provider_id: String,
name: String,
auth_type: String,
capabilities: Option<serde_json::Value>,
is_active: bool,
) -> Result<Self, crate::DataLayerError> {
if name.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"provider_api_keys.name is empty".to_string(),
));
}
if auth_type.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"provider_api_keys.auth_type is empty".to_string(),
));
}
Ok(Self {
id,
provider_id,
name,
auth_type,
capabilities,
is_active,
})
}
}
#[async_trait]
pub trait ProviderCatalogReadRepository: Send + Sync {
async fn list_providers_by_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogProvider>, crate::DataLayerError>;
async fn list_endpoints_by_ids(
&self,
endpoint_ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, crate::DataLayerError>;
async fn list_keys_by_ids(
&self,
key_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, crate::DataLayerError>;
}
#[cfg(test)]
mod tests {
use super::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
#[test]
fn rejects_empty_provider_name() {
assert!(StoredProviderCatalogProvider::new(
"provider-1".to_string(),
"".to_string(),
None,
"custom".to_string(),
)
.is_err());
}
#[test]
fn rejects_empty_endpoint_api_format() {
assert!(StoredProviderCatalogEndpoint::new(
"endpoint-1".to_string(),
"provider-1".to_string(),
"".to_string(),
None,
None,
true,
)
.is_err());
}
#[test]
fn rejects_empty_key_auth_type() {
assert!(StoredProviderCatalogKey::new(
"key-1".to_string(),
"provider-1".to_string(),
"default".to_string(),
"".to_string(),
None,
true,
)
.is_err());
}
}

170
crates/aether-data/src/repository/shadow_results/memory.rs Normal file
View File

@@ -0,0 +1,170 @@
use std::collections::BTreeMap;
use std::sync::RwLock;
use async_trait::async_trait;
use super::types::{
ShadowResultLookupKey, ShadowResultReadRepository, ShadowResultWriteRepository,
StoredShadowResult, UpsertShadowResult,
};
use crate::DataLayerError;
#[derive(Debug, Default)]
pub struct InMemoryShadowResultRepository {
results: RwLock<BTreeMap<(String, String), StoredShadowResult>>,
}
#[async_trait]
impl ShadowResultReadRepository for InMemoryShadowResultRepository {
async fn find(
&self,
key: ShadowResultLookupKey<'_>,
) -> Result<Option<StoredShadowResult>, DataLayerError> {
let results = self.results.read().expect("shadow result repository lock");
Ok(match key {
ShadowResultLookupKey::TraceFingerprint {
trace_id,
request_fingerprint,
} => results
.get(&(trace_id.to_string(), request_fingerprint.to_string()))
.cloned(),
})
}
async fn list_recent(&self, limit: usize) -> Result<Vec<StoredShadowResult>, DataLayerError> {
if limit == 0 {
return Ok(Vec::new());
}
let mut results = self
.results
.read()
.expect("shadow result repository lock")
.values()
.cloned()
.collect::<Vec<_>>();
results.sort_by(|left, right| right.updated_at_unix_secs.cmp(&left.updated_at_unix_secs));
results.truncate(limit);
Ok(results)
}
}
#[async_trait]
impl ShadowResultWriteRepository for InMemoryShadowResultRepository {
async fn upsert(
&self,
result: UpsertShadowResult,
) -> Result<StoredShadowResult, DataLayerError> {
let stored = result.into_stored();
let mut results = self.results.write().expect("shadow result repository lock");
results.insert(
(stored.trace_id.clone(), stored.request_fingerprint.clone()),
stored.clone(),
);
Ok(stored)
}
}
#[cfg(test)]
mod tests {
use super::InMemoryShadowResultRepository;
use crate::repository::shadow_results::{
ShadowResultLookupKey, ShadowResultMatchStatus, ShadowResultReadRepository,
ShadowResultWriteRepository, UpsertShadowResult,
};
fn sample_result(
trace_id: &str,
request_fingerprint: &str,
updated_at_unix_secs: u64,
) -> UpsertShadowResult {
UpsertShadowResult {
trace_id: trace_id.to_string(),
request_fingerprint: request_fingerprint.to_string(),
request_id: Some(format!("req-{trace_id}")),
route_family: Some("openai".to_string()),
route_kind: Some("chat".to_string()),
candidate_id: Some("cand-1".to_string()),
rust_result_digest: Some("rust-digest".to_string()),
python_result_digest: Some("python-digest".to_string()),
match_status: ShadowResultMatchStatus::Match,
status_code: Some(200),
error_message: None,
created_at_unix_secs: updated_at_unix_secs.saturating_sub(10),
updated_at_unix_secs,
}
}
#[tokio::test]
async fn reads_result_by_trace_and_fingerprint() {
let repo = InMemoryShadowResultRepository::default();
repo.upsert(sample_result("trace-1", "fp-1", 100))
.await
.expect("upsert should succeed");
assert!(repo
.find(ShadowResultLookupKey::TraceFingerprint {
trace_id: "trace-1",
request_fingerprint: "fp-1",
})
.await
.expect("find should succeed")
.is_some());
}
#[tokio::test]
async fn list_recent_returns_results_in_descending_update_order() {
let repo = InMemoryShadowResultRepository::default();
repo.upsert(sample_result("trace-1", "fp-1", 100))
.await
.expect("upsert should succeed");
repo.upsert(sample_result("trace-2", "fp-2", 200))
.await
.expect("upsert should succeed");
let recent = repo
.list_recent(10)
.await
.expect("list recent should succeed");
assert_eq!(recent.len(), 2);
assert_eq!(recent[0].trace_id, "trace-2");
assert_eq!(recent[1].trace_id, "trace-1");
}
#[tokio::test]
async fn upsert_replaces_existing_shadow_result() {
let repo = InMemoryShadowResultRepository::default();
repo.upsert(sample_result("trace-1", "fp-1", 100))
.await
.expect("upsert should succeed");
repo.upsert(UpsertShadowResult {
trace_id: "trace-1".to_string(),
request_fingerprint: "fp-1".to_string(),
request_id: Some("req-trace-1".to_string()),
route_family: Some("openai".to_string()),
route_kind: Some("chat".to_string()),
candidate_id: Some("cand-2".to_string()),
rust_result_digest: Some("rust-digest-2".to_string()),
python_result_digest: Some("python-digest-2".to_string()),
match_status: ShadowResultMatchStatus::Mismatch,
status_code: Some(502),
error_message: Some("mismatch".to_string()),
created_at_unix_secs: 100,
updated_at_unix_secs: 200,
})
.await
.expect("upsert should succeed");
let stored = repo
.find(ShadowResultLookupKey::TraceFingerprint {
trace_id: "trace-1",
request_fingerprint: "fp-1",
})
.await
.expect("find should succeed")
.expect("stored result should exist");
assert_eq!(stored.request_id.as_deref(), Some("req-trace-1"));
assert_eq!(stored.candidate_id.as_deref(), Some("cand-2"));
assert_eq!(stored.match_status, ShadowResultMatchStatus::Mismatch);
}
}

12
crates/aether-data/src/repository/shadow_results/mod.rs Normal file
View File

@@ -0,0 +1,12 @@
mod memory;
mod record;
mod sql;
mod types;
pub use memory::InMemoryShadowResultRepository;
pub use record::{merge_shadow_result_sample, RecordShadowResultSample, ShadowResultSampleOrigin};
pub use sql::SqlxShadowResultRepository;
pub use types::{
ShadowResultLookupKey, ShadowResultMatchStatus, ShadowResultReadRepository,
ShadowResultRepository, ShadowResultWriteRepository, StoredShadowResult, UpsertShadowResult,
};

188
crates/aether-data/src/repository/shadow_results/record.rs Normal file
View File

@@ -0,0 +1,188 @@
use super::types::{ShadowResultMatchStatus, StoredShadowResult, UpsertShadowResult};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ShadowResultSampleOrigin {
Rust,
Python,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RecordShadowResultSample {
pub trace_id: String,
pub request_fingerprint: String,
pub request_id: Option<String>,
pub route_family: Option<String>,
pub route_kind: Option<String>,
pub candidate_id: Option<String>,
pub origin: ShadowResultSampleOrigin,
pub result_digest: String,
pub status_code: Option<u16>,
pub error_message: Option<String>,
pub recorded_at_unix_secs: u64,
}
pub fn merge_shadow_result_sample(
existing: Option<&StoredShadowResult>,
sample: RecordShadowResultSample,
) -> UpsertShadowResult {
let RecordShadowResultSample {
trace_id,
request_fingerprint,
request_id,
route_family,
route_kind,
candidate_id,
origin,
result_digest,
status_code,
error_message,
recorded_at_unix_secs,
} = sample;
let (rust_result_digest, python_result_digest) = match origin {
ShadowResultSampleOrigin::Rust => (
Some(result_digest),
existing.and_then(|row| row.python_result_digest.clone()),
),
ShadowResultSampleOrigin::Python => (
existing.and_then(|row| row.rust_result_digest.clone()),
Some(result_digest),
),
};
let match_status = resolve_match_status(
rust_result_digest.as_deref(),
python_result_digest.as_deref(),
);
UpsertShadowResult {
trace_id,
request_fingerprint,
request_id: request_id.or_else(|| existing.and_then(|row| row.request_id.clone())),
route_family: route_family.or_else(|| existing.and_then(|row| row.route_family.clone())),
route_kind: route_kind.or_else(|| existing.and_then(|row| row.route_kind.clone())),
candidate_id: candidate_id.or_else(|| existing.and_then(|row| row.candidate_id.clone())),
rust_result_digest,
python_result_digest,
match_status,
status_code: status_code.or(existing.and_then(|row| row.status_code)),
error_message: resolve_error_message(existing, error_message, match_status),
created_at_unix_secs: existing
.map(|row| row.created_at_unix_secs)
.unwrap_or(recorded_at_unix_secs),
updated_at_unix_secs: recorded_at_unix_secs,
}
}
fn resolve_match_status(
rust_result_digest: Option<&str>,
python_result_digest: Option<&str>,
) -> ShadowResultMatchStatus {
match (rust_result_digest, python_result_digest) {
(Some(rust_digest), Some(python_digest)) if rust_digest == python_digest => {
ShadowResultMatchStatus::Match
}
(Some(_), Some(_)) => ShadowResultMatchStatus::Mismatch,
_ => ShadowResultMatchStatus::Pending,
}
}
fn resolve_error_message(
existing: Option<&StoredShadowResult>,
error_message: Option<String>,
match_status: ShadowResultMatchStatus,
) -> Option<String> {
if match_status == ShadowResultMatchStatus::Mismatch {
error_message
.or_else(|| existing.and_then(|row| row.error_message.clone()))
.or_else(|| Some("shadow result digest mismatch".to_string()))
} else {
error_message.or_else(|| existing.and_then(|row| row.error_message.clone()))
}
}
#[cfg(test)]
mod tests {
use super::{merge_shadow_result_sample, RecordShadowResultSample, ShadowResultSampleOrigin};
use crate::repository::shadow_results::{ShadowResultMatchStatus, UpsertShadowResult};
fn rust_sample(result_digest: &str, recorded_at_unix_secs: u64) -> RecordShadowResultSample {
RecordShadowResultSample {
trace_id: "trace-1".to_string(),
request_fingerprint: "fp-1".to_string(),
request_id: Some("req-1".to_string()),
route_family: Some("openai".to_string()),
route_kind: Some("chat".to_string()),
candidate_id: None,
origin: ShadowResultSampleOrigin::Rust,
result_digest: result_digest.to_string(),
status_code: Some(200),
error_message: None,
recorded_at_unix_secs,
}
}
fn python_sample(result_digest: &str, recorded_at_unix_secs: u64) -> RecordShadowResultSample {
RecordShadowResultSample {
trace_id: "trace-1".to_string(),
request_fingerprint: "fp-1".to_string(),
request_id: Some("req-1".to_string()),
route_family: Some("openai".to_string()),
route_kind: Some("chat".to_string()),
candidate_id: None,
origin: ShadowResultSampleOrigin::Python,
result_digest: result_digest.to_string(),
status_code: Some(200),
error_message: None,
recorded_at_unix_secs,
}
}
fn stored(upsert: UpsertShadowResult) -> crate::repository::shadow_results::StoredShadowResult {
upsert.into_stored()
}
#[test]
fn keeps_pending_until_both_samples_exist() {
let merged = merge_shadow_result_sample(None, rust_sample("digest-1", 100));
assert_eq!(merged.match_status, ShadowResultMatchStatus::Pending);
assert_eq!(merged.request_id.as_deref(), Some("req-1"));
assert_eq!(merged.rust_result_digest.as_deref(), Some("digest-1"));
assert!(merged.python_result_digest.is_none());
}
#[test]
fn marks_match_when_rust_and_python_digests_are_equal() {
let existing = stored(merge_shadow_result_sample(
None,
rust_sample("digest-1", 100),
));
let merged = merge_shadow_result_sample(Some(&existing), python_sample("digest-1", 200));
assert_eq!(merged.match_status, ShadowResultMatchStatus::Match);
assert_eq!(merged.created_at_unix_secs, 100);
assert_eq!(merged.updated_at_unix_secs, 200);
assert_eq!(merged.request_id.as_deref(), Some("req-1"));
assert_eq!(merged.rust_result_digest.as_deref(), Some("digest-1"));
assert_eq!(merged.python_result_digest.as_deref(), Some("digest-1"));
}
#[test]
fn marks_mismatch_when_rust_and_python_digests_differ() {
let existing = stored(merge_shadow_result_sample(
None,
rust_sample("digest-1", 100),
));
let merged = merge_shadow_result_sample(Some(&existing), python_sample("digest-2", 200));
assert_eq!(merged.match_status, ShadowResultMatchStatus::Mismatch);
assert_eq!(
merged.error_message.as_deref(),
Some("shadow result digest mismatch")
);
assert_eq!(merged.request_id.as_deref(), Some("req-1"));
assert_eq!(merged.rust_result_digest.as_deref(), Some("digest-1"));
assert_eq!(merged.python_result_digest.as_deref(), Some("digest-2"));
}
}

283
crates/aether-data/src/repository/shadow_results/sql.rs Normal file
View File

@@ -0,0 +1,283 @@
use async_trait::async_trait;
use futures_util::future::BoxFuture;
use sqlx::{PgPool, Row};
use super::types::{
ShadowResultLookupKey, ShadowResultMatchStatus, ShadowResultReadRepository,
ShadowResultWriteRepository, StoredShadowResult, UpsertShadowResult,
};
use crate::postgres::PostgresTransactionRunner;
use crate::DataLayerError;
const FIND_BY_TRACE_FINGERPRINT_SQL: &str = r#"
SELECT
trace_id,
request_fingerprint,
NULL::TEXT AS request_id,
route_family,
route_kind,
candidate_id,
rust_result_digest,
python_result_digest,
match_status,
status_code,
error_message,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
FROM gateway_shadow_results
WHERE trace_id = $1 AND request_fingerprint = $2
LIMIT 1
"#;
const LIST_RECENT_SQL: &str = r#"
SELECT
trace_id,
request_fingerprint,
NULL::TEXT AS request_id,
route_family,
route_kind,
candidate_id,
rust_result_digest,
python_result_digest,
match_status,
status_code,
error_message,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
FROM gateway_shadow_results
ORDER BY updated_at DESC
LIMIT $1
"#;
const UPSERT_SQL: &str = r#"
INSERT INTO gateway_shadow_results (
trace_id,
request_fingerprint,
route_family,
route_kind,
candidate_id,
rust_result_digest,
python_result_digest,
match_status,
status_code,
error_message,
created_at,
updated_at
) VALUES (
$1,
$2,
$3,
$4,
$5,
$6,
$7,
$8,
$9,
$10,
TO_TIMESTAMP($11::double precision),
TO_TIMESTAMP($12::double precision)
)
ON CONFLICT (trace_id, request_fingerprint)
DO UPDATE SET
route_family = EXCLUDED.route_family,
route_kind = EXCLUDED.route_kind,
candidate_id = EXCLUDED.candidate_id,
rust_result_digest = EXCLUDED.rust_result_digest,
python_result_digest = EXCLUDED.python_result_digest,
match_status = EXCLUDED.match_status,
status_code = EXCLUDED.status_code,
error_message = EXCLUDED.error_message,
updated_at = EXCLUDED.updated_at
RETURNING
trace_id,
request_fingerprint,
NULL::TEXT AS request_id,
route_family,
route_kind,
candidate_id,
rust_result_digest,
python_result_digest,
match_status,
status_code,
error_message,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
"#;
#[derive(Debug, Clone)]
pub struct SqlxShadowResultRepository {
pool: PgPool,
tx_runner: PostgresTransactionRunner,
}
impl SqlxShadowResultRepository {
pub fn new(pool: PgPool) -> Self {
let tx_runner = PostgresTransactionRunner::new(pool.clone());
Self { pool, tx_runner }
}
pub fn pool(&self) -> &PgPool {
&self.pool
}
pub fn transaction_runner(&self) -> &PostgresTransactionRunner {
&self.tx_runner
}
pub async fn find(
&self,
key: ShadowResultLookupKey<'_>,
) -> Result<Option<StoredShadowResult>, DataLayerError> {
match key {
ShadowResultLookupKey::TraceFingerprint {
trace_id,
request_fingerprint,
} => {
self.find_by_trace_fingerprint(trace_id, request_fingerprint)
.await
}
}
}
pub async fn find_by_trace_fingerprint(
&self,
trace_id: &str,
request_fingerprint: &str,
) -> Result<Option<StoredShadowResult>, DataLayerError> {
let row = sqlx::query(FIND_BY_TRACE_FINGERPRINT_SQL)
.bind(trace_id)
.bind(request_fingerprint)
.fetch_optional(&self.pool)
.await?;
row.as_ref().map(map_shadow_result_row).transpose()
}
pub async fn list_recent(
&self,
limit: usize,
) -> Result<Vec<StoredShadowResult>, DataLayerError> {
if limit == 0 {
return Ok(Vec::new());
}
let rows = sqlx::query(LIST_RECENT_SQL)
.bind(i64::try_from(limit).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid recent shadow result limit: {limit}"
))
})?)
.fetch_all(&self.pool)
.await?;
rows.iter().map(map_shadow_result_row).collect()
}
pub async fn upsert(
&self,
result: UpsertShadowResult,
) -> Result<StoredShadowResult, DataLayerError> {
self.tx_runner
.run_read_write(|tx| {
Box::pin(async move {
let row = sqlx::query(UPSERT_SQL)
.bind(&result.trace_id)
.bind(&result.request_fingerprint)
.bind(&result.route_family)
.bind(&result.route_kind)
.bind(&result.candidate_id)
.bind(&result.rust_result_digest)
.bind(&result.python_result_digest)
.bind(match_status_to_database(result.match_status))
.bind(result.status_code.map(i32::from))
.bind(&result.error_message)
.bind(result.created_at_unix_secs as f64)
.bind(result.updated_at_unix_secs as f64)
.fetch_one(&mut **tx)
.await?;
map_shadow_result_row(&row)
}) as BoxFuture<'_, Result<StoredShadowResult, DataLayerError>>
})
.await
}
}
#[async_trait]
impl ShadowResultReadRepository for SqlxShadowResultRepository {
async fn find(
&self,
key: ShadowResultLookupKey<'_>,
) -> Result<Option<StoredShadowResult>, DataLayerError> {
Self::find(self, key).await
}
async fn list_recent(&self, limit: usize) -> Result<Vec<StoredShadowResult>, DataLayerError> {
Self::list_recent(self, limit).await
}
}
#[async_trait]
impl ShadowResultWriteRepository for SqlxShadowResultRepository {
async fn upsert(
&self,
result: UpsertShadowResult,
) -> Result<StoredShadowResult, DataLayerError> {
Self::upsert(self, result).await
}
}
fn match_status_to_database(status: ShadowResultMatchStatus) -> &'static str {
match status {
ShadowResultMatchStatus::Pending => "pending",
ShadowResultMatchStatus::Match => "match",
ShadowResultMatchStatus::Mismatch => "mismatch",
ShadowResultMatchStatus::Error => "error",
}
}
fn map_shadow_result_row(
row: &sqlx::postgres::PgRow,
) -> Result<StoredShadowResult, DataLayerError> {
let match_status =
ShadowResultMatchStatus::from_database(row.try_get::<String, _>("match_status")?.as_str())?;
StoredShadowResult::new(
row.try_get("trace_id")?,
row.try_get("request_fingerprint")?,
row.try_get("request_id")?,
row.try_get("route_family")?,
row.try_get("route_kind")?,
row.try_get("candidate_id")?,
row.try_get("rust_result_digest")?,
row.try_get("python_result_digest")?,
match_status,
row.try_get("status_code")?,
row.try_get("error_message")?,
row.try_get("created_at_unix_secs")?,
row.try_get("updated_at_unix_secs")?,
)
}
#[cfg(test)]
mod tests {
use super::SqlxShadowResultRepository;
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let pool = factory.connect_lazy().expect("pool should build");
let repository = SqlxShadowResultRepository::new(pool);
let _ = repository.pool();
let _ = repository.transaction_runner();
}
}

227
crates/aether-data/src/repository/shadow_results/types.rs Normal file
View File

@@ -0,0 +1,227 @@
use async_trait::async_trait;
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum ShadowResultMatchStatus {
Pending,
Match,
Mismatch,
Error,
}
impl ShadowResultMatchStatus {
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
match value.trim().to_ascii_lowercase().as_str() {
"pending" => Ok(Self::Pending),
"match" => Ok(Self::Match),
"mismatch" => Ok(Self::Mismatch),
"error" => Ok(Self::Error),
other => Err(crate::DataLayerError::UnexpectedValue(format!(
"unsupported gateway_shadow_results.match_status: {other}"
))),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredShadowResult {
pub trace_id: String,
pub request_fingerprint: String,
pub request_id: Option<String>,
pub route_family: Option<String>,
pub route_kind: Option<String>,
pub candidate_id: Option<String>,
pub rust_result_digest: Option<String>,
pub python_result_digest: Option<String>,
pub match_status: ShadowResultMatchStatus,
pub status_code: Option<u16>,
pub error_message: Option<String>,
pub created_at_unix_secs: u64,
pub updated_at_unix_secs: u64,
}
impl StoredShadowResult {
#[allow(clippy::too_many_arguments)]
pub fn new(
trace_id: String,
request_fingerprint: String,
request_id: Option<String>,
route_family: Option<String>,
route_kind: Option<String>,
candidate_id: Option<String>,
rust_result_digest: Option<String>,
python_result_digest: Option<String>,
match_status: ShadowResultMatchStatus,
status_code: Option<i32>,
error_message: Option<String>,
created_at_unix_secs: i64,
updated_at_unix_secs: i64,
) -> Result<Self, crate::DataLayerError> {
let status_code = status_code
.map(|value| {
u16::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("invalid status_code: {value}"))
})
})
.transpose()?;
let created_at_unix_secs = u64::try_from(created_at_unix_secs).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid created_at_unix_secs: {created_at_unix_secs}"
))
})?;
let updated_at_unix_secs = u64::try_from(updated_at_unix_secs).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid updated_at_unix_secs: {updated_at_unix_secs}"
))
})?;
Ok(Self {
trace_id,
request_fingerprint,
request_id,
route_family,
route_kind,
candidate_id,
rust_result_digest,
python_result_digest,
match_status,
status_code,
error_message,
created_at_unix_secs,
updated_at_unix_secs,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpsertShadowResult {
pub trace_id: String,
pub request_fingerprint: String,
pub request_id: Option<String>,
pub route_family: Option<String>,
pub route_kind: Option<String>,
pub candidate_id: Option<String>,
pub rust_result_digest: Option<String>,
pub python_result_digest: Option<String>,
pub match_status: ShadowResultMatchStatus,
pub status_code: Option<u16>,
pub error_message: Option<String>,
pub created_at_unix_secs: u64,
pub updated_at_unix_secs: u64,
}
impl UpsertShadowResult {
pub fn into_stored(self) -> StoredShadowResult {
StoredShadowResult {
trace_id: self.trace_id,
request_fingerprint: self.request_fingerprint,
request_id: self.request_id,
route_family: self.route_family,
route_kind: self.route_kind,
candidate_id: self.candidate_id,
rust_result_digest: self.rust_result_digest,
python_result_digest: self.python_result_digest,
match_status: self.match_status,
status_code: self.status_code,
error_message: self.error_message,
created_at_unix_secs: self.created_at_unix_secs,
updated_at_unix_secs: self.updated_at_unix_secs,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ShadowResultLookupKey<'a> {
TraceFingerprint {
trace_id: &'a str,
request_fingerprint: &'a str,
},
}
#[async_trait]
pub trait ShadowResultReadRepository: Send + Sync {
async fn find(
&self,
key: ShadowResultLookupKey<'_>,
) -> Result<Option<StoredShadowResult>, crate::DataLayerError>;
async fn list_recent(
&self,
limit: usize,
) -> Result<Vec<StoredShadowResult>, crate::DataLayerError>;
}
#[async_trait]
pub trait ShadowResultWriteRepository: Send + Sync {
async fn upsert(
&self,
result: UpsertShadowResult,
) -> Result<StoredShadowResult, crate::DataLayerError>;
}
pub trait ShadowResultRepository:
ShadowResultReadRepository + ShadowResultWriteRepository + Send + Sync
{
}
impl<T> ShadowResultRepository for T where
T: ShadowResultReadRepository + ShadowResultWriteRepository + Send + Sync
{
}
#[cfg(test)]
mod tests {
use super::{ShadowResultMatchStatus, StoredShadowResult};
#[test]
fn parses_match_status_from_database_text() {
assert_eq!(
ShadowResultMatchStatus::from_database("match").expect("status should parse"),
ShadowResultMatchStatus::Match
);
}
#[test]
fn rejects_invalid_database_status() {
assert!(ShadowResultMatchStatus::from_database("mystery").is_err());
}
#[test]
fn rejects_invalid_numeric_fields() {
assert!(StoredShadowResult::new(
"trace-1".to_string(),
"fp-1".to_string(),
None,
None,
None,
None,
None,
None,
ShadowResultMatchStatus::Pending,
Some(-1),
None,
1,
1,
)
.is_err());
}
#[test]
fn rejects_negative_updated_at_values() {
assert!(StoredShadowResult::new(
"trace-1".to_string(),
"fp-1".to_string(),
None,
None,
None,
None,
None,
None,
ShadowResultMatchStatus::Pending,
Some(200),
None,
1,
-1,
)
.is_err());
}
}

107
crates/aether-data/src/repository/usage/memory.rs Normal file
View File

@@ -0,0 +1,107 @@
use std::collections::BTreeMap;
use std::sync::RwLock;
use async_trait::async_trait;
use super::types::{StoredRequestUsageAudit, UsageReadRepository};
use crate::DataLayerError;
#[derive(Debug, Default)]
pub struct InMemoryUsageReadRepository {
by_request_id: RwLock<BTreeMap<String, StoredRequestUsageAudit>>,
}
impl InMemoryUsageReadRepository {
pub fn seed<I>(items: I) -> Self
where
I: IntoIterator<Item = StoredRequestUsageAudit>,
{
let mut by_request_id = BTreeMap::new();
for item in items {
by_request_id.insert(item.request_id.clone(), item);
}
Self {
by_request_id: RwLock::new(by_request_id),
}
}
}
#[async_trait]
impl UsageReadRepository for InMemoryUsageReadRepository {
async fn find_by_request_id(
&self,
request_id: &str,
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
Ok(self
.by_request_id
.read()
.expect("usage repository lock")
.get(request_id)
.cloned())
}
}
#[cfg(test)]
mod tests {
use super::InMemoryUsageReadRepository;
use crate::repository::usage::{StoredRequestUsageAudit, UsageReadRepository};
fn sample_usage(request_id: &str, created_at_unix_secs: i64) -> StoredRequestUsageAudit {
StoredRequestUsageAudit::new(
"usage-1".to_string(),
request_id.to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
"OpenAI".to_string(),
"gpt-4.1".to_string(),
Some("gpt-4.1-mini".to_string()),
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
true,
false,
100,
50,
150,
0.12,
0.18,
Some(200),
None,
None,
Some(420),
Some(120),
"completed".to_string(),
"settled".to_string(),
created_at_unix_secs,
created_at_unix_secs + 1,
Some(created_at_unix_secs + 2),
)
.expect("usage should build")
}
#[tokio::test]
async fn finds_usage_by_request_id() {
let repository = InMemoryUsageReadRepository::seed(vec![
sample_usage("req-1", 100),
sample_usage("req-2", 200),
]);
let usage = repository
.find_by_request_id("req-2")
.await
.expect("find should succeed")
.expect("usage should exist");
assert_eq!(usage.request_id, "req-2");
assert_eq!(usage.total_tokens, 150);
}
}

7
crates/aether-data/src/repository/usage/mod.rs Normal file
View File

@@ -0,0 +1,7 @@
mod memory;
mod sql;
mod types;
pub use memory::InMemoryUsageReadRepository;
pub use sql::SqlxUsageReadRepository;
pub use types::{StoredRequestUsageAudit, UsageReadRepository, UsageRepository};

150
crates/aether-data/src/repository/usage/sql.rs Normal file
View File

@@ -0,0 +1,150 @@
use async_trait::async_trait;
use sqlx::{PgPool, Row};
use super::types::{StoredRequestUsageAudit, UsageReadRepository};
use crate::DataLayerError;
const FIND_BY_REQUEST_ID_SQL: &str = r#"
SELECT
id,
request_id,
user_id,
api_key_id,
username,
api_key_name,
provider_name,
model,
target_model,
provider_id,
provider_endpoint_id,
provider_api_key_id,
request_type,
api_format,
api_family,
endpoint_kind,
endpoint_api_format,
provider_api_family,
provider_endpoint_kind,
COALESCE(has_format_conversion, FALSE) AS has_format_conversion,
COALESCE(is_stream, FALSE) AS is_stream,
input_tokens,
output_tokens,
total_tokens,
COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) AS total_cost_usd,
COALESCE(CAST(actual_total_cost_usd AS DOUBLE PRECISION), 0) AS actual_total_cost_usd,
status_code,
error_message,
error_category,
response_time_ms,
first_byte_time_ms,
status,
billing_status,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs,
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
FROM "usage"
WHERE request_id = $1
LIMIT 1
"#;
#[derive(Debug, Clone)]
pub struct SqlxUsageReadRepository {
pool: PgPool,
}
impl SqlxUsageReadRepository {
pub fn new(pool: PgPool) -> Self {
Self { pool }
}
pub fn pool(&self) -> &PgPool {
&self.pool
}
pub async fn find_by_request_id(
&self,
request_id: &str,
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
let row = sqlx::query(FIND_BY_REQUEST_ID_SQL)
.bind(request_id)
.fetch_optional(&self.pool)
.await?;
row.as_ref().map(map_usage_row).transpose()
}
}
#[async_trait]
impl UsageReadRepository for SqlxUsageReadRepository {
async fn find_by_request_id(
&self,
request_id: &str,
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
Self::find_by_request_id(self, request_id).await
}
}
fn map_usage_row(row: &sqlx::postgres::PgRow) -> Result<StoredRequestUsageAudit, DataLayerError> {
StoredRequestUsageAudit::new(
row.try_get("id")?,
row.try_get("request_id")?,
row.try_get("user_id")?,
row.try_get("api_key_id")?,
row.try_get("username")?,
row.try_get("api_key_name")?,
row.try_get("provider_name")?,
row.try_get("model")?,
row.try_get("target_model")?,
row.try_get("provider_id")?,
row.try_get("provider_endpoint_id")?,
row.try_get("provider_api_key_id")?,
row.try_get("request_type")?,
row.try_get("api_format")?,
row.try_get("api_family")?,
row.try_get("endpoint_kind")?,
row.try_get("endpoint_api_format")?,
row.try_get("provider_api_family")?,
row.try_get("provider_endpoint_kind")?,
row.try_get("has_format_conversion")?,
row.try_get("is_stream")?,
row.try_get("input_tokens")?,
row.try_get("output_tokens")?,
row.try_get("total_tokens")?,
row.try_get("total_cost_usd")?,
row.try_get("actual_total_cost_usd")?,
row.try_get("status_code")?,
row.try_get("error_message")?,
row.try_get("error_category")?,
row.try_get("response_time_ms")?,
row.try_get("first_byte_time_ms")?,
row.try_get("status")?,
row.try_get("billing_status")?,
row.try_get("created_at_unix_secs")?,
row.try_get("updated_at_unix_secs")?,
row.try_get("finalized_at_unix_secs")?,
)
}
#[cfg(test)]
mod tests {
use super::SqlxUsageReadRepository;
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let pool = factory.connect_lazy().expect("pool should build");
let repository = SqlxUsageReadRepository::new(pool);
let _ = repository.pool();
}
}

304
crates/aether-data/src/repository/usage/types.rs Normal file
View File

@@ -0,0 +1,304 @@
use async_trait::async_trait;
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredRequestUsageAudit {
pub id: String,
pub request_id: String,
pub user_id: Option<String>,
pub api_key_id: Option<String>,
pub username: Option<String>,
pub api_key_name: Option<String>,
pub provider_name: String,
pub model: String,
pub target_model: Option<String>,
pub provider_id: Option<String>,
pub provider_endpoint_id: Option<String>,
pub provider_api_key_id: Option<String>,
pub request_type: Option<String>,
pub api_format: Option<String>,
pub api_family: Option<String>,
pub endpoint_kind: Option<String>,
pub endpoint_api_format: Option<String>,
pub provider_api_family: Option<String>,
pub provider_endpoint_kind: Option<String>,
pub has_format_conversion: bool,
pub is_stream: bool,
pub input_tokens: u64,
pub output_tokens: u64,
pub total_tokens: u64,
pub total_cost_usd: f64,
pub actual_total_cost_usd: f64,
pub status_code: Option<u16>,
pub error_message: Option<String>,
pub error_category: Option<String>,
pub response_time_ms: Option<u64>,
pub first_byte_time_ms: Option<u64>,
pub status: String,
pub billing_status: String,
pub created_at_unix_secs: u64,
pub updated_at_unix_secs: u64,
pub finalized_at_unix_secs: Option<u64>,
}
impl StoredRequestUsageAudit {
#[allow(clippy::too_many_arguments)]
pub fn new(
id: String,
request_id: String,
user_id: Option<String>,
api_key_id: Option<String>,
username: Option<String>,
api_key_name: Option<String>,
provider_name: String,
model: String,
target_model: Option<String>,
provider_id: Option<String>,
provider_endpoint_id: Option<String>,
provider_api_key_id: Option<String>,
request_type: Option<String>,
api_format: Option<String>,
api_family: Option<String>,
endpoint_kind: Option<String>,
endpoint_api_format: Option<String>,
provider_api_family: Option<String>,
provider_endpoint_kind: Option<String>,
has_format_conversion: bool,
is_stream: bool,
input_tokens: i32,
output_tokens: i32,
total_tokens: i32,
total_cost_usd: f64,
actual_total_cost_usd: f64,
status_code: Option<i32>,
error_message: Option<String>,
error_category: Option<String>,
response_time_ms: Option<i32>,
first_byte_time_ms: Option<i32>,
status: String,
billing_status: String,
created_at_unix_secs: i64,
updated_at_unix_secs: i64,
finalized_at_unix_secs: Option<i64>,
) -> Result<Self, crate::DataLayerError> {
if request_id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"usage.request_id is empty".to_string(),
));
}
if provider_name.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"usage.provider_name is empty".to_string(),
));
}
if model.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"usage.model is empty".to_string(),
));
}
if status.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"usage.status is empty".to_string(),
));
}
if billing_status.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"usage.billing_status is empty".to_string(),
));
}
if !total_cost_usd.is_finite() {
return Err(crate::DataLayerError::UnexpectedValue(
"usage.total_cost_usd is not finite".to_string(),
));
}
if !actual_total_cost_usd.is_finite() {
return Err(crate::DataLayerError::UnexpectedValue(
"usage.actual_total_cost_usd is not finite".to_string(),
));
}
Ok(Self {
id,
request_id,
user_id,
api_key_id,
username,
api_key_name,
provider_name,
model,
target_model,
provider_id,
provider_endpoint_id,
provider_api_key_id,
request_type,
api_format,
api_family,
endpoint_kind,
endpoint_api_format,
provider_api_family,
provider_endpoint_kind,
has_format_conversion,
is_stream,
input_tokens: parse_u64(input_tokens, "usage.input_tokens")?,
output_tokens: parse_u64(output_tokens, "usage.output_tokens")?,
total_tokens: parse_u64(total_tokens, "usage.total_tokens")?,
total_cost_usd,
actual_total_cost_usd,
status_code: parse_u16(status_code, "usage.status_code")?,
error_message,
error_category,
response_time_ms: parse_optional_u64(response_time_ms, "usage.response_time_ms")?,
first_byte_time_ms: parse_optional_u64(first_byte_time_ms, "usage.first_byte_time_ms")?,
status,
billing_status,
created_at_unix_secs: parse_timestamp(
created_at_unix_secs,
"usage.created_at_unix_secs",
)?,
updated_at_unix_secs: parse_timestamp(
updated_at_unix_secs,
"usage.updated_at_unix_secs",
)?,
finalized_at_unix_secs: finalized_at_unix_secs
.map(|value| parse_timestamp(value, "usage.finalized_at_unix_secs"))
.transpose()?,
})
}
}
#[async_trait]
pub trait UsageReadRepository: Send + Sync {
async fn find_by_request_id(
&self,
request_id: &str,
) -> Result<Option<StoredRequestUsageAudit>, crate::DataLayerError>;
}
pub trait UsageRepository: UsageReadRepository + Send + Sync {}
impl<T> UsageRepository for T where T: UsageReadRepository + Send + Sync {}
fn parse_u64(value: i32, field_name: &str) -> Result<u64, crate::DataLayerError> {
u64::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}"))
})
}
fn parse_optional_u64(
value: Option<i32>,
field_name: &str,
) -> Result<Option<u64>, crate::DataLayerError> {
value
.map(|value| {
u64::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}"))
})
})
.transpose()
}
fn parse_u16(value: Option<i32>, field_name: &str) -> Result<Option<u16>, crate::DataLayerError> {
value
.map(|value| {
u16::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}"))
})
})
.transpose()
}
fn parse_timestamp(value: i64, field_name: &str) -> Result<u64, crate::DataLayerError> {
u64::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}"))
})
}
#[cfg(test)]
mod tests {
use super::StoredRequestUsageAudit;
#[test]
fn rejects_empty_request_id() {
assert!(StoredRequestUsageAudit::new(
"usage-1".to_string(),
"".to_string(),
None,
None,
None,
None,
"OpenAI".to_string(),
"gpt-4.1".to_string(),
None,
None,
None,
None,
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
false,
false,
10,
20,
30,
0.1,
0.1,
Some(200),
None,
None,
Some(120),
Some(80),
"completed".to_string(),
"settled".to_string(),
100,
101,
Some(102),
)
.is_err());
}
#[test]
fn rejects_negative_token_count() {
assert!(StoredRequestUsageAudit::new(
"usage-1".to_string(),
"req-1".to_string(),
None,
None,
None,
None,
"OpenAI".to_string(),
"gpt-4.1".to_string(),
None,
None,
None,
None,
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
false,
false,
-1,
20,
30,
0.1,
0.1,
Some(200),
None,
None,
Some(120),
Some(80),
"completed".to_string(),
"settled".to_string(),
100,
101,
Some(102),
)
.is_err());
}
}

232
crates/aether-data/src/repository/video_tasks/memory.rs Normal file
View File

@@ -0,0 +1,232 @@
use std::collections::BTreeMap;
use std::sync::RwLock;
use async_trait::async_trait;
use super::types::{
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository,
VideoTaskWriteRepository,
};
use crate::DataLayerError;
#[derive(Debug, Default)]
struct MemoryVideoTaskIndex {
by_id: BTreeMap<String, StoredVideoTask>,
short_to_id: BTreeMap<String, String>,
user_external_to_id: BTreeMap<(String, String), String>,
}
#[derive(Debug, Default)]
pub struct InMemoryVideoTaskRepository {
index: RwLock<MemoryVideoTaskIndex>,
}
impl InMemoryVideoTaskRepository {
fn store_locked(index: &mut MemoryVideoTaskIndex, task: StoredVideoTask) -> StoredVideoTask {
if let Some(previous) = index.by_id.insert(task.id.clone(), task.clone()) {
if let Some(short_id) = previous.short_id {
index.short_to_id.remove(&short_id);
}
if let (Some(user_id), Some(external_task_id)) =
(previous.user_id, previous.external_task_id)
{
index
.user_external_to_id
.remove(&(user_id, external_task_id));
}
}
if let Some(short_id) = &task.short_id {
index.short_to_id.insert(short_id.clone(), task.id.clone());
}
if let (Some(user_id), Some(external_task_id)) = (&task.user_id, &task.external_task_id) {
index
.user_external_to_id
.insert((user_id.clone(), external_task_id.clone()), task.id.clone());
}
task
}
}
#[async_trait]
impl VideoTaskReadRepository for InMemoryVideoTaskRepository {
async fn find(
&self,
key: VideoTaskLookupKey<'_>,
) -> Result<Option<StoredVideoTask>, DataLayerError> {
let index = self.index.read().expect("video task repository lock");
Ok(match key {
VideoTaskLookupKey::Id(id) => index.by_id.get(id).cloned(),
VideoTaskLookupKey::ShortId(short_id) => index
.short_to_id
.get(short_id)
.and_then(|id| index.by_id.get(id))
.cloned(),
VideoTaskLookupKey::UserExternal {
user_id,
external_task_id,
} => index
.user_external_to_id
.get(&(user_id.to_string(), external_task_id.to_string()))
.and_then(|id| index.by_id.get(id))
.cloned(),
})
}
async fn list_active(&self, limit: usize) -> Result<Vec<StoredVideoTask>, DataLayerError> {
if limit == 0 {
return Ok(Vec::new());
}
let mut tasks = self
.index
.read()
.expect("video task repository lock")
.by_id
.values()
.filter(|task| task.status.is_active())
.cloned()
.collect::<Vec<_>>();
tasks.sort_by(|left, right| right.updated_at_unix_secs.cmp(&left.updated_at_unix_secs));
tasks.truncate(limit);
Ok(tasks)
}
}
#[async_trait]
impl VideoTaskWriteRepository for InMemoryVideoTaskRepository {
async fn upsert(&self, task: UpsertVideoTask) -> Result<StoredVideoTask, DataLayerError> {
let mut index = self.index.write().expect("video task repository lock");
Ok(Self::store_locked(&mut index, task.into_stored()))
}
}
#[cfg(test)]
mod tests {
use super::InMemoryVideoTaskRepository;
use crate::repository::video_tasks::{
UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository, VideoTaskStatus,
VideoTaskWriteRepository,
};
fn sample_task(
id: &str,
status: VideoTaskStatus,
updated_at_unix_secs: u64,
) -> UpsertVideoTask {
UpsertVideoTask {
id: id.to_string(),
short_id: Some(format!("short-{id}")),
user_id: Some("user-1".to_string()),
external_task_id: Some(format!("ext-{id}")),
provider_api_format: Some("openai:video".to_string()),
model: Some("sora-2".to_string()),
prompt: Some("hello".to_string()),
size: Some("1280x720".to_string()),
status,
progress_percent: 0,
created_at_unix_secs: updated_at_unix_secs.saturating_sub(10),
updated_at_unix_secs,
error_code: None,
error_message: None,
video_url: None,
}
}
#[tokio::test]
async fn reads_task_by_all_supported_lookup_keys() {
let repo = InMemoryVideoTaskRepository::default();
repo.upsert(sample_task("task-1", VideoTaskStatus::Submitted, 100))
.await
.expect("upsert should succeed");
assert!(repo
.find(VideoTaskLookupKey::Id("task-1"))
.await
.expect("find by id should succeed")
.is_some());
assert!(repo
.find(VideoTaskLookupKey::ShortId("short-task-1"))
.await
.expect("find by short id should succeed")
.is_some());
assert!(repo
.find(VideoTaskLookupKey::UserExternal {
user_id: "user-1",
external_task_id: "ext-task-1",
})
.await
.expect("find by user/external should succeed")
.is_some());
}
#[tokio::test]
async fn list_active_only_returns_active_tasks_in_descending_update_order() {
let repo = InMemoryVideoTaskRepository::default();
repo.upsert(sample_task("task-1", VideoTaskStatus::Completed, 100))
.await
.expect("upsert should succeed");
repo.upsert(sample_task("task-2", VideoTaskStatus::Processing, 200))
.await
.expect("upsert should succeed");
repo.upsert(sample_task("task-3", VideoTaskStatus::Queued, 150))
.await
.expect("upsert should succeed");
let active = repo
.list_active(10)
.await
.expect("list active should succeed");
assert_eq!(active.len(), 2);
assert_eq!(active[0].id, "task-2");
assert_eq!(active[1].id, "task-3");
}
#[tokio::test]
async fn upsert_replaces_secondary_indexes() {
let repo = InMemoryVideoTaskRepository::default();
repo.upsert(sample_task("task-1", VideoTaskStatus::Submitted, 100))
.await
.expect("upsert should succeed");
repo.upsert(UpsertVideoTask {
id: "task-1".to_string(),
short_id: Some("short-task-1b".to_string()),
user_id: Some("user-2".to_string()),
external_task_id: Some("ext-task-1b".to_string()),
provider_api_format: Some("gemini:video".to_string()),
model: Some("veo-3".to_string()),
prompt: Some("remix".to_string()),
size: Some("720p".to_string()),
status: VideoTaskStatus::Processing,
progress_percent: 50,
created_at_unix_secs: 150,
updated_at_unix_secs: 200,
error_code: None,
error_message: None,
video_url: None,
})
.await
.expect("upsert should succeed");
assert!(repo
.find(VideoTaskLookupKey::ShortId("short-task-1"))
.await
.expect("find should succeed")
.is_none());
assert!(repo
.find(VideoTaskLookupKey::UserExternal {
user_id: "user-1",
external_task_id: "ext-task-1",
})
.await
.expect("find should succeed")
.is_none());
assert!(repo
.find(VideoTaskLookupKey::ShortId("short-task-1b"))
.await
.expect("find should succeed")
.is_some());
}
}

10
crates/aether-data/src/repository/video_tasks/mod.rs Normal file
View File

@@ -0,0 +1,10 @@
mod memory;
mod sql;
mod types;
pub use memory::InMemoryVideoTaskRepository;
pub use sql::SqlxVideoTaskReadRepository;
pub use types::{
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository,
VideoTaskRepository, VideoTaskStatus, VideoTaskWriteRepository,
};

259
crates/aether-data/src/repository/video_tasks/sql.rs Normal file
View File

@@ -0,0 +1,259 @@
use sqlx::{PgPool, Row};
use async_trait::async_trait;
use crate::repository::video_tasks::{
StoredVideoTask, VideoTaskLookupKey, VideoTaskReadRepository, VideoTaskStatus,
};
use crate::DataLayerError;
const FIND_BY_ID_SQL: &str = r#"
SELECT
id,
short_id,
user_id,
external_task_id,
provider_api_format,
model,
prompt,
size,
status,
progress_percent,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
,
error_code,
error_message,
video_url
FROM video_tasks
WHERE id = $1
LIMIT 1
"#;
const FIND_BY_SHORT_ID_SQL: &str = r#"
SELECT
id,
short_id,
user_id,
external_task_id,
provider_api_format,
model,
prompt,
size,
status,
progress_percent,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
,
error_code,
error_message,
video_url
FROM video_tasks
WHERE short_id = $1
LIMIT 1
"#;
const FIND_BY_USER_EXTERNAL_SQL: &str = r#"
SELECT
id,
short_id,
user_id,
external_task_id,
provider_api_format,
model,
prompt,
size,
status,
progress_percent,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
,
error_code,
error_message,
video_url
FROM video_tasks
WHERE user_id = $1 AND external_task_id = $2
LIMIT 1
"#;
const LIST_ACTIVE_SQL: &str = r#"
SELECT
id,
short_id,
user_id,
external_task_id,
provider_api_format,
model,
prompt,
size,
status,
progress_percent,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
,
error_code,
error_message,
video_url
FROM video_tasks
WHERE status = ANY($1)
ORDER BY updated_at DESC
LIMIT $2
"#;
#[derive(Debug, Clone)]
pub struct SqlxVideoTaskReadRepository {
pool: PgPool,
}
impl SqlxVideoTaskReadRepository {
pub fn new(pool: PgPool) -> Self {
Self { pool }
}
pub fn pool(&self) -> &PgPool {
&self.pool
}
pub async fn find(
&self,
key: VideoTaskLookupKey<'_>,
) -> Result<Option<StoredVideoTask>, DataLayerError> {
match key {
VideoTaskLookupKey::Id(id) => self.find_by_id(id).await,
VideoTaskLookupKey::ShortId(short_id) => self.find_by_short_id(short_id).await,
VideoTaskLookupKey::UserExternal {
user_id,
external_task_id,
} => self.find_by_user_external(user_id, external_task_id).await,
}
}
pub async fn find_by_id(&self, id: &str) -> Result<Option<StoredVideoTask>, DataLayerError> {
let row = sqlx::query(FIND_BY_ID_SQL)
.bind(id)
.fetch_optional(&self.pool)
.await?;
row.as_ref().map(map_video_task_row).transpose()
}
pub async fn find_by_short_id(
&self,
short_id: &str,
) -> Result<Option<StoredVideoTask>, DataLayerError> {
let row = sqlx::query(FIND_BY_SHORT_ID_SQL)
.bind(short_id)
.fetch_optional(&self.pool)
.await?;
row.as_ref().map(map_video_task_row).transpose()
}
pub async fn find_by_user_external(
&self,
user_id: &str,
external_task_id: &str,
) -> Result<Option<StoredVideoTask>, DataLayerError> {
let row = sqlx::query(FIND_BY_USER_EXTERNAL_SQL)
.bind(user_id)
.bind(external_task_id)
.fetch_optional(&self.pool)
.await?;
row.as_ref().map(map_video_task_row).transpose()
}
pub async fn list_active(&self, limit: usize) -> Result<Vec<StoredVideoTask>, DataLayerError> {
if limit == 0 {
return Ok(Vec::new());
}
let active_statuses = vec!["pending", "submitted", "queued", "processing"];
let rows = sqlx::query(LIST_ACTIVE_SQL)
.bind(active_statuses)
.bind(i64::try_from(limit).map_err(|_| {
DataLayerError::UnexpectedValue(format!("invalid active task limit: {limit}"))
})?)
.fetch_all(&self.pool)
.await?;
rows.iter().map(map_video_task_row).collect()
}
}
#[async_trait]
impl VideoTaskReadRepository for SqlxVideoTaskReadRepository {
async fn find(
&self,
key: VideoTaskLookupKey<'_>,
) -> Result<Option<StoredVideoTask>, DataLayerError> {
Self::find(self, key).await
}
async fn list_active(&self, limit: usize) -> Result<Vec<StoredVideoTask>, DataLayerError> {
Self::list_active(self, limit).await
}
}
fn map_video_task_row(row: &sqlx::postgres::PgRow) -> Result<StoredVideoTask, DataLayerError> {
let status = VideoTaskStatus::from_database(row.try_get::<String, _>("status")?.as_str())?;
StoredVideoTask::new(
row.try_get("id")?,
row.try_get("short_id")?,
row.try_get("user_id")?,
row.try_get("external_task_id")?,
row.try_get("provider_api_format")?,
row.try_get("model")?,
row.try_get("prompt")?,
row.try_get("size")?,
status,
row.try_get("progress_percent")?,
row.try_get("created_at_unix_secs")?,
row.try_get("updated_at_unix_secs")?,
row.try_get("error_code")?,
row.try_get("error_message")?,
row.try_get("video_url")?,
)
}
#[cfg(test)]
mod tests {
use super::SqlxVideoTaskReadRepository;
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
use crate::repository::video_tasks::{VideoTaskLookupKey, VideoTaskReadRepository};
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let pool = factory.connect_lazy().expect("pool should build");
let repository = SqlxVideoTaskReadRepository::new(pool);
let _ = repository.pool();
}
#[tokio::test]
async fn read_trait_delegates_to_sqlx_repository() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 1,
max_connections: 4,
acquire_timeout_ms: 1_000,
idle_timeout_ms: 5_000,
max_lifetime_ms: 30_000,
statement_cache_capacity: 64,
require_ssl: false,
})
.expect("factory should build");
let pool = factory.connect_lazy().expect("pool should build");
let repository = SqlxVideoTaskReadRepository::new(pool);
let _ = VideoTaskReadRepository::find(&repository, VideoTaskLookupKey::Id("task-1")).await;
}
}

277
crates/aether-data/src/repository/video_tasks/types.rs Normal file
View File

@@ -0,0 +1,277 @@
use async_trait::async_trait;
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum VideoTaskStatus {
Pending,
Submitted,
Queued,
Processing,
Completed,
Failed,
Cancelled,
Expired,
Deleted,
}
impl VideoTaskStatus {
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
match value.trim().to_ascii_lowercase().as_str() {
"pending" => Ok(Self::Pending),
"submitted" => Ok(Self::Submitted),
"queued" => Ok(Self::Queued),
"processing" => Ok(Self::Processing),
"completed" => Ok(Self::Completed),
"failed" => Ok(Self::Failed),
"cancelled" => Ok(Self::Cancelled),
"expired" => Ok(Self::Expired),
"deleted" => Ok(Self::Deleted),
other => Err(crate::DataLayerError::UnexpectedValue(format!(
"unsupported video_tasks.status: {other}"
))),
}
}
pub fn is_active(self) -> bool {
matches!(
self,
Self::Pending | Self::Submitted | Self::Queued | Self::Processing
)
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredVideoTask {
pub id: String,
pub short_id: Option<String>,
pub user_id: Option<String>,
pub external_task_id: Option<String>,
pub provider_api_format: Option<String>,
pub model: Option<String>,
pub prompt: Option<String>,
pub size: Option<String>,
pub status: VideoTaskStatus,
pub progress_percent: u16,
pub created_at_unix_secs: u64,
pub updated_at_unix_secs: u64,
pub error_code: Option<String>,
pub error_message: Option<String>,
pub video_url: Option<String>,
}
impl StoredVideoTask {
pub fn new(
id: String,
short_id: Option<String>,
user_id: Option<String>,
external_task_id: Option<String>,
provider_api_format: Option<String>,
model: Option<String>,
prompt: Option<String>,
size: Option<String>,
status: VideoTaskStatus,
progress_percent: i32,
created_at_unix_secs: i64,
updated_at_unix_secs: i64,
error_code: Option<String>,
error_message: Option<String>,
video_url: Option<String>,
) -> Result<Self, crate::DataLayerError> {
let progress_percent = u16::try_from(progress_percent).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid progress_percent: {progress_percent}"
))
})?;
let created_at_unix_secs = u64::try_from(created_at_unix_secs).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid created_at_unix_secs: {created_at_unix_secs}"
))
})?;
let updated_at_unix_secs = u64::try_from(updated_at_unix_secs).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid updated_at_unix_secs: {updated_at_unix_secs}"
))
})?;
Ok(Self {
id,
short_id,
user_id,
external_task_id,
provider_api_format,
model,
prompt,
size,
status,
progress_percent,
created_at_unix_secs,
updated_at_unix_secs,
error_code,
error_message,
video_url,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpsertVideoTask {
pub id: String,
pub short_id: Option<String>,
pub user_id: Option<String>,
pub external_task_id: Option<String>,
pub provider_api_format: Option<String>,
pub model: Option<String>,
pub prompt: Option<String>,
pub size: Option<String>,
pub status: VideoTaskStatus,
pub progress_percent: u16,
pub created_at_unix_secs: u64,
pub updated_at_unix_secs: u64,
pub error_code: Option<String>,
pub error_message: Option<String>,
pub video_url: Option<String>,
}
impl UpsertVideoTask {
pub fn into_stored(self) -> StoredVideoTask {
StoredVideoTask {
id: self.id,
short_id: self.short_id,
user_id: self.user_id,
external_task_id: self.external_task_id,
provider_api_format: self.provider_api_format,
model: self.model,
prompt: self.prompt,
size: self.size,
status: self.status,
progress_percent: self.progress_percent,
created_at_unix_secs: self.created_at_unix_secs,
updated_at_unix_secs: self.updated_at_unix_secs,
error_code: self.error_code,
error_message: self.error_message,
video_url: self.video_url,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VideoTaskLookupKey<'a> {
Id(&'a str),
ShortId(&'a str),
UserExternal {
user_id: &'a str,
external_task_id: &'a str,
},
}
#[async_trait]
pub trait VideoTaskReadRepository: Send + Sync {
async fn find(
&self,
key: VideoTaskLookupKey<'_>,
) -> Result<Option<StoredVideoTask>, crate::DataLayerError>;
async fn list_active(
&self,
limit: usize,
) -> Result<Vec<StoredVideoTask>, crate::DataLayerError>;
}
#[async_trait]
pub trait VideoTaskWriteRepository: Send + Sync {
async fn upsert(&self, task: UpsertVideoTask)
-> Result<StoredVideoTask, crate::DataLayerError>;
}
pub trait VideoTaskRepository:
VideoTaskReadRepository + VideoTaskWriteRepository + Send + Sync
{
}
impl<T> VideoTaskRepository for T where
T: VideoTaskReadRepository + VideoTaskWriteRepository + Send + Sync
{
}
#[cfg(test)]
mod tests {
use super::{StoredVideoTask, VideoTaskStatus};
#[test]
fn parses_status_from_database_text() {
assert_eq!(
VideoTaskStatus::from_database("processing").expect("status should parse"),
VideoTaskStatus::Processing
);
}
#[test]
fn rejects_invalid_database_status() {
assert!(VideoTaskStatus::from_database("mystery").is_err());
}
#[test]
fn rejects_invalid_numeric_fields() {
assert!(StoredVideoTask::new(
"task-1".to_string(),
None,
None,
None,
None,
None,
None,
None,
VideoTaskStatus::Submitted,
-1,
1,
1,
None,
None,
None
)
.is_err());
}
#[test]
fn rejects_negative_updated_at_values() {
assert!(StoredVideoTask::new(
"task-1".to_string(),
None,
None,
None,
None,
None,
None,
None,
VideoTaskStatus::Submitted,
10,
1,
-1,
None,
None,
None
)
.is_err());
}
#[test]
fn rejects_negative_created_at_values() {
assert!(StoredVideoTask::new(
"task-1".to_string(),
None,
None,
None,
None,
None,
None,
None,
VideoTaskStatus::Submitted,
10,
-1,
1,
None,
None,
None
)
.is_err());
}
}

5
crates/aether-executor/Cargo.toml
View File

@@ -8,6 +8,8 @@ description = "Rust executor scaffold for Aether request execution"
[dependencies] [dependencies]
aether-contracts.workspace = true aether-contracts.workspace = true
aether-http.workspace = true
aether-runtime.workspace = true
async-stream.workspace = true async-stream.workspace = true
axum = { version = "0.8" } axum = { version = "0.8" }
base64.workspace = true base64.workspace = true
@@ -22,6 +24,5 @@ serde.workspace = true
serde_json.workspace = true serde_json.workspace = true
thiserror.workspace = true thiserror.workspace = true
tokio.workspace = true tokio.workspace = true
tracing = "0.1" tracing.workspace = true
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
webpki-roots.workspace = true webpki-roots.workspace = true

6
crates/aether-executor/src/error.rs
View File

@@ -39,6 +39,12 @@ pub enum ExecutorServiceError {
BodyEncode(serde_json::Error), BodyEncode(serde_json::Error),
#[error("failed to build HTTP client: {0}")] #[error("failed to build HTTP client: {0}")]
ClientBuild(reqwest::Error), ClientBuild(reqwest::Error),
#[error("failed to read executor request body: {0}")]
RequestRead(String),
#[error("executor request body is not valid JSON: {0}")]
InvalidRequestJson(serde_json::Error),
#[error("executor overloaded: gate {gate} saturated at {limit}")]
Overloaded { gate: &'static str, limit: usize },
#[error("failed to execute upstream request: {0}")] #[error("failed to execute upstream request: {0}")]
UpstreamRequest(reqwest::Error), UpstreamRequest(reqwest::Error),
#[error("hub relay request failed: {0}")] #[error("hub relay request failed: {0}")]

100
crates/aether-executor/src/main.rs
View File

@@ -4,6 +4,10 @@ use clap::Parser;
use tracing::info; use tracing::info;
use aether_executor::server; use aether_executor::server;
use aether_runtime::{
init_service_runtime, DistributedConcurrencyGate, RedisDistributedConcurrencyConfig,
ServiceRuntimeConfig,
};
#[derive(Parser, Debug)] #[derive(Parser, Debug)]
#[command(name = "aether-executor", about = "Internal Rust executor for Aether")] #[command(name = "aether-executor", about = "Internal Rust executor for Aether")]
@@ -20,28 +24,106 @@ struct Args {
default_value = "/tmp/aether-executor.sock" default_value = "/tmp/aether-executor.sock"
)] )]
unix_socket: PathBuf, unix_socket: PathBuf,
#[arg(long, env = "AETHER_EXECUTOR_MAX_IN_FLIGHT_REQUESTS")]
max_in_flight_requests: Option<usize>,
#[arg(long, env = "AETHER_EXECUTOR_DISTRIBUTED_REQUEST_LIMIT")]
distributed_request_limit: Option<usize>,
#[arg(long, env = "AETHER_EXECUTOR_DISTRIBUTED_REQUEST_REDIS_URL")]
distributed_request_redis_url: Option<String>,
#[arg(long, env = "AETHER_EXECUTOR_DISTRIBUTED_REQUEST_REDIS_KEY_PREFIX")]
distributed_request_redis_key_prefix: Option<String>,
#[arg(
long,
env = "AETHER_EXECUTOR_DISTRIBUTED_REQUEST_LEASE_TTL_MS",
default_value_t = 30_000
)]
distributed_request_lease_ttl_ms: u64,
#[arg(
long,
env = "AETHER_EXECUTOR_DISTRIBUTED_REQUEST_RENEW_INTERVAL_MS",
default_value_t = 10_000
)]
distributed_request_renew_interval_ms: u64,
#[arg(
long,
env = "AETHER_EXECUTOR_DISTRIBUTED_REQUEST_COMMAND_TIMEOUT_MS",
default_value_t = 1_000
)]
distributed_request_command_timeout_ms: u64,
} }
#[tokio::main] #[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> { async fn main() -> Result<(), Box<dyn std::error::Error>> {
let _ = rustls::crypto::ring::default_provider().install_default(); let _ = rustls::crypto::ring::default_provider().install_default();
tracing_subscriber::fmt() init_service_runtime(ServiceRuntimeConfig::new(
.with_env_filter( "aether-executor",
tracing_subscriber::EnvFilter::try_from_default_env() "aether_executor=info",
.unwrap_or_else(|_| "aether_executor=info".into()), ))?;
)
.init();
let args = Args::parse(); let args = Args::parse();
let distributed_request_gate = match args.distributed_request_limit.filter(|limit| *limit > 0) {
Some(limit) => {
let redis_url = args
.distributed_request_redis_url
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"AETHER_EXECUTOR_DISTRIBUTED_REQUEST_REDIS_URL is required when distributed request limit is enabled",
)
})?;
Some(DistributedConcurrencyGate::new_redis(
"executor_requests_distributed",
limit,
RedisDistributedConcurrencyConfig {
url: redis_url.to_string(),
key_prefix: args
.distributed_request_redis_key_prefix
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
lease_ttl_ms: args.distributed_request_lease_ttl_ms.max(1),
renew_interval_ms: args.distributed_request_renew_interval_ms.max(1),
command_timeout_ms: Some(args.distributed_request_command_timeout_ms.max(1)),
},
)?)
}
None => None,
};
match args.transport.trim().to_ascii_lowercase().as_str() { match args.transport.trim().to_ascii_lowercase().as_str() {
"unix_socket" | "unix" | "uds" => { "unix_socket" | "unix" | "uds" => {
info!(socket = %args.unix_socket.display(), "aether-executor started"); info!(socket = %args.unix_socket.display(), "aether-executor started");
server::serve_unix(&args.unix_socket).await?; server::serve_unix(
&args.unix_socket,
args.max_in_flight_requests,
distributed_request_gate.clone(),
)
.await?;
} }
"tcp" => { "tcp" => {
info!(bind = %args.bind, "aether-executor started"); info!(
server::serve_tcp(&args.bind).await?; bind = %args.bind,
max_in_flight_requests = args.max_in_flight_requests.unwrap_or_default(),
distributed_request_limit = args.distributed_request_limit.unwrap_or_default(),
"aether-executor started"
);
server::serve_tcp(
&args.bind,
args.max_in_flight_requests,
distributed_request_gate,
)
.await?;
} }
other => { other => {
return Err(format!("unsupported executor transport: {other}").into()); return Err(format!("unsupported executor transport: {other}").into());

491
crates/aether-executor/src/server.rs
View File

@@ -1,13 +1,20 @@
use std::convert::Infallible; use std::convert::Infallible;
use std::path::Path; use std::path::Path;
use std::sync::Arc;
use aether_contracts::{ use aether_contracts::{
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionResult, ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionTelemetry,
ExecutionTelemetry, StreamFrame, StreamFramePayload, StreamFrameType, StreamFrame, StreamFramePayload, StreamFrameType,
};
use aether_runtime::{
maybe_hold_axum_response_permit, prometheus_response, service_up_sample, AdmissionPermit,
ConcurrencyError, ConcurrencyGate, ConcurrencySnapshot, DistributedConcurrencyError,
DistributedConcurrencyGate, DistributedConcurrencySnapshot, MetricKind, MetricLabel,
MetricSample,
}; };
use async_stream::stream; use async_stream::stream;
use axum::body::Body; use axum::body::{to_bytes, Body};
use axum::extract::State; use axum::extract::{Request, State};
use axum::http::StatusCode; use axum::http::StatusCode;
use axum::response::{IntoResponse, Response}; use axum::response::{IntoResponse, Response};
use axum::routing::{get, post}; use axum::routing::{get, post};
@@ -22,25 +29,132 @@ use crate::{encode_frame, ExecutorServiceError, SyncExecutor};
#[derive(Debug, Clone, Default)] #[derive(Debug, Clone, Default)]
pub struct AppState { pub struct AppState {
executor: SyncExecutor, executor: SyncExecutor,
request_gate: Option<Arc<ConcurrencyGate>>,
distributed_request_gate: Option<Arc<DistributedConcurrencyGate>>,
}
impl AppState {
fn with_request_concurrency_limit(limit: Option<usize>) -> Self {
Self {
executor: SyncExecutor::new(),
request_gate: limit
.filter(|limit| *limit > 0)
.map(|limit| Arc::new(ConcurrencyGate::new("executor_requests", limit))),
distributed_request_gate: None,
}
}
fn with_distributed_request_gate(mut self, gate: DistributedConcurrencyGate) -> Self {
self.distributed_request_gate = Some(Arc::new(gate));
self
}
fn request_concurrency_snapshot(&self) -> Option<ConcurrencySnapshot> {
self.request_gate.as_ref().map(|gate| gate.snapshot())
}
async fn distributed_request_concurrency_snapshot(
&self,
) -> Result<Option<DistributedConcurrencySnapshot>, DistributedConcurrencyError> {
match self.distributed_request_gate.as_ref() {
Some(gate) => gate.snapshot().await.map(Some),
None => Ok(None),
}
}
async fn metric_samples(&self) -> Vec<MetricSample> {
let mut samples = vec![service_up_sample("aether-executor")];
if let Some(snapshot) = self.request_concurrency_snapshot() {
samples.extend(snapshot.to_metric_samples("executor_requests"));
}
if let Some(gate) = self.distributed_request_gate.as_ref() {
match gate.snapshot().await {
Ok(snapshot) => {
samples.extend(snapshot.to_metric_samples("executor_requests_distributed"));
}
Err(_) => samples.push(
MetricSample::new(
"concurrency_unavailable",
"Whether the distributed concurrency gate is currently unavailable.",
MetricKind::Gauge,
1,
)
.with_labels(vec![MetricLabel::new(
"gate",
"executor_requests_distributed",
)]),
),
}
}
samples
}
async fn try_acquire_request_permit(
&self,
) -> Result<Option<AdmissionPermit>, RequestAdmissionError> {
let local = self
.request_gate
.as_ref()
.map(|gate| gate.try_acquire())
.transpose()
.map_err(RequestAdmissionError::Local)?;
let distributed = match self.distributed_request_gate.as_ref() {
Some(gate) => Some(
gate.try_acquire()
.await
.map_err(RequestAdmissionError::Distributed)?,
),
None => None,
};
Ok(AdmissionPermit::from_parts(local, distributed))
}
} }
pub fn build_router() -> Router { pub fn build_router() -> Router {
Router::new() build_router_with_request_concurrency_limit(None)
.route("/health", get(health))
.route("/v1/execute/sync", post(execute_sync))
.route("/v1/execute/stream", post(execute_stream))
.with_state(AppState {
executor: SyncExecutor::new(),
})
} }
pub async fn serve_tcp(bind: &str) -> Result<(), Box<dyn std::error::Error>> { pub fn build_router_with_request_concurrency_limit(limit: Option<usize>) -> Router {
build_router_with_request_gates(limit, None)
}
pub fn build_router_with_request_gates(
limit: Option<usize>,
distributed_gate: Option<DistributedConcurrencyGate>,
) -> Router {
let state = match distributed_gate {
Some(gate) => {
AppState::with_request_concurrency_limit(limit).with_distributed_request_gate(gate)
}
None => AppState::with_request_concurrency_limit(limit),
};
Router::new()
.route("/health", get(health))
.route("/metrics", get(metrics))
.route("/v1/execute/sync", post(execute_sync))
.route("/v1/execute/stream", post(execute_stream))
.with_state(state)
}
pub async fn serve_tcp(
bind: &str,
max_in_flight_requests: Option<usize>,
distributed_request_gate: Option<DistributedConcurrencyGate>,
) -> Result<(), Box<dyn std::error::Error>> {
let listener = tokio::net::TcpListener::bind(bind).await?; let listener = tokio::net::TcpListener::bind(bind).await?;
axum::serve(listener, build_router()).await?; axum::serve(
listener,
build_router_with_request_gates(max_in_flight_requests, distributed_request_gate),
)
.await?;
Ok(()) Ok(())
} }
pub async fn serve_unix(socket_path: &Path) -> Result<(), Box<dyn std::error::Error>> { pub async fn serve_unix(
socket_path: &Path,
max_in_flight_requests: Option<usize>,
distributed_request_gate: Option<DistributedConcurrencyGate>,
) -> Result<(), Box<dyn std::error::Error>> {
if let Some(parent) = socket_path.parent() { if let Some(parent) = socket_path.parent() {
std::fs::create_dir_all(parent)?; std::fs::create_dir_all(parent)?;
} }
@@ -49,30 +163,69 @@ pub async fn serve_unix(socket_path: &Path) -> Result<(), Box<dyn std::error::Er
} }
let listener = tokio::net::UnixListener::bind(socket_path)?; let listener = tokio::net::UnixListener::bind(socket_path)?;
axum::serve(listener, build_router()).await?; axum::serve(
listener,
build_router_with_request_gates(max_in_flight_requests, distributed_request_gate),
)
.await?;
Ok(()) Ok(())
} }
async fn health() -> impl IntoResponse { async fn health(State(state): State<AppState>) -> impl IntoResponse {
Json(json!({"status": "ok"})) let request_concurrency = state.request_concurrency_snapshot().map(|snapshot| {
json!({
"limit": snapshot.limit,
"in_flight": snapshot.in_flight,
"available_permits": snapshot.available_permits,
"high_watermark": snapshot.high_watermark,
"rejected": snapshot.rejected,
})
});
let distributed_request_concurrency = state
.distributed_request_concurrency_snapshot()
.await
.ok()
.flatten()
.map(|snapshot| {
json!({
"limit": snapshot.limit,
"in_flight": snapshot.in_flight,
"available_permits": snapshot.available_permits,
"high_watermark": snapshot.high_watermark,
"rejected": snapshot.rejected,
})
});
Json(json!({
"status": "ok",
"component": "aether-executor",
"request_concurrency": request_concurrency,
"distributed_request_concurrency": distributed_request_concurrency,
}))
}
async fn metrics(State(state): State<AppState>) -> Response {
prometheus_response(&state.metric_samples().await)
} }
async fn execute_sync( async fn execute_sync(
State(state): State<AppState>, State(state): State<AppState>,
Json(plan): Json<ExecutionPlan>, request: Request,
) -> Result<Json<ExecutionResult>, AppError> { ) -> Result<Response, AppError> {
state let request_permit = acquire_request_permit(&state).await?;
.executor let plan = parse_request_json::<ExecutionPlan>(request).await?;
.execute_sync(plan) let result = state.executor.execute_sync(plan).await.map_err(AppError)?;
.await Ok(maybe_hold_axum_response_permit(
.map(Json) Json(result).into_response(),
.map_err(AppError) request_permit,
))
} }
async fn execute_stream( async fn execute_stream(
State(state): State<AppState>, State(state): State<AppState>,
Json(plan): Json<ExecutionPlan>, request: Request,
) -> Result<Response, AppError> { ) -> Result<Response, AppError> {
let request_permit = acquire_request_permit(&state).await?;
let plan = parse_request_json::<ExecutionPlan>(request).await?;
let execution = state let execution = state
.executor .executor
.execute_stream(plan) .execute_stream(plan)
@@ -149,7 +302,61 @@ async fn execute_stream(
axum::http::header::CONTENT_TYPE, axum::http::header::CONTENT_TYPE,
axum::http::HeaderValue::from_static("application/x-ndjson"), axum::http::HeaderValue::from_static("application/x-ndjson"),
); );
Ok(response) Ok(maybe_hold_axum_response_permit(response, request_permit))
}
async fn acquire_request_permit(state: &AppState) -> Result<Option<AdmissionPermit>, AppError> {
match state.try_acquire_request_permit().await {
Ok(permit) => Ok(permit),
Err(RequestAdmissionError::Local(ConcurrencyError::Saturated { gate, limit }))
| Err(RequestAdmissionError::Distributed(DistributedConcurrencyError::Saturated {
gate,
limit,
}))
| Err(RequestAdmissionError::Distributed(DistributedConcurrencyError::Unavailable {
gate,
limit,
..
})) => Err(AppError(ExecutorServiceError::Overloaded { gate, limit })),
Err(RequestAdmissionError::Local(ConcurrencyError::Closed { gate })) => {
Err(AppError(ExecutorServiceError::RequestRead(format!(
"executor request concurrency gate {gate} is closed"
))))
}
Err(RequestAdmissionError::Distributed(
DistributedConcurrencyError::InvalidConfiguration(message),
)) => Err(AppError(ExecutorServiceError::RequestRead(message))),
}
}
#[derive(Debug)]
enum RequestAdmissionError {
Local(ConcurrencyError),
Distributed(DistributedConcurrencyError),
}
async fn parse_request_json<T>(request: Request) -> Result<T, AppError>
where
T: serde::de::DeserializeOwned,
{
let body = to_bytes(request.into_body(), usize::MAX)
.await
.map_err(|err| AppError(ExecutorServiceError::RequestRead(err.to_string())))?;
serde_json::from_slice(&body)
.map_err(|err| AppError(ExecutorServiceError::InvalidRequestJson(err)))
}
fn build_overloaded_response(message: &str) -> Response {
(
StatusCode::SERVICE_UNAVAILABLE,
Json(json!({
"error": {
"type": "overloaded",
"message": message,
}
})),
)
.into_response()
} }
#[derive(Debug)] #[derive(Debug)]
@@ -158,6 +365,12 @@ struct AppError(ExecutorServiceError);
impl IntoResponse for AppError { impl IntoResponse for AppError {
fn into_response(self) -> Response { fn into_response(self) -> Response {
let status_code = match self.0 { let status_code = match self.0 {
ExecutorServiceError::RequestRead(_) | ExecutorServiceError::InvalidRequestJson(_) => {
StatusCode::BAD_REQUEST
}
ExecutorServiceError::Overloaded { .. } => {
return build_overloaded_response(&self.0.to_string());
}
ExecutorServiceError::StreamUnsupported ExecutorServiceError::StreamUnsupported
| ExecutorServiceError::RequestBodyRequired | ExecutorServiceError::RequestBodyRequired
| ExecutorServiceError::BodyDecode(_) | ExecutorServiceError::BodyDecode(_)
@@ -185,3 +398,227 @@ impl IntoResponse for AppError {
.into_response() .into_response()
} }
} }
#[cfg(test)]
mod tests {
use super::{build_router_with_request_concurrency_limit, build_router_with_request_gates};
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
use axum::body::{Body, Bytes};
use axum::response::Response;
use axum::routing::any;
use axum::{extract::Request, Router};
use http::StatusCode;
use std::convert::Infallible;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
async fn start_server(app: Router) -> (String, tokio::task::JoinHandle<()>) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("listener should bind");
let addr = listener.local_addr().expect("local addr should resolve");
let handle = tokio::spawn(async move {
axum::serve(listener, app).await.expect("server should run");
});
(format!("http://{addr}"), handle)
}
fn stream_plan(url: String) -> ExecutionPlan {
ExecutionPlan {
request_id: "req-1".into(),
candidate_id: Some("cand-1".into()),
provider_name: Some("openai".into()),
provider_id: "prov-1".into(),
endpoint_id: "ep-1".into(),
key_id: "key-1".into(),
method: "GET".into(),
url,
headers: std::collections::BTreeMap::new(),
content_type: None,
content_encoding: None,
body: RequestBody {
json_body: None,
body_bytes_b64: None,
body_ref: None,
},
stream: true,
client_api_format: "openai:chat".into(),
provider_api_format: "openai:chat".into(),
model_name: Some("gpt-4.1".into()),
proxy: None,
tls_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(5_000),
total_ms: Some(30_000),
..ExecutionTimeouts::default()
}),
}
}
#[tokio::test]
async fn executor_rejects_second_in_flight_stream_request_with_overload() {
let upstream_hits = Arc::new(AtomicUsize::new(0));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/slow",
any(move |_request: Request| {
let upstream_hits = Arc::clone(&upstream_hits_clone);
async move {
upstream_hits.fetch_add(1, Ordering::SeqCst);
let stream = async_stream::stream! {
yield Ok::<_, Infallible>(Bytes::from_static(b"chunk-1"));
futures_util::future::pending::<()>().await;
};
Response::builder()
.status(StatusCode::OK)
.body(Body::from_stream(stream))
.expect("response should build")
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let executor = build_router_with_request_concurrency_limit(Some(1));
let (executor_url, executor_handle) = start_server(executor).await;
let client = reqwest::Client::new();
let first_response = client
.post(format!("{executor_url}/v1/execute/stream"))
.json(&stream_plan(format!("{upstream_url}/slow")))
.send()
.await
.expect("first request should succeed");
for _ in 0..50 {
if upstream_hits.load(Ordering::SeqCst) == 1 {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
assert_eq!(upstream_hits.load(Ordering::SeqCst), 1);
let second_response = client
.post(format!("{executor_url}/v1/execute/stream"))
.json(&stream_plan(format!("{upstream_url}/slow")))
.send()
.await
.expect("second request should complete");
assert_eq!(second_response.status(), StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(
second_response
.json::<serde_json::Value>()
.await
.expect("json body should decode")["error"]["type"],
"overloaded"
);
assert_eq!(upstream_hits.load(Ordering::SeqCst), 1);
drop(first_response);
executor_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn executor_rejects_second_in_flight_stream_request_with_distributed_overload() {
let upstream_hits = Arc::new(AtomicUsize::new(0));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/slow",
any(move |_request: Request| {
let upstream_hits = Arc::clone(&upstream_hits_clone);
async move {
upstream_hits.fetch_add(1, Ordering::SeqCst);
let stream = async_stream::stream! {
yield Ok::<_, Infallible>(Bytes::from_static(b"chunk-1"));
futures_util::future::pending::<()>().await;
};
Response::builder()
.status(StatusCode::OK)
.body(Body::from_stream(stream))
.expect("response should build")
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let distributed_gate = aether_runtime::DistributedConcurrencyGate::new_in_memory(
"executor_requests_distributed",
1,
);
let executor_a = build_router_with_request_gates(None, Some(distributed_gate.clone()));
let executor_b = build_router_with_request_gates(None, Some(distributed_gate));
let (executor_a_url, executor_a_handle) = start_server(executor_a).await;
let (executor_b_url, executor_b_handle) = start_server(executor_b).await;
let client = reqwest::Client::new();
let first_response = client
.post(format!("{executor_a_url}/v1/execute/stream"))
.json(&stream_plan(format!("{upstream_url}/slow")))
.send()
.await
.expect("first request should succeed");
for _ in 0..50 {
if upstream_hits.load(Ordering::SeqCst) == 1 {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
assert_eq!(upstream_hits.load(Ordering::SeqCst), 1);
let second_response = client
.post(format!("{executor_b_url}/v1/execute/stream"))
.json(&stream_plan(format!("{upstream_url}/slow")))
.send()
.await
.expect("second request should complete");
assert_eq!(second_response.status(), StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(
second_response
.json::<serde_json::Value>()
.await
.expect("json body should decode")["error"]["type"],
"overloaded"
);
assert_eq!(upstream_hits.load(Ordering::SeqCst), 1);
drop(first_response);
executor_a_handle.abort();
executor_b_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn executor_exposes_request_concurrency_metrics() {
let executor = build_router_with_request_gates(
Some(4),
Some(aether_runtime::DistributedConcurrencyGate::new_in_memory(
"executor_requests_distributed",
6,
)),
);
let (executor_url, executor_handle) = start_server(executor).await;
let response = reqwest::Client::new()
.get(format!("{executor_url}/metrics"))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok()),
Some("text/plain; version=0.0.4; charset=utf-8")
);
let body = response.text().await.expect("body should read");
assert!(body.contains("service_up{service=\"aether-executor\"} 1"));
assert!(body.contains("concurrency_available_permits{gate=\"executor_requests\"} 4"));
assert!(body
.contains("concurrency_available_permits{gate=\"executor_requests_distributed\"} 6"));
executor_handle.abort();
}
}

24
crates/aether-executor/src/service.rs
View File

@@ -5,6 +5,7 @@ use std::time::{Duration, Instant};
use aether_contracts::{ use aether_contracts::{
ExecutionPlan, ExecutionResult, ExecutionTelemetry, ProxySnapshot, ResponseBody, ExecutionPlan, ExecutionResult, ExecutionTelemetry, ProxySnapshot, ResponseBody,
}; };
use aether_http::{apply_http_client_config, HttpClientConfig};
use base64::Engine as _; use base64::Engine as _;
use flate2::write::GzEncoder; use flate2::write::GzEncoder;
use flate2::Compression; use flate2::Compression;
@@ -256,10 +257,14 @@ fn gzip_bytes(body_bytes: &[u8]) -> Result<Vec<u8>, ExecutorServiceError> {
fn build_relay_client( fn build_relay_client(
timeouts: Option<&aether_contracts::ExecutionTimeouts>, timeouts: Option<&aether_contracts::ExecutionTimeouts>,
) -> Result<reqwest::Client, ExecutorServiceError> { ) -> Result<reqwest::Client, ExecutorServiceError> {
let mut builder = reqwest::Client::builder(); let builder = apply_http_client_config(
if let Some(connect_ms) = timeouts.and_then(|timeouts| timeouts.connect_ms) { reqwest::Client::builder(),
builder = builder.connect_timeout(Duration::from_millis(connect_ms)); &HttpClientConfig {
} connect_timeout_ms: timeouts.and_then(|timeouts| timeouts.connect_ms),
use_rustls_tls: false,
..HttpClientConfig::default()
},
);
builder.build().map_err(ExecutorServiceError::ClientBuild) builder.build().map_err(ExecutorServiceError::ClientBuild)
} }
@@ -339,10 +344,13 @@ fn build_client(
proxy: Option<&ProxySnapshot>, proxy: Option<&ProxySnapshot>,
tls_profile: Option<&str>, tls_profile: Option<&str>,
) -> Result<reqwest::Client, ExecutorServiceError> { ) -> Result<reqwest::Client, ExecutorServiceError> {
let mut builder = reqwest::Client::builder().use_rustls_tls(); let mut builder = apply_http_client_config(
if let Some(connect_ms) = timeouts.and_then(|timeouts| timeouts.connect_ms) { reqwest::Client::builder(),
builder = builder.connect_timeout(Duration::from_millis(connect_ms)); &HttpClientConfig {
} connect_timeout_ms: timeouts.and_then(|timeouts| timeouts.connect_ms),
..HttpClientConfig::default()
},
);
builder = apply_tls_profile(builder, tls_profile); builder = apply_tls_profile(builder, tls_profile);
if let Some(proxy_url) = resolve_proxy_url(proxy)? { if let Some(proxy_url) = resolve_proxy_url(proxy)? {
let proxy = reqwest::Proxy::all(&proxy_url).map_err(ExecutorServiceError::InvalidProxy)?; let proxy = reqwest::Proxy::all(&proxy_url).map_err(ExecutorServiceError::InvalidProxy)?;

7
crates/aether-gateway/Cargo.toml
View File

@@ -7,7 +7,11 @@ repository.workspace = true
description = "Rust ingress gateway for Aether phase 3a transparent proxy" description = "Rust ingress gateway for Aether phase 3a transparent proxy"
[dependencies] [dependencies]
aether-cache.workspace = true
aether-contracts.workspace = true aether-contracts.workspace = true
aether-data.workspace = true
aether-http.workspace = true
aether-runtime.workspace = true
async-stream.workspace = true async-stream.workspace = true
axum = { version = "0.8" } axum = { version = "0.8" }
base64.workspace = true base64.workspace = true
@@ -21,7 +25,6 @@ serde_json.workspace = true
thiserror.workspace = true thiserror.workspace = true
tokio.workspace = true tokio.workspace = true
tokio-util.workspace = true tokio-util.workspace = true
tracing = "0.1" tracing.workspace = true
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
url.workspace = true url.workspace = true
uuid.workspace = true uuid.workspace = true

205
crates/aether-gateway/src/audit/http.rs Normal file
View File

@@ -0,0 +1,205 @@
use std::time::{SystemTime, UNIX_EPOCH};
use axum::extract::{Path, Query, State};
use axum::response::IntoResponse;
use axum::Json;
use serde::{Deserialize, Serialize};
use serde_json::json;
use crate::gateway::{AppState, GatewayError};
const DEFAULT_RECENT_LIMIT: usize = 20;
const MAX_RECENT_LIMIT: usize = 200;
#[derive(Debug, Deserialize)]
pub(crate) struct ListRecentShadowResultsQuery {
pub(crate) limit: Option<usize>,
}
#[derive(Debug, Serialize)]
pub(crate) struct ShadowResultStatusCounts {
pub(crate) pending: usize,
pub(crate) r#match: usize,
pub(crate) mismatch: usize,
pub(crate) error: usize,
}
#[derive(Debug, Serialize)]
pub(crate) struct ListRecentShadowResultsResponse {
pub(crate) items: Vec<aether_data::repository::shadow_results::StoredShadowResult>,
pub(crate) limit_applied: usize,
pub(crate) counts: ShadowResultStatusCounts,
}
pub(crate) async fn list_recent_shadow_results(
State(state): State<AppState>,
Query(query): Query<ListRecentShadowResultsQuery>,
) -> Result<Json<ListRecentShadowResultsResponse>, GatewayError> {
let limit = query
.limit
.unwrap_or(DEFAULT_RECENT_LIMIT)
.clamp(1, MAX_RECENT_LIMIT);
let items = state.list_recent_shadow_results(limit).await?;
let mut counts = ShadowResultStatusCounts {
pending: 0,
r#match: 0,
mismatch: 0,
error: 0,
};
for item in &items {
match item.match_status {
aether_data::repository::shadow_results::ShadowResultMatchStatus::Pending => {
counts.pending += 1
}
aether_data::repository::shadow_results::ShadowResultMatchStatus::Match => {
counts.r#match += 1
}
aether_data::repository::shadow_results::ShadowResultMatchStatus::Mismatch => {
counts.mismatch += 1
}
aether_data::repository::shadow_results::ShadowResultMatchStatus::Error => {
counts.error += 1
}
}
}
Ok(Json(ListRecentShadowResultsResponse {
items,
limit_applied: limit,
counts,
}))
}
#[derive(Debug, Deserialize)]
pub(crate) struct GetRequestCandidateTraceQuery {
pub(crate) attempted_only: Option<bool>,
}
pub(crate) async fn get_request_candidate_trace(
State(state): State<AppState>,
Path(request_id): Path<String>,
Query(query): Query<GetRequestCandidateTraceQuery>,
) -> Result<Json<crate::gateway::data::RequestCandidateTrace>, axum::response::Response> {
let attempted_only = query.attempted_only.unwrap_or(false);
let trace = state
.read_request_candidate_trace(&request_id, attempted_only)
.await
.map_err(IntoResponse::into_response)?;
match trace {
Some(trace) => Ok(Json(trace)),
None => Err((
axum::http::StatusCode::NOT_FOUND,
Json(json!({
"error": {
"message": "Request not found",
}
})),
)
.into_response()),
}
}
pub(crate) async fn get_decision_trace(
State(state): State<AppState>,
Path(request_id): Path<String>,
Query(query): Query<GetRequestCandidateTraceQuery>,
) -> Result<Json<crate::gateway::data::DecisionTrace>, axum::response::Response> {
let attempted_only = query.attempted_only.unwrap_or(false);
let trace = state
.read_decision_trace(&request_id, attempted_only)
.await
.map_err(IntoResponse::into_response)?;
match trace {
Some(trace) => Ok(Json(trace)),
None => Err((
axum::http::StatusCode::NOT_FOUND,
Json(json!({
"error": {
"message": "Decision trace not found",
}
})),
)
.into_response()),
}
}
pub(crate) async fn get_request_usage_audit(
State(state): State<AppState>,
Path(request_id): Path<String>,
) -> Result<Json<crate::gateway::data::RequestUsageAudit>, axum::response::Response> {
let usage = state
.read_request_usage_audit(&request_id)
.await
.map_err(IntoResponse::into_response)?;
match usage {
Some(usage) => Ok(Json(usage)),
None => Err((
axum::http::StatusCode::NOT_FOUND,
Json(json!({
"error": {
"message": "Request usage not found",
}
})),
)
.into_response()),
}
}
pub(crate) async fn get_request_audit_bundle(
State(state): State<AppState>,
Path(request_id): Path<String>,
Query(query): Query<GetRequestCandidateTraceQuery>,
) -> Result<Json<crate::gateway::data::RequestAuditBundle>, axum::response::Response> {
let attempted_only = query.attempted_only.unwrap_or(false);
let bundle = state
.read_request_audit_bundle(&request_id, attempted_only, current_unix_secs())
.await
.map_err(IntoResponse::into_response)?;
match bundle {
Some(bundle) => Ok(Json(bundle)),
None => Err((
axum::http::StatusCode::NOT_FOUND,
Json(json!({
"error": {
"message": "Request audit bundle not found",
}
})),
)
.into_response()),
}
}
pub(crate) async fn get_auth_api_key_snapshot(
State(state): State<AppState>,
Path((user_id, api_key_id)): Path<(String, String)>,
) -> Result<Json<crate::gateway::data::StoredGatewayAuthApiKeySnapshot>, axum::response::Response> {
let snapshot = state
.read_auth_api_key_snapshot(&user_id, &api_key_id, current_unix_secs())
.await
.map_err(IntoResponse::into_response)?;
match snapshot {
Some(snapshot) => Ok(Json(snapshot)),
None => Err((
axum::http::StatusCode::NOT_FOUND,
Json(json!({
"error": {
"message": "Auth snapshot not found",
}
})),
)
.into_response()),
}
}
fn current_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}

10
crates/aether-gateway/src/audit/mod.rs Normal file
View File

@@ -0,0 +1,10 @@
mod http;
mod shadow;
pub(crate) use http::get_auth_api_key_snapshot;
pub(crate) use http::get_decision_trace;
pub(crate) use http::get_request_audit_bundle;
pub(crate) use http::get_request_candidate_trace;
pub(crate) use http::get_request_usage_audit;
pub(crate) use http::list_recent_shadow_results;
pub(crate) use shadow::record_shadow_result_non_blocking;

271
crates/aether-gateway/src/audit/shadow.rs Normal file
View File

@@ -0,0 +1,271 @@
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::time::{SystemTime, UNIX_EPOCH};
use aether_data::repository::shadow_results::{RecordShadowResultSample, ShadowResultSampleOrigin};
use axum::body::Body;
use axum::http::header::CONTENT_TYPE;
use axum::http::Response;
use tracing::warn;
use crate::gateway::constants::{
CONTROL_CANDIDATE_ID_HEADER, CONTROL_REQUEST_ID_HEADER, EXECUTION_PATH_CONTROL_EXECUTE_STREAM,
EXECUTION_PATH_CONTROL_EXECUTE_SYNC,
};
use crate::gateway::{AppState, GatewayControlDecision};
pub(crate) fn record_shadow_result_non_blocking(
state: AppState,
trace_id: &str,
method: &http::Method,
path_and_query: &str,
control_decision: Option<&GatewayControlDecision>,
execution_path: &'static str,
response: &Response<Body>,
) {
let Some(decision) =
control_decision.filter(|decision| decision.route_class.as_deref() == Some("ai_public"))
else {
return;
};
if !state.has_shadow_result_data_writer() {
return;
}
let route_family = decision.route_family.clone();
let route_kind = decision.route_kind.clone();
let status_code = response.status().as_u16();
let content_type = response
.headers()
.get(CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string();
let candidate_id = response
.headers()
.get(CONTROL_CANDIDATE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let request_id = response
.headers()
.get(CONTROL_REQUEST_ID_HEADER)
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
let sample = RecordShadowResultSample {
trace_id: trace_id.to_string(),
request_fingerprint: build_request_fingerprint(
method,
path_and_query,
route_family.as_deref(),
route_kind.as_deref(),
),
request_id,
route_family,
route_kind,
candidate_id,
origin: sample_origin_for_execution_path(execution_path),
result_digest: build_result_digest(status_code, &content_type),
status_code: Some(status_code),
error_message: (status_code >= 400)
.then(|| format!("gateway response status {status_code}")),
recorded_at_unix_secs: now_unix_secs,
};
let trace_id = trace_id.to_string();
tokio::spawn(async move {
if let Err(err) = state.record_shadow_result_sample(sample).await {
warn!(trace_id = %trace_id, error = ?err, "gateway failed to record shadow result");
}
});
}
fn build_request_fingerprint(
method: &http::Method,
path_and_query: &str,
route_family: Option<&str>,
route_kind: Option<&str>,
) -> String {
let mut hasher = DefaultHasher::new();
method.as_str().hash(&mut hasher);
path_and_query.hash(&mut hasher);
route_family.unwrap_or_default().hash(&mut hasher);
route_kind.unwrap_or_default().hash(&mut hasher);
format!("{:x}", hasher.finish())
}
fn build_result_digest(status_code: u16, content_type: &str) -> String {
let mut hasher = DefaultHasher::new();
status_code.hash(&mut hasher);
content_type.hash(&mut hasher);
format!("{:x}", hasher.finish())
}
fn sample_origin_for_execution_path(execution_path: &str) -> ShadowResultSampleOrigin {
match execution_path {
EXECUTION_PATH_CONTROL_EXECUTE_SYNC | EXECUTION_PATH_CONTROL_EXECUTE_STREAM => {
ShadowResultSampleOrigin::Python
}
_ => ShadowResultSampleOrigin::Rust,
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use aether_data::repository::shadow_results::{
InMemoryShadowResultRepository, ShadowResultMatchStatus, ShadowResultReadRepository,
};
use axum::body::Body;
use axum::http::header::CONTENT_TYPE;
use axum::http::{Method, Response, StatusCode};
use super::record_shadow_result_non_blocking;
use crate::gateway::constants::{
CONTROL_REQUEST_ID_HEADER, EXECUTION_PATH_CONTROL_EXECUTE_SYNC,
EXECUTION_PATH_EXECUTOR_SYNC,
};
use crate::gateway::{AppState, GatewayControlDecision};
fn sample_decision() -> GatewayControlDecision {
GatewayControlDecision {
public_path: "/v1/chat/completions".to_string(),
public_query_string: Some("stream=true".to_string()),
route_class: Some("ai_public".to_string()),
route_family: Some("openai".to_string()),
route_kind: Some("chat".to_string()),
auth_endpoint_signature: Some("openai:chat".to_string()),
executor_candidate: true,
auth_context: None,
}
}
#[tokio::test]
async fn records_shadow_result_for_ai_public_response() {
let repository = Arc::new(InMemoryShadowResultRepository::default());
let state = AppState::new_with_executor(
"http://127.0.0.1:18084",
Some("http://127.0.0.1:18085".to_string()),
Some("http://127.0.0.1:18086".to_string()),
)
.expect("app state should build")
.with_shadow_result_data_writer_for_tests(repository.clone());
let response = Response::builder()
.status(StatusCode::OK)
.header(CONTENT_TYPE, "application/json")
.header(CONTROL_REQUEST_ID_HEADER, "req-shadow-123")
.body(Body::from("{}"))
.expect("response should build");
record_shadow_result_non_blocking(
state,
"trace-shadow-123",
&Method::POST,
"/v1/chat/completions?stream=true",
Some(&sample_decision()),
EXECUTION_PATH_EXECUTOR_SYNC,
&response,
);
for _ in 0..30 {
if repository
.list_recent(1)
.await
.map(|rows| !rows.is_empty())
.unwrap_or(false)
{
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let stored = repository
.list_recent(1)
.await
.expect("list should succeed")
.into_iter()
.next()
.expect("stored result should exist");
assert_eq!(stored.trace_id, "trace-shadow-123");
assert_eq!(stored.request_id.as_deref(), Some("req-shadow-123"));
assert_eq!(stored.route_family.as_deref(), Some("openai"));
assert_eq!(stored.route_kind.as_deref(), Some("chat"));
assert_eq!(stored.match_status, ShadowResultMatchStatus::Pending);
assert_eq!(stored.status_code, Some(200));
assert!(stored.rust_result_digest.is_some());
}
#[tokio::test]
async fn merges_rust_and_python_shadow_samples_into_match() {
let repository = Arc::new(InMemoryShadowResultRepository::default());
let state = AppState::new_with_executor(
"http://127.0.0.1:18084",
Some("http://127.0.0.1:18085".to_string()),
Some("http://127.0.0.1:18086".to_string()),
)
.expect("app state should build")
.with_shadow_result_data_repository_for_tests(repository.clone());
let response = Response::builder()
.status(StatusCode::OK)
.header(CONTENT_TYPE, "application/json")
.header(CONTROL_REQUEST_ID_HEADER, "req-shadow-compare-123")
.body(Body::from("{}"))
.expect("response should build");
record_shadow_result_non_blocking(
state.clone(),
"trace-shadow-compare-123",
&Method::POST,
"/v1/chat/completions?stream=true",
Some(&sample_decision()),
EXECUTION_PATH_EXECUTOR_SYNC,
&response,
);
record_shadow_result_non_blocking(
state,
"trace-shadow-compare-123",
&Method::POST,
"/v1/chat/completions?stream=true",
Some(&sample_decision()),
EXECUTION_PATH_CONTROL_EXECUTE_SYNC,
&response,
);
for _ in 0..30 {
if repository
.list_recent(1)
.await
.map(|rows| {
rows.first()
.map(|row| row.match_status == ShadowResultMatchStatus::Match)
.unwrap_or(false)
})
.unwrap_or(false)
{
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let stored = repository
.list_recent(1)
.await
.expect("list should succeed")
.into_iter()
.next()
.expect("stored result should exist");
assert_eq!(stored.request_id.as_deref(), Some("req-shadow-compare-123"));
assert_eq!(stored.match_status, ShadowResultMatchStatus::Match);
assert!(stored.rust_result_digest.is_some());
assert!(stored.python_result_digest.is_some());
}
}

31
crates/aether-gateway/src/cache/auth_context.rs vendored Normal file
View File

@@ -0,0 +1,31 @@
use std::time::Duration;
use aether_cache::ExpiringMap;
use crate::gateway::GatewayControlAuthContext;
#[derive(Debug, Default)]
pub(crate) struct AuthContextCache {
entries: ExpiringMap<String, GatewayControlAuthContext>,
}
impl AuthContextCache {
pub(crate) fn get_fresh(
&self,
cache_key: &str,
ttl: Duration,
) -> Option<GatewayControlAuthContext> {
self.entries.get_fresh(&cache_key.to_string(), ttl)
}
pub(crate) fn insert(
&self,
cache_key: String,
auth_context: GatewayControlAuthContext,
ttl: Duration,
max_entries: usize,
) {
self.entries
.insert(cache_key, auth_context, ttl, max_entries);
}
}

18
crates/aether-gateway/src/cache/direct_plan_bypass.rs vendored Normal file
View File

@@ -0,0 +1,18 @@
use std::time::Duration;
use aether_cache::ExpiringMap;
#[derive(Debug, Default)]
pub(crate) struct DirectPlanBypassCache {
entries: ExpiringMap<String, ()>,
}
impl DirectPlanBypassCache {
pub(crate) fn should_skip(&self, cache_key: &str, ttl: Duration) -> bool {
self.entries.contains_fresh(&cache_key.to_string(), ttl)
}
pub(crate) fn mark(&self, cache_key: String, ttl: Duration, max_entries: usize) {
self.entries.insert(cache_key, (), ttl, max_entries);
}
}

5
crates/aether-gateway/src/cache/mod.rs vendored Normal file
View File

@@ -0,0 +1,5 @@
mod auth_context;
mod direct_plan_bypass;
pub(crate) use auth_context::AuthContextCache;
pub(crate) use direct_plan_bypass::DirectPlanBypassCache;

5
crates/aether-gateway/src/constants.rs
View File

@@ -11,10 +11,15 @@ pub(crate) const EXECUTION_PATH_EXECUTOR_SYNC: &str = "executor_sync";
pub(crate) const EXECUTION_PATH_EXECUTOR_STREAM: &str = "executor_stream"; pub(crate) const EXECUTION_PATH_EXECUTOR_STREAM: &str = "executor_stream";
pub(crate) const EXECUTION_PATH_CONTROL_EXECUTE_SYNC: &str = "control_execute_sync"; pub(crate) const EXECUTION_PATH_CONTROL_EXECUTE_SYNC: &str = "control_execute_sync";
pub(crate) const EXECUTION_PATH_CONTROL_EXECUTE_STREAM: &str = "control_execute_stream"; pub(crate) const EXECUTION_PATH_CONTROL_EXECUTE_STREAM: &str = "control_execute_stream";
pub(crate) const EXECUTION_PATH_LOCAL_AUTH_DENIED: &str = "local_auth_denied";
pub(crate) const EXECUTION_PATH_LOCAL_OVERLOADED: &str = "local_overloaded";
pub(crate) const EXECUTION_PATH_DISTRIBUTED_OVERLOADED: &str = "distributed_overloaded";
pub(crate) const CONTROL_ROUTE_CLASS_HEADER: &str = "x-aether-control-route-class"; pub(crate) const CONTROL_ROUTE_CLASS_HEADER: &str = "x-aether-control-route-class";
pub(crate) const CONTROL_ROUTE_FAMILY_HEADER: &str = "x-aether-control-route-family"; pub(crate) const CONTROL_ROUTE_FAMILY_HEADER: &str = "x-aether-control-route-family";
pub(crate) const CONTROL_ROUTE_KIND_HEADER: &str = "x-aether-control-route-kind"; pub(crate) const CONTROL_ROUTE_KIND_HEADER: &str = "x-aether-control-route-kind";
pub(crate) const CONTROL_EXECUTOR_HEADER: &str = "x-aether-control-executor-candidate"; pub(crate) const CONTROL_EXECUTOR_HEADER: &str = "x-aether-control-executor-candidate";
pub(crate) const CONTROL_REQUEST_ID_HEADER: &str = "x-aether-control-request-id";
pub(crate) const CONTROL_CANDIDATE_ID_HEADER: &str = "x-aether-control-candidate-id";
pub(crate) const CONTROL_ENDPOINT_SIGNATURE_HEADER: &str = "x-aether-control-endpoint-signature"; pub(crate) const CONTROL_ENDPOINT_SIGNATURE_HEADER: &str = "x-aether-control-endpoint-signature";
pub(crate) const CONTROL_EXECUTED_HEADER: &str = "x-aether-control-executed"; pub(crate) const CONTROL_EXECUTED_HEADER: &str = "x-aether-control-executed";
pub(crate) const CONTROL_ACTION_HEADER: &str = "x-aether-control-action"; pub(crate) const CONTROL_ACTION_HEADER: &str = "x-aether-control-action";

174
crates/aether-gateway/src/control.rs
View File

@@ -1,4 +1,4 @@
use std::time::{Duration, Instant}; use std::time::{Duration, SystemTime, UNIX_EPOCH};
use axum::body::{Body, Bytes}; use axum::body::{Body, Bytes};
use axum::http::{Response, StatusCode, Uri}; use axum::http::{Response, StatusCode, Uri};
@@ -10,7 +10,7 @@ use crate::gateway::constants::*;
use crate::gateway::headers::{ use crate::gateway::headers::{
collect_control_headers, header_equals, header_value_str, header_value_u64, is_json_request, collect_control_headers, header_equals, header_value_str, header_value_u64, is_json_request,
}; };
use crate::gateway::{build_client_response, AppState, CachedAuthContextEntry, GatewayError}; use crate::gateway::{build_client_response, AppState, GatewayError};
const AUTH_CONTEXT_CACHE_TTL: Duration = Duration::from_secs(60); const AUTH_CONTEXT_CACHE_TTL: Duration = Duration::from_secs(60);
const AUTH_CONTEXT_CACHE_MAX_ENTRIES: usize = 256; const AUTH_CONTEXT_CACHE_MAX_ENTRIES: usize = 256;
@@ -72,6 +72,15 @@ pub(crate) struct GatewayControlAuthContext {
pub(crate) api_key_id: String, pub(crate) api_key_id: String,
pub(crate) balance_remaining: Option<f64>, pub(crate) balance_remaining: Option<f64>,
pub(crate) access_allowed: bool, pub(crate) access_allowed: bool,
#[serde(skip)]
pub(crate) local_rejection: Option<GatewayLocalAuthRejection>,
}
#[derive(Debug, Clone)]
pub(crate) enum GatewayLocalAuthRejection {
InvalidApiKey,
LockedApiKey,
BalanceDenied { remaining: Option<f64> },
} }
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
@@ -116,10 +125,31 @@ pub(crate) async fn resolve_control_route(
}; };
decision.public_query_string = uri.query().map(ToOwned::to_owned); decision.public_query_string = uri.query().map(ToOwned::to_owned);
if let Some(auth_context) = resolve_data_backed_auth_context(
state,
headers,
decision.auth_endpoint_signature.as_deref(),
)
.await?
{
if let Some(cache_key) = decision
.auth_endpoint_signature
.as_deref()
.and_then(|signature| build_auth_context_cache_key(headers, uri, signature))
{
put_cached_auth_context(state, cache_key, auth_context.clone());
}
decision.auth_context = Some(auth_context);
}
if state.executor_base_url.is_some() && decision.executor_candidate { if state.executor_base_url.is_some() && decision.executor_candidate {
return Ok(Some(decision)); return Ok(Some(decision));
} }
if decision.auth_context.is_some() {
return Ok(Some(decision));
}
match fetch_auth_context( match fetch_auth_context(
state, state,
control_base_url, control_base_url,
@@ -170,6 +200,13 @@ pub(crate) async fn resolve_executor_auth_context(
return Ok(Some(auth_context)); return Ok(Some(auth_context));
} }
if let Some(auth_context) =
resolve_data_backed_auth_context(state, headers, Some(auth_endpoint_signature)).await?
{
put_cached_auth_context(state, cache_key, auth_context.clone());
return Ok(Some(auth_context));
}
Ok(None) Ok(None)
} }
@@ -190,6 +227,19 @@ pub(crate) fn cache_executor_auth_context(
put_cached_auth_context(state, cache_key, auth_context); put_cached_auth_context(state, cache_key, auth_context);
} }
pub(crate) fn trusted_auth_local_rejection(
decision: Option<&GatewayControlDecision>,
_headers: &http::HeaderMap,
) -> Option<GatewayLocalAuthRejection> {
let decision = decision?;
if decision.route_class.as_deref() != Some("ai_public") {
return None;
}
let auth_context = decision.auth_context.as_ref()?;
auth_context.local_rejection.clone()
}
async fn fetch_auth_context( async fn fetch_auth_context(
state: &AppState, state: &AppState,
control_base_url: &str, control_base_url: &str,
@@ -334,13 +384,9 @@ fn build_auth_context_cache_key(
} }
fn get_cached_auth_context(state: &AppState, cache_key: &str) -> Option<GatewayControlAuthContext> { fn get_cached_auth_context(state: &AppState, cache_key: &str) -> Option<GatewayControlAuthContext> {
let mut cache = state.auth_context_cache.lock().ok()?; state
let entry = cache.get(cache_key)?.clone(); .auth_context_cache
if entry.cached_at.elapsed() > AUTH_CONTEXT_CACHE_TTL { .get_fresh(cache_key, AUTH_CONTEXT_CACHE_TTL)
cache.remove(cache_key);
return None;
}
Some(entry.auth_context)
} }
fn put_cached_auth_context( fn put_cached_auth_context(
@@ -348,28 +394,100 @@ fn put_cached_auth_context(
cache_key: String, cache_key: String,
auth_context: GatewayControlAuthContext, auth_context: GatewayControlAuthContext,
) { ) {
let Ok(mut cache) = state.auth_context_cache.lock() else { state.auth_context_cache.insert(
return; cache_key,
auth_context,
AUTH_CONTEXT_CACHE_TTL,
AUTH_CONTEXT_CACHE_MAX_ENTRIES,
);
}
async fn resolve_data_backed_auth_context(
state: &AppState,
headers: &http::HeaderMap,
auth_endpoint_signature: Option<&str>,
) -> Result<Option<GatewayControlAuthContext>, GatewayError> {
let Some(signature) = auth_endpoint_signature
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return Ok(None);
};
let _ = signature;
let Some(user_id) =
header_value_str(headers, TRUSTED_AUTH_USER_ID_HEADER).filter(|value| !value.is_empty())
else {
return Ok(None);
};
let Some(api_key_id) =
header_value_str(headers, TRUSTED_AUTH_API_KEY_ID_HEADER).filter(|value| !value.is_empty())
else {
return Ok(None);
}; };
cache.retain(|_, entry| entry.cached_at.elapsed() <= AUTH_CONTEXT_CACHE_TTL); let snapshot = state
if cache.len() >= AUTH_CONTEXT_CACHE_MAX_ENTRIES { .read_auth_api_key_snapshot(&user_id, &api_key_id, current_unix_secs())
if let Some(oldest_key) = cache .await?;
.iter() let Some(snapshot) = snapshot else {
.min_by_key(|(_, entry)| entry.cached_at) return Ok(None);
.map(|(key, _)| key.clone()) };
{
cache.remove(&oldest_key);
}
}
cache.insert( let header_access_allowed = header_value_str(headers, TRUSTED_AUTH_ACCESS_ALLOWED_HEADER)
cache_key, .as_deref()
CachedAuthContextEntry { .and_then(parse_bool_header);
auth_context, let invalid_api_key = !snapshot.user_is_active
cached_at: Instant::now(), || snapshot.user_is_deleted
}, || !snapshot.api_key_is_active
); || snapshot
.api_key_expires_at_unix_secs
.is_some_and(|expires_at| expires_at < current_unix_secs());
let locked_api_key = snapshot.api_key_is_locked && !snapshot.api_key_is_standalone;
let access_allowed = header_access_allowed
.map(|value| value && snapshot.currently_usable)
.unwrap_or(snapshot.currently_usable);
let local_rejection = if invalid_api_key {
Some(GatewayLocalAuthRejection::InvalidApiKey)
} else if locked_api_key {
Some(GatewayLocalAuthRejection::LockedApiKey)
} else if header_access_allowed.is_some_and(|value| !value) && snapshot.currently_usable {
Some(GatewayLocalAuthRejection::BalanceDenied {
remaining: header_value_str(headers, TRUSTED_AUTH_BALANCE_HEADER)
.as_deref()
.and_then(parse_f64_header),
})
} else {
None
};
Ok(Some(GatewayControlAuthContext {
user_id: snapshot.user_id,
api_key_id: snapshot.api_key_id,
balance_remaining: header_value_str(headers, TRUSTED_AUTH_BALANCE_HEADER)
.as_deref()
.and_then(parse_f64_header),
access_allowed,
local_rejection,
}))
}
fn parse_bool_header(value: &str) -> Option<bool> {
match value.trim().to_ascii_lowercase().as_str() {
"true" | "1" | "yes" => Some(true),
"false" | "0" | "no" => Some(false),
_ => None,
}
}
fn parse_f64_header(value: &str) -> Option<f64> {
value.trim().parse::<f64>().ok()
}
fn current_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
} }
fn classify_control_route( fn classify_control_route(

155
crates/aether-gateway/src/data/auth.rs Normal file
View File

@@ -0,0 +1,155 @@
use aether_data::repository::auth::{AuthApiKeyLookupKey, StoredAuthApiKeySnapshot};
use aether_data::DataLayerError;
use super::state::GatewayDataState;
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub(crate) struct StoredGatewayAuthApiKeySnapshot {
pub(crate) user_id: String,
pub(crate) username: String,
pub(crate) email: Option<String>,
pub(crate) user_role: String,
pub(crate) user_auth_source: String,
pub(crate) user_is_active: bool,
pub(crate) user_is_deleted: bool,
pub(crate) user_allowed_providers: Option<Vec<String>>,
pub(crate) user_allowed_api_formats: Option<Vec<String>>,
pub(crate) user_allowed_models: Option<Vec<String>>,
pub(crate) api_key_id: String,
pub(crate) api_key_name: Option<String>,
pub(crate) api_key_is_active: bool,
pub(crate) api_key_is_locked: bool,
pub(crate) api_key_is_standalone: bool,
pub(crate) api_key_rate_limit: Option<i32>,
pub(crate) api_key_concurrent_limit: Option<i32>,
pub(crate) api_key_expires_at_unix_secs: Option<u64>,
pub(crate) api_key_allowed_providers: Option<Vec<String>>,
pub(crate) api_key_allowed_api_formats: Option<Vec<String>>,
pub(crate) api_key_allowed_models: Option<Vec<String>>,
pub(crate) currently_usable: bool,
}
impl StoredGatewayAuthApiKeySnapshot {
fn from_stored(snapshot: StoredAuthApiKeySnapshot, now_unix_secs: u64) -> Self {
let currently_usable = snapshot.is_currently_usable(now_unix_secs);
Self {
user_id: snapshot.user_id,
username: snapshot.username,
email: snapshot.email,
user_role: snapshot.user_role,
user_auth_source: snapshot.user_auth_source,
user_is_active: snapshot.user_is_active,
user_is_deleted: snapshot.user_is_deleted,
user_allowed_providers: snapshot.user_allowed_providers,
user_allowed_api_formats: snapshot.user_allowed_api_formats,
user_allowed_models: snapshot.user_allowed_models,
api_key_id: snapshot.api_key_id,
api_key_name: snapshot.api_key_name,
api_key_is_active: snapshot.api_key_is_active,
api_key_is_locked: snapshot.api_key_is_locked,
api_key_is_standalone: snapshot.api_key_is_standalone,
api_key_rate_limit: snapshot.api_key_rate_limit,
api_key_concurrent_limit: snapshot.api_key_concurrent_limit,
api_key_expires_at_unix_secs: snapshot.api_key_expires_at_unix_secs,
api_key_allowed_providers: snapshot.api_key_allowed_providers,
api_key_allowed_api_formats: snapshot.api_key_allowed_api_formats,
api_key_allowed_models: snapshot.api_key_allowed_models,
currently_usable,
}
}
}
pub(crate) async fn read_auth_api_key_snapshot(
state: &GatewayDataState,
user_id: &str,
api_key_id: &str,
now_unix_secs: u64,
) -> Result<Option<StoredGatewayAuthApiKeySnapshot>, DataLayerError> {
let snapshot = state
.find_auth_api_key_snapshot(AuthApiKeyLookupKey::UserApiKeyIds {
user_id,
api_key_id,
})
.await?;
Ok(snapshot
.map(|snapshot| StoredGatewayAuthApiKeySnapshot::from_stored(snapshot, now_unix_secs)))
}
#[cfg(test)]
mod tests {
use super::super::GatewayDataState;
use super::{read_auth_api_key_snapshot, StoredGatewayAuthApiKeySnapshot};
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
};
use std::sync::Arc;
fn sample_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-4.1"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(200),
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-4.1"])),
)
.expect("snapshot should build")
}
#[tokio::test]
async fn reads_trusted_auth_snapshot_and_derives_usability() {
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-1".to_string()),
sample_snapshot("key-1", "user-1"),
)]));
let state = GatewayDataState::with_auth_api_key_reader_for_tests(repository);
let snapshot = read_auth_api_key_snapshot(&state, "user-1", "key-1", 150)
.await
.expect("read should succeed")
.expect("snapshot should exist");
assert_eq!(
snapshot,
StoredGatewayAuthApiKeySnapshot {
user_id: "user-1".to_string(),
username: "alice".to_string(),
email: Some("alice@example.com".to_string()),
user_role: "user".to_string(),
user_auth_source: "local".to_string(),
user_is_active: true,
user_is_deleted: false,
user_allowed_providers: Some(vec!["openai".to_string()]),
user_allowed_api_formats: Some(vec!["openai:chat".to_string()]),
user_allowed_models: Some(vec!["gpt-4.1".to_string()]),
api_key_id: "key-1".to_string(),
api_key_name: Some("default".to_string()),
api_key_is_active: true,
api_key_is_locked: false,
api_key_is_standalone: false,
api_key_rate_limit: Some(60),
api_key_concurrent_limit: Some(5),
api_key_expires_at_unix_secs: Some(200),
api_key_allowed_providers: Some(vec!["openai".to_string()]),
api_key_allowed_api_formats: Some(vec!["openai:chat".to_string()]),
api_key_allowed_models: Some(vec!["gpt-4.1".to_string()]),
currently_usable: true,
}
);
}
}

213
crates/aether-gateway/src/data/candidates.rs Normal file
View File

@@ -0,0 +1,213 @@
use aether_data::repository::candidates::{RequestCandidateStatus, StoredRequestCandidate};
use aether_data::DataLayerError;
use super::state::GatewayDataState;
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "snake_case")]
pub(crate) enum RequestCandidateFinalStatus {
Success,
Failed,
Cancelled,
Streaming,
Pending,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub(crate) struct RequestCandidateTrace {
pub(crate) request_id: String,
pub(crate) total_candidates: usize,
pub(crate) final_status: RequestCandidateFinalStatus,
pub(crate) total_latency_ms: u64,
pub(crate) candidates: Vec<StoredRequestCandidate>,
}
pub(crate) async fn read_request_candidate_trace(
state: &GatewayDataState,
request_id: &str,
attempted_only: bool,
) -> Result<Option<RequestCandidateTrace>, DataLayerError> {
let all_candidates = state
.list_request_candidates_by_request_id(request_id)
.await?;
if all_candidates.is_empty() {
return Ok(None);
}
let candidates = if attempted_only {
all_candidates
.iter()
.filter(|candidate| {
candidate
.status
.is_attempted(candidate.started_at_unix_secs)
})
.cloned()
.collect::<Vec<_>>()
} else {
all_candidates.clone()
};
let total_latency_ms = candidates
.iter()
.filter(|candidate| {
matches!(
candidate.status,
RequestCandidateStatus::Success
| RequestCandidateStatus::Failed
| RequestCandidateStatus::Cancelled
) && candidate.latency_ms.is_some()
})
.map(|candidate| candidate.latency_ms.unwrap_or(0))
.sum();
let final_status_source = if attempted_only && candidates.is_empty() {
&all_candidates
} else {
&candidates
};
Ok(Some(RequestCandidateTrace {
request_id: request_id.to_string(),
total_candidates: candidates.len(),
final_status: derive_final_status(final_status_source),
total_latency_ms,
candidates,
}))
}
fn derive_final_status(candidates: &[StoredRequestCandidate]) -> RequestCandidateFinalStatus {
let has_success = candidates.iter().any(|candidate| {
candidate.status == RequestCandidateStatus::Success
|| matches!(candidate.status_code, Some(status_code) if (200..300).contains(&status_code))
});
if has_success {
return RequestCandidateFinalStatus::Success;
}
if candidates
.iter()
.any(|candidate| candidate.status == RequestCandidateStatus::Streaming)
{
return RequestCandidateFinalStatus::Streaming;
}
if candidates
.iter()
.any(|candidate| candidate.status == RequestCandidateStatus::Pending)
{
return RequestCandidateFinalStatus::Pending;
}
let has_cancelled = candidates
.iter()
.any(|candidate| candidate.status == RequestCandidateStatus::Cancelled);
let has_failed = candidates
.iter()
.any(|candidate| candidate.status == RequestCandidateStatus::Failed);
if has_cancelled && !has_failed {
return RequestCandidateFinalStatus::Cancelled;
}
RequestCandidateFinalStatus::Failed
}
#[cfg(test)]
mod tests {
use super::super::GatewayDataState;
use super::{derive_final_status, read_request_candidate_trace, RequestCandidateFinalStatus};
use aether_data::repository::candidates::{
InMemoryRequestCandidateRepository, RequestCandidateStatus, StoredRequestCandidate,
};
use std::sync::Arc;
fn sample_candidate(
id: &str,
request_id: &str,
candidate_index: i32,
status: RequestCandidateStatus,
started_at_unix_secs: Option<i64>,
latency_ms: Option<i32>,
status_code: Option<i32>,
) -> StoredRequestCandidate {
StoredRequestCandidate::new(
id.to_string(),
request_id.to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
candidate_index,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
status,
None,
false,
status_code,
None,
None,
latency_ms,
Some(1),
None,
None,
100 + i64::from(candidate_index),
started_at_unix_secs,
started_at_unix_secs.map(|value| value + 1),
)
.expect("candidate should build")
}
#[test]
fn derive_final_status_prefers_success() {
let candidates = vec![sample_candidate(
"cand-1",
"req-1",
0,
RequestCandidateStatus::Success,
Some(100),
Some(25),
Some(200),
)];
assert_eq!(
derive_final_status(&candidates),
RequestCandidateFinalStatus::Success
);
}
#[tokio::test]
async fn read_request_candidate_trace_filters_attempted_rows() {
let repository = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
sample_candidate(
"cand-1",
"req-1",
0,
RequestCandidateStatus::Pending,
None,
None,
None,
),
sample_candidate(
"cand-2",
"req-1",
1,
RequestCandidateStatus::Failed,
Some(101),
Some(33),
Some(502),
),
]));
let state = GatewayDataState::with_request_candidate_reader_for_tests(repository);
let trace = read_request_candidate_trace(&state, "req-1", true)
.await
.expect("trace should succeed")
.expect("trace should exist");
assert_eq!(trace.total_candidates, 1);
assert_eq!(trace.candidates[0].id, "cand-2");
assert_eq!(trace.final_status, RequestCandidateFinalStatus::Failed);
assert_eq!(trace.total_latency_ms, 33);
}
}

41
crates/aether-gateway/src/data/config.rs Normal file
View File

@@ -0,0 +1,41 @@
use aether_data::postgres::PostgresPoolConfig;
use aether_data::DataLayerConfig;
#[derive(Debug, Clone, Default)]
pub struct GatewayDataConfig {
postgres: Option<PostgresPoolConfig>,
}
impl GatewayDataConfig {
pub fn disabled() -> Self {
Self::default()
}
pub fn from_postgres_config(postgres: PostgresPoolConfig) -> Self {
Self {
postgres: Some(postgres),
}
}
pub fn from_postgres_url(database_url: impl Into<String>, require_ssl: bool) -> Self {
let mut postgres = PostgresPoolConfig::default();
postgres.database_url = database_url.into();
postgres.require_ssl = require_ssl;
Self::from_postgres_config(postgres)
}
pub fn postgres(&self) -> Option<&PostgresPoolConfig> {
self.postgres.as_ref()
}
pub fn is_enabled(&self) -> bool {
self.postgres.is_some()
}
pub fn to_data_layer_config(&self) -> DataLayerConfig {
DataLayerConfig {
postgres: self.postgres.clone(),
redis: None,
}
}
}

267
crates/aether-gateway/src/data/decision_trace.rs Normal file
View File

@@ -0,0 +1,267 @@
use std::collections::{BTreeMap, BTreeSet};
use aether_data::repository::candidates::StoredRequestCandidate;
use aether_data::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data::DataLayerError;
use super::candidates::RequestCandidateFinalStatus;
use super::state::GatewayDataState;
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub(crate) struct DecisionTraceCandidate {
#[serde(flatten)]
pub(crate) candidate: StoredRequestCandidate,
pub(crate) provider_name: Option<String>,
pub(crate) provider_website: Option<String>,
pub(crate) provider_type: Option<String>,
pub(crate) endpoint_api_format: Option<String>,
pub(crate) endpoint_api_family: Option<String>,
pub(crate) endpoint_kind: Option<String>,
pub(crate) provider_key_name: Option<String>,
pub(crate) provider_key_auth_type: Option<String>,
pub(crate) provider_key_capabilities: Option<serde_json::Value>,
pub(crate) provider_key_is_active: Option<bool>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub(crate) struct DecisionTrace {
pub(crate) request_id: String,
pub(crate) total_candidates: usize,
pub(crate) final_status: RequestCandidateFinalStatus,
pub(crate) total_latency_ms: u64,
pub(crate) candidates: Vec<DecisionTraceCandidate>,
}
pub(crate) async fn read_decision_trace(
state: &GatewayDataState,
request_id: &str,
attempted_only: bool,
) -> Result<Option<DecisionTrace>, DataLayerError> {
let Some(trace) = state
.read_request_candidate_trace(request_id, attempted_only)
.await?
else {
return Ok(None);
};
let provider_ids = unique_ids(
trace
.candidates
.iter()
.filter_map(|item| item.provider_id.as_ref()),
);
let endpoint_ids = unique_ids(
trace
.candidates
.iter()
.filter_map(|item| item.endpoint_id.as_ref()),
);
let key_ids = unique_ids(
trace
.candidates
.iter()
.filter_map(|item| item.key_id.as_ref()),
);
let provider_map = state
.list_provider_catalog_providers_by_ids(&provider_ids)
.await?
.into_iter()
.map(|item| (item.id.clone(), item))
.collect::<BTreeMap<_, _>>();
let endpoint_map = state
.list_provider_catalog_endpoints_by_ids(&endpoint_ids)
.await?
.into_iter()
.map(|item| (item.id.clone(), item))
.collect::<BTreeMap<_, _>>();
let key_map = state
.list_provider_catalog_keys_by_ids(&key_ids)
.await?
.into_iter()
.map(|item| (item.id.clone(), item))
.collect::<BTreeMap<_, _>>();
Ok(Some(DecisionTrace {
request_id: trace.request_id,
total_candidates: trace.total_candidates,
final_status: trace.final_status,
total_latency_ms: trace.total_latency_ms,
candidates: trace
.candidates
.into_iter()
.map(|candidate| enrich_candidate(candidate, &provider_map, &endpoint_map, &key_map))
.collect(),
}))
}
fn enrich_candidate(
candidate: StoredRequestCandidate,
provider_map: &BTreeMap<String, StoredProviderCatalogProvider>,
endpoint_map: &BTreeMap<String, StoredProviderCatalogEndpoint>,
key_map: &BTreeMap<String, StoredProviderCatalogKey>,
) -> DecisionTraceCandidate {
let provider = candidate
.provider_id
.as_ref()
.and_then(|provider_id| provider_map.get(provider_id));
let endpoint = candidate
.endpoint_id
.as_ref()
.and_then(|endpoint_id| endpoint_map.get(endpoint_id));
let provider_key = candidate
.key_id
.as_ref()
.and_then(|key_id| key_map.get(key_id));
DecisionTraceCandidate {
provider_name: provider.map(|item| item.name.clone()),
provider_website: provider.and_then(|item| item.website.clone()),
provider_type: provider.map(|item| item.provider_type.clone()),
endpoint_api_format: endpoint.map(|item| item.api_format.clone()),
endpoint_api_family: endpoint.and_then(|item| item.api_family.clone()),
endpoint_kind: endpoint.and_then(|item| item.endpoint_kind.clone()),
provider_key_name: provider_key
.map(|item| item.name.clone())
.or_else(|| candidate.api_key_name.clone()),
provider_key_auth_type: provider_key.map(|item| item.auth_type.clone()),
provider_key_capabilities: provider_key.and_then(|item| item.capabilities.clone()),
provider_key_is_active: provider_key.map(|item| item.is_active),
candidate,
}
}
fn unique_ids<'a>(items: impl Iterator<Item = &'a String>) -> Vec<String> {
items
.cloned()
.collect::<BTreeSet<_>>()
.into_iter()
.collect()
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use aether_data::repository::candidates::{
InMemoryRequestCandidateRepository, RequestCandidateStatus, StoredRequestCandidate,
};
use aether_data::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint,
StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use super::{read_decision_trace, DecisionTrace, DecisionTraceCandidate};
use crate::gateway::data::candidates::RequestCandidateFinalStatus;
use crate::gateway::data::GatewayDataState;
fn sample_candidate(request_id: &str) -> StoredRequestCandidate {
StoredRequestCandidate::new(
"cand-1".to_string(),
request_id.to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
RequestCandidateStatus::Failed,
None,
false,
Some(502),
Some("bad_gateway".to_string()),
Some("upstream failed".to_string()),
Some(37),
Some(1),
None,
Some(serde_json::json!({"cache_1h": true})),
100,
Some(101),
Some(102),
)
.expect("candidate should build")
}
fn sample_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-1".to_string(),
"OpenAI".to_string(),
Some("https://openai.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
}
fn sample_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-1".to_string(),
"provider-1".to_string(),
"openai:chat".to_string(),
Some("openai".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
}
fn sample_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"provider-key-1".to_string(),
"provider-1".to_string(),
"prod-key".to_string(),
"api_key".to_string(),
Some(serde_json::json!({"cache_1h": true})),
true,
)
.expect("key should build")
}
#[tokio::test]
async fn enriches_request_candidate_trace_with_provider_catalog_metadata() {
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
sample_candidate("req-1"),
]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider()],
vec![sample_endpoint()],
vec![sample_key()],
));
let state = GatewayDataState::with_decision_trace_readers_for_tests(
request_candidates,
provider_catalog,
);
let trace = read_decision_trace(&state, "req-1", true)
.await
.expect("trace should read")
.expect("trace should exist");
assert_eq!(
trace,
DecisionTrace {
request_id: "req-1".to_string(),
total_candidates: 1,
final_status: RequestCandidateFinalStatus::Failed,
total_latency_ms: 37,
candidates: vec![DecisionTraceCandidate {
candidate: sample_candidate("req-1"),
provider_name: Some("OpenAI".to_string()),
provider_website: Some("https://openai.com".to_string()),
provider_type: Some("custom".to_string()),
endpoint_api_format: Some("openai:chat".to_string()),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
provider_key_name: Some("prod-key".to_string()),
provider_key_auth_type: Some("api_key".to_string()),
provider_key_capabilities: Some(serde_json::json!({"cache_1h": true})),
provider_key_is_active: Some(true),
}],
}
);
}
}

80
crates/aether-gateway/src/data/gemini.rs Normal file
View File

@@ -0,0 +1,80 @@
use aether_data::repository::video_tasks::{StoredVideoTask, VideoTaskStatus};
use serde_json::json;
use crate::gateway::video_tasks::LocalVideoTaskReadResponse;
pub(super) fn map_gemini_video_task_to_read_response(
task: StoredVideoTask,
) -> LocalVideoTaskReadResponse {
match task.status {
VideoTaskStatus::Cancelled => LocalVideoTaskReadResponse {
status_code: 404,
body_json: json!({"detail": "Video task was cancelled"}),
},
VideoTaskStatus::Deleted => LocalVideoTaskReadResponse {
status_code: 404,
body_json: json!({"detail": "Video task not found"}),
},
VideoTaskStatus::Completed => LocalVideoTaskReadResponse {
status_code: 200,
body_json: build_gemini_completed_body(task),
},
VideoTaskStatus::Failed | VideoTaskStatus::Expired => LocalVideoTaskReadResponse {
status_code: 200,
body_json: build_gemini_failed_body(task),
},
_ => LocalVideoTaskReadResponse {
status_code: 200,
body_json: build_gemini_pending_body(task),
},
}
}
fn build_gemini_completed_body(task: StoredVideoTask) -> serde_json::Value {
let operation_name = operation_name(&task);
let short_id = task.short_id.unwrap_or_default();
json!({
"name": operation_name,
"done": true,
"response": {
"generateVideoResponse": {
"generatedSamples": [
{
"video": {
"uri": format!("/v1beta/files/aev_{short_id}:download?alt=media"),
"mimeType": "video/mp4"
}
}
]
}
}
})
}
fn build_gemini_failed_body(task: StoredVideoTask) -> serde_json::Value {
json!({
"name": operation_name(&task),
"done": true,
"error": {
"code": task.error_code.unwrap_or_else(|| "UNKNOWN".to_string()),
"message": task
.error_message
.unwrap_or_else(|| "Video generation failed".to_string()),
}
})
}
fn build_gemini_pending_body(task: StoredVideoTask) -> serde_json::Value {
json!({
"name": operation_name(&task),
"done": false,
"metadata": {}
})
}
fn operation_name(task: &StoredVideoTask) -> String {
let model = task.model.clone().unwrap_or_else(|| "unknown".to_string());
let short_id = task.short_id.clone().unwrap_or_else(|| task.id.clone());
format!("models/{model}/operations/{short_id}")
}

21
crates/aether-gateway/src/data/mod.rs Normal file
View File

@@ -0,0 +1,21 @@
mod auth;
mod candidates;
mod config;
mod decision_trace;
mod gemini;
mod openai;
mod request_audit;
mod state;
mod usage;
mod video_tasks;
#[cfg(test)]
mod tests;
pub(crate) use auth::StoredGatewayAuthApiKeySnapshot;
pub(crate) use candidates::RequestCandidateTrace;
pub use config::GatewayDataConfig;
pub(crate) use decision_trace::DecisionTrace;
pub(crate) use request_audit::RequestAuditBundle;
pub(crate) use state::GatewayDataState;
pub(crate) use usage::RequestUsageAudit;

69
crates/aether-gateway/src/data/openai.rs Normal file
View File

@@ -0,0 +1,69 @@
use aether_data::repository::video_tasks::{StoredVideoTask, VideoTaskStatus};
use serde_json::{json, Value};
use crate::gateway::video_tasks::LocalVideoTaskReadResponse;
pub(super) fn map_openai_video_task_to_read_response(
task: StoredVideoTask,
) -> LocalVideoTaskReadResponse {
match task.status {
VideoTaskStatus::Cancelled => LocalVideoTaskReadResponse {
status_code: 404,
body_json: json!({"detail": "Video task was cancelled"}),
},
VideoTaskStatus::Deleted => LocalVideoTaskReadResponse {
status_code: 404,
body_json: json!({"detail": "Video task not found"}),
},
status => LocalVideoTaskReadResponse {
status_code: 200,
body_json: build_openai_video_task_body(task, status),
},
}
}
fn build_openai_video_task_body(task: StoredVideoTask, status: VideoTaskStatus) -> Value {
let mut body = json!({
"id": task.id,
"object": "video",
"status": map_openai_video_status(status),
"progress": task.progress_percent,
"created_at": task.created_at_unix_secs,
});
if let Some(model) = task.model {
body["model"] = Value::String(model);
}
if let Some(prompt) = task.prompt {
body["prompt"] = Value::String(prompt);
}
if let Some(size) = task.size {
body["size"] = Value::String(size);
}
if let Some(video_url) = task.video_url {
body["video_url"] = Value::String(video_url);
}
if matches!(
status,
VideoTaskStatus::Failed | VideoTaskStatus::Expired | VideoTaskStatus::Cancelled
) {
body["error"] = json!({
"code": task.error_code.unwrap_or_else(|| "unknown".to_string()),
"message": task
.error_message
.unwrap_or_else(|| "Video generation failed".to_string()),
});
}
body
}
fn map_openai_video_status(status: VideoTaskStatus) -> &'static str {
match status {
VideoTaskStatus::Pending | VideoTaskStatus::Submitted | VideoTaskStatus::Queued => "queued",
VideoTaskStatus::Processing => "processing",
VideoTaskStatus::Completed => "completed",
VideoTaskStatus::Failed | VideoTaskStatus::Cancelled | VideoTaskStatus::Expired => "failed",
VideoTaskStatus::Deleted => "deleted",
}
}

56
crates/aether-gateway/src/data/request_audit.rs Normal file
View File

@@ -0,0 +1,56 @@
use aether_data::DataLayerError;
use super::auth::StoredGatewayAuthApiKeySnapshot;
use super::decision_trace::DecisionTrace;
use super::state::GatewayDataState;
use super::usage::RequestUsageAudit;
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
pub(crate) struct RequestAuditBundle {
pub(crate) request_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) usage: Option<RequestUsageAudit>,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) decision_trace: Option<DecisionTrace>,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) auth_snapshot: Option<StoredGatewayAuthApiKeySnapshot>,
}
pub(crate) async fn read_request_audit_bundle(
state: &GatewayDataState,
request_id: &str,
attempted_only: bool,
now_unix_secs: u64,
) -> Result<Option<RequestAuditBundle>, DataLayerError> {
let usage = state.read_request_usage_audit(request_id).await?;
let decision_trace = state
.read_decision_trace(request_id, attempted_only)
.await?;
let auth_snapshot = if let Some(usage) = usage.as_ref() {
match (
usage.usage.user_id.as_deref(),
usage.usage.api_key_id.as_deref(),
) {
(Some(user_id), Some(api_key_id)) => {
state
.read_auth_api_key_snapshot(user_id, api_key_id, now_unix_secs)
.await?
}
_ => None,
}
} else {
None
};
if usage.is_none() && decision_trace.is_none() && auth_snapshot.is_none() {
return Ok(None);
}
Ok(Some(RequestAuditBundle {
request_id: request_id.to_string(),
usage,
decision_trace,
auth_snapshot,
}))
}

454
crates/aether-gateway/src/data/state.rs Normal file
View File

@@ -0,0 +1,454 @@
use std::fmt;
use std::sync::Arc;
use aether_data::repository::auth::{
AuthApiKeyLookupKey, AuthApiKeyReadRepository, StoredAuthApiKeySnapshot,
};
use aether_data::repository::candidates::{RequestCandidateReadRepository, StoredRequestCandidate};
use aether_data::repository::provider_catalog::{
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogProvider,
};
use aether_data::repository::shadow_results::{
merge_shadow_result_sample, RecordShadowResultSample, ShadowResultLookupKey,
ShadowResultReadRepository, ShadowResultWriteRepository, StoredShadowResult,
};
use aether_data::repository::usage::{StoredRequestUsageAudit, UsageReadRepository};
use aether_data::repository::video_tasks::{
StoredVideoTask, VideoTaskLookupKey, VideoTaskReadRepository,
};
use aether_data::{DataBackends, DataLayerError};
use super::auth::{read_auth_api_key_snapshot, StoredGatewayAuthApiKeySnapshot};
use super::candidates::{read_request_candidate_trace, RequestCandidateTrace};
use super::config::GatewayDataConfig;
use super::decision_trace::{read_decision_trace, DecisionTrace};
use super::request_audit::{read_request_audit_bundle, RequestAuditBundle};
use super::usage::{read_request_usage_audit, RequestUsageAudit};
use super::video_tasks::read_video_task_response;
use crate::gateway::video_tasks::LocalVideoTaskReadResponse;
#[derive(Clone, Default)]
pub(crate) struct GatewayDataState {
config: GatewayDataConfig,
backends: Option<DataBackends>,
auth_api_key_reader: Option<Arc<dyn AuthApiKeyReadRepository>>,
request_candidate_reader: Option<Arc<dyn RequestCandidateReadRepository>>,
provider_catalog_reader: Option<Arc<dyn ProviderCatalogReadRepository>>,
usage_reader: Option<Arc<dyn UsageReadRepository>>,
video_task_reader: Option<Arc<dyn VideoTaskReadRepository>>,
shadow_result_reader: Option<Arc<dyn ShadowResultReadRepository>>,
shadow_result_writer: Option<Arc<dyn ShadowResultWriteRepository>>,
}
impl fmt::Debug for GatewayDataState {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("GatewayDataState")
.field("config", &self.config)
.field("has_backends", &self.backends.is_some())
.field(
"has_auth_api_key_reader",
&self.auth_api_key_reader.is_some(),
)
.field(
"has_request_candidate_reader",
&self.request_candidate_reader.is_some(),
)
.field(
"has_provider_catalog_reader",
&self.provider_catalog_reader.is_some(),
)
.field("has_usage_reader", &self.usage_reader.is_some())
.field("has_video_task_reader", &self.video_task_reader.is_some())
.field(
"has_shadow_result_reader",
&self.shadow_result_reader.is_some(),
)
.field(
"has_shadow_result_writer",
&self.shadow_result_writer.is_some(),
)
.finish()
}
}
impl GatewayDataState {
pub(crate) fn disabled() -> Self {
Self::default()
}
pub(crate) fn from_config(config: GatewayDataConfig) -> Result<Self, DataLayerError> {
if !config.is_enabled() {
return Ok(Self {
config,
backends: None,
auth_api_key_reader: None,
request_candidate_reader: None,
provider_catalog_reader: None,
usage_reader: None,
video_task_reader: None,
shadow_result_reader: None,
shadow_result_writer: None,
});
}
let backends = DataBackends::from_config(config.to_data_layer_config())?;
let auth_api_key_reader = backends.read().auth_api_keys();
let request_candidate_reader = backends.read().request_candidates();
let provider_catalog_reader = backends.read().provider_catalog();
let usage_reader = backends.read().usage();
let video_task_reader = backends.read().video_tasks();
let shadow_result_reader = backends.read().shadow_results();
let shadow_result_writer = backends.write().shadow_results();
Ok(Self {
config,
backends: Some(backends),
auth_api_key_reader,
request_candidate_reader,
provider_catalog_reader,
usage_reader,
video_task_reader,
shadow_result_reader,
shadow_result_writer,
})
}
pub(crate) fn has_backends(&self) -> bool {
self.backends.is_some()
}
pub(crate) fn has_auth_api_key_reader(&self) -> bool {
self.auth_api_key_reader.is_some()
}
pub(crate) fn has_request_candidate_reader(&self) -> bool {
self.request_candidate_reader.is_some()
}
pub(crate) fn has_provider_catalog_reader(&self) -> bool {
self.provider_catalog_reader.is_some()
}
pub(crate) fn has_usage_reader(&self) -> bool {
self.usage_reader.is_some()
}
pub(crate) fn has_video_task_reader(&self) -> bool {
self.video_task_reader.is_some()
}
pub(crate) fn has_shadow_result_writer(&self) -> bool {
self.shadow_result_writer.is_some()
}
pub(crate) fn has_shadow_result_reader(&self) -> bool {
self.shadow_result_reader.is_some()
}
pub(super) async fn find_video_task(
&self,
key: VideoTaskLookupKey<'_>,
) -> Result<Option<StoredVideoTask>, DataLayerError> {
match &self.video_task_reader {
Some(repository) => repository.find(key).await,
None => Ok(None),
}
}
pub(super) async fn find_auth_api_key_snapshot(
&self,
key: AuthApiKeyLookupKey<'_>,
) -> Result<Option<StoredAuthApiKeySnapshot>, DataLayerError> {
match &self.auth_api_key_reader {
Some(repository) => repository.find_api_key_snapshot(key).await,
None => Ok(None),
}
}
pub(super) async fn list_request_candidates_by_request_id(
&self,
request_id: &str,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
match &self.request_candidate_reader {
Some(repository) => repository.list_by_request_id(request_id).await,
None => Ok(Vec::new()),
}
}
pub(super) async fn list_provider_catalog_providers_by_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
match &self.provider_catalog_reader {
Some(repository) => repository.list_providers_by_ids(provider_ids).await,
None => Ok(Vec::new()),
}
}
pub(super) async fn list_provider_catalog_endpoints_by_ids(
&self,
endpoint_ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
match &self.provider_catalog_reader {
Some(repository) => repository.list_endpoints_by_ids(endpoint_ids).await,
None => Ok(Vec::new()),
}
}
pub(super) async fn list_provider_catalog_keys_by_ids(
&self,
key_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
match &self.provider_catalog_reader {
Some(repository) => repository.list_keys_by_ids(key_ids).await,
None => Ok(Vec::new()),
}
}
pub(super) async fn find_request_usage_by_request_id(
&self,
request_id: &str,
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
match &self.usage_reader {
Some(repository) => repository.find_by_request_id(request_id).await,
None => Ok(None),
}
}
pub(crate) async fn read_request_candidate_trace(
&self,
request_id: &str,
attempted_only: bool,
) -> Result<Option<RequestCandidateTrace>, DataLayerError> {
read_request_candidate_trace(self, request_id, attempted_only).await
}
pub(crate) async fn read_decision_trace(
&self,
request_id: &str,
attempted_only: bool,
) -> Result<Option<DecisionTrace>, DataLayerError> {
read_decision_trace(self, request_id, attempted_only).await
}
pub(crate) async fn read_request_usage_audit(
&self,
request_id: &str,
) -> Result<Option<RequestUsageAudit>, DataLayerError> {
read_request_usage_audit(self, request_id).await
}
pub(crate) async fn read_request_audit_bundle(
&self,
request_id: &str,
attempted_only: bool,
now_unix_secs: u64,
) -> Result<Option<RequestAuditBundle>, DataLayerError> {
read_request_audit_bundle(self, request_id, attempted_only, now_unix_secs).await
}
pub(crate) async fn read_auth_api_key_snapshot(
&self,
user_id: &str,
api_key_id: &str,
now_unix_secs: u64,
) -> Result<Option<StoredGatewayAuthApiKeySnapshot>, DataLayerError> {
read_auth_api_key_snapshot(self, user_id, api_key_id, now_unix_secs).await
}
pub(crate) async fn read_video_task_response(
&self,
route_family: Option<&str>,
request_path: &str,
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
read_video_task_response(self, route_family, request_path).await
}
#[cfg(test)]
pub(crate) async fn write_shadow_result(
&self,
result: aether_data::repository::shadow_results::UpsertShadowResult,
) -> Result<Option<StoredShadowResult>, DataLayerError> {
match &self.shadow_result_writer {
Some(repository) => repository.upsert(result).await.map(Some),
None => Ok(None),
}
}
pub(crate) async fn record_shadow_result_sample(
&self,
sample: RecordShadowResultSample,
) -> Result<Option<StoredShadowResult>, DataLayerError> {
let Some(writer) = &self.shadow_result_writer else {
return Ok(None);
};
let existing = match &self.shadow_result_reader {
Some(reader) => {
reader
.find(ShadowResultLookupKey::TraceFingerprint {
trace_id: &sample.trace_id,
request_fingerprint: &sample.request_fingerprint,
})
.await?
}
None => None,
};
let merged = merge_shadow_result_sample(existing.as_ref(), sample);
writer.upsert(merged).await.map(Some)
}
pub(crate) async fn list_recent_shadow_results(
&self,
limit: usize,
) -> Result<Vec<StoredShadowResult>, DataLayerError> {
match &self.shadow_result_reader {
Some(repository) => repository.list_recent(limit).await,
None => Ok(Vec::new()),
}
}
#[cfg(test)]
pub(crate) fn with_video_task_reader_for_tests(
repository: Arc<dyn VideoTaskReadRepository>,
) -> Self {
Self {
config: GatewayDataConfig::disabled(),
backends: None,
auth_api_key_reader: None,
request_candidate_reader: None,
provider_catalog_reader: None,
usage_reader: None,
video_task_reader: Some(repository),
shadow_result_reader: None,
shadow_result_writer: None,
}
}
#[cfg(test)]
pub(crate) fn with_request_candidate_reader_for_tests(
repository: Arc<dyn RequestCandidateReadRepository>,
) -> Self {
Self {
config: GatewayDataConfig::disabled(),
backends: None,
auth_api_key_reader: None,
request_candidate_reader: Some(repository),
provider_catalog_reader: None,
usage_reader: None,
video_task_reader: None,
shadow_result_reader: None,
shadow_result_writer: None,
}
}
#[cfg(test)]
pub(crate) fn with_usage_reader_for_tests(repository: Arc<dyn UsageReadRepository>) -> Self {
Self {
config: GatewayDataConfig::disabled(),
backends: None,
auth_api_key_reader: None,
request_candidate_reader: None,
provider_catalog_reader: None,
usage_reader: Some(repository),
video_task_reader: None,
shadow_result_reader: None,
shadow_result_writer: None,
}
}
#[cfg(test)]
pub(crate) fn with_auth_api_key_reader_for_tests(
repository: Arc<dyn AuthApiKeyReadRepository>,
) -> Self {
Self {
config: GatewayDataConfig::disabled(),
backends: None,
auth_api_key_reader: Some(repository),
request_candidate_reader: None,
provider_catalog_reader: None,
usage_reader: None,
video_task_reader: None,
shadow_result_reader: None,
shadow_result_writer: None,
}
}
#[cfg(test)]
pub(crate) fn with_decision_trace_readers_for_tests(
request_candidate_repository: Arc<dyn RequestCandidateReadRepository>,
provider_catalog_repository: Arc<dyn ProviderCatalogReadRepository>,
) -> Self {
Self {
config: GatewayDataConfig::disabled(),
backends: None,
auth_api_key_reader: None,
request_candidate_reader: Some(request_candidate_repository),
provider_catalog_reader: Some(provider_catalog_repository),
usage_reader: None,
video_task_reader: None,
shadow_result_reader: None,
shadow_result_writer: None,
}
}
#[cfg(test)]
pub(crate) fn with_request_audit_readers_for_tests(
auth_api_key_repository: Arc<dyn AuthApiKeyReadRepository>,
request_candidate_repository: Arc<dyn RequestCandidateReadRepository>,
provider_catalog_repository: Arc<dyn ProviderCatalogReadRepository>,
usage_repository: Arc<dyn UsageReadRepository>,
) -> Self {
Self {
config: GatewayDataConfig::disabled(),
backends: None,
auth_api_key_reader: Some(auth_api_key_repository),
request_candidate_reader: Some(request_candidate_repository),
provider_catalog_reader: Some(provider_catalog_repository),
usage_reader: Some(usage_repository),
video_task_reader: None,
shadow_result_reader: None,
shadow_result_writer: None,
}
}
#[cfg(test)]
pub(crate) fn with_shadow_result_writer_for_tests(
repository: Arc<dyn ShadowResultWriteRepository>,
) -> Self {
Self {
config: GatewayDataConfig::disabled(),
backends: None,
auth_api_key_reader: None,
request_candidate_reader: None,
provider_catalog_reader: None,
usage_reader: None,
video_task_reader: None,
shadow_result_reader: None,
shadow_result_writer: Some(repository),
}
}
#[cfg(test)]
pub(crate) fn with_shadow_result_repository_for_tests<T>(repository: Arc<T>) -> Self
where
T: aether_data::repository::shadow_results::ShadowResultRepository + 'static,
{
let shadow_result_reader: Arc<dyn ShadowResultReadRepository> = repository.clone();
let shadow_result_writer: Arc<dyn ShadowResultWriteRepository> = repository;
Self {
config: GatewayDataConfig::disabled(),
backends: None,
auth_api_key_reader: None,
request_candidate_reader: None,
provider_catalog_reader: None,
usage_reader: None,
video_task_reader: None,
shadow_result_reader: Some(shadow_result_reader),
shadow_result_writer: Some(shadow_result_writer),
}
}
}

683
crates/aether-gateway/src/data/tests.rs Normal file
View File

@@ -0,0 +1,683 @@
use std::sync::Arc;
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
};
use aether_data::repository::candidates::{
InMemoryRequestCandidateRepository, RequestCandidateStatus, StoredRequestCandidate,
};
use aether_data::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogProvider,
};
use aether_data::repository::shadow_results::{
InMemoryShadowResultRepository, RecordShadowResultSample, ShadowResultLookupKey,
ShadowResultMatchStatus, ShadowResultReadRepository, ShadowResultSampleOrigin,
UpsertShadowResult,
};
use aether_data::repository::usage::{InMemoryUsageReadRepository, StoredRequestUsageAudit};
use aether_data::repository::video_tasks::{
InMemoryVideoTaskRepository, UpsertVideoTask, VideoTaskLookupKey, VideoTaskStatus,
VideoTaskWriteRepository,
};
use super::{GatewayDataConfig, GatewayDataState};
use crate::gateway::AppState;
#[test]
fn disabled_gateway_data_state_has_no_backends() {
let state = GatewayDataState::from_config(GatewayDataConfig::disabled())
.expect("disabled config should build");
assert!(!state.has_backends());
assert!(!state.has_auth_api_key_reader());
assert!(!state.has_request_candidate_reader());
assert!(!state.has_provider_catalog_reader());
assert!(!state.has_usage_reader());
assert!(!state.has_video_task_reader());
assert!(!state.has_shadow_result_reader());
assert!(!state.has_shadow_result_writer());
}
#[tokio::test]
async fn postgres_gateway_data_state_builds_video_task_reader() {
let state = GatewayDataState::from_config(GatewayDataConfig::from_postgres_url(
"postgres://localhost/aether",
false,
))
.expect("postgres-backed state should build");
assert!(state.has_backends());
assert!(state.has_auth_api_key_reader());
assert!(state.has_request_candidate_reader());
assert!(state.has_provider_catalog_reader());
assert!(state.has_usage_reader());
assert!(state.has_video_task_reader());
assert!(state.has_shadow_result_reader());
assert!(state.has_shadow_result_writer());
}
#[tokio::test]
async fn data_state_find_uses_configured_read_repository() {
let repository = Arc::new(InMemoryVideoTaskRepository::default());
repository
.upsert(UpsertVideoTask {
id: "task-1".to_string(),
short_id: Some("short-task-1".to_string()),
user_id: Some("user-1".to_string()),
external_task_id: Some("ext-task-1".to_string()),
provider_api_format: Some("openai:video".to_string()),
model: Some("sora-2".to_string()),
prompt: Some("hello".to_string()),
size: Some("1280x720".to_string()),
status: VideoTaskStatus::Queued,
progress_percent: 0,
created_at_unix_secs: 100,
updated_at_unix_secs: 100,
error_code: None,
error_message: None,
video_url: None,
})
.await
.expect("upsert should succeed");
let state = GatewayDataState::with_video_task_reader_for_tests(repository);
let task = state
.find_video_task(VideoTaskLookupKey::Id("task-1"))
.await
.expect("find should succeed");
assert_eq!(task.expect("task should exist").id, "task-1");
}
#[tokio::test]
async fn app_state_wires_gateway_data_state_from_config() {
let state = AppState::new_with_executor(
"http://127.0.0.1:18084",
Some("http://127.0.0.1:18085".to_string()),
Some("http://127.0.0.1:18086".to_string()),
)
.expect("app state should build")
.with_data_config(GatewayDataConfig::from_postgres_url(
"postgres://localhost/aether",
false,
))
.expect("data config should wire");
assert!(state.data.has_backends());
assert!(state.data.has_auth_api_key_reader());
assert!(state.data.has_request_candidate_reader());
assert!(state.data.has_provider_catalog_reader());
assert!(state.data.has_usage_reader());
assert!(state.data.has_video_task_reader());
assert!(state.data.has_shadow_result_reader());
assert!(state.data.has_shadow_result_writer());
}
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-4.1"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(200),
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-4.1"])),
)
.expect("auth snapshot should build")
}
#[tokio::test]
async fn data_state_reads_auth_api_key_snapshot_from_reader() {
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-1".to_string()),
sample_auth_snapshot("key-1", "user-1"),
)]));
let state = GatewayDataState::with_auth_api_key_reader_for_tests(repository);
let snapshot = state
.read_auth_api_key_snapshot("user-1", "key-1", 150)
.await
.expect("read should succeed")
.expect("snapshot should exist");
assert_eq!(snapshot.user_id, "user-1");
assert_eq!(snapshot.api_key_id, "key-1");
assert_eq!(snapshot.username, "alice");
assert_eq!(
snapshot.api_key_allowed_models,
Some(vec!["gpt-4.1".to_string()])
);
assert!(snapshot.currently_usable);
}
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-1".to_string(),
"OpenAI".to_string(),
Some("https://openai.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
}
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-1".to_string(),
"provider-1".to_string(),
"openai:chat".to_string(),
Some("openai".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
}
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"provider-key-1".to_string(),
"provider-1".to_string(),
"prod-key".to_string(),
"api_key".to_string(),
Some(serde_json::json!({"cache_1h": true})),
true,
)
.expect("key should build")
}
fn sample_request_usage(request_id: &str) -> StoredRequestUsageAudit {
StoredRequestUsageAudit::new(
"usage-1".to_string(),
request_id.to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
"OpenAI".to_string(),
"gpt-4.1".to_string(),
Some("gpt-4.1-mini".to_string()),
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
true,
false,
120,
40,
160,
0.24,
0.36,
Some(200),
None,
None,
Some(450),
Some(120),
"completed".to_string(),
"settled".to_string(),
100,
101,
Some(102),
)
.expect("usage should build")
}
#[tokio::test]
async fn data_state_reads_decision_trace_with_provider_catalog_metadata() {
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
StoredRequestCandidate::new(
"cand-1".to_string(),
"req-1".to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
RequestCandidateStatus::Failed,
None,
false,
Some(502),
None,
None,
Some(37),
Some(1),
None,
Some(serde_json::json!({"cache_1h": true})),
100,
Some(101),
Some(102),
)
.expect("candidate should build"),
]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_catalog_provider()],
vec![sample_provider_catalog_endpoint()],
vec![sample_provider_catalog_key()],
));
let state = GatewayDataState::with_decision_trace_readers_for_tests(
request_candidates,
provider_catalog,
);
let trace = state
.read_decision_trace("req-1", true)
.await
.expect("trace should read")
.expect("trace should exist");
assert_eq!(trace.request_id, "req-1");
assert_eq!(trace.total_candidates, 1);
assert_eq!(trace.candidates[0].provider_name.as_deref(), Some("OpenAI"));
assert_eq!(
trace.candidates[0].endpoint_api_format.as_deref(),
Some("openai:chat")
);
assert_eq!(
trace.candidates[0].provider_key_auth_type.as_deref(),
Some("api_key")
);
assert_eq!(
trace.candidates[0].provider_key_capabilities,
Some(serde_json::json!({"cache_1h": true}))
);
}
#[tokio::test]
async fn data_state_reads_request_usage_audit_from_reader() {
let repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
sample_request_usage("req-usage-1"),
]));
let state = GatewayDataState::with_usage_reader_for_tests(repository);
let usage = state
.read_request_usage_audit("req-usage-1")
.await
.expect("read should succeed")
.expect("usage should exist");
assert_eq!(usage.usage.request_id, "req-usage-1");
assert_eq!(usage.usage.provider_name, "OpenAI");
assert_eq!(usage.usage.total_tokens, 160);
assert_eq!(usage.usage.total_cost_usd, 0.24);
assert!(usage.usage.has_format_conversion);
}
#[tokio::test]
async fn data_state_reads_request_audit_bundle_from_multiple_readers() {
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-1".to_string()),
sample_auth_snapshot("api-key-1", "user-1"),
)]));
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
StoredRequestCandidate::new(
"cand-1".to_string(),
"req-usage-1".to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
RequestCandidateStatus::Success,
None,
false,
Some(200),
None,
None,
Some(37),
Some(1),
None,
Some(serde_json::json!({"cache_1h": true})),
100,
Some(101),
Some(102),
)
.expect("candidate should build"),
]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_catalog_provider()],
vec![sample_provider_catalog_endpoint()],
vec![sample_provider_catalog_key()],
));
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
sample_request_usage("req-usage-1"),
]));
let state = GatewayDataState::with_request_audit_readers_for_tests(
auth_repository,
request_candidates,
provider_catalog,
usage_repository,
);
let bundle = state
.read_request_audit_bundle("req-usage-1", true, 150)
.await
.expect("bundle should read")
.expect("bundle should exist");
assert_eq!(bundle.request_id, "req-usage-1");
assert_eq!(
bundle
.usage
.as_ref()
.and_then(|usage| usage.usage.target_model.as_deref()),
Some("gpt-4.1-mini")
);
assert_eq!(
bundle
.decision_trace
.as_ref()
.and_then(|trace| trace.candidates.first())
.and_then(|candidate| candidate.provider_name.as_deref()),
Some("OpenAI")
);
assert_eq!(
bundle
.auth_snapshot
.as_ref()
.map(|snapshot| snapshot.currently_usable),
Some(true)
);
}
#[tokio::test]
async fn maps_openai_video_task_repository_row_into_read_response() {
let repository = Arc::new(InMemoryVideoTaskRepository::default());
repository
.upsert(UpsertVideoTask {
id: "task-1".to_string(),
short_id: Some("short-task-1".to_string()),
user_id: Some("user-1".to_string()),
external_task_id: Some("ext-task-1".to_string()),
provider_api_format: Some("openai:video".to_string()),
model: Some("sora-2".to_string()),
prompt: Some("hello".to_string()),
size: Some("1280x720".to_string()),
status: VideoTaskStatus::Processing,
progress_percent: 45,
created_at_unix_secs: 100,
updated_at_unix_secs: 120,
error_code: None,
error_message: None,
video_url: None,
})
.await
.expect("upsert should succeed");
let state = GatewayDataState::with_video_task_reader_for_tests(repository);
let response = state
.read_video_task_response(Some("openai"), "/v1/videos/task-1")
.await
.expect("read should succeed")
.expect("read response should exist");
assert_eq!(response.status_code, 200);
assert_eq!(response.body_json["id"], "task-1");
assert_eq!(response.body_json["status"], "processing");
assert_eq!(response.body_json["created_at"], 100);
}
#[tokio::test]
async fn maps_gemini_video_task_repository_row_into_read_response() {
let repository = Arc::new(InMemoryVideoTaskRepository::default());
repository
.upsert(UpsertVideoTask {
id: "task-1".to_string(),
short_id: Some("localshort123".to_string()),
user_id: Some("user-1".to_string()),
external_task_id: Some("operations/ext-task-1".to_string()),
provider_api_format: Some("gemini:video".to_string()),
model: Some("veo-3".to_string()),
prompt: Some("hello".to_string()),
size: Some("720p".to_string()),
status: VideoTaskStatus::Completed,
progress_percent: 100,
created_at_unix_secs: 100,
updated_at_unix_secs: 120,
error_code: None,
error_message: None,
video_url: None,
})
.await
.expect("upsert should succeed");
let state = GatewayDataState::with_video_task_reader_for_tests(repository);
let response = state
.read_video_task_response(
Some("gemini"),
"/v1beta/models/veo-3/operations/localshort123",
)
.await
.expect("read should succeed")
.expect("read response should exist");
assert_eq!(response.status_code, 200);
assert_eq!(
response.body_json["name"],
"models/veo-3/operations/localshort123"
);
assert_eq!(response.body_json["done"], true);
}
#[tokio::test]
async fn data_state_write_uses_configured_shadow_result_writer() {
let repository = Arc::new(InMemoryShadowResultRepository::default());
let state = GatewayDataState::with_shadow_result_writer_for_tests(repository.clone());
let written = state
.write_shadow_result(UpsertShadowResult {
trace_id: "trace-1".to_string(),
request_fingerprint: "fp-1".to_string(),
request_id: Some("req-1".to_string()),
route_family: Some("openai".to_string()),
route_kind: Some("chat".to_string()),
candidate_id: None,
rust_result_digest: Some("rust-digest".to_string()),
python_result_digest: None,
match_status: ShadowResultMatchStatus::Pending,
status_code: Some(200),
error_message: None,
created_at_unix_secs: 100,
updated_at_unix_secs: 100,
})
.await
.expect("write should succeed");
assert!(written.is_some());
let stored = repository
.find(ShadowResultLookupKey::TraceFingerprint {
trace_id: "trace-1",
request_fingerprint: "fp-1",
})
.await
.expect("find should succeed");
assert_eq!(
stored.expect("stored result should exist").match_status,
ShadowResultMatchStatus::Pending
);
}
#[tokio::test]
async fn data_state_records_shadow_result_samples_and_merges_match_status() {
let repository = Arc::new(InMemoryShadowResultRepository::default());
let state = GatewayDataState::with_shadow_result_repository_for_tests(repository);
let first = state
.record_shadow_result_sample(RecordShadowResultSample {
trace_id: "trace-1".to_string(),
request_fingerprint: "fp-1".to_string(),
request_id: Some("req-1".to_string()),
route_family: Some("openai".to_string()),
route_kind: Some("chat".to_string()),
candidate_id: None,
origin: ShadowResultSampleOrigin::Rust,
result_digest: "digest-1".to_string(),
status_code: Some(200),
error_message: None,
recorded_at_unix_secs: 100,
})
.await
.expect("first record should succeed")
.expect("first stored result should exist");
assert_eq!(first.match_status, ShadowResultMatchStatus::Pending);
let second = state
.record_shadow_result_sample(RecordShadowResultSample {
trace_id: "trace-1".to_string(),
request_fingerprint: "fp-1".to_string(),
request_id: Some("req-1".to_string()),
route_family: Some("openai".to_string()),
route_kind: Some("chat".to_string()),
candidate_id: None,
origin: ShadowResultSampleOrigin::Python,
result_digest: "digest-1".to_string(),
status_code: Some(200),
error_message: None,
recorded_at_unix_secs: 200,
})
.await
.expect("second record should succeed")
.expect("second stored result should exist");
assert_eq!(second.match_status, ShadowResultMatchStatus::Match);
assert_eq!(second.created_at_unix_secs, 100);
assert_eq!(second.updated_at_unix_secs, 200);
assert_eq!(second.request_id.as_deref(), Some("req-1"));
}
#[tokio::test]
async fn data_state_lists_recent_shadow_results_from_reader() {
let repository = Arc::new(InMemoryShadowResultRepository::default());
let state = GatewayDataState::with_shadow_result_repository_for_tests(repository.clone());
state
.record_shadow_result_sample(RecordShadowResultSample {
trace_id: "trace-1".to_string(),
request_fingerprint: "fp-1".to_string(),
request_id: Some("req-shadow-1".to_string()),
route_family: Some("openai".to_string()),
route_kind: Some("chat".to_string()),
candidate_id: None,
origin: ShadowResultSampleOrigin::Rust,
result_digest: "digest-1".to_string(),
status_code: Some(200),
error_message: None,
recorded_at_unix_secs: 100,
})
.await
.expect("record should succeed");
let recent = state
.list_recent_shadow_results(5)
.await
.expect("list recent should succeed");
assert_eq!(recent.len(), 1);
assert_eq!(recent[0].trace_id, "trace-1");
assert_eq!(recent[0].request_id.as_deref(), Some("req-shadow-1"));
}
fn sample_request_candidate(
id: &str,
request_id: &str,
candidate_index: i32,
status: RequestCandidateStatus,
started_at_unix_secs: Option<i64>,
latency_ms: Option<i32>,
status_code: Option<i32>,
) -> StoredRequestCandidate {
StoredRequestCandidate::new(
id.to_string(),
request_id.to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
candidate_index,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
status,
None,
false,
status_code,
None,
None,
latency_ms,
Some(1),
None,
None,
100 + i64::from(candidate_index),
started_at_unix_secs,
started_at_unix_secs.map(|value| value + 1),
)
.expect("candidate should build")
}
#[tokio::test]
async fn data_state_reads_request_candidate_trace_from_reader() {
let repository = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
sample_request_candidate(
"cand-1",
"req-1",
0,
RequestCandidateStatus::Pending,
None,
None,
None,
),
sample_request_candidate(
"cand-2",
"req-1",
1,
RequestCandidateStatus::Success,
Some(101),
Some(42),
Some(200),
),
]));
let state = GatewayDataState::with_request_candidate_reader_for_tests(repository);
let trace = state
.read_request_candidate_trace("req-1", true)
.await
.expect("trace should succeed")
.expect("trace should exist");
assert_eq!(trace.request_id, "req-1");
assert_eq!(trace.total_candidates, 1);
assert_eq!(
trace.final_status,
super::candidates::RequestCandidateFinalStatus::Success
);
assert_eq!(trace.total_latency_ms, 42);
assert_eq!(trace.candidates[0].id, "cand-2");
}

20
crates/aether-gateway/src/data/usage.rs Normal file
View File

@@ -0,0 +1,20 @@
use aether_data::repository::usage::StoredRequestUsageAudit;
use aether_data::DataLayerError;
use super::state::GatewayDataState;
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
pub(crate) struct RequestUsageAudit {
#[serde(flatten)]
pub(crate) usage: StoredRequestUsageAudit,
}
pub(crate) async fn read_request_usage_audit(
state: &GatewayDataState,
request_id: &str,
) -> Result<Option<RequestUsageAudit>, DataLayerError> {
Ok(state
.find_request_usage_by_request_id(request_id)
.await?
.map(|usage| RequestUsageAudit { usage }))
}

65
crates/aether-gateway/src/data/video_tasks.rs Normal file
View File

@@ -0,0 +1,65 @@
use aether_data::repository::video_tasks::VideoTaskLookupKey;
use aether_data::DataLayerError;
use super::gemini::map_gemini_video_task_to_read_response;
use super::openai::map_openai_video_task_to_read_response;
use super::state::GatewayDataState;
use crate::gateway::video_tasks::{
extract_gemini_short_id_from_path, extract_openai_task_id_from_path, LocalVideoTaskReadResponse,
};
pub(super) async fn read_video_task_response(
state: &GatewayDataState,
route_family: Option<&str>,
request_path: &str,
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
match route_family {
Some("openai") => read_openai_video_task_response(state, request_path).await,
Some("gemini") => read_gemini_video_task_response(state, request_path).await,
_ => Ok(None),
}
}
async fn read_openai_video_task_response(
state: &GatewayDataState,
request_path: &str,
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
let Some(task_id) = extract_openai_task_id_from_path(request_path) else {
return Ok(None);
};
let Some(task) = state
.find_video_task(VideoTaskLookupKey::Id(task_id))
.await?
else {
return Ok(None);
};
if !matches!(task.provider_api_format.as_deref(), Some("openai:video")) {
return Ok(None);
}
Ok(Some(map_openai_video_task_to_read_response(task)))
}
async fn read_gemini_video_task_response(
state: &GatewayDataState,
request_path: &str,
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
let Some(short_id) = extract_gemini_short_id_from_path(request_path) else {
return Ok(None);
};
let Some(task) = state
.find_video_task(VideoTaskLookupKey::ShortId(short_id))
.await?
else {
return Ok(None);
};
if !matches!(task.provider_api_format.as_deref(), Some("gemini:video")) {
return Ok(None);
}
Ok(Some(map_gemini_video_task_to_read_response(task)))
}

6
crates/aether-gateway/src/executor.rs
View File

@@ -2,7 +2,7 @@ use std::collections::hash_map::DefaultHasher;
use std::collections::BTreeMap; use std::collections::BTreeMap;
use std::hash::{Hash, Hasher}; use std::hash::{Hash, Hasher};
use std::io::Error as IoError; use std::io::Error as IoError;
use std::time::{Duration, Instant}; use std::time::Duration;
use aether_contracts::{ use aether_contracts::{
ExecutionPlan, ExecutionResult, ExecutionTelemetry, ExecutionTimeouts, ProxySnapshot, ExecutionPlan, ExecutionResult, ExecutionTelemetry, ExecutionTimeouts, ProxySnapshot,
@@ -24,8 +24,8 @@ use crate::gateway::headers::{
should_skip_upstream_passthrough_header, should_skip_upstream_passthrough_header,
}; };
use crate::gateway::{ use crate::gateway::{
build_client_response, build_client_response_from_parts, cache_executor_auth_context, attach_control_metadata_headers, build_client_response, build_client_response_from_parts,
local_finalize::maybe_build_local_core_sync_finalize_response, cache_executor_auth_context, local_finalize::maybe_build_local_core_sync_finalize_response,
local_stream::maybe_build_local_stream_rewriter, resolve_executor_auth_context, AppState, local_stream::maybe_build_local_stream_rewriter, resolve_executor_auth_context, AppState,
GatewayControlAuthContext, GatewayControlDecision, GatewayError, GatewayControlAuthContext, GatewayControlDecision, GatewayError,
}; };

Some files were not shown because too many files have changed in this diff Show More