mirror of
https://github.com/telemt/telemt.git
synced 2026-07-19 07:09:49 +03:00
Compare commits
14 Commits
3.4.23
...
flow-runti
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
fabd98ce89 | ||
|
|
7df3dab5e8 | ||
|
|
c6f40e3717 | ||
|
|
91e05265be | ||
|
|
991d5b2c38 | ||
|
|
1f9c82c924 | ||
|
|
73afeccae1 | ||
|
|
feb51cbf57 | ||
|
|
8c65cd868c | ||
|
|
ea296bbdc8 | ||
|
|
fb042f826e | ||
|
|
96425f15c8 | ||
|
|
d4c4980e5a | ||
|
|
893ce0cf36 |
3
Cargo.lock
generated
3
Cargo.lock
generated
@@ -21,6 +21,7 @@ dependencies = [
|
||||
"cfg-if",
|
||||
"cipher",
|
||||
"cpufeatures 0.2.17",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2899,7 +2900,7 @@ checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417"
|
||||
|
||||
[[package]]
|
||||
name = "telemt"
|
||||
version = "3.4.23"
|
||||
version = "3.4.24"
|
||||
dependencies = [
|
||||
"aes",
|
||||
"anyhow",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "telemt"
|
||||
version = "3.4.23"
|
||||
version = "3.4.24"
|
||||
edition = "2024"
|
||||
|
||||
[features]
|
||||
@@ -15,8 +15,8 @@ tokio = { version = "1.52.3", features = ["full", "tracing"] }
|
||||
tokio-util = { version = "0.7.18", features = ["full"] }
|
||||
|
||||
# Crypto
|
||||
aes = "0.8.4"
|
||||
ctr = "0.9.2"
|
||||
aes = { version = "0.8.4", features = ["zeroize"] }
|
||||
ctr = { version = "0.9.2", features = ["zeroize"] }
|
||||
cbc = "0.1.2"
|
||||
sha2 = "0.10.9"
|
||||
sha1 = "0.10.6"
|
||||
|
||||
@@ -107,7 +107,9 @@ Notes:
|
||||
| `GET` | `/v1/stats/users/active-ips` | none | `200` | `UserActiveIps[]` |
|
||||
| `GET` | `/v1/stats/users` | none | `200` | `UserInfo[]` |
|
||||
| `GET` | `/v1/config` | none | `200` | `ConfigData` |
|
||||
| `PATCH` | `/v1/config` | sparse JSON object | `200` | `PatchConfigResponse` |
|
||||
| `PATCH` | `/v1/config` | sparse JSON object; optional reload query | `200` or `202` | `PatchConfigResponse` |
|
||||
| `POST` | `/v1/system/reload` | `ReloadRequest` or empty body | `202` | `ReloadAccepted` |
|
||||
| `GET` | `/v1/system/reload/{id}` | none | `200` | `ReloadStatus` |
|
||||
| `GET` | `/v1/users` | none | `200` | `UserInfo[]` |
|
||||
| `POST` | `/v1/users` | `CreateUserRequest` | `201` or `202` | `CreateUserResponse` |
|
||||
| `GET` | `/v1/users/{username}` | none | `200` | `UserInfo` |
|
||||
@@ -146,7 +148,9 @@ Notes:
|
||||
| `GET /v1/stats/users/active-ips` | Returns users that currently have non-empty active source-IP lists. |
|
||||
| `GET /v1/stats/users` | Alias of `GET /v1/users`; returns disk-first user views with runtime lag flag. |
|
||||
| `GET /v1/config` | Returns the current editable config sections as JSON (no `access.*`) plus the revision. |
|
||||
| `PATCH /v1/config` | Applies a sparse patch to editable config sections; validates, writes, and reports restart impact. |
|
||||
| `PATCH /v1/config` | Applies a sparse patch and optionally submits an in-process runtime reload to Maestro. |
|
||||
| `POST /v1/system/reload` | Loads and validates the current on-disk config, then asks Maestro to prepare and activate a new runtime generation. |
|
||||
| `GET /v1/system/reload/{id}` | Returns one retained reload status; the coordinator retains the most recent 32 operations. |
|
||||
| `GET /v1/users` | Returns disk-first user views sorted by username. |
|
||||
| `POST /v1/users` | Creates a user and returns the effective user view plus secret. |
|
||||
| `GET /v1/users/{username}` | Returns one disk-first user view or `404` when absent. |
|
||||
@@ -170,11 +174,13 @@ Notes:
|
||||
| `404` | `not_found` | Unknown route, unknown user, or unsupported sub-route. |
|
||||
| `405` | `method_not_allowed` | Unsupported method for `/v1/users/{username}` route shape. |
|
||||
| `409` | `revision_conflict` | `If-Match` revision mismatch. |
|
||||
| `409` | `reload_in_progress` | Another reload operation is non-terminal. |
|
||||
| `409` | `user_exists` | User already exists on create. |
|
||||
| `409` | `last_user_forbidden` | Attempt to delete last configured user. |
|
||||
| `413` | `payload_too_large` | Body exceeds `request_body_limit_bytes`. |
|
||||
| `500` | `internal_error` | Internal error (I/O, serialization, config load/save). |
|
||||
| `503` | `api_disabled` | API disabled in config. |
|
||||
| `503` | `maestro_unavailable` | Maestro's reload command channel is unavailable. |
|
||||
|
||||
## Routing and Method Edge Cases
|
||||
|
||||
@@ -292,13 +298,17 @@ Sections absent from the config file are absent from the response (not `null`).
|
||||
|
||||
### `PatchConfigResponse`
|
||||
|
||||
Returned by `PATCH /v1/config` on success (`200`).
|
||||
Returned by `PATCH /v1/config` on success (`200`, or `202` when a reload was accepted).
|
||||
|
||||
| Field | Type | Description |
|
||||
| --- | --- | --- |
|
||||
| `revision` | `string` | SHA-256 hex of the config file after the patch was written. |
|
||||
| `restart_required` | `bool` | `true` when one or more changed fields require a process restart to take effect. Hot-reloadable fields (e.g. `general.log_level`) are applied automatically by the config file watcher; restart-required fields (e.g. any `censorship.*`, `timeouts.*`, `upstreams`, or `general.modes` change) are written to disk but only take effect after the Telemt process is restarted. The caller is responsible for triggering a restart when this flag is `true`. |
|
||||
| `restart_required` | `bool` | Legacy classifier result: `true` when the old file watcher alone cannot apply every changed field. Use `runtime_reload_required` and `process_restart_required` for new integrations. |
|
||||
| `runtime_reload_required` | `bool` | `true` when full effect requires a Maestro runtime-generation reload rather than the legacy hot-field overlay. |
|
||||
| `process_restart_required` | `bool` | `true` when process-owned sockets or paths changed and remain deferred after an in-process reload. |
|
||||
| `deferred_process_fields` | `string[]` | Process-owned fields that the active process cannot rebind during generation activation. |
|
||||
| `changed` | `string[]` | Top-level section names that differed between the old and new config (e.g. `["censorship"]`). |
|
||||
| `reload` | `ReloadAccepted?` | Present only when the patch included a valid reload query and Maestro accepted the operation. |
|
||||
|
||||
### `HealthData`
|
||||
| Field | Type | Description |
|
||||
@@ -1376,24 +1386,50 @@ Applies a sparse patch to the editable config sections. The merged config is ful
|
||||
|
||||
**Read-only mode:** returns `403 read_only` when the API runs with `read_only = true`.
|
||||
|
||||
**Success `200` response body** (`data` field of the standard envelope):
|
||||
**Optional in-process reload query:**
|
||||
|
||||
| Query | Required | Description |
|
||||
| --- | --- | --- |
|
||||
| `reload=instant` | no | Activates a new generation and cancels sessions owned by the previous generation. |
|
||||
| `reload=drain` | no | Activates a new generation and lets old sessions finish until `timeout_secs`. |
|
||||
| `timeout_secs=1..3600` | for `reload=drain` | Bounded old-generation drain interval. Invalid with `reload=instant`. |
|
||||
| `failure_policy=keep_new\|rollback` | no | Defaults to `keep_new`. `rollback` applies only through the activation barrier, before old-generation teardown. |
|
||||
|
||||
Without a `reload` query parameter, the endpoint preserves the legacy behavior: it writes the patch and the file watcher applies only supported hot fields.
|
||||
|
||||
**Success `200` or `202` response body** (`data` field of the standard envelope):
|
||||
```json
|
||||
{
|
||||
"revision": "<new-sha256-hex>",
|
||||
"restart_required": true,
|
||||
"changed": ["censorship"]
|
||||
"runtime_reload_required": true,
|
||||
"process_restart_required": false,
|
||||
"deferred_process_fields": [],
|
||||
"changed": ["censorship"],
|
||||
"reload": {
|
||||
"reload_id": 7,
|
||||
"target_generation": 2,
|
||||
"config_revision": "<new-sha256-hex>",
|
||||
"state": "accepted",
|
||||
"mode": "instant",
|
||||
"failure_policy": "keep_new"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- `revision` — SHA-256 hex of the config file after the write.
|
||||
- `restart_required` — `true` when the change affects a field that Telemt cannot hot-reload (e.g. `censorship.*`, `timeouts.*`, `upstreams`, `general.modes`). Hot-reloadable fields (e.g. `general.log_level`) are applied automatically by the config file watcher. Restart-required fields are written to disk but only take effect after the Telemt process is restarted; the caller is responsible for triggering the restart.
|
||||
- `restart_required` — legacy file-watcher classification retained for compatibility.
|
||||
- `runtime_reload_required` — reports whether a full Maestro generation reload is needed for runtime effect.
|
||||
- `process_restart_required` and `deferred_process_fields` — report process-owned sockets or paths that remain unchanged by an in-process reload.
|
||||
- `changed` — list of top-level section names that differed.
|
||||
- `reload` — accepted operation metadata; omitted when no reload query was supplied.
|
||||
|
||||
**Status codes:**
|
||||
|
||||
| HTTP | `error.code` | Condition |
|
||||
| --- | --- | --- |
|
||||
| `200` | — | Patch applied successfully. |
|
||||
| `202` | — | Patch applied and runtime reload accepted. |
|
||||
| `400` | `bad_request` | Invalid JSON, empty patch, or config validation/deserialization failure. |
|
||||
| `400` | `access_not_editable` | Patch contains an `access` key. |
|
||||
| `400` | `section_not_editable` | Patch contains `server`, `network`, or an unknown top-level key. |
|
||||
@@ -1401,6 +1437,7 @@ Applies a sparse patch to the editable config sections. The merged config is ful
|
||||
| `403` | `read_only` | API is in read-only mode. |
|
||||
| `405` | `method_not_allowed` | Method other than `GET` or `PATCH` used on `/v1/config`. |
|
||||
| `409` | `revision_conflict` | `If-Match` header supplied but does not match current revision. |
|
||||
| `409` | `reload_in_progress` | Another runtime reload is active; the patch is not written. |
|
||||
| `500` | `internal_error` | I/O or serialization failure. |
|
||||
|
||||
**curl example:**
|
||||
@@ -1412,14 +1449,40 @@ curl -s -H "Authorization: <token>" http://127.0.0.1:<api>/v1/system/info | jq -
|
||||
curl -s -X PATCH -H "Authorization: <token>" -H "If-Match: <revision>" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"censorship":{"tls_domain":"front.example.com"}}' \
|
||||
http://127.0.0.1:<api>/v1/config
|
||||
'http://127.0.0.1:<api>/v1/config?reload=instant'
|
||||
```
|
||||
|
||||
## Runtime Reload Endpoints
|
||||
|
||||
### `POST /v1/system/reload`
|
||||
|
||||
Loads the current on-disk config under the API mutation lock and submits an immutable config snapshot to Maestro. `If-Match` is optional and uses the same revision contract as `PATCH /v1/config`. An empty body defaults to `{"mode":"instant","failure_policy":"keep_new"}`.
|
||||
|
||||
```json
|
||||
{
|
||||
"mode": "drain",
|
||||
"timeout_secs": 30,
|
||||
"failure_policy": "rollback"
|
||||
}
|
||||
```
|
||||
|
||||
The endpoint returns `202` with `ReloadAccepted`. A concurrent non-terminal reload returns `409 reload_in_progress`. Config parsing or validation failure is reported before a command is submitted.
|
||||
|
||||
### `GET /v1/system/reload/{id}`
|
||||
|
||||
Returns `ReloadStatus` with `state` equal to `accepted`, `preparing`, `activating`, `draining`, `succeeded`, `rolled_back`, or `failed`. Terminal statuses include `finished_at_epoch_secs`; failures include `error`. Successful activation may include `warnings` for old-generation cleanup failures and `deferred_process_fields` for process-owned settings.
|
||||
|
||||
Runtime generation activation rebuilds statistics, upstream routing, replay and buffer state, TLS-front cache, IP tracking, admission/route state, and Middle-End orchestration. Per-user quota accounting is process-scoped and remains continuous across generations. API, metrics, client TCP/Unix listeners, PID ownership, and logging remain process-scoped; changed bind/path fields are reported as deferred and do not cause Maestro to invoke systemd, containerd, or another process supervisor.
|
||||
|
||||
Reload preparation requires every configured TLS-front domain to have a non-default cached profile and requires a ready Middle-End pool when direct fallback is disabled. A candidate that does not satisfy either readiness condition fails without replacing the active generation.
|
||||
|
||||
The revision is verified again after preparation. With `failure_policy=rollback`, a changed revision or revision read failure rolls the candidate back; with `failure_policy=keep_new`, the condition is reported in `warnings` and activation continues.
|
||||
|
||||
## Mutation Semantics
|
||||
|
||||
| Endpoint | Notes |
|
||||
| --- | --- |
|
||||
| `PATCH /v1/config` | Deep-merges the patch into editable config sections (tables merged per-field; arrays/scalars replaced wholesale). Validates the merged result before writing. Writes only the touched sections via atomic `tmp + rename`. Returns the new revision and which sections changed. |
|
||||
| `PATCH /v1/config` | Deep-merges and validates the patch, writes touched sections via atomic `tmp + rename`, and optionally submits the exact written revision for an in-process Maestro reload. |
|
||||
| `POST /v1/users` | Creates user, validates config, then atomically updates only affected `access.*` TOML tables (`access.users` always, plus optional per-user tables present in request). |
|
||||
| `PATCH /v1/users/{username}` | Partial update of provided fields only. Missing fields remain unchanged; explicit `null` removes optional per-user entries. The write path updates only affected `access.*` TOML tables. |
|
||||
| `POST /v1/users/{username}/rotate-secret` | Replaces the user's secret with a provided valid 32-hex value or a generated value, then returns the effective secret in `CreateUserResponse`. |
|
||||
|
||||
@@ -6,19 +6,28 @@ use toml::Value as Toml;
|
||||
|
||||
use super::ApiShared;
|
||||
use super::config_store::{
|
||||
EDITABLE_SECTIONS, compute_revision, current_revision, save_sections_to_disk,
|
||||
EDITABLE_SECTIONS, compute_revision, current_revision, load_config_from_disk,
|
||||
save_sections_to_disk,
|
||||
};
|
||||
use super::model::ApiFailure;
|
||||
use crate::config::ProxyConfig;
|
||||
use crate::config::hot_reload::classify_config_changes;
|
||||
use crate::maestro::reload::{ReloadAccepted, ReloadRequest, ReloadSubmitError};
|
||||
use crate::maestro::runtime_build::deferred_process_fields;
|
||||
use serde::Serialize;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(super) struct PatchConfigResponse {
|
||||
pub revision: String,
|
||||
pub restart_required: bool,
|
||||
pub runtime_reload_required: bool,
|
||||
pub process_restart_required: bool,
|
||||
pub deferred_process_fields: Vec<String>,
|
||||
pub changed: Vec<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub reload: Option<ReloadAccepted>,
|
||||
}
|
||||
|
||||
/// Shared-state wrapper around [`apply_patch_to_path`]: serializes config
|
||||
@@ -27,10 +36,40 @@ pub(super) struct PatchConfigResponse {
|
||||
pub(super) async fn patch_config(
|
||||
patch_json: Json,
|
||||
expected_revision: Option<String>,
|
||||
reload_request: Option<ReloadRequest>,
|
||||
shared: &ApiShared,
|
||||
) -> Result<PatchConfigResponse, ApiFailure> {
|
||||
let _guard = shared.mutation_lock.lock().await;
|
||||
let resp = apply_patch_to_path(&shared.config_path, &patch_json, expected_revision).await?;
|
||||
if reload_request.is_some()
|
||||
&& let Some(reload_id) = shared.reload_control.in_progress().await
|
||||
{
|
||||
return Err(ApiFailure::new(
|
||||
hyper::StatusCode::CONFLICT,
|
||||
"reload_in_progress",
|
||||
format!("Reload {} is already in progress", reload_id),
|
||||
));
|
||||
}
|
||||
let mut resp = apply_patch_to_path(&shared.config_path, &patch_json, expected_revision).await?;
|
||||
if let Some(request) = reload_request {
|
||||
let config = Arc::new(load_config_from_disk(&shared.config_path).await?);
|
||||
let accepted = shared
|
||||
.reload_control
|
||||
.submit(config, resp.revision.clone(), request)
|
||||
.await
|
||||
.map_err(|error| match error {
|
||||
ReloadSubmitError::InProgress(reload_id) => ApiFailure::new(
|
||||
hyper::StatusCode::CONFLICT,
|
||||
"reload_in_progress",
|
||||
format!("Reload {} is already in progress", reload_id),
|
||||
),
|
||||
ReloadSubmitError::MaestroUnavailable => ApiFailure::new(
|
||||
hyper::StatusCode::SERVICE_UNAVAILABLE,
|
||||
"maestro_unavailable",
|
||||
"Maestro reload coordinator is unavailable",
|
||||
),
|
||||
})?;
|
||||
resp.reload = Some(accepted);
|
||||
}
|
||||
drop(_guard);
|
||||
shared
|
||||
.runtime_events
|
||||
@@ -114,6 +153,7 @@ pub(super) async fn apply_patch_to_path(
|
||||
|
||||
// 4. classify changes (Telemt's own hot/restart rule)
|
||||
let class = classify_config_changes(&old_cfg, &new_cfg);
|
||||
let deferred_process_fields = deferred_process_fields(&old_cfg, &new_cfg);
|
||||
|
||||
// 5. write only the touched top-level sections
|
||||
let revision = save_sections_to_disk(config_path, &new_cfg, &touched).await?;
|
||||
@@ -121,7 +161,11 @@ pub(super) async fn apply_patch_to_path(
|
||||
Ok(PatchConfigResponse {
|
||||
revision,
|
||||
restart_required: class.restart_required,
|
||||
runtime_reload_required: class.restart_required,
|
||||
process_restart_required: !deferred_process_fields.is_empty(),
|
||||
deferred_process_fields,
|
||||
changed: class.changed,
|
||||
reload: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -266,6 +310,9 @@ mod tests {
|
||||
let patch: Json = serde_json::json!({"censorship": {"tls_domain": "b.com"}});
|
||||
let resp = apply_patch_to_path(&path, &patch, None).await.unwrap();
|
||||
assert!(resp.restart_required);
|
||||
assert!(resp.runtime_reload_required);
|
||||
assert!(!resp.process_restart_required);
|
||||
assert!(resp.deferred_process_fields.is_empty());
|
||||
assert!(resp.changed.iter().any(|c| c == "censorship"));
|
||||
let written = std::fs::read_to_string(&path).unwrap();
|
||||
assert!(written.contains("tls_domain = \"b.com\""));
|
||||
@@ -406,6 +453,8 @@ mod tests {
|
||||
let patch: Json = serde_json::json!({"general": {"log_level": "debug"}});
|
||||
let resp = apply_patch_to_path(&path, &patch, None).await.unwrap();
|
||||
assert!(!resp.restart_required);
|
||||
assert!(!resp.runtime_reload_required);
|
||||
assert!(!resp.process_restart_required);
|
||||
assert!(resp.changed.iter().any(|c| c == "general"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -72,6 +72,13 @@ pub(super) async fn current_revision(config_path: &Path) -> Result<String, ApiFa
|
||||
Ok(compute_revision(&content))
|
||||
}
|
||||
|
||||
pub(crate) async fn current_revision_for_maestro(config_path: &Path) -> Result<String, String> {
|
||||
let content = tokio::fs::read_to_string(config_path)
|
||||
.await
|
||||
.map_err(|error| format!("failed to read config: {}", error))?;
|
||||
Ok(compute_revision(&content))
|
||||
}
|
||||
|
||||
pub(super) fn compute_revision(content: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(content.as_bytes());
|
||||
@@ -86,6 +93,14 @@ pub(super) async fn load_config_from_disk(config_path: &Path) -> Result<ProxyCon
|
||||
.map_err(|e| ApiFailure::internal(format!("failed to load config: {}", e)))
|
||||
}
|
||||
|
||||
pub(super) async fn load_config_for_reload(config_path: &Path) -> Result<ProxyConfig, ApiFailure> {
|
||||
let config_path = config_path.to_path_buf();
|
||||
tokio::task::spawn_blocking(move || ProxyConfig::load(config_path))
|
||||
.await
|
||||
.map_err(|error| ApiFailure::internal(format!("failed to join config loader: {}", error)))?
|
||||
.map_err(|error| ApiFailure::bad_request(format!("invalid runtime config: {}", error)))
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub(super) async fn save_config_to_disk(
|
||||
config_path: &Path,
|
||||
|
||||
164
src/api/mod.rs
164
src/api/mod.rs
@@ -7,6 +7,7 @@ use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use arc_swap::ArcSwap;
|
||||
use http_body_util::Full;
|
||||
use hyper::body::{Bytes, Incoming};
|
||||
use hyper::header::AUTHORIZATION;
|
||||
@@ -19,8 +20,10 @@ use tokio::sync::{Mutex, RwLock, Semaphore, watch};
|
||||
use tokio::time::timeout;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use crate::config::{ApiGrayAction, ProxyConfig};
|
||||
use crate::config::ApiGrayAction;
|
||||
use crate::ip_tracker::UserIpTracker;
|
||||
use crate::maestro::generation::{RuntimeGeneration, RuntimeWatchState};
|
||||
use crate::maestro::reload::{ReloadControl, ReloadRequest, ReloadSubmitError};
|
||||
use crate::proxy::route_mode::RouteRuntimeController;
|
||||
use crate::proxy::shared_state::ProxySharedState;
|
||||
use crate::startup::StartupTracker;
|
||||
@@ -29,7 +32,7 @@ use crate::transport::UpstreamManager;
|
||||
use crate::transport::middle_proxy::MePool;
|
||||
|
||||
mod config_edit;
|
||||
mod config_store;
|
||||
pub(crate) mod config_store;
|
||||
mod events;
|
||||
mod http_utils;
|
||||
mod model;
|
||||
@@ -44,7 +47,8 @@ mod runtime_zero;
|
||||
mod users;
|
||||
|
||||
use config_store::{
|
||||
current_revision, ensure_expected_revision, load_config_from_disk, parse_if_match,
|
||||
current_revision, ensure_expected_revision, load_config_for_reload, load_config_from_disk,
|
||||
parse_if_match,
|
||||
};
|
||||
use events::ApiEventStore;
|
||||
use http_utils::{error_response, read_json, read_optional_json, success_response};
|
||||
@@ -107,12 +111,15 @@ pub(super) struct ApiShared {
|
||||
pub(super) minimal_cache: Arc<Mutex<Option<MinimalCacheEntry>>>,
|
||||
pub(super) runtime_edge_connections_cache: Arc<Mutex<Option<EdgeConnectionsCacheEntry>>>,
|
||||
pub(super) runtime_edge_recompute_lock: Arc<Mutex<()>>,
|
||||
pub(super) cache_generation: Arc<AtomicU64>,
|
||||
pub(super) runtime_events: Arc<ApiEventStore>,
|
||||
pub(super) request_id: Arc<AtomicU64>,
|
||||
pub(super) runtime_state: Arc<ApiRuntimeState>,
|
||||
pub(super) startup_tracker: Arc<StartupTracker>,
|
||||
pub(super) route_runtime: Arc<RouteRuntimeController>,
|
||||
pub(super) proxy_shared: Arc<ProxySharedState>,
|
||||
pub(super) reload_control: ReloadControl,
|
||||
pub(super) active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
||||
}
|
||||
|
||||
impl ApiShared {
|
||||
@@ -123,6 +130,31 @@ impl ApiShared {
|
||||
fn detected_link_ips(&self) -> (Option<IpAddr>, Option<IpAddr>) {
|
||||
*self.detected_ips_rx.borrow()
|
||||
}
|
||||
|
||||
fn for_runtime(&self, runtime: &RuntimeGeneration) -> Self {
|
||||
Self {
|
||||
stats: runtime.stats.clone(),
|
||||
ip_tracker: runtime.ip_tracker.clone(),
|
||||
me_pool: runtime.me_pool_runtime.clone(),
|
||||
upstream_manager: runtime.upstream_manager.clone(),
|
||||
config_path: self.config_path.clone(),
|
||||
quota_state_path: self.quota_state_path.clone(),
|
||||
detected_ips_rx: self.detected_ips_rx.clone(),
|
||||
mutation_lock: self.mutation_lock.clone(),
|
||||
minimal_cache: self.minimal_cache.clone(),
|
||||
runtime_edge_connections_cache: self.runtime_edge_connections_cache.clone(),
|
||||
runtime_edge_recompute_lock: self.runtime_edge_recompute_lock.clone(),
|
||||
cache_generation: self.cache_generation.clone(),
|
||||
runtime_events: self.runtime_events.clone(),
|
||||
request_id: self.request_id.clone(),
|
||||
runtime_state: self.runtime_state.clone(),
|
||||
startup_tracker: self.startup_tracker.clone(),
|
||||
route_runtime: runtime.route_runtime.clone(),
|
||||
proxy_shared: runtime.proxy_shared.clone(),
|
||||
reload_control: self.reload_control.clone(),
|
||||
active_runtime: self.active_runtime.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn auth_header_matches(actual: &str, expected: &str) -> bool {
|
||||
@@ -144,6 +176,12 @@ fn user_action_route_matches(path: &str, suffix: &str) -> bool {
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn reload_status_route_id(path: &str) -> Option<u64> {
|
||||
path.strip_prefix("/v1/system/reload/")
|
||||
.filter(|id| !id.is_empty() && !id.contains('/'))
|
||||
.and_then(|id| id.parse().ok())
|
||||
}
|
||||
|
||||
fn allowed_methods_for_path(path: &str) -> Option<&'static str> {
|
||||
match path {
|
||||
"/v1/health"
|
||||
@@ -175,12 +213,14 @@ fn allowed_methods_for_path(path: &str) -> Option<&'static str> {
|
||||
| "/v1/stats/users/active-ips"
|
||||
| "/v1/stats/users/quota"
|
||||
| "/v1/stats/users" => Some(ALLOW_GET),
|
||||
"/v1/system/reload" => Some(ALLOW_POST),
|
||||
"/v1/users" => Some(ALLOW_GET_POST),
|
||||
"/v1/config" => Some(ALLOW_GET_PATCH),
|
||||
_ if user_action_route_matches(path, "/reset-quota") => Some(ALLOW_POST),
|
||||
_ if user_action_route_matches(path, "/rotate-secret") => Some(ALLOW_POST),
|
||||
_ if user_action_route_matches(path, "/enable") => Some(ALLOW_POST),
|
||||
_ if user_action_route_matches(path, "/disable") => Some(ALLOW_POST),
|
||||
_ if reload_status_route_id(path).is_some() => Some(ALLOW_GET),
|
||||
_ if path
|
||||
.strip_prefix("/v1/users/")
|
||||
.map(|user| !user.is_empty() && !user.contains('/'))
|
||||
@@ -200,14 +240,35 @@ pub async fn serve(
|
||||
route_runtime: Arc<RouteRuntimeController>,
|
||||
proxy_shared: Arc<ProxySharedState>,
|
||||
upstream_manager: Arc<UpstreamManager>,
|
||||
config_rx: watch::Receiver<Arc<ProxyConfig>>,
|
||||
admission_rx: watch::Receiver<bool>,
|
||||
config_path: PathBuf,
|
||||
quota_state_path: PathBuf,
|
||||
detected_ips_rx: watch::Receiver<(Option<IpAddr>, Option<IpAddr>)>,
|
||||
process_started_at_epoch_secs: u64,
|
||||
startup_tracker: Arc<StartupTracker>,
|
||||
reload_control: ReloadControl,
|
||||
mut active_runtime_rx: watch::Receiver<Option<Arc<ArcSwap<RuntimeGeneration>>>>,
|
||||
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
||||
) {
|
||||
let active_runtime = loop {
|
||||
if let Some(active_runtime) = active_runtime_rx.borrow().clone() {
|
||||
break active_runtime;
|
||||
}
|
||||
if active_runtime_rx.changed().await.is_err() {
|
||||
warn!("Runtime generation channel closed before API bootstrap");
|
||||
return;
|
||||
}
|
||||
};
|
||||
let initial_watch_state = loop {
|
||||
if let Some(watch_state) = runtime_watch_rx.borrow().clone() {
|
||||
break watch_state;
|
||||
}
|
||||
if runtime_watch_rx.changed().await.is_err() {
|
||||
warn!("Runtime watch channel closed before API bootstrap");
|
||||
return;
|
||||
}
|
||||
};
|
||||
let config_rx = initial_watch_state.config_rx.clone();
|
||||
let admission_rx = initial_watch_state.admission_rx.clone();
|
||||
let listener = match TcpListener::bind(listen).await {
|
||||
Ok(listener) => listener,
|
||||
Err(error) => {
|
||||
@@ -241,6 +302,7 @@ pub async fn serve(
|
||||
minimal_cache: Arc::new(Mutex::new(None)),
|
||||
runtime_edge_connections_cache: Arc::new(Mutex::new(None)),
|
||||
runtime_edge_recompute_lock: Arc::new(Mutex::new(())),
|
||||
cache_generation: Arc::new(AtomicU64::new(1)),
|
||||
runtime_events: Arc::new(ApiEventStore::new(
|
||||
config_rx.borrow().server.api.runtime_edge_events_capacity,
|
||||
)),
|
||||
@@ -249,11 +311,12 @@ pub async fn serve(
|
||||
startup_tracker,
|
||||
route_runtime,
|
||||
proxy_shared,
|
||||
reload_control,
|
||||
active_runtime,
|
||||
});
|
||||
|
||||
spawn_runtime_watchers(
|
||||
config_rx.clone(),
|
||||
admission_rx.clone(),
|
||||
runtime_watch_rx,
|
||||
runtime_state.clone(),
|
||||
shared.runtime_events.clone(),
|
||||
);
|
||||
@@ -282,13 +345,11 @@ pub async fn serve(
|
||||
};
|
||||
|
||||
let shared_conn = shared.clone();
|
||||
let config_rx_conn = config_rx.clone();
|
||||
tokio::spawn(async move {
|
||||
let _connection_permit = connection_permit;
|
||||
let svc = service_fn(move |req: Request<Incoming>| {
|
||||
let shared_req = shared_conn.clone();
|
||||
let config_rx_req = config_rx_conn.clone();
|
||||
async move { handle(req, peer, shared_req, config_rx_req).await }
|
||||
async move { handle(req, peer, shared_req).await }
|
||||
});
|
||||
match timeout(
|
||||
API_HTTP_CONNECTION_TIMEOUT,
|
||||
@@ -318,8 +379,19 @@ async fn handle(
|
||||
req: Request<Incoming>,
|
||||
peer: SocketAddr,
|
||||
shared: Arc<ApiShared>,
|
||||
config_rx: watch::Receiver<Arc<ProxyConfig>>,
|
||||
) -> Result<Response<Full<Bytes>>, IoError> {
|
||||
let runtime = shared.active_runtime.load_full();
|
||||
let previous_cache_generation = shared.cache_generation.swap(runtime.id, Ordering::AcqRel);
|
||||
if previous_cache_generation != runtime.id {
|
||||
*shared.minimal_cache.lock().await = None;
|
||||
*shared.runtime_edge_connections_cache.lock().await = None;
|
||||
}
|
||||
let shared = Arc::new(shared.for_runtime(runtime.as_ref()));
|
||||
let config_rx = runtime.config_rx.clone();
|
||||
shared
|
||||
.runtime_state
|
||||
.admission_open
|
||||
.store(*runtime.admission_rx.borrow(), Ordering::Relaxed);
|
||||
let request_id = shared.next_request_id();
|
||||
let cfg = config_rx.borrow().clone();
|
||||
let api_cfg = &cfg.server.api;
|
||||
@@ -651,6 +723,45 @@ async fn handle(
|
||||
config_edit::read_managed_config(&shared.config_path).await?;
|
||||
Ok(success_response(StatusCode::OK, value, revision))
|
||||
}
|
||||
("POST", "/v1/system/reload") => {
|
||||
if api_cfg.read_only {
|
||||
return Ok(error_response(
|
||||
request_id,
|
||||
ApiFailure::new(
|
||||
StatusCode::FORBIDDEN,
|
||||
"read_only",
|
||||
"API runs in read-only mode",
|
||||
),
|
||||
));
|
||||
}
|
||||
let expected_revision = parse_if_match(req.headers());
|
||||
let request = read_optional_json::<ReloadRequest>(req.into_body(), body_limit)
|
||||
.await?
|
||||
.unwrap_or_default();
|
||||
request.validate().map_err(ApiFailure::bad_request)?;
|
||||
|
||||
let _guard = shared.mutation_lock.lock().await;
|
||||
ensure_expected_revision(&shared.config_path, expected_revision.as_deref()).await?;
|
||||
let revision = current_revision(&shared.config_path).await?;
|
||||
let config = Arc::new(load_config_for_reload(&shared.config_path).await?);
|
||||
let accepted = shared
|
||||
.reload_control
|
||||
.submit(config, revision.clone(), request)
|
||||
.await
|
||||
.map_err(|error| match error {
|
||||
ReloadSubmitError::InProgress(reload_id) => ApiFailure::new(
|
||||
StatusCode::CONFLICT,
|
||||
"reload_in_progress",
|
||||
format!("Reload {} is already in progress", reload_id),
|
||||
),
|
||||
ReloadSubmitError::MaestroUnavailable => ApiFailure::new(
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"maestro_unavailable",
|
||||
"Maestro reload coordinator is unavailable",
|
||||
),
|
||||
})?;
|
||||
Ok(success_response(StatusCode::ACCEPTED, accepted, revision))
|
||||
}
|
||||
("PATCH", "/v1/config") => {
|
||||
if api_cfg.read_only {
|
||||
return Ok(error_response(
|
||||
@@ -663,11 +774,20 @@ async fn handle(
|
||||
));
|
||||
}
|
||||
let expected_revision = parse_if_match(req.headers());
|
||||
let reload_request =
|
||||
ReloadRequest::from_query(query.as_deref()).map_err(ApiFailure::bad_request)?;
|
||||
let body = read_json::<serde_json::Value>(req.into_body(), body_limit).await?;
|
||||
match config_edit::patch_config(body, expected_revision, &shared).await {
|
||||
match config_edit::patch_config(body, expected_revision, reload_request, &shared)
|
||||
.await
|
||||
{
|
||||
Ok(resp) => {
|
||||
let revision = resp.revision.clone();
|
||||
Ok(success_response(StatusCode::OK, resp, revision))
|
||||
let status = if resp.reload.is_some() {
|
||||
StatusCode::ACCEPTED
|
||||
} else {
|
||||
StatusCode::OK
|
||||
};
|
||||
Ok(success_response(status, resp, revision))
|
||||
}
|
||||
Err(error) => {
|
||||
shared
|
||||
@@ -678,6 +798,24 @@ async fn handle(
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
if method == Method::GET
|
||||
&& let Some(reload_id) = reload_status_route_id(normalized_path)
|
||||
{
|
||||
let revision = current_revision(&shared.config_path).await?;
|
||||
let status =
|
||||
shared
|
||||
.reload_control
|
||||
.status(reload_id)
|
||||
.await
|
||||
.ok_or_else(|| {
|
||||
ApiFailure::new(
|
||||
StatusCode::NOT_FOUND,
|
||||
"reload_not_found",
|
||||
format!("Reload {} was not found", reload_id),
|
||||
)
|
||||
})?;
|
||||
return Ok(success_response(StatusCode::OK, status, revision));
|
||||
}
|
||||
if method == Method::POST
|
||||
&& let Some(base_user) = normalized_path
|
||||
.strip_prefix("/v1/users/")
|
||||
|
||||
@@ -4,63 +4,283 @@ use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use tokio::sync::watch;
|
||||
|
||||
use crate::config::ProxyConfig;
|
||||
use crate::maestro::generation::RuntimeWatchState;
|
||||
|
||||
use super::ApiRuntimeState;
|
||||
use super::events::ApiEventStore;
|
||||
|
||||
pub(super) fn spawn_runtime_watchers(
|
||||
config_rx: watch::Receiver<Arc<ProxyConfig>>,
|
||||
admission_rx: watch::Receiver<bool>,
|
||||
runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
||||
runtime_state: Arc<ApiRuntimeState>,
|
||||
runtime_events: Arc<ApiEventStore>,
|
||||
) {
|
||||
let mut config_rx_reload = config_rx;
|
||||
let runtime_state_reload = runtime_state.clone();
|
||||
let runtime_events_reload = runtime_events.clone();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
if config_rx_reload.changed().await.is_err() {
|
||||
break;
|
||||
}
|
||||
runtime_state_reload
|
||||
.config_reload_count
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
runtime_state_reload
|
||||
.last_config_reload_epoch_secs
|
||||
.store(now_epoch_secs(), Ordering::Relaxed);
|
||||
runtime_events_reload.record("config.reload.applied", "config receiver updated");
|
||||
}
|
||||
});
|
||||
spawn_config_watcher(
|
||||
runtime_watch_rx.clone(),
|
||||
runtime_state.clone(),
|
||||
runtime_events.clone(),
|
||||
);
|
||||
spawn_admission_watcher(runtime_watch_rx, runtime_state, runtime_events);
|
||||
}
|
||||
|
||||
let mut admission_rx_watch = admission_rx;
|
||||
fn spawn_config_watcher(
|
||||
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
||||
runtime_state: Arc<ApiRuntimeState>,
|
||||
runtime_events: Arc<ApiEventStore>,
|
||||
) {
|
||||
tokio::spawn(async move {
|
||||
runtime_state
|
||||
.admission_open
|
||||
.store(*admission_rx_watch.borrow(), Ordering::Relaxed);
|
||||
runtime_events.record(
|
||||
"admission.state",
|
||||
format!("accepting_new_connections={}", *admission_rx_watch.borrow()),
|
||||
);
|
||||
let Some(mut current) = runtime_watch_rx.borrow().clone() else {
|
||||
return;
|
||||
};
|
||||
loop {
|
||||
if admission_rx_watch.changed().await.is_err() {
|
||||
break;
|
||||
tokio::select! {
|
||||
biased;
|
||||
changed = runtime_watch_rx.changed() => {
|
||||
if changed.is_err() {
|
||||
break;
|
||||
}
|
||||
let Some(next) = runtime_watch_rx.borrow().clone() else {
|
||||
continue;
|
||||
};
|
||||
if next.generation_id != current.generation_id {
|
||||
current = next;
|
||||
record_config_reload(
|
||||
&runtime_state,
|
||||
&runtime_events,
|
||||
format!("runtime generation {} activated", current.generation_id),
|
||||
);
|
||||
}
|
||||
}
|
||||
changed = current.config_rx.changed() => {
|
||||
if changed.is_err() {
|
||||
let Some(next) = wait_for_new_generation(
|
||||
&mut runtime_watch_rx,
|
||||
current.generation_id,
|
||||
).await else {
|
||||
break;
|
||||
};
|
||||
current = next;
|
||||
record_config_reload(
|
||||
&runtime_state,
|
||||
&runtime_events,
|
||||
format!("runtime generation {} activated", current.generation_id),
|
||||
);
|
||||
continue;
|
||||
}
|
||||
if active_generation_id(&runtime_watch_rx) != Some(current.generation_id) {
|
||||
continue;
|
||||
}
|
||||
record_config_reload(
|
||||
&runtime_state,
|
||||
&runtime_events,
|
||||
format!("generation {} config receiver updated", current.generation_id),
|
||||
);
|
||||
}
|
||||
}
|
||||
let admission_open = *admission_rx_watch.borrow();
|
||||
runtime_state
|
||||
.admission_open
|
||||
.store(admission_open, Ordering::Relaxed);
|
||||
runtime_events.record(
|
||||
"admission.state",
|
||||
format!("accepting_new_connections={}", admission_open),
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn spawn_admission_watcher(
|
||||
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
||||
runtime_state: Arc<ApiRuntimeState>,
|
||||
runtime_events: Arc<ApiEventStore>,
|
||||
) {
|
||||
tokio::spawn(async move {
|
||||
let Some(mut current) = runtime_watch_rx.borrow().clone() else {
|
||||
return;
|
||||
};
|
||||
record_admission_state(&runtime_state, &runtime_events, ¤t);
|
||||
loop {
|
||||
tokio::select! {
|
||||
biased;
|
||||
changed = runtime_watch_rx.changed() => {
|
||||
if changed.is_err() {
|
||||
break;
|
||||
}
|
||||
let Some(next) = runtime_watch_rx.borrow().clone() else {
|
||||
continue;
|
||||
};
|
||||
if next.generation_id != current.generation_id {
|
||||
current = next;
|
||||
record_admission_state(&runtime_state, &runtime_events, ¤t);
|
||||
}
|
||||
}
|
||||
changed = current.admission_rx.changed() => {
|
||||
if changed.is_err() {
|
||||
let Some(next) = wait_for_new_generation(
|
||||
&mut runtime_watch_rx,
|
||||
current.generation_id,
|
||||
).await else {
|
||||
break;
|
||||
};
|
||||
current = next;
|
||||
record_admission_state(&runtime_state, &runtime_events, ¤t);
|
||||
continue;
|
||||
}
|
||||
if active_generation_id(&runtime_watch_rx) == Some(current.generation_id) {
|
||||
record_admission_state(&runtime_state, &runtime_events, ¤t);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn active_generation_id(
|
||||
runtime_watch_rx: &watch::Receiver<Option<RuntimeWatchState>>,
|
||||
) -> Option<u64> {
|
||||
runtime_watch_rx
|
||||
.borrow()
|
||||
.as_ref()
|
||||
.map(|state| state.generation_id)
|
||||
}
|
||||
|
||||
async fn wait_for_new_generation(
|
||||
runtime_watch_rx: &mut watch::Receiver<Option<RuntimeWatchState>>,
|
||||
previous_generation_id: u64,
|
||||
) -> Option<RuntimeWatchState> {
|
||||
loop {
|
||||
if let Some(state) = runtime_watch_rx.borrow().clone()
|
||||
&& state.generation_id != previous_generation_id
|
||||
{
|
||||
return Some(state);
|
||||
}
|
||||
if runtime_watch_rx.changed().await.is_err() {
|
||||
return None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn record_config_reload(
|
||||
runtime_state: &ApiRuntimeState,
|
||||
runtime_events: &ApiEventStore,
|
||||
context: String,
|
||||
) {
|
||||
runtime_state
|
||||
.config_reload_count
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
runtime_state
|
||||
.last_config_reload_epoch_secs
|
||||
.store(now_epoch_secs(), Ordering::Relaxed);
|
||||
runtime_events.record("config.reload.applied", context);
|
||||
}
|
||||
|
||||
fn record_admission_state(
|
||||
runtime_state: &ApiRuntimeState,
|
||||
runtime_events: &ApiEventStore,
|
||||
current: &RuntimeWatchState,
|
||||
) {
|
||||
let admission_open = *current.admission_rx.borrow();
|
||||
runtime_state
|
||||
.admission_open
|
||||
.store(admission_open, Ordering::Relaxed);
|
||||
runtime_events.record(
|
||||
"admission.state",
|
||||
format!(
|
||||
"generation={} accepting_new_connections={}",
|
||||
current.generation_id, admission_open
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
fn now_epoch_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::config::ProxyConfig;
|
||||
use std::sync::atomic::{AtomicBool, AtomicU64};
|
||||
use std::time::Duration;
|
||||
|
||||
fn state(
|
||||
generation_id: u64,
|
||||
) -> (
|
||||
RuntimeWatchState,
|
||||
watch::Sender<Arc<ProxyConfig>>,
|
||||
watch::Sender<bool>,
|
||||
) {
|
||||
let (config_tx, config_rx) = watch::channel(Arc::new(ProxyConfig::default()));
|
||||
let (admission_tx, admission_rx) = watch::channel(true);
|
||||
(
|
||||
RuntimeWatchState {
|
||||
generation_id,
|
||||
config_rx,
|
||||
admission_rx,
|
||||
},
|
||||
config_tx,
|
||||
admission_tx,
|
||||
)
|
||||
}
|
||||
|
||||
fn runtime_state() -> Arc<ApiRuntimeState> {
|
||||
Arc::new(ApiRuntimeState {
|
||||
process_started_at_epoch_secs: 1,
|
||||
config_reload_count: AtomicU64::new(0),
|
||||
last_config_reload_epoch_secs: AtomicU64::new(0),
|
||||
admission_open: AtomicBool::new(false),
|
||||
})
|
||||
}
|
||||
|
||||
async fn wait_for_count(runtime_state: &ApiRuntimeState, expected: u64) {
|
||||
tokio::time::timeout(Duration::from_secs(1), async {
|
||||
loop {
|
||||
if runtime_state.config_reload_count.load(Ordering::Relaxed) == expected {
|
||||
break;
|
||||
}
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn watchers_follow_only_the_active_generation() {
|
||||
let (initial, initial_config_tx, initial_admission_tx) = state(1);
|
||||
let (runtime_watch_tx, runtime_watch_rx) = watch::channel(Some(initial));
|
||||
let runtime_state = runtime_state();
|
||||
let events = Arc::new(ApiEventStore::new(16));
|
||||
spawn_runtime_watchers(runtime_watch_rx, runtime_state.clone(), events.clone());
|
||||
tokio::task::yield_now().await;
|
||||
|
||||
assert_eq!(runtime_state.config_reload_count.load(Ordering::Relaxed), 0);
|
||||
initial_config_tx.send_replace(Arc::new(ProxyConfig::default()));
|
||||
wait_for_count(&runtime_state, 1).await;
|
||||
|
||||
let (next, next_config_tx, next_admission_tx) = state(2);
|
||||
runtime_watch_tx.send_replace(Some(next));
|
||||
wait_for_count(&runtime_state, 2).await;
|
||||
|
||||
initial_config_tx.send_replace(Arc::new(ProxyConfig::default()));
|
||||
initial_admission_tx.send_replace(false);
|
||||
tokio::task::yield_now().await;
|
||||
assert_eq!(runtime_state.config_reload_count.load(Ordering::Relaxed), 2);
|
||||
assert!(runtime_state.admission_open.load(Ordering::Relaxed));
|
||||
|
||||
next_config_tx.send_replace(Arc::new(ProxyConfig::default()));
|
||||
next_admission_tx.send_replace(false);
|
||||
wait_for_count(&runtime_state, 3).await;
|
||||
tokio::time::timeout(Duration::from_secs(1), async {
|
||||
while runtime_state.admission_open.load(Ordering::Relaxed) {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let snapshot = events.snapshot(16);
|
||||
assert_eq!(
|
||||
snapshot
|
||||
.events
|
||||
.iter()
|
||||
.filter(|event| event.event_type == "config.reload.applied")
|
||||
.count(),
|
||||
3
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use serde::Serialize;
|
||||
|
||||
@@ -162,7 +163,7 @@ pub(super) fn build_system_info_data(
|
||||
build_time_utc,
|
||||
rustc_version,
|
||||
process_started_at_epoch_secs: shared.runtime_state.process_started_at_epoch_secs,
|
||||
uptime_seconds: shared.stats.uptime_secs(),
|
||||
uptime_seconds: process_uptime_seconds(shared.runtime_state.process_started_at_epoch_secs),
|
||||
config_path: shared.config_path.display().to_string(),
|
||||
config_hash: revision.to_string(),
|
||||
config_reload_count: shared
|
||||
@@ -173,6 +174,18 @@ pub(super) fn build_system_info_data(
|
||||
}
|
||||
}
|
||||
|
||||
fn process_uptime_seconds(process_started_at_epoch_secs: u64) -> f64 {
|
||||
let now_epoch_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
process_uptime_seconds_at(process_started_at_epoch_secs, now_epoch_secs)
|
||||
}
|
||||
|
||||
fn process_uptime_seconds_at(process_started_at_epoch_secs: u64, now_epoch_secs: u64) -> f64 {
|
||||
now_epoch_secs.saturating_sub(process_started_at_epoch_secs) as f64
|
||||
}
|
||||
|
||||
pub(super) async fn build_runtime_gates_data(
|
||||
shared: &ApiShared,
|
||||
cfg: &ProxyConfig,
|
||||
@@ -339,3 +352,14 @@ fn me_writer_pick_mode_label(mode: MeWriterPickMode) -> &'static str {
|
||||
MeWriterPickMode::P2c => "p2c",
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::process_uptime_seconds_at;
|
||||
|
||||
#[test]
|
||||
fn process_uptime_is_monotonic_and_saturating() {
|
||||
assert_eq!(process_uptime_seconds_at(100, 135), 35.0);
|
||||
assert_eq!(process_uptime_seconds_at(135, 100), 0.0);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -24,6 +24,10 @@ const DEFAULT_ME_ADAPTIVE_FLOOR_MAX_WARM_WRITERS_GLOBAL: u32 = 256;
|
||||
const DEFAULT_ME_ROUTE_BACKPRESSURE_ENABLED: bool = false;
|
||||
const DEFAULT_ME_ROUTE_FAIRSHARE_ENABLED: bool = false;
|
||||
const DEFAULT_ME_WRITER_CMD_CHANNEL_CAPACITY: usize = 4096;
|
||||
pub(crate) const ME_WRITER_BYTE_PERMIT_UNIT_BYTES: usize = 16 * 1024;
|
||||
pub(crate) const ME_WRITER_FRAME_OVERHEAD_RESERVE_BYTES: usize = 256;
|
||||
const DEFAULT_ME_WRITER_BYTE_BUDGET_BYTES: usize =
|
||||
32 * 1024 * 1024 + ME_WRITER_BYTE_PERMIT_UNIT_BYTES;
|
||||
const DEFAULT_ME_ROUTE_CHANNEL_CAPACITY: usize = 768;
|
||||
const DEFAULT_ME_C2ME_CHANNEL_CAPACITY: usize = 1024;
|
||||
const DEFAULT_ME_READER_ROUTE_DATA_WAIT_MS: u64 = 2;
|
||||
@@ -35,6 +39,8 @@ const DEFAULT_ME_QUOTA_SOFT_OVERSHOOT_BYTES: u64 = 64 * 1024;
|
||||
const DEFAULT_ME_D2C_FRAME_BUF_SHRINK_THRESHOLD_BYTES: usize = 256 * 1024;
|
||||
const DEFAULT_DIRECT_RELAY_COPY_BUF_C2S_BYTES: usize = 64 * 1024;
|
||||
const DEFAULT_DIRECT_RELAY_COPY_BUF_S2C_BYTES: usize = 256 * 1024;
|
||||
pub(crate) const DIRECT_RELAY_BUFFER_BUDGET_UNIT_BYTES: usize = 4 * 1024;
|
||||
const DEFAULT_DIRECT_RELAY_BUFFER_BUDGET_MAX_BYTES: usize = 0;
|
||||
const DEFAULT_ME_WRITER_PICK_SAMPLE_SIZE: u8 = 3;
|
||||
const DEFAULT_ME_HEALTH_INTERVAL_MS_UNHEALTHY: u64 = 1000;
|
||||
const DEFAULT_ME_HEALTH_INTERVAL_MS_HEALTHY: u64 = 3000;
|
||||
@@ -455,6 +461,18 @@ pub(crate) fn default_me_writer_cmd_channel_capacity() -> usize {
|
||||
DEFAULT_ME_WRITER_CMD_CHANNEL_CAPACITY
|
||||
}
|
||||
|
||||
pub(crate) fn default_me_writer_byte_budget_bytes() -> usize {
|
||||
DEFAULT_ME_WRITER_BYTE_BUDGET_BYTES
|
||||
}
|
||||
|
||||
pub(crate) fn minimum_me_writer_byte_budget_bytes(max_client_frame: usize) -> usize {
|
||||
max_client_frame
|
||||
.saturating_mul(2)
|
||||
.saturating_add(ME_WRITER_FRAME_OVERHEAD_RESERVE_BYTES)
|
||||
.div_ceil(ME_WRITER_BYTE_PERMIT_UNIT_BYTES)
|
||||
.saturating_mul(ME_WRITER_BYTE_PERMIT_UNIT_BYTES)
|
||||
}
|
||||
|
||||
pub(crate) fn default_me_route_channel_capacity() -> usize {
|
||||
DEFAULT_ME_ROUTE_CHANNEL_CAPACITY
|
||||
}
|
||||
@@ -499,6 +517,10 @@ pub(crate) fn default_direct_relay_copy_buf_s2c_bytes() -> usize {
|
||||
DEFAULT_DIRECT_RELAY_COPY_BUF_S2C_BYTES
|
||||
}
|
||||
|
||||
pub(crate) fn default_direct_relay_buffer_budget_max_bytes() -> usize {
|
||||
DEFAULT_DIRECT_RELAY_BUFFER_BUDGET_MAX_BYTES
|
||||
}
|
||||
|
||||
pub(crate) fn default_me_writer_pick_sample_size() -> u8 {
|
||||
DEFAULT_ME_WRITER_PICK_SAMPLE_SIZE
|
||||
}
|
||||
|
||||
@@ -1403,11 +1403,13 @@ fn reload_config(
|
||||
/// `detected_ip_v4` / `detected_ip_v6` are the IPs discovered during the
|
||||
/// startup probe — used when generating proxy links for newly added users,
|
||||
/// matching the same logic as the startup output.
|
||||
/// The watcher releases its notify and signal resources when `cancellation` fires.
|
||||
pub fn spawn_config_watcher(
|
||||
config_path: PathBuf,
|
||||
initial: Arc<ProxyConfig>,
|
||||
detected_ip_v4: Option<IpAddr>,
|
||||
detected_ip_v6: Option<IpAddr>,
|
||||
cancellation: tokio_util::sync::CancellationToken,
|
||||
) -> (watch::Receiver<Arc<ProxyConfig>>, watch::Receiver<LogLevel>) {
|
||||
let initial_level = initial.general.log_level.clone();
|
||||
let (config_tx, config_rx) = watch::channel(initial);
|
||||
@@ -1515,10 +1517,14 @@ pub fn spawn_config_watcher(
|
||||
_ = sighup.recv() => {
|
||||
info!("SIGHUP received — reloading {:?}", config_path);
|
||||
}
|
||||
_ = cancellation.cancelled() => break,
|
||||
}
|
||||
#[cfg(not(unix))]
|
||||
if notify_rx.recv().await.is_none() {
|
||||
break;
|
||||
tokio::select! {
|
||||
msg = notify_rx.recv() => {
|
||||
if msg.is_none() { break; }
|
||||
}
|
||||
_ = cancellation.cancelled() => break,
|
||||
}
|
||||
|
||||
// Debounce: drain extra events that arrive within a short quiet window.
|
||||
|
||||
@@ -39,8 +39,11 @@ use self::validation::{
|
||||
};
|
||||
|
||||
const MAX_ME_WRITER_CMD_CHANNEL_CAPACITY: usize = 16_384;
|
||||
const MAX_ME_WRITER_BYTE_BUDGET_BYTES: usize = 256 * 1024 * 1024;
|
||||
const MAX_ME_ROUTE_CHANNEL_CAPACITY: usize = 8_192;
|
||||
const MAX_ME_C2ME_CHANNEL_CAPACITY: usize = 8_192;
|
||||
const MIN_DIRECT_RELAY_BUFFER_BUDGET_BYTES: usize = 16 * 1024 * 1024;
|
||||
const MAX_DIRECT_RELAY_BUFFER_BUDGET_BYTES: usize = 2 * 1024 * 1024 * 1024;
|
||||
const MIN_MAX_CLIENT_FRAME_BYTES: usize = 4 * 1024;
|
||||
const MAX_MAX_CLIENT_FRAME_BYTES: usize = 16 * 1024 * 1024;
|
||||
const MAX_API_REQUEST_BODY_LIMIT_BYTES: usize = 1024 * 1024;
|
||||
@@ -540,6 +543,22 @@ impl ProxyConfig {
|
||||
)));
|
||||
}
|
||||
|
||||
let min_writer_byte_budget =
|
||||
minimum_me_writer_byte_budget_bytes(config.general.max_client_frame);
|
||||
if config.general.me_writer_byte_budget_bytes % ME_WRITER_BYTE_PERMIT_UNIT_BYTES != 0 {
|
||||
return Err(ProxyError::Config(format!(
|
||||
"general.me_writer_byte_budget_bytes must be a multiple of {ME_WRITER_BYTE_PERMIT_UNIT_BYTES}"
|
||||
)));
|
||||
}
|
||||
if !(min_writer_byte_budget..=MAX_ME_WRITER_BYTE_BUDGET_BYTES)
|
||||
.contains(&config.general.me_writer_byte_budget_bytes)
|
||||
{
|
||||
return Err(ProxyError::Config(format!(
|
||||
"general.me_writer_byte_budget_bytes must be within [{min_writer_byte_budget}, {MAX_ME_WRITER_BYTE_BUDGET_BYTES}] for general.max_client_frame={}",
|
||||
config.general.max_client_frame
|
||||
)));
|
||||
}
|
||||
|
||||
if config.general.me_c2me_send_timeout_ms > 60_000 {
|
||||
return Err(ProxyError::Config(
|
||||
"general.me_c2me_send_timeout_ms must be within [0, 60000]".to_string(),
|
||||
@@ -599,6 +618,24 @@ impl ProxyConfig {
|
||||
));
|
||||
}
|
||||
|
||||
if config.general.direct_relay_buffer_budget_max_bytes != 0 {
|
||||
if config.general.direct_relay_buffer_budget_max_bytes
|
||||
% DIRECT_RELAY_BUFFER_BUDGET_UNIT_BYTES
|
||||
!= 0
|
||||
{
|
||||
return Err(ProxyError::Config(format!(
|
||||
"general.direct_relay_buffer_budget_max_bytes must be 0 or a multiple of {DIRECT_RELAY_BUFFER_BUDGET_UNIT_BYTES}"
|
||||
)));
|
||||
}
|
||||
if !(MIN_DIRECT_RELAY_BUFFER_BUDGET_BYTES..=MAX_DIRECT_RELAY_BUFFER_BUDGET_BYTES)
|
||||
.contains(&config.general.direct_relay_buffer_budget_max_bytes)
|
||||
{
|
||||
return Err(ProxyError::Config(format!(
|
||||
"general.direct_relay_buffer_budget_max_bytes must be 0 or within [{MIN_DIRECT_RELAY_BUFFER_BUDGET_BYTES}, {MAX_DIRECT_RELAY_BUFFER_BUDGET_BYTES}]"
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
if config.general.me_health_interval_ms_unhealthy == 0 {
|
||||
return Err(ProxyError::Config(
|
||||
"general.me_health_interval_ms_unhealthy must be > 0".to_string(),
|
||||
|
||||
@@ -52,6 +52,7 @@ const GENERAL_CONFIG_KEYS: &[&str] = &[
|
||||
"me_keepalive_payload_random",
|
||||
"rpc_proxy_req_every",
|
||||
"me_writer_cmd_channel_capacity",
|
||||
"me_writer_byte_budget_bytes",
|
||||
"me_route_channel_capacity",
|
||||
"me_c2me_channel_capacity",
|
||||
"me_c2me_send_timeout_ms",
|
||||
@@ -64,6 +65,7 @@ const GENERAL_CONFIG_KEYS: &[&str] = &[
|
||||
"me_d2c_frame_buf_shrink_threshold_bytes",
|
||||
"direct_relay_copy_buf_c2s_bytes",
|
||||
"direct_relay_copy_buf_s2c_bytes",
|
||||
"direct_relay_buffer_budget_max_bytes",
|
||||
"crypto_pending_buffer",
|
||||
"max_client_frame",
|
||||
"desync_all_full",
|
||||
|
||||
@@ -95,6 +95,101 @@ max_client_frame = 16777217
|
||||
remove_temp_config(&path);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_rejects_writer_byte_budget_below_frame_residency_minimum() {
|
||||
let path = write_temp_config(
|
||||
r#"
|
||||
[general]
|
||||
max_client_frame = 16777216
|
||||
me_writer_byte_budget_bytes = 33554432
|
||||
"#,
|
||||
);
|
||||
|
||||
let err = ProxyConfig::load(&path)
|
||||
.expect_err("writer byte budget below frame residency minimum must fail");
|
||||
let msg = err.to_string();
|
||||
assert!(
|
||||
msg.contains("general.me_writer_byte_budget_bytes must be within [33570816, 268435456]"),
|
||||
"error must explain writer byte budget minimum, got: {msg}"
|
||||
);
|
||||
|
||||
remove_temp_config(&path);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_rejects_unaligned_writer_byte_budget() {
|
||||
let path = write_temp_config(
|
||||
r#"
|
||||
[general]
|
||||
me_writer_byte_budget_bytes = 33570817
|
||||
"#,
|
||||
);
|
||||
|
||||
let err = ProxyConfig::load(&path)
|
||||
.expect_err("writer byte budget outside permit granularity must fail");
|
||||
let msg = err.to_string();
|
||||
assert!(
|
||||
msg.contains("general.me_writer_byte_budget_bytes must be a multiple of 16384"),
|
||||
"error must explain writer byte budget alignment, got: {msg}"
|
||||
);
|
||||
|
||||
remove_temp_config(&path);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_rejects_writer_byte_budget_above_hard_cap() {
|
||||
let path = write_temp_config(
|
||||
r#"
|
||||
[general]
|
||||
me_writer_byte_budget_bytes = 268451840
|
||||
"#,
|
||||
);
|
||||
|
||||
let err = ProxyConfig::load(&path).expect_err("writer byte budget above hard cap must fail");
|
||||
let msg = err.to_string();
|
||||
assert!(
|
||||
msg.contains("general.me_writer_byte_budget_bytes must be within [33570816, 268435456]"),
|
||||
"error must explain writer byte budget hard cap, got: {msg}"
|
||||
);
|
||||
|
||||
remove_temp_config(&path);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_rejects_unaligned_direct_relay_buffer_budget() {
|
||||
let path = write_temp_config(
|
||||
r#"
|
||||
[general]
|
||||
direct_relay_buffer_budget_max_bytes = 16777217
|
||||
"#,
|
||||
);
|
||||
|
||||
let err = ProxyConfig::load(&path).expect_err("unaligned direct relay buffer budget must fail");
|
||||
assert!(
|
||||
err.to_string().contains(
|
||||
"general.direct_relay_buffer_budget_max_bytes must be 0 or a multiple of 4096"
|
||||
)
|
||||
);
|
||||
remove_temp_config(&path);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_rejects_direct_relay_buffer_budget_above_hard_cap() {
|
||||
let path = write_temp_config(
|
||||
r#"
|
||||
[general]
|
||||
direct_relay_buffer_budget_max_bytes = 2147487744
|
||||
"#,
|
||||
);
|
||||
|
||||
let err =
|
||||
ProxyConfig::load(&path).expect_err("direct relay buffer budget above hard cap must fail");
|
||||
assert!(err.to_string().contains(
|
||||
"general.direct_relay_buffer_budget_max_bytes must be 0 or within [16777216, 2147483648]"
|
||||
));
|
||||
remove_temp_config(&path);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_rejects_listen_backlog_above_i32_upper_bound() {
|
||||
let path = write_temp_config(
|
||||
@@ -139,16 +234,23 @@ fn load_accepts_memory_limits_at_hard_upper_bounds() {
|
||||
r#"
|
||||
[general]
|
||||
me_writer_cmd_channel_capacity = 16384
|
||||
me_writer_byte_budget_bytes = 268435456
|
||||
me_route_channel_capacity = 8192
|
||||
me_c2me_channel_capacity = 8192
|
||||
direct_relay_buffer_budget_max_bytes = 2147483648
|
||||
max_client_frame = 16777216
|
||||
"#,
|
||||
);
|
||||
|
||||
let cfg = ProxyConfig::load(&path).expect("hard upper bound values must be accepted");
|
||||
assert_eq!(cfg.general.me_writer_cmd_channel_capacity, 16384);
|
||||
assert_eq!(cfg.general.me_writer_byte_budget_bytes, 256 * 1024 * 1024);
|
||||
assert_eq!(cfg.general.me_route_channel_capacity, 8192);
|
||||
assert_eq!(cfg.general.me_c2me_channel_capacity, 8192);
|
||||
assert_eq!(
|
||||
cfg.general.direct_relay_buffer_budget_max_bytes,
|
||||
2 * 1024 * 1024 * 1024
|
||||
);
|
||||
assert_eq!(cfg.general.max_client_frame, 16 * 1024 * 1024);
|
||||
|
||||
remove_temp_config(&path);
|
||||
|
||||
@@ -579,6 +579,10 @@ pub struct GeneralConfig {
|
||||
#[serde(default = "default_me_writer_cmd_channel_capacity")]
|
||||
pub me_writer_cmd_channel_capacity: usize,
|
||||
|
||||
/// Resident-memory budget in bytes for each ME writer data queue.
|
||||
#[serde(default = "default_me_writer_byte_budget_bytes")]
|
||||
pub me_writer_byte_budget_bytes: usize,
|
||||
|
||||
/// Capacity of per-connection ME response route channel.
|
||||
#[serde(default = "default_me_route_channel_capacity")]
|
||||
pub me_route_channel_capacity: usize,
|
||||
@@ -622,7 +626,7 @@ pub struct GeneralConfig {
|
||||
#[serde(default = "default_me_d2c_frame_buf_shrink_threshold_bytes")]
|
||||
pub me_d2c_frame_buf_shrink_threshold_bytes: usize,
|
||||
|
||||
/// Copy buffer size for client->DC direction in direct relay.
|
||||
/// Copy buffer ceiling for client->DC direction in direct relay.
|
||||
///
|
||||
/// This is also the upper bound for one amortized upload rate-limit burst:
|
||||
/// upload debt is settled before the next relay read instead of blocking
|
||||
@@ -630,13 +634,18 @@ pub struct GeneralConfig {
|
||||
#[serde(default = "default_direct_relay_copy_buf_c2s_bytes")]
|
||||
pub direct_relay_copy_buf_c2s_bytes: usize,
|
||||
|
||||
/// Copy buffer size for DC->client direction in direct relay.
|
||||
/// Copy buffer ceiling for DC->client direction in direct relay.
|
||||
///
|
||||
/// This bounds one direct download rate-limit grant because writes are
|
||||
/// clipped to the currently available shaper budget.
|
||||
#[serde(default = "default_direct_relay_copy_buf_s2c_bytes")]
|
||||
pub direct_relay_copy_buf_s2c_bytes: usize,
|
||||
|
||||
/// Process-wide hard ceiling for Direct relay copy buffers.
|
||||
/// `0` derives the ceiling from host and cgroup memory limits.
|
||||
#[serde(default = "default_direct_relay_buffer_budget_max_bytes")]
|
||||
pub direct_relay_buffer_budget_max_bytes: usize,
|
||||
|
||||
/// Max pending ciphertext buffer per client writer (bytes).
|
||||
/// Controls FakeTLS backpressure vs throughput.
|
||||
#[serde(default = "default_crypto_pending_buffer")]
|
||||
@@ -1103,6 +1112,7 @@ impl Default for GeneralConfig {
|
||||
me_keepalive_payload_random: default_true(),
|
||||
rpc_proxy_req_every: default_rpc_proxy_req_every(),
|
||||
me_writer_cmd_channel_capacity: default_me_writer_cmd_channel_capacity(),
|
||||
me_writer_byte_budget_bytes: default_me_writer_byte_budget_bytes(),
|
||||
me_route_channel_capacity: default_me_route_channel_capacity(),
|
||||
me_c2me_channel_capacity: default_me_c2me_channel_capacity(),
|
||||
me_c2me_send_timeout_ms: default_me_c2me_send_timeout_ms(),
|
||||
@@ -1116,6 +1126,7 @@ impl Default for GeneralConfig {
|
||||
default_me_d2c_frame_buf_shrink_threshold_bytes(),
|
||||
direct_relay_copy_buf_c2s_bytes: default_direct_relay_copy_buf_c2s_bytes(),
|
||||
direct_relay_copy_buf_s2c_bytes: default_direct_relay_copy_buf_s2c_bytes(),
|
||||
direct_relay_buffer_budget_max_bytes: default_direct_relay_buffer_budget_max_bytes(),
|
||||
me_warmup_stagger_enabled: default_true(),
|
||||
me_warmup_step_delay_ms: default_warmup_step_delay_ms(),
|
||||
me_warmup_step_jitter_ms: default_warmup_step_jitter_ms(),
|
||||
|
||||
@@ -7,6 +7,7 @@ use std::time::Duration;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::process::Command;
|
||||
use tokio::sync::{mpsc, watch};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use crate::config::{ConntrackBackend, ConntrackMode, ProxyConfig};
|
||||
@@ -57,10 +58,11 @@ impl PressureState {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn spawn_conntrack_controller(
|
||||
pub(crate) async fn run_conntrack_controller(
|
||||
config_rx: watch::Receiver<Arc<ProxyConfig>>,
|
||||
stats: Arc<Stats>,
|
||||
shared: Arc<ProxySharedState>,
|
||||
cancellation: CancellationToken,
|
||||
) {
|
||||
if !cfg!(target_os = "linux") {
|
||||
let cfg = config_rx.borrow();
|
||||
@@ -87,16 +89,15 @@ pub(crate) fn spawn_conntrack_controller(
|
||||
|
||||
let (tx, rx) = mpsc::channel(CONNTRACK_EVENT_QUEUE_CAPACITY);
|
||||
shared.set_conntrack_close_sender(tx);
|
||||
tokio::spawn(async move {
|
||||
run_conntrack_controller(config_rx, stats, shared, rx).await;
|
||||
});
|
||||
run_conntrack_controller_worker(config_rx, stats, shared, rx, cancellation).await;
|
||||
}
|
||||
|
||||
async fn run_conntrack_controller(
|
||||
async fn run_conntrack_controller_worker(
|
||||
mut config_rx: watch::Receiver<Arc<ProxyConfig>>,
|
||||
stats: Arc<Stats>,
|
||||
shared: Arc<ProxySharedState>,
|
||||
mut close_rx: mpsc::Receiver<ConntrackCloseEvent>,
|
||||
cancellation: CancellationToken,
|
||||
) {
|
||||
let mut cfg = config_rx.borrow().clone();
|
||||
let mut pressure_state = PressureState::new(stats.as_ref());
|
||||
@@ -115,6 +116,7 @@ async fn run_conntrack_controller(
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = cancellation.cancelled() => break,
|
||||
changed = config_rx.changed() => {
|
||||
if changed.is_err() {
|
||||
break;
|
||||
|
||||
@@ -4,12 +4,12 @@
|
||||
//!
|
||||
//! ## Zeroize policy
|
||||
//!
|
||||
//! - `AesCbc` stores raw key/IV bytes and zeroizes them on drop.
|
||||
//! - `AesCtr` wraps an opaque `Aes256Ctr` cipher from the `ctr` crate.
|
||||
//! The expanded key schedule lives inside that type and cannot be
|
||||
//! zeroized from outside. Callers that hold raw key material (e.g.
|
||||
//! `HandshakeSuccess`, `ObfuscationParams`) are responsible for
|
||||
//! zeroizing their own copies.
|
||||
//! - `AesCbc` stores raw key/IV bytes and zeroizes them on drop. Temporary
|
||||
//! expanded key schedules are also zeroized by the RustCrypto backend.
|
||||
//! - `AesCtr` uses the RustCrypto `zeroize` contract to clear its expanded
|
||||
//! key schedule, counter, and buffered keystream on drop.
|
||||
//! - Callers that hold raw key material (e.g. `HandshakeSuccess`,
|
||||
//! `ObfuscationParams`) remain responsible for zeroizing their own copies.
|
||||
|
||||
#![allow(dead_code)]
|
||||
|
||||
@@ -19,24 +19,28 @@ use ctr::{
|
||||
Ctr128BE,
|
||||
cipher::{KeyIvInit, StreamCipher},
|
||||
};
|
||||
use zeroize::Zeroize;
|
||||
use zeroize::{Zeroize, ZeroizeOnDrop};
|
||||
|
||||
type Aes256Ctr = Ctr128BE<Aes256>;
|
||||
|
||||
static_assertions::assert_impl_all!(Aes256: ZeroizeOnDrop);
|
||||
static_assertions::assert_impl_all!(Aes256Ctr: ZeroizeOnDrop);
|
||||
|
||||
// ============= AES-256-CTR =============
|
||||
|
||||
/// AES-256-CTR encryptor/decryptor
|
||||
///
|
||||
/// CTR mode is symmetric — encryption and decryption are the same operation.
|
||||
///
|
||||
/// **Zeroize note:** The inner `Aes256Ctr` cipher state (expanded key schedule
|
||||
/// + counter) is opaque and cannot be zeroized. If you need to protect key
|
||||
/// material, zeroize the `[u8; 32]` key and `u128` IV at the call site
|
||||
/// before dropping them.
|
||||
/// **Zeroize note:** The inner `Aes256Ctr` zeroizes its expanded key schedule,
|
||||
/// counter, and buffered keystream on drop. Callers remain responsible for
|
||||
/// zeroizing their own raw key and IV copies.
|
||||
pub struct AesCtr {
|
||||
cipher: Aes256Ctr,
|
||||
}
|
||||
|
||||
impl ZeroizeOnDrop for AesCtr {}
|
||||
|
||||
impl AesCtr {
|
||||
/// Create new AES-CTR cipher with key and IV
|
||||
pub fn new(key: &[u8; 32], iv: u128) -> Self {
|
||||
@@ -92,7 +96,7 @@ impl AesCtr {
|
||||
/// are different operations. This implementation handles CBC chaining
|
||||
/// correctly across multiple blocks.
|
||||
///
|
||||
/// Key and IV are zeroized on drop.
|
||||
/// Key, IV, and temporary expanded key schedules are zeroized on drop.
|
||||
pub struct AesCbc {
|
||||
key: [u8; 32],
|
||||
iv: [u8; 16],
|
||||
@@ -105,6 +109,8 @@ impl Drop for AesCbc {
|
||||
}
|
||||
}
|
||||
|
||||
impl ZeroizeOnDrop for AesCbc {}
|
||||
|
||||
impl AesCbc {
|
||||
/// AES block size
|
||||
const BLOCK_SIZE: usize = 16;
|
||||
|
||||
@@ -8,6 +8,8 @@ use crate::config::ProxyConfig;
|
||||
use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController};
|
||||
use crate::transport::middle_proxy::MePool;
|
||||
|
||||
use super::generation::RuntimeTaskScope;
|
||||
|
||||
const STARTUP_FALLBACK_AFTER: Duration = Duration::from_secs(80);
|
||||
const RUNTIME_FALLBACK_AFTER: Duration = Duration::from_secs(6);
|
||||
|
||||
@@ -19,6 +21,7 @@ pub(crate) async fn configure_admission_gate(
|
||||
admission_tx: &watch::Sender<bool>,
|
||||
config_rx: watch::Receiver<Arc<ProxyConfig>>,
|
||||
me_ready_rx: watch::Receiver<u64>,
|
||||
task_scope: RuntimeTaskScope,
|
||||
) {
|
||||
if config.general.use_middle_proxy {
|
||||
if me_pool.is_some() || config.general.me2dc_fallback {
|
||||
@@ -64,7 +67,7 @@ pub(crate) async fn configure_admission_gate(
|
||||
let mut config_rx_gate = config_rx.clone();
|
||||
let mut me_ready_rx_gate = me_ready_rx;
|
||||
let mut admission_poll_ms = config.general.me_admission_poll_ms.max(1);
|
||||
tokio::spawn(async move {
|
||||
task_scope.spawn(async move {
|
||||
let mut gate_open = initial_gate_open;
|
||||
let mut route_mode = initial_route_mode;
|
||||
let mut ready_observed = initial_ready;
|
||||
|
||||
238
src/maestro/generation.rs
Normal file
238
src/maestro/generation.rs
Normal file
@@ -0,0 +1,238 @@
|
||||
use std::future::Future;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::sync::{RwLock, Semaphore, watch};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tokio_util::task::TaskTracker;
|
||||
|
||||
use crate::config::ProxyConfig;
|
||||
use crate::crypto::SecureRandom;
|
||||
use crate::ip_tracker::UserIpTracker;
|
||||
use crate::proxy::route_mode::RouteRuntimeController;
|
||||
use crate::proxy::shared_state::ProxySharedState;
|
||||
use crate::stats::beobachten::BeobachtenStore;
|
||||
use crate::stats::{ReplayChecker, Stats};
|
||||
use crate::stream::BufferPool;
|
||||
use crate::tls_front::TlsFrontCache;
|
||||
use crate::transport::UpstreamManager;
|
||||
use crate::transport::middle_proxy::MePool;
|
||||
|
||||
const SESSION_STOP_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
const BACKGROUND_STOP_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
|
||||
/// Process-visible control-plane receivers for one active runtime generation.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct RuntimeWatchState {
|
||||
pub(crate) generation_id: u64,
|
||||
pub(crate) config_rx: watch::Receiver<Arc<ProxyConfig>>,
|
||||
pub(crate) admission_rx: watch::Receiver<bool>,
|
||||
}
|
||||
|
||||
/// Cancellation and join ownership for one generation's background tasks.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct RuntimeTaskScope {
|
||||
tracker: TaskTracker,
|
||||
cancel: CancellationToken,
|
||||
}
|
||||
|
||||
impl RuntimeTaskScope {
|
||||
pub(crate) fn new() -> Self {
|
||||
Self {
|
||||
tracker: TaskTracker::new(),
|
||||
cancel: CancellationToken::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn spawn<F>(&self, future: F)
|
||||
where
|
||||
F: Future<Output = ()> + Send + 'static,
|
||||
{
|
||||
let cancel = self.cancel.clone();
|
||||
self.tracker.spawn(async move {
|
||||
tokio::select! {
|
||||
_ = cancel.cancelled() => {}
|
||||
_ = future => {}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
pub(crate) fn cancellation_token(&self) -> CancellationToken {
|
||||
self.cancel.clone()
|
||||
}
|
||||
|
||||
pub(crate) async fn stop(&self) {
|
||||
self.cancel.cancel();
|
||||
self.tracker.close();
|
||||
let _ = tokio::time::timeout(BACKGROUND_STOP_TIMEOUT, self.tracker.wait()).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Runtime-owned data plane and control-plane dependencies for one generation.
|
||||
pub(crate) struct RuntimeGeneration {
|
||||
pub(crate) id: u64,
|
||||
pub(crate) config_rx: watch::Receiver<Arc<ProxyConfig>>,
|
||||
pub(crate) admission_rx: watch::Receiver<bool>,
|
||||
pub(crate) stats: Arc<Stats>,
|
||||
pub(crate) upstream_manager: Arc<UpstreamManager>,
|
||||
pub(crate) replay_checker: Arc<ReplayChecker>,
|
||||
pub(crate) buffer_pool: Arc<BufferPool>,
|
||||
pub(crate) rng: Arc<SecureRandom>,
|
||||
pub(crate) me_pool: Option<Arc<MePool>>,
|
||||
pub(crate) me_pool_runtime: Arc<RwLock<Option<Arc<MePool>>>>,
|
||||
pub(crate) route_runtime: Arc<RouteRuntimeController>,
|
||||
pub(crate) tls_cache: Option<Arc<TlsFrontCache>>,
|
||||
pub(crate) ip_tracker: Arc<UserIpTracker>,
|
||||
pub(crate) beobachten: Arc<BeobachtenStore>,
|
||||
pub(crate) proxy_shared: Arc<ProxySharedState>,
|
||||
pub(crate) max_connections: Arc<Semaphore>,
|
||||
background_tasks: RuntimeTaskScope,
|
||||
sessions: TaskTracker,
|
||||
session_cancel: CancellationToken,
|
||||
accepting_sessions: AtomicBool,
|
||||
}
|
||||
|
||||
impl RuntimeGeneration {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) fn new(
|
||||
id: u64,
|
||||
config_rx: watch::Receiver<Arc<ProxyConfig>>,
|
||||
admission_rx: watch::Receiver<bool>,
|
||||
stats: Arc<Stats>,
|
||||
upstream_manager: Arc<UpstreamManager>,
|
||||
replay_checker: Arc<ReplayChecker>,
|
||||
buffer_pool: Arc<BufferPool>,
|
||||
rng: Arc<SecureRandom>,
|
||||
me_pool: Option<Arc<MePool>>,
|
||||
me_pool_runtime: Arc<RwLock<Option<Arc<MePool>>>>,
|
||||
route_runtime: Arc<RouteRuntimeController>,
|
||||
tls_cache: Option<Arc<TlsFrontCache>>,
|
||||
ip_tracker: Arc<UserIpTracker>,
|
||||
beobachten: Arc<BeobachtenStore>,
|
||||
proxy_shared: Arc<ProxySharedState>,
|
||||
max_connections: Arc<Semaphore>,
|
||||
background_tasks: RuntimeTaskScope,
|
||||
) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
id,
|
||||
config_rx,
|
||||
admission_rx,
|
||||
stats,
|
||||
upstream_manager,
|
||||
replay_checker,
|
||||
buffer_pool,
|
||||
rng,
|
||||
me_pool,
|
||||
me_pool_runtime,
|
||||
route_runtime,
|
||||
tls_cache,
|
||||
ip_tracker,
|
||||
beobachten,
|
||||
proxy_shared,
|
||||
max_connections,
|
||||
background_tasks,
|
||||
sessions: TaskTracker::new(),
|
||||
session_cancel: CancellationToken::new(),
|
||||
accepting_sessions: AtomicBool::new(true),
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn config(&self) -> Arc<ProxyConfig> {
|
||||
self.config_rx.borrow().clone()
|
||||
}
|
||||
|
||||
/// Returns receivers used by process-scoped observers of this generation.
|
||||
pub(crate) fn watch_state(&self) -> RuntimeWatchState {
|
||||
RuntimeWatchState {
|
||||
generation_id: self.id,
|
||||
config_rx: self.config_rx.clone(),
|
||||
admission_rx: self.admission_rx.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn current_me_pool(&self) -> Option<Arc<MePool>> {
|
||||
if let Some(pool) = &self.me_pool {
|
||||
return Some(pool.clone());
|
||||
}
|
||||
self.me_pool_runtime.read().await.clone()
|
||||
}
|
||||
|
||||
pub(crate) fn spawn_session<F>(&self, future: F) -> bool
|
||||
where
|
||||
F: Future<Output = ()> + Send + 'static,
|
||||
{
|
||||
if !self.accepting_sessions.load(Ordering::Acquire) {
|
||||
return false;
|
||||
}
|
||||
let cancel = self.session_cancel.clone();
|
||||
self.sessions.spawn(async move {
|
||||
tokio::select! {
|
||||
_ = cancel.cancelled() => {}
|
||||
_ = future => {}
|
||||
}
|
||||
});
|
||||
true
|
||||
}
|
||||
|
||||
pub(crate) fn stop_accepting_sessions(&self) {
|
||||
self.accepting_sessions.store(false, Ordering::Release);
|
||||
}
|
||||
|
||||
pub(crate) fn resume_accepting_sessions(&self) {
|
||||
self.accepting_sessions.store(true, Ordering::Release);
|
||||
}
|
||||
|
||||
pub(crate) async fn drain_sessions(&self, timeout: Duration) -> bool {
|
||||
self.stop_accepting_sessions();
|
||||
self.sessions.close();
|
||||
if tokio::time::timeout(timeout, self.sessions.wait())
|
||||
.await
|
||||
.is_ok()
|
||||
{
|
||||
return true;
|
||||
}
|
||||
self.stop_sessions().await;
|
||||
false
|
||||
}
|
||||
|
||||
pub(crate) async fn stop_sessions(&self) {
|
||||
self.stop_accepting_sessions();
|
||||
self.session_cancel.cancel();
|
||||
self.sessions.close();
|
||||
let _ = tokio::time::timeout(SESSION_STOP_TIMEOUT, self.sessions.wait()).await;
|
||||
}
|
||||
|
||||
pub(crate) async fn stop_background_tasks(&self) {
|
||||
self.background_tasks.stop().await;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn stop_sessions_cancels_tracked_future() {
|
||||
let tracker = TaskTracker::new();
|
||||
let cancel = CancellationToken::new();
|
||||
let child_cancel = cancel.clone();
|
||||
tracker.spawn(async move {
|
||||
child_cancel.cancelled().await;
|
||||
});
|
||||
cancel.cancel();
|
||||
tracker.close();
|
||||
tokio::time::timeout(Duration::from_secs(1), tracker.wait())
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn runtime_task_scope_joins_cancelled_background_task() {
|
||||
let scope = RuntimeTaskScope::new();
|
||||
scope.spawn(std::future::pending());
|
||||
tokio::time::timeout(Duration::from_secs(1), scope.stop())
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
@@ -3,35 +3,33 @@ use std::net::{IpAddr, SocketAddr};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use arc_swap::ArcSwap;
|
||||
use tokio::net::TcpListener;
|
||||
#[cfg(unix)]
|
||||
use tokio::net::UnixListener;
|
||||
use tokio::sync::{RwLock, Semaphore, watch};
|
||||
use tracing::{debug, error, info, warn};
|
||||
|
||||
use crate::config::{ProxyConfig, RstOnCloseMode};
|
||||
use crate::crypto::SecureRandom;
|
||||
use crate::ip_tracker::UserIpTracker;
|
||||
use crate::proxy::ClientHandler;
|
||||
use crate::proxy::route_mode::RouteRuntimeController;
|
||||
use crate::proxy::shared_state::ProxySharedState;
|
||||
use crate::startup::{COMPONENT_LISTENERS_BIND, StartupTracker};
|
||||
use crate::stats::beobachten::BeobachtenStore;
|
||||
use crate::stats::{ReplayChecker, Stats};
|
||||
use crate::stream::BufferPool;
|
||||
use crate::tls_front::TlsFrontCache;
|
||||
use crate::transport::middle_proxy::MePool;
|
||||
use crate::transport::socket::set_linger_zero;
|
||||
use crate::transport::{ListenOptions, UpstreamManager, create_listener, find_listener_processes};
|
||||
use crate::transport::{ListenOptions, create_listener, find_listener_processes};
|
||||
|
||||
use super::generation::RuntimeGeneration;
|
||||
use super::helpers::{
|
||||
expected_handshake_close_description, is_expected_handshake_eof, peer_close_description,
|
||||
print_proxy_links,
|
||||
};
|
||||
|
||||
#[cfg(unix)]
|
||||
mod unix;
|
||||
#[cfg(unix)]
|
||||
pub(crate) use unix::spawn_unix_accept_loop;
|
||||
|
||||
pub(crate) struct BoundListeners {
|
||||
pub(crate) listeners: Vec<(TcpListener, bool)>,
|
||||
pub(crate) has_unix_listener: bool,
|
||||
#[cfg(unix)]
|
||||
pub(crate) unix_listener: Option<UnixListener>,
|
||||
}
|
||||
|
||||
fn listener_port_or_legacy(listener: &crate::config::ListenerConfig, config: &ProxyConfig) -> u16 {
|
||||
@@ -59,21 +57,6 @@ pub(crate) async fn bind_listeners(
|
||||
detected_ip_v4: Option<IpAddr>,
|
||||
detected_ip_v6: Option<IpAddr>,
|
||||
startup_tracker: &Arc<StartupTracker>,
|
||||
config_rx: watch::Receiver<Arc<ProxyConfig>>,
|
||||
admission_rx: watch::Receiver<bool>,
|
||||
stats: Arc<Stats>,
|
||||
upstream_manager: Arc<UpstreamManager>,
|
||||
replay_checker: Arc<ReplayChecker>,
|
||||
buffer_pool: Arc<BufferPool>,
|
||||
rng: Arc<SecureRandom>,
|
||||
me_pool: Option<Arc<MePool>>,
|
||||
me_pool_runtime: Arc<RwLock<Option<Arc<MePool>>>>,
|
||||
route_runtime: Arc<RouteRuntimeController>,
|
||||
tls_cache: Option<Arc<TlsFrontCache>>,
|
||||
ip_tracker: Arc<UserIpTracker>,
|
||||
beobachten: Arc<BeobachtenStore>,
|
||||
shared: Arc<ProxySharedState>,
|
||||
max_connections: Arc<Semaphore>,
|
||||
) -> Result<BoundListeners, Box<dyn Error>> {
|
||||
startup_tracker
|
||||
.start_component(
|
||||
@@ -218,7 +201,8 @@ pub(crate) async fn bind_listeners(
|
||||
print_proxy_links(&host, port, config);
|
||||
}
|
||||
|
||||
let mut has_unix_listener = false;
|
||||
#[cfg(unix)]
|
||||
let mut unix_listener_out = None;
|
||||
#[cfg(unix)]
|
||||
if let Some(ref unix_path) = config.server.listen_unix_sock {
|
||||
let _ = tokio::fs::remove_file(unix_path).await;
|
||||
@@ -251,123 +235,14 @@ pub(crate) async fn bind_listeners(
|
||||
info!("Listening on unix:{}", unix_path);
|
||||
}
|
||||
|
||||
has_unix_listener = true;
|
||||
|
||||
let mut config_rx_unix: watch::Receiver<Arc<ProxyConfig>> = config_rx.clone();
|
||||
let admission_rx_unix = admission_rx.clone();
|
||||
let stats = stats.clone();
|
||||
let upstream_manager = upstream_manager.clone();
|
||||
let replay_checker = replay_checker.clone();
|
||||
let buffer_pool = buffer_pool.clone();
|
||||
let rng = rng.clone();
|
||||
let me_pool = me_pool.clone();
|
||||
let me_pool_runtime = me_pool_runtime.clone();
|
||||
let route_runtime = route_runtime.clone();
|
||||
let tls_cache = tls_cache.clone();
|
||||
let ip_tracker = ip_tracker.clone();
|
||||
let beobachten = beobachten.clone();
|
||||
let shared = shared.clone();
|
||||
let max_connections_unix = max_connections.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let unix_conn_counter = Arc::new(std::sync::atomic::AtomicU64::new(1));
|
||||
|
||||
loop {
|
||||
match unix_listener.accept().await {
|
||||
Ok((stream, _)) => {
|
||||
if !*admission_rx_unix.borrow() {
|
||||
drop(stream);
|
||||
continue;
|
||||
}
|
||||
let accept_permit_timeout_ms =
|
||||
config_rx_unix.borrow().server.accept_permit_timeout_ms;
|
||||
let permit = if accept_permit_timeout_ms == 0 {
|
||||
match max_connections_unix.clone().acquire_owned().await {
|
||||
Ok(permit) => permit,
|
||||
Err(_) => {
|
||||
error!("Connection limiter is closed");
|
||||
break;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
match tokio::time::timeout(
|
||||
Duration::from_millis(accept_permit_timeout_ms),
|
||||
max_connections_unix.clone().acquire_owned(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(permit)) => permit,
|
||||
Ok(Err(_)) => {
|
||||
error!("Connection limiter is closed");
|
||||
break;
|
||||
}
|
||||
Err(_) => {
|
||||
stats.increment_accept_permit_timeout_total();
|
||||
debug!(
|
||||
timeout_ms = accept_permit_timeout_ms,
|
||||
"Dropping accepted unix connection: permit wait timeout"
|
||||
);
|
||||
drop(stream);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
};
|
||||
let conn_id =
|
||||
unix_conn_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||
let fake_peer =
|
||||
SocketAddr::from(([127, 0, 0, 1], (conn_id % 65535) as u16));
|
||||
|
||||
let config = config_rx_unix.borrow_and_update().clone();
|
||||
let stats = stats.clone();
|
||||
let upstream_manager = upstream_manager.clone();
|
||||
let replay_checker = replay_checker.clone();
|
||||
let buffer_pool = buffer_pool.clone();
|
||||
let rng = rng.clone();
|
||||
let me_pool = me_pool.clone();
|
||||
let me_pool_runtime = me_pool_runtime.clone();
|
||||
let route_runtime = route_runtime.clone();
|
||||
let tls_cache = tls_cache.clone();
|
||||
let ip_tracker = ip_tracker.clone();
|
||||
let beobachten = beobachten.clone();
|
||||
let shared = shared.clone();
|
||||
let proxy_protocol_enabled = config.server.proxy_protocol;
|
||||
|
||||
tokio::spawn(async move {
|
||||
let _permit = permit;
|
||||
if let Err(e) =
|
||||
crate::proxy::client::handle_client_stream_with_shared_and_pool_runtime(
|
||||
stream,
|
||||
fake_peer,
|
||||
config,
|
||||
stats,
|
||||
upstream_manager,
|
||||
replay_checker,
|
||||
buffer_pool,
|
||||
rng,
|
||||
me_pool,
|
||||
Some(me_pool_runtime),
|
||||
route_runtime,
|
||||
tls_cache,
|
||||
ip_tracker,
|
||||
beobachten,
|
||||
shared,
|
||||
proxy_protocol_enabled,
|
||||
)
|
||||
.await
|
||||
{
|
||||
debug!(error = %e, "Unix socket connection error");
|
||||
}
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
error!("Unix socket accept error: {}", e);
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
unix_listener_out = Some(unix_listener);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
let has_unix_listener = unix_listener_out.is_some();
|
||||
#[cfg(not(unix))]
|
||||
let has_unix_listener = false;
|
||||
|
||||
startup_tracker
|
||||
.complete_component(
|
||||
COMPONENT_LISTENERS_BIND,
|
||||
@@ -381,51 +256,25 @@ pub(crate) async fn bind_listeners(
|
||||
|
||||
Ok(BoundListeners {
|
||||
listeners,
|
||||
has_unix_listener,
|
||||
#[cfg(unix)]
|
||||
unix_listener: unix_listener_out,
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) fn spawn_tcp_accept_loops(
|
||||
listeners: Vec<(TcpListener, bool)>,
|
||||
config_rx: watch::Receiver<Arc<ProxyConfig>>,
|
||||
admission_rx: watch::Receiver<bool>,
|
||||
stats: Arc<Stats>,
|
||||
upstream_manager: Arc<UpstreamManager>,
|
||||
replay_checker: Arc<ReplayChecker>,
|
||||
buffer_pool: Arc<BufferPool>,
|
||||
rng: Arc<SecureRandom>,
|
||||
me_pool: Option<Arc<MePool>>,
|
||||
me_pool_runtime: Arc<RwLock<Option<Arc<MePool>>>>,
|
||||
route_runtime: Arc<RouteRuntimeController>,
|
||||
tls_cache: Option<Arc<TlsFrontCache>>,
|
||||
ip_tracker: Arc<UserIpTracker>,
|
||||
beobachten: Arc<BeobachtenStore>,
|
||||
shared: Arc<ProxySharedState>,
|
||||
max_connections: Arc<Semaphore>,
|
||||
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
||||
) {
|
||||
for (listener, listener_proxy_protocol) in listeners {
|
||||
let mut config_rx: watch::Receiver<Arc<ProxyConfig>> = config_rx.clone();
|
||||
let admission_rx_tcp = admission_rx.clone();
|
||||
let stats = stats.clone();
|
||||
let upstream_manager = upstream_manager.clone();
|
||||
let replay_checker = replay_checker.clone();
|
||||
let buffer_pool = buffer_pool.clone();
|
||||
let rng = rng.clone();
|
||||
let me_pool = me_pool.clone();
|
||||
let me_pool_runtime = me_pool_runtime.clone();
|
||||
let route_runtime = route_runtime.clone();
|
||||
let tls_cache = tls_cache.clone();
|
||||
let ip_tracker = ip_tracker.clone();
|
||||
let beobachten = beobachten.clone();
|
||||
let shared = shared.clone();
|
||||
let max_connections_tcp = max_connections.clone();
|
||||
let active_runtime = active_runtime.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
match listener.accept().await {
|
||||
Ok((stream, peer_addr)) => {
|
||||
let rst_mode = config_rx.borrow().general.rst_on_close;
|
||||
let runtime = active_runtime.load_full();
|
||||
let config = runtime.config();
|
||||
let rst_mode = config.general.rst_on_close;
|
||||
#[cfg(unix)]
|
||||
let raw_fd = {
|
||||
use std::os::unix::io::AsRawFd;
|
||||
@@ -434,15 +283,14 @@ pub(crate) fn spawn_tcp_accept_loops(
|
||||
if matches!(rst_mode, RstOnCloseMode::Errors | RstOnCloseMode::Always) {
|
||||
let _ = set_linger_zero(&stream);
|
||||
}
|
||||
if !*admission_rx_tcp.borrow() {
|
||||
if !*runtime.admission_rx.borrow() {
|
||||
debug!(peer = %peer_addr, "Admission gate closed, dropping connection");
|
||||
drop(stream);
|
||||
continue;
|
||||
}
|
||||
let accept_permit_timeout_ms =
|
||||
config_rx.borrow().server.accept_permit_timeout_ms;
|
||||
let accept_permit_timeout_ms = config.server.accept_permit_timeout_ms;
|
||||
let permit = if accept_permit_timeout_ms == 0 {
|
||||
match max_connections_tcp.clone().acquire_owned().await {
|
||||
match runtime.max_connections.clone().acquire_owned().await {
|
||||
Ok(permit) => permit,
|
||||
Err(_) => {
|
||||
error!("Connection limiter is closed");
|
||||
@@ -452,7 +300,7 @@ pub(crate) fn spawn_tcp_accept_loops(
|
||||
} else {
|
||||
match tokio::time::timeout(
|
||||
Duration::from_millis(accept_permit_timeout_ms),
|
||||
max_connections_tcp.clone().acquire_owned(),
|
||||
runtime.max_connections.clone().acquire_owned(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
@@ -462,7 +310,7 @@ pub(crate) fn spawn_tcp_accept_loops(
|
||||
break;
|
||||
}
|
||||
Err(_) => {
|
||||
stats.increment_accept_permit_timeout_total();
|
||||
runtime.stats.increment_accept_permit_timeout_total();
|
||||
debug!(
|
||||
peer = %peer_addr,
|
||||
timeout_ms = accept_permit_timeout_ms,
|
||||
@@ -473,24 +321,23 @@ pub(crate) fn spawn_tcp_accept_loops(
|
||||
}
|
||||
}
|
||||
};
|
||||
let config = config_rx.borrow_and_update().clone();
|
||||
let stats = stats.clone();
|
||||
let upstream_manager = upstream_manager.clone();
|
||||
let replay_checker = replay_checker.clone();
|
||||
let buffer_pool = buffer_pool.clone();
|
||||
let rng = rng.clone();
|
||||
let me_pool = me_pool.clone();
|
||||
let me_pool_runtime = me_pool_runtime.clone();
|
||||
let route_runtime = route_runtime.clone();
|
||||
let tls_cache = tls_cache.clone();
|
||||
let ip_tracker = ip_tracker.clone();
|
||||
let beobachten = beobachten.clone();
|
||||
let shared = shared.clone();
|
||||
let stats = runtime.stats.clone();
|
||||
let upstream_manager = runtime.upstream_manager.clone();
|
||||
let replay_checker = runtime.replay_checker.clone();
|
||||
let buffer_pool = runtime.buffer_pool.clone();
|
||||
let rng = runtime.rng.clone();
|
||||
let me_pool = runtime.me_pool.clone();
|
||||
let me_pool_runtime = runtime.me_pool_runtime.clone();
|
||||
let route_runtime = runtime.route_runtime.clone();
|
||||
let tls_cache = runtime.tls_cache.clone();
|
||||
let ip_tracker = runtime.ip_tracker.clone();
|
||||
let beobachten = runtime.beobachten.clone();
|
||||
let shared = runtime.proxy_shared.clone();
|
||||
let proxy_protocol_enabled = listener_proxy_protocol;
|
||||
let real_peer_report = Arc::new(std::sync::Mutex::new(None));
|
||||
let real_peer_report_for_handler = real_peer_report.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let _ = runtime.spawn_session(async move {
|
||||
let _permit = permit;
|
||||
if let Err(e) = ClientHandler::new_with_shared(
|
||||
stream,
|
||||
|
||||
117
src/maestro/listeners/unix.rs
Normal file
117
src/maestro/listeners/unix.rs
Normal file
@@ -0,0 +1,117 @@
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use arc_swap::ArcSwap;
|
||||
use tokio::net::UnixListener;
|
||||
use tracing::{debug, error};
|
||||
|
||||
use super::RuntimeGeneration;
|
||||
|
||||
pub(crate) fn spawn_unix_accept_loop(
|
||||
listener: Option<UnixListener>,
|
||||
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
||||
) {
|
||||
let Some(listener) = listener else {
|
||||
return;
|
||||
};
|
||||
|
||||
tokio::spawn(async move {
|
||||
let connection_counter = AtomicU64::new(1);
|
||||
|
||||
loop {
|
||||
match listener.accept().await {
|
||||
Ok((stream, _)) => {
|
||||
let runtime = active_runtime.load_full();
|
||||
if !*runtime.admission_rx.borrow() {
|
||||
drop(stream);
|
||||
continue;
|
||||
}
|
||||
|
||||
let config = runtime.config();
|
||||
let timeout_ms = config.server.accept_permit_timeout_ms;
|
||||
let permit = if timeout_ms == 0 {
|
||||
match runtime.max_connections.clone().acquire_owned().await {
|
||||
Ok(permit) => permit,
|
||||
Err(_) => {
|
||||
error!("Connection limiter is closed");
|
||||
break;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
match tokio::time::timeout(
|
||||
Duration::from_millis(timeout_ms),
|
||||
runtime.max_connections.clone().acquire_owned(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(permit)) => permit,
|
||||
Ok(Err(_)) => {
|
||||
error!("Connection limiter is closed");
|
||||
break;
|
||||
}
|
||||
Err(_) => {
|
||||
runtime.stats.increment_accept_permit_timeout_total();
|
||||
debug!(
|
||||
timeout_ms,
|
||||
"Dropping accepted unix connection: permit wait timeout"
|
||||
);
|
||||
drop(stream);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let connection_id = connection_counter.fetch_add(1, Ordering::Relaxed);
|
||||
let fake_peer =
|
||||
SocketAddr::from(([127, 0, 0, 1], (connection_id % 65535) as u16));
|
||||
let stats = runtime.stats.clone();
|
||||
let upstream_manager = runtime.upstream_manager.clone();
|
||||
let replay_checker = runtime.replay_checker.clone();
|
||||
let buffer_pool = runtime.buffer_pool.clone();
|
||||
let rng = runtime.rng.clone();
|
||||
let me_pool = runtime.me_pool.clone();
|
||||
let me_pool_runtime = runtime.me_pool_runtime.clone();
|
||||
let route_runtime = runtime.route_runtime.clone();
|
||||
let tls_cache = runtime.tls_cache.clone();
|
||||
let ip_tracker = runtime.ip_tracker.clone();
|
||||
let beobachten = runtime.beobachten.clone();
|
||||
let shared = runtime.proxy_shared.clone();
|
||||
let proxy_protocol_enabled = config.server.proxy_protocol;
|
||||
|
||||
let _ = runtime.spawn_session(async move {
|
||||
let _permit = permit;
|
||||
if let Err(error) =
|
||||
crate::proxy::client::handle_client_stream_with_shared_and_pool_runtime(
|
||||
stream,
|
||||
fake_peer,
|
||||
config,
|
||||
stats,
|
||||
upstream_manager,
|
||||
replay_checker,
|
||||
buffer_pool,
|
||||
rng,
|
||||
me_pool,
|
||||
Some(me_pool_runtime),
|
||||
route_runtime,
|
||||
tls_cache,
|
||||
ip_tracker,
|
||||
beobachten,
|
||||
shared,
|
||||
proxy_protocol_enabled,
|
||||
)
|
||||
.await
|
||||
{
|
||||
debug!(error = %error, "Unix socket connection error");
|
||||
}
|
||||
});
|
||||
}
|
||||
Err(error) => {
|
||||
error!(error = %error, "Unix socket accept error");
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -1,9 +1,11 @@
|
||||
#![allow(clippy::too_many_arguments)]
|
||||
|
||||
use std::future::Future;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::sync::{RwLock, watch};
|
||||
use tokio_util::task::AbortOnDropHandle;
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
use crate::config::ProxyConfig;
|
||||
@@ -17,8 +19,58 @@ use crate::stats::Stats;
|
||||
use crate::transport::UpstreamManager;
|
||||
use crate::transport::middle_proxy::MePool;
|
||||
|
||||
use super::generation::RuntimeTaskScope;
|
||||
use super::helpers::load_startup_proxy_config_snapshot;
|
||||
|
||||
async fn supervise_me_task<F, Fut>(task_name: &'static str, mut task: F)
|
||||
where
|
||||
F: FnMut() -> Fut,
|
||||
Fut: Future<Output = ()> + Send + 'static,
|
||||
{
|
||||
loop {
|
||||
let result = AbortOnDropHandle::new(tokio::spawn(task())).await;
|
||||
match result {
|
||||
Ok(()) => warn!(task = task_name, "Middle-End supervisor task exited unexpectedly, restarting"),
|
||||
Err(error) => {
|
||||
error!(task = task_name, error = %error, "Middle-End supervisor task panicked, restarting in 1s");
|
||||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn spawn_me_supervisors(
|
||||
task_scope: RuntimeTaskScope,
|
||||
pool: Arc<MePool>,
|
||||
rng: Arc<SecureRandom>,
|
||||
min_connections: usize,
|
||||
) {
|
||||
let health_pool = pool.clone();
|
||||
let health_rng = rng;
|
||||
task_scope.spawn(supervise_me_task("health_monitor", move || {
|
||||
let pool = health_pool.clone();
|
||||
let rng = health_rng.clone();
|
||||
async move {
|
||||
crate::transport::middle_proxy::me_health_monitor(pool, rng, min_connections).await;
|
||||
}
|
||||
}));
|
||||
|
||||
let drain_pool = pool.clone();
|
||||
task_scope.spawn(supervise_me_task("drain_timeout_enforcer", move || {
|
||||
let pool = drain_pool.clone();
|
||||
async move {
|
||||
crate::transport::middle_proxy::me_drain_timeout_enforcer(pool).await;
|
||||
}
|
||||
}));
|
||||
|
||||
task_scope.spawn(supervise_me_task("zombie_writer_watchdog", move || {
|
||||
let pool = pool.clone();
|
||||
async move {
|
||||
crate::transport::middle_proxy::me_zombie_writer_watchdog(pool).await;
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
pub(crate) async fn initialize_me_pool(
|
||||
use_middle_proxy: bool,
|
||||
config: &ProxyConfig,
|
||||
@@ -30,6 +82,7 @@ pub(crate) async fn initialize_me_pool(
|
||||
stats: Arc<Stats>,
|
||||
api_me_pool: Arc<RwLock<Option<Arc<MePool>>>>,
|
||||
me_ready_tx: watch::Sender<u64>,
|
||||
task_scope: RuntimeTaskScope,
|
||||
) -> Option<Arc<MePool>> {
|
||||
if !use_middle_proxy {
|
||||
return None;
|
||||
@@ -52,15 +105,8 @@ pub(crate) async fn initialize_me_pool(
|
||||
.as_ref()
|
||||
.map(|tag| hex::decode(tag).expect("general.ad_tag must be validated before startup"));
|
||||
|
||||
// =============================================================
|
||||
// CRITICAL: Download Telegram proxy-secret (NOT user secret!)
|
||||
//
|
||||
// C MTProxy uses TWO separate secrets:
|
||||
// -S flag = 16-byte user secret for client obfuscation
|
||||
// --aes-pwd = 32-512 byte binary file for ME RPC auth
|
||||
//
|
||||
// proxy-secret is from: https://core.telegram.org/getProxySecret
|
||||
// =============================================================
|
||||
// The Telegram proxy-secret authenticates ME RPC and is distinct from client secrets.
|
||||
// It corresponds to the C MTProxy --aes-pwd input and may be fetched from Telegram.
|
||||
let proxy_secret_path = config.general.proxy_secret_path.as_deref();
|
||||
let pool_size = config.general.middle_proxy_pool_size.max(1);
|
||||
let proxy_secret = loop {
|
||||
@@ -279,6 +325,7 @@ pub(crate) async fn initialize_me_pool(
|
||||
config.general.me_writer_pick_sample_size,
|
||||
config.general.me_socks_kdf_policy,
|
||||
config.general.me_writer_cmd_channel_capacity,
|
||||
config.general.me_writer_byte_budget_bytes,
|
||||
config.general.me_route_channel_capacity,
|
||||
config.general.me_route_backpressure_enabled,
|
||||
config.general.me_route_fairshare_enabled,
|
||||
@@ -318,143 +365,70 @@ pub(crate) async fn initialize_me_pool(
|
||||
let rng_bg = rng.clone();
|
||||
let startup_tracker_bg = startup_tracker.clone();
|
||||
let me_ready_tx_bg = me_ready_tx.clone();
|
||||
let task_scope_bg = task_scope.clone();
|
||||
let retry_limit = if me_init_retry_attempts == 0 {
|
||||
String::from("unlimited")
|
||||
} else {
|
||||
me_init_retry_attempts.to_string()
|
||||
};
|
||||
std::thread::spawn(move || {
|
||||
let runtime = match tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
{
|
||||
Ok(runtime) => runtime,
|
||||
Err(error) => {
|
||||
error!(error = %error, "Failed to build background runtime for ME initialization");
|
||||
return;
|
||||
}
|
||||
};
|
||||
runtime.block_on(async move {
|
||||
let mut init_attempt: u32 = 0;
|
||||
loop {
|
||||
init_attempt = init_attempt.saturating_add(1);
|
||||
startup_tracker_bg.set_me_init_attempt(init_attempt).await;
|
||||
match pool_bg.init(pool_size, &rng_bg).await {
|
||||
Ok(()) => {
|
||||
startup_tracker_bg.set_me_last_error(None).await;
|
||||
startup_tracker_bg
|
||||
.complete_component(
|
||||
COMPONENT_ME_POOL_INIT_STAGE1,
|
||||
Some("ME pool initialized".to_string()),
|
||||
)
|
||||
.await;
|
||||
startup_tracker_bg
|
||||
.set_me_status(StartupMeStatus::Ready, "ready")
|
||||
.await;
|
||||
me_ready_tx_bg.send_modify(|version| {
|
||||
*version = version.saturating_add(1);
|
||||
});
|
||||
info!(
|
||||
task_scope.spawn(async move {
|
||||
let mut init_attempt: u32 = 0;
|
||||
loop {
|
||||
init_attempt = init_attempt.saturating_add(1);
|
||||
startup_tracker_bg.set_me_init_attempt(init_attempt).await;
|
||||
match pool_bg.init(pool_size, &rng_bg).await {
|
||||
Ok(()) => {
|
||||
startup_tracker_bg.set_me_last_error(None).await;
|
||||
startup_tracker_bg
|
||||
.complete_component(
|
||||
COMPONENT_ME_POOL_INIT_STAGE1,
|
||||
Some("ME pool initialized".to_string()),
|
||||
)
|
||||
.await;
|
||||
startup_tracker_bg
|
||||
.set_me_status(StartupMeStatus::Ready, "ready")
|
||||
.await;
|
||||
me_ready_tx_bg.send_modify(|version| {
|
||||
*version = version.saturating_add(1);
|
||||
});
|
||||
info!(
|
||||
attempt = init_attempt,
|
||||
"Middle-End pool initialized successfully"
|
||||
);
|
||||
spawn_me_supervisors(
|
||||
task_scope_bg,
|
||||
pool_bg.clone(),
|
||||
rng_bg.clone(),
|
||||
pool_size,
|
||||
);
|
||||
break;
|
||||
}
|
||||
Err(e) => {
|
||||
startup_tracker_bg
|
||||
.set_me_last_error(Some(e.to_string()))
|
||||
.await;
|
||||
if init_attempt >= me_init_warn_after_attempts {
|
||||
warn!(
|
||||
error = %e,
|
||||
attempt = init_attempt,
|
||||
"Middle-End pool initialized successfully"
|
||||
retry_limit = %retry_limit,
|
||||
retry_in_secs = 2,
|
||||
"ME pool is not ready yet; retrying background initialization"
|
||||
);
|
||||
} else {
|
||||
info!(
|
||||
error = %e,
|
||||
attempt = init_attempt,
|
||||
retry_limit = %retry_limit,
|
||||
retry_in_secs = 2,
|
||||
"ME pool startup warmup: retrying background initialization"
|
||||
);
|
||||
|
||||
// ── Supervised background tasks ──────────────────
|
||||
// Each task runs inside a nested tokio::spawn so
|
||||
// that a panic is caught via JoinHandle and the
|
||||
// outer loop restarts the task automatically.
|
||||
let pool_health = pool_bg.clone();
|
||||
let rng_health = rng_bg.clone();
|
||||
let min_conns = pool_size;
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let p = pool_health.clone();
|
||||
let r = rng_health.clone();
|
||||
let res = tokio::spawn(async move {
|
||||
crate::transport::middle_proxy::me_health_monitor(
|
||||
p, r, min_conns,
|
||||
)
|
||||
.await;
|
||||
})
|
||||
.await;
|
||||
match res {
|
||||
Ok(()) => warn!("me_health_monitor exited unexpectedly, restarting"),
|
||||
Err(e) => {
|
||||
error!(error = %e, "me_health_monitor panicked, restarting in 1s");
|
||||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
let pool_drain_enforcer = pool_bg.clone();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let p = pool_drain_enforcer.clone();
|
||||
let res = tokio::spawn(async move {
|
||||
crate::transport::middle_proxy::me_drain_timeout_enforcer(p).await;
|
||||
})
|
||||
.await;
|
||||
match res {
|
||||
Ok(()) => warn!("me_drain_timeout_enforcer exited unexpectedly, restarting"),
|
||||
Err(e) => {
|
||||
error!(error = %e, "me_drain_timeout_enforcer panicked, restarting in 1s");
|
||||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
let pool_watchdog = pool_bg.clone();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let p = pool_watchdog.clone();
|
||||
let res = tokio::spawn(async move {
|
||||
crate::transport::middle_proxy::me_zombie_writer_watchdog(p).await;
|
||||
})
|
||||
.await;
|
||||
match res {
|
||||
Ok(()) => warn!("me_zombie_writer_watchdog exited unexpectedly, restarting"),
|
||||
Err(e) => {
|
||||
error!(error = %e, "me_zombie_writer_watchdog panicked, restarting in 1s");
|
||||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
// CRITICAL: keep the current-thread runtime
|
||||
// alive. Without this, block_on() returns,
|
||||
// the Runtime is dropped, and ALL spawned
|
||||
// background tasks (health monitor, drain
|
||||
// enforcer, zombie watchdog) are silently
|
||||
// cancelled — causing the draining-writer
|
||||
// leak that brought us here.
|
||||
std::future::pending::<()>().await;
|
||||
unreachable!();
|
||||
}
|
||||
Err(e) => {
|
||||
startup_tracker_bg.set_me_last_error(Some(e.to_string())).await;
|
||||
if init_attempt >= me_init_warn_after_attempts {
|
||||
warn!(
|
||||
error = %e,
|
||||
attempt = init_attempt,
|
||||
retry_limit = %retry_limit,
|
||||
retry_in_secs = 2,
|
||||
"ME pool is not ready yet; retrying background initialization"
|
||||
);
|
||||
} else {
|
||||
info!(
|
||||
error = %e,
|
||||
attempt = init_attempt,
|
||||
retry_limit = %retry_limit,
|
||||
retry_in_secs = 2,
|
||||
"ME pool startup warmup: retrying background initialization"
|
||||
);
|
||||
}
|
||||
pool_bg.reset_stun_state();
|
||||
tokio::time::sleep(Duration::from_secs(2)).await;
|
||||
}
|
||||
pool_bg.reset_stun_state();
|
||||
tokio::time::sleep(Duration::from_secs(2)).await;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
startup_tracker
|
||||
.set_me_status(StartupMeStatus::Initializing, "background_init")
|
||||
@@ -489,70 +463,12 @@ pub(crate) async fn initialize_me_pool(
|
||||
"Middle-End pool initialized successfully"
|
||||
);
|
||||
|
||||
// ── Supervised background tasks ──────────────────
|
||||
let pool_clone = pool.clone();
|
||||
let rng_clone = rng.clone();
|
||||
let min_conns = pool_size;
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let p = pool_clone.clone();
|
||||
let r = rng_clone.clone();
|
||||
let res = tokio::spawn(async move {
|
||||
crate::transport::middle_proxy::me_health_monitor(
|
||||
p, r, min_conns,
|
||||
)
|
||||
.await;
|
||||
})
|
||||
.await;
|
||||
match res {
|
||||
Ok(()) => warn!(
|
||||
"me_health_monitor exited unexpectedly, restarting"
|
||||
),
|
||||
Err(e) => {
|
||||
error!(error = %e, "me_health_monitor panicked, restarting in 1s");
|
||||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
let pool_drain_enforcer = pool.clone();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let p = pool_drain_enforcer.clone();
|
||||
let res = tokio::spawn(async move {
|
||||
crate::transport::middle_proxy::me_drain_timeout_enforcer(p).await;
|
||||
})
|
||||
.await;
|
||||
match res {
|
||||
Ok(()) => warn!(
|
||||
"me_drain_timeout_enforcer exited unexpectedly, restarting"
|
||||
),
|
||||
Err(e) => {
|
||||
error!(error = %e, "me_drain_timeout_enforcer panicked, restarting in 1s");
|
||||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
let pool_watchdog = pool.clone();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let p = pool_watchdog.clone();
|
||||
let res = tokio::spawn(async move {
|
||||
crate::transport::middle_proxy::me_zombie_writer_watchdog(p).await;
|
||||
})
|
||||
.await;
|
||||
match res {
|
||||
Ok(()) => warn!(
|
||||
"me_zombie_writer_watchdog exited unexpectedly, restarting"
|
||||
),
|
||||
Err(e) => {
|
||||
error!(error = %e, "me_zombie_writer_watchdog panicked, restarting in 1s");
|
||||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
spawn_me_supervisors(
|
||||
task_scope.clone(),
|
||||
pool.clone(),
|
||||
rng.clone(),
|
||||
pool_size,
|
||||
);
|
||||
|
||||
break Some(pool);
|
||||
}
|
||||
@@ -665,3 +581,38 @@ pub(crate) async fn initialize_me_pool(
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tokio::sync::Notify;
|
||||
|
||||
struct DropSignal(Arc<Notify>);
|
||||
|
||||
impl Drop for DropSignal {
|
||||
fn drop(&mut self) {
|
||||
self.0.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn scoped_supervisor_aborts_its_current_child() {
|
||||
let scope = RuntimeTaskScope::new();
|
||||
let dropped = Arc::new(Notify::new());
|
||||
let dropped_for_task = dropped.clone();
|
||||
scope.spawn(supervise_me_task("test", move || {
|
||||
let dropped = dropped_for_task.clone();
|
||||
async move {
|
||||
let _signal = DropSignal(dropped);
|
||||
std::future::pending::<()>().await;
|
||||
}
|
||||
}));
|
||||
tokio::task::yield_now().await;
|
||||
|
||||
scope.stop().await;
|
||||
|
||||
tokio::time::timeout(Duration::from_secs(1), dropped.notified())
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -13,19 +13,24 @@
|
||||
// - shutdown: graceful shutdown sequence and uptime logging.
|
||||
mod admission;
|
||||
mod connectivity;
|
||||
pub(crate) mod generation;
|
||||
mod helpers;
|
||||
mod listeners;
|
||||
mod me_startup;
|
||||
pub(crate) mod reload;
|
||||
mod reload_supervisor;
|
||||
pub(crate) mod runtime_build;
|
||||
mod runtime_tasks;
|
||||
mod shutdown;
|
||||
mod tls_bootstrap;
|
||||
|
||||
use arc_swap::ArcSwap;
|
||||
use std::net::{IpAddr, SocketAddr};
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
|
||||
use tokio::sync::{RwLock, Semaphore, watch};
|
||||
use tracing::{error, info, warn};
|
||||
use tracing_subscriber::{EnvFilter, fmt, prelude::*, reload};
|
||||
use tracing_subscriber::{EnvFilter, fmt, prelude::*, reload as tracing_reload};
|
||||
|
||||
use crate::api;
|
||||
use crate::config::{LogLevel, ProxyConfig};
|
||||
@@ -33,6 +38,9 @@ use crate::conntrack_control;
|
||||
use crate::crypto::SecureRandom;
|
||||
use crate::ip_tracker::UserIpTracker;
|
||||
use crate::network::probe::{decide_network_capabilities, log_probe_result, run_probe};
|
||||
use crate::proxy::direct_buffer_budget::{
|
||||
DirectBufferBudget, resolve_direct_buffer_hard_limit, run_direct_buffer_budget_controller,
|
||||
};
|
||||
use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController};
|
||||
use crate::proxy::shared_state::ProxySharedState;
|
||||
use crate::startup::{
|
||||
@@ -43,7 +51,7 @@ use crate::startup::{
|
||||
};
|
||||
use crate::stats::beobachten::BeobachtenStore;
|
||||
use crate::stats::telemetry::TelemetryPolicy;
|
||||
use crate::stats::{ReplayChecker, Stats};
|
||||
use crate::stats::{QuotaStore, ReplayChecker, Stats};
|
||||
use crate::stream::BufferPool;
|
||||
use crate::synlimit_control;
|
||||
use crate::transport::UpstreamManager;
|
||||
@@ -340,7 +348,7 @@ async fn run_telemt_core(
|
||||
}
|
||||
};
|
||||
let (filter_layer, filter_handle) =
|
||||
reload::Layer::new(EnvFilter::new(initial_filter_spec.clone()));
|
||||
tracing_reload::Layer::new(EnvFilter::new(initial_filter_spec.clone()));
|
||||
startup_tracker
|
||||
.start_component(
|
||||
COMPONENT_TRACING_INIT,
|
||||
@@ -384,6 +392,7 @@ async fn run_telemt_core(
|
||||
_logging_guard = Some(guard);
|
||||
}
|
||||
}
|
||||
let runtime_log_filter = runtime_tasks::RuntimeLogFilter::new(filter_handle);
|
||||
|
||||
startup_tracker
|
||||
.complete_component(
|
||||
@@ -430,21 +439,26 @@ async fn run_telemt_core(
|
||||
warn!("Using default tls_domain. Consider setting a custom domain.");
|
||||
}
|
||||
|
||||
let stats = Arc::new(Stats::new());
|
||||
let quota_store = Arc::new(QuotaStore::default());
|
||||
let stats = Arc::new(Stats::with_quota_store(quota_store.clone()));
|
||||
let runtime_task_scope = generation::RuntimeTaskScope::new();
|
||||
stats.apply_telemetry_policy(TelemetryPolicy::from_config(&config.general.telemetry));
|
||||
let quota_state_path = config.general.quota_state_path.clone();
|
||||
crate::quota_state::load_quota_state("a_state_path, stats.as_ref()).await;
|
||||
|
||||
let upstream_manager = Arc::new(UpstreamManager::new(
|
||||
config.upstreams.clone(),
|
||||
config.general.upstream_connect_retry_attempts,
|
||||
config.general.upstream_connect_retry_backoff_ms,
|
||||
config.general.upstream_connect_budget_ms,
|
||||
config.general.tg_connect,
|
||||
config.general.upstream_unhealthy_fail_threshold,
|
||||
config.general.upstream_connect_failfast_hard_errors,
|
||||
stats.clone(),
|
||||
));
|
||||
let upstream_manager = Arc::new(
|
||||
UpstreamManager::new(
|
||||
config.upstreams.clone(),
|
||||
config.general.upstream_connect_retry_attempts,
|
||||
config.general.upstream_connect_retry_backoff_ms,
|
||||
config.general.upstream_connect_budget_ms,
|
||||
config.general.tg_connect,
|
||||
config.general.upstream_unhealthy_fail_threshold,
|
||||
config.general.upstream_connect_failfast_hard_errors,
|
||||
stats.clone(),
|
||||
)
|
||||
.with_dns_overrides(&config.network.dns_overrides)?,
|
||||
);
|
||||
let ip_tracker = Arc::new(UserIpTracker::new());
|
||||
ip_tracker
|
||||
.load_limits(
|
||||
@@ -473,18 +487,31 @@ async fn run_telemt_core(
|
||||
config.network.dns_overrides.len()
|
||||
);
|
||||
}
|
||||
let shared_state = ProxySharedState::new();
|
||||
let direct_buffer_hard_limit =
|
||||
resolve_direct_buffer_hard_limit(config.general.direct_relay_buffer_budget_max_bytes).await;
|
||||
let direct_buffer_budget = DirectBufferBudget::new(direct_buffer_hard_limit);
|
||||
info!(
|
||||
hard_limit_bytes = direct_buffer_hard_limit,
|
||||
configured_override_bytes = config.general.direct_relay_buffer_budget_max_bytes,
|
||||
"Direct relay buffer budget initialized"
|
||||
);
|
||||
let shared_state =
|
||||
ProxySharedState::new_with_direct_buffer_budget(direct_buffer_budget.clone());
|
||||
shared_state.apply_user_enabled_config(&config.access.user_enabled);
|
||||
shared_state.traffic_limiter.apply_policy(
|
||||
config.access.user_rate_limits.clone(),
|
||||
config.access.cidr_rate_limits.clone(),
|
||||
);
|
||||
|
||||
let (api_config_tx, api_config_rx) = watch::channel(Arc::new(config.clone()));
|
||||
let (detected_ips_tx, detected_ips_rx) = watch::channel((None::<IpAddr>, None::<IpAddr>));
|
||||
let initial_direct_first = config.general.use_middle_proxy && config.general.me2dc_fallback;
|
||||
let initial_admission_open = !config.general.use_middle_proxy || initial_direct_first;
|
||||
let (admission_tx, admission_rx) = watch::channel(initial_admission_open);
|
||||
let (reload_control, reload_commands) = reload::ReloadControl::channel(1);
|
||||
let (active_runtime_tx, active_runtime_rx) =
|
||||
watch::channel(None::<Arc<ArcSwap<generation::RuntimeGeneration>>>);
|
||||
let (runtime_watch_tx, runtime_watch_rx) =
|
||||
watch::channel(None::<generation::RuntimeWatchState>);
|
||||
let initial_route_mode = if !config.general.use_middle_proxy || initial_direct_first {
|
||||
RelayRouteMode::Direct
|
||||
} else {
|
||||
@@ -518,12 +545,13 @@ async fn run_telemt_core(
|
||||
let upstream_manager_api = upstream_manager.clone();
|
||||
let route_runtime_api = route_runtime.clone();
|
||||
let proxy_shared_api = shared_state.clone();
|
||||
let config_rx_api = api_config_rx.clone();
|
||||
let admission_rx_api = admission_rx.clone();
|
||||
let config_path_api = config_path.clone();
|
||||
let quota_state_path_api = quota_state_path.clone();
|
||||
let startup_tracker_api = startup_tracker.clone();
|
||||
let detected_ips_rx_api = detected_ips_rx.clone();
|
||||
let reload_control_api = reload_control.clone();
|
||||
let active_runtime_rx_api = active_runtime_rx.clone();
|
||||
let runtime_watch_rx_api = runtime_watch_rx.clone();
|
||||
tokio::spawn(async move {
|
||||
api::serve(
|
||||
listen,
|
||||
@@ -533,13 +561,14 @@ async fn run_telemt_core(
|
||||
route_runtime_api,
|
||||
proxy_shared_api,
|
||||
upstream_manager_api,
|
||||
config_rx_api,
|
||||
admission_rx_api,
|
||||
config_path_api,
|
||||
quota_state_path_api,
|
||||
detected_ips_rx_api,
|
||||
process_started_at_epoch_secs,
|
||||
startup_tracker_api,
|
||||
reload_control_api,
|
||||
active_runtime_rx_api,
|
||||
runtime_watch_rx_api,
|
||||
)
|
||||
.await;
|
||||
});
|
||||
@@ -579,8 +608,10 @@ async fn run_telemt_core(
|
||||
&tls_domains,
|
||||
upstream_manager.clone(),
|
||||
&startup_tracker,
|
||||
runtime_task_scope.clone(),
|
||||
tls_bootstrap::TlsBootstrapPolicy::BestEffort,
|
||||
)
|
||||
.await;
|
||||
.await?;
|
||||
|
||||
startup_tracker
|
||||
.start_component(
|
||||
@@ -706,6 +737,7 @@ async fn run_telemt_core(
|
||||
stats.clone(),
|
||||
api_me_pool.clone(),
|
||||
me_ready_tx.clone(),
|
||||
runtime_task_scope.clone(),
|
||||
)
|
||||
.await
|
||||
};
|
||||
@@ -793,16 +825,22 @@ async fn run_telemt_core(
|
||||
rng.clone(),
|
||||
ip_tracker.clone(),
|
||||
beobachten.clone(),
|
||||
api_config_tx.clone(),
|
||||
me_pool.clone(),
|
||||
shared_state.clone(),
|
||||
me_ready_tx.clone(),
|
||||
runtime_task_scope.clone(),
|
||||
)
|
||||
.await;
|
||||
let config_rx = runtime_watches.config_rx;
|
||||
let log_level_rx = runtime_watches.log_level_rx;
|
||||
let detected_ip_v4 = runtime_watches.detected_ip_v4;
|
||||
let detected_ip_v6 = runtime_watches.detected_ip_v6;
|
||||
runtime_log_filter.start(
|
||||
has_rust_log,
|
||||
&effective_log_level,
|
||||
log_level_rx,
|
||||
runtime_task_scope.clone(),
|
||||
);
|
||||
|
||||
if direct_first_startup {
|
||||
let config_bg = config.clone();
|
||||
@@ -815,7 +853,8 @@ async fn run_telemt_core(
|
||||
let api_me_pool_bg = api_me_pool.clone();
|
||||
let me_ready_tx_bg = me_ready_tx.clone();
|
||||
let config_rx_bg = config_rx.clone();
|
||||
tokio::spawn(async move {
|
||||
let task_scope_bg = runtime_task_scope.clone();
|
||||
runtime_task_scope.spawn(async move {
|
||||
let mut bootstrap_attempt: u32 = 0;
|
||||
loop {
|
||||
bootstrap_attempt = bootstrap_attempt.saturating_add(1);
|
||||
@@ -830,6 +869,7 @@ async fn run_telemt_core(
|
||||
stats_bg.clone(),
|
||||
api_me_pool_bg.clone(),
|
||||
me_ready_tx_bg.clone(),
|
||||
task_scope_bg.clone(),
|
||||
)
|
||||
.await;
|
||||
if let Some(pool) = pool {
|
||||
@@ -839,6 +879,7 @@ async fn run_telemt_core(
|
||||
pool,
|
||||
rng_bg,
|
||||
me_ready_tx_bg,
|
||||
task_scope_bg,
|
||||
);
|
||||
break;
|
||||
}
|
||||
@@ -852,7 +893,7 @@ async fn run_telemt_core(
|
||||
let startup_tracker_ready = startup_tracker.clone();
|
||||
let api_me_pool_ready = api_me_pool.clone();
|
||||
let mut me_ready_rx_transport = me_ready_tx.subscribe();
|
||||
tokio::spawn(async move {
|
||||
runtime_task_scope.spawn(async move {
|
||||
if me_ready_rx_transport.changed().await.is_ok() {
|
||||
if let Some(pool) = api_me_pool_ready.read().await.as_ref() {
|
||||
pool.set_runtime_ready(true);
|
||||
@@ -874,13 +915,54 @@ async fn run_telemt_core(
|
||||
&admission_tx,
|
||||
config_rx.clone(),
|
||||
me_ready_rx,
|
||||
runtime_task_scope.clone(),
|
||||
)
|
||||
.await;
|
||||
let _admission_tx_hold = admission_tx;
|
||||
conntrack_control::spawn_conntrack_controller(
|
||||
let conntrack_scope = runtime_task_scope.clone();
|
||||
runtime_task_scope.spawn(conntrack_control::run_conntrack_controller(
|
||||
config_rx.clone(),
|
||||
stats.clone(),
|
||||
shared_state.clone(),
|
||||
conntrack_scope.cancellation_token(),
|
||||
));
|
||||
runtime_task_scope.spawn(run_direct_buffer_budget_controller(
|
||||
direct_buffer_budget,
|
||||
buffer_pool.clone(),
|
||||
stats.clone(),
|
||||
shared_state.clone(),
|
||||
config.server.max_connections,
|
||||
));
|
||||
|
||||
let runtime_generation = generation::RuntimeGeneration::new(
|
||||
1,
|
||||
config_rx.clone(),
|
||||
admission_rx.clone(),
|
||||
stats.clone(),
|
||||
upstream_manager.clone(),
|
||||
replay_checker.clone(),
|
||||
buffer_pool.clone(),
|
||||
rng.clone(),
|
||||
me_pool.clone(),
|
||||
api_me_pool.clone(),
|
||||
route_runtime.clone(),
|
||||
tls_cache.clone(),
|
||||
ip_tracker.clone(),
|
||||
beobachten.clone(),
|
||||
shared_state.clone(),
|
||||
max_connections.clone(),
|
||||
runtime_task_scope.clone(),
|
||||
);
|
||||
let active_runtime = Arc::new(ArcSwap::from(runtime_generation));
|
||||
reload_supervisor::ReloadSupervisor::spawn(
|
||||
active_runtime.clone(),
|
||||
reload_control,
|
||||
reload_commands,
|
||||
config_path.clone(),
|
||||
quota_store,
|
||||
detected_ips_tx,
|
||||
runtime_log_filter,
|
||||
runtime_watch_tx.clone(),
|
||||
);
|
||||
|
||||
let bound = listeners::bind_listeners(
|
||||
@@ -890,25 +972,15 @@ async fn run_telemt_core(
|
||||
detected_ip_v4,
|
||||
detected_ip_v6,
|
||||
&startup_tracker,
|
||||
config_rx.clone(),
|
||||
admission_rx.clone(),
|
||||
stats.clone(),
|
||||
upstream_manager.clone(),
|
||||
replay_checker.clone(),
|
||||
buffer_pool.clone(),
|
||||
rng.clone(),
|
||||
me_pool.clone(),
|
||||
api_me_pool.clone(),
|
||||
route_runtime.clone(),
|
||||
tls_cache.clone(),
|
||||
ip_tracker.clone(),
|
||||
beobachten.clone(),
|
||||
shared_state.clone(),
|
||||
max_connections.clone(),
|
||||
)
|
||||
.await?;
|
||||
let listeners = bound.listeners;
|
||||
let has_unix_listener = bound.has_unix_listener;
|
||||
#[cfg(unix)]
|
||||
let unix_listener = bound.unix_listener;
|
||||
#[cfg(unix)]
|
||||
let has_unix_listener = unix_listener.is_some();
|
||||
#[cfg(not(unix))]
|
||||
let has_unix_listener = false;
|
||||
|
||||
if listeners.is_empty() && !has_unix_listener {
|
||||
error!("No listeners. Exiting.");
|
||||
@@ -918,54 +990,29 @@ async fn run_telemt_core(
|
||||
// On Unix, caller supplies privilege drop after bind (may require root for port < 1024).
|
||||
drop_after_bind();
|
||||
|
||||
synlimit_control::reconcile_synlimit_rules(&config).await;
|
||||
synlimit_control::spawn_synlimit_controller(config_rx.clone());
|
||||
let synlimit_controller = synlimit_control::spawn_synlimit_controller(runtime_watch_rx);
|
||||
|
||||
runtime_tasks::apply_runtime_log_filter(
|
||||
has_rust_log,
|
||||
&effective_log_level,
|
||||
filter_handle,
|
||||
log_level_rx,
|
||||
)
|
||||
.await;
|
||||
|
||||
runtime_tasks::spawn_metrics_if_configured(
|
||||
&config,
|
||||
&startup_tracker,
|
||||
stats.clone(),
|
||||
beobachten.clone(),
|
||||
shared_state.clone(),
|
||||
ip_tracker.clone(),
|
||||
tls_cache.clone(),
|
||||
config_rx.clone(),
|
||||
)
|
||||
.await;
|
||||
runtime_tasks::spawn_metrics_if_configured(&config, &startup_tracker, active_runtime.clone())
|
||||
.await;
|
||||
|
||||
runtime_watch_tx.send_replace(Some(active_runtime.load_full().watch_state()));
|
||||
active_runtime_tx.send_replace(Some(active_runtime.clone()));
|
||||
runtime_tasks::mark_runtime_ready(&startup_tracker).await;
|
||||
|
||||
// Spawn signal handlers for SIGUSR1/SIGUSR2 (non-shutdown signals)
|
||||
shutdown::spawn_signal_handlers(stats.clone(), process_started_at);
|
||||
shutdown::spawn_signal_handlers(active_runtime.clone(), process_started_at);
|
||||
|
||||
listeners::spawn_tcp_accept_loops(
|
||||
listeners,
|
||||
config_rx.clone(),
|
||||
admission_rx.clone(),
|
||||
stats.clone(),
|
||||
upstream_manager.clone(),
|
||||
replay_checker.clone(),
|
||||
buffer_pool.clone(),
|
||||
rng.clone(),
|
||||
me_pool.clone(),
|
||||
api_me_pool.clone(),
|
||||
route_runtime.clone(),
|
||||
tls_cache.clone(),
|
||||
ip_tracker.clone(),
|
||||
beobachten.clone(),
|
||||
shared_state,
|
||||
max_connections.clone(),
|
||||
);
|
||||
listeners::spawn_tcp_accept_loops(listeners, active_runtime.clone());
|
||||
#[cfg(unix)]
|
||||
listeners::spawn_unix_accept_loop(unix_listener, active_runtime.clone());
|
||||
|
||||
shutdown::wait_for_shutdown(process_started_at, me_pool, stats, quota_state_path).await;
|
||||
shutdown::wait_for_shutdown(
|
||||
process_started_at,
|
||||
active_runtime,
|
||||
quota_state_path,
|
||||
synlimit_controller,
|
||||
)
|
||||
.await;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
493
src/maestro/reload.rs
Normal file
493
src/maestro/reload.rs
Normal file
@@ -0,0 +1,493 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::sync::{Mutex, mpsc};
|
||||
|
||||
use crate::config::ProxyConfig;
|
||||
|
||||
const RELOAD_HISTORY_CAPACITY: usize = 32;
|
||||
const RELOAD_COMMAND_CAPACITY: usize = 1;
|
||||
const MAX_DRAIN_TIMEOUT_SECS: u64 = 3_600;
|
||||
|
||||
/// Session handling policy for an in-process runtime reload.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub(crate) enum ReloadMode {
|
||||
#[default]
|
||||
Instant,
|
||||
Drain,
|
||||
}
|
||||
|
||||
/// Failure policy applied during the activation barrier.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub(crate) enum ReloadFailurePolicy {
|
||||
#[default]
|
||||
KeepNew,
|
||||
Rollback,
|
||||
}
|
||||
|
||||
/// Request body accepted by the maestro reload endpoint.
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize, Serialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub(crate) struct ReloadRequest {
|
||||
#[serde(default)]
|
||||
pub(crate) mode: ReloadMode,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) timeout_secs: Option<u64>,
|
||||
#[serde(default)]
|
||||
pub(crate) failure_policy: ReloadFailurePolicy,
|
||||
}
|
||||
|
||||
impl ReloadRequest {
|
||||
pub(crate) fn validate(&self) -> Result<(), &'static str> {
|
||||
match (self.mode, self.timeout_secs) {
|
||||
(ReloadMode::Instant, None) => Ok(()),
|
||||
(ReloadMode::Instant, Some(_)) => Err("timeout_secs is only valid when mode is drain"),
|
||||
(ReloadMode::Drain, Some(1..=MAX_DRAIN_TIMEOUT_SECS)) => Ok(()),
|
||||
(ReloadMode::Drain, Some(_)) => Err("timeout_secs must be within 1..=3600"),
|
||||
(ReloadMode::Drain, None) => Err("timeout_secs is required when mode is drain"),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn from_query(query: Option<&str>) -> Result<Option<Self>, String> {
|
||||
let Some(query) = query.filter(|query| !query.is_empty()) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let mut mode = None;
|
||||
let mut timeout_secs = None;
|
||||
let mut failure_policy = None;
|
||||
for (key, value) in url::form_urlencoded::parse(query.as_bytes()) {
|
||||
match key.as_ref() {
|
||||
"reload" if mode.is_none() => {
|
||||
mode = Some(match value.as_ref() {
|
||||
"instant" => ReloadMode::Instant,
|
||||
"drain" => ReloadMode::Drain,
|
||||
_ => return Err("reload must be instant or drain".to_string()),
|
||||
});
|
||||
}
|
||||
"timeout_secs" if timeout_secs.is_none() => {
|
||||
timeout_secs = Some(
|
||||
value
|
||||
.parse::<u64>()
|
||||
.map_err(|_| "timeout_secs must be an integer".to_string())?,
|
||||
);
|
||||
}
|
||||
"failure_policy" if failure_policy.is_none() => {
|
||||
failure_policy = Some(match value.as_ref() {
|
||||
"keep_new" => ReloadFailurePolicy::KeepNew,
|
||||
"rollback" => ReloadFailurePolicy::Rollback,
|
||||
_ => {
|
||||
return Err("failure_policy must be keep_new or rollback".to_string());
|
||||
}
|
||||
});
|
||||
}
|
||||
"reload" | "timeout_secs" | "failure_policy" => {
|
||||
return Err(format!("duplicate query parameter: {}", key));
|
||||
}
|
||||
_ => return Err(format!("unknown query parameter: {}", key)),
|
||||
}
|
||||
}
|
||||
let mode = mode.ok_or_else(|| "reload query parameter is required".to_string())?;
|
||||
let request = Self {
|
||||
mode,
|
||||
timeout_secs,
|
||||
failure_policy: failure_policy.unwrap_or_default(),
|
||||
};
|
||||
request.validate().map_err(str::to_string)?;
|
||||
Ok(Some(request))
|
||||
}
|
||||
}
|
||||
|
||||
/// Observable phase of one reload operation.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub(crate) enum ReloadPhase {
|
||||
Accepted,
|
||||
Preparing,
|
||||
Activating,
|
||||
Draining,
|
||||
Succeeded,
|
||||
RolledBack,
|
||||
Failed,
|
||||
}
|
||||
|
||||
impl ReloadPhase {
|
||||
fn is_terminal(self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
ReloadPhase::Succeeded | ReloadPhase::RolledBack | ReloadPhase::Failed
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Bounded public status for one reload operation.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub(crate) struct ReloadStatus {
|
||||
pub(crate) reload_id: u64,
|
||||
pub(crate) target_generation: u64,
|
||||
pub(crate) config_revision: String,
|
||||
pub(crate) state: ReloadPhase,
|
||||
pub(crate) mode: ReloadMode,
|
||||
pub(crate) failure_policy: ReloadFailurePolicy,
|
||||
pub(crate) requested_at_epoch_secs: u64,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) started_at_epoch_secs: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) finished_at_epoch_secs: Option<u64>,
|
||||
#[serde(
|
||||
rename = "deferred_process_fields",
|
||||
default,
|
||||
skip_serializing_if = "Vec::is_empty"
|
||||
)]
|
||||
pub(crate) deferred_fields: Vec<String>,
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub(crate) warnings: Vec<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) error: Option<String>,
|
||||
}
|
||||
|
||||
/// Accepted operation metadata returned before asynchronous preparation starts.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
|
||||
pub(crate) struct ReloadAccepted {
|
||||
pub(crate) reload_id: u64,
|
||||
pub(crate) target_generation: u64,
|
||||
pub(crate) config_revision: String,
|
||||
pub(crate) state: ReloadPhase,
|
||||
pub(crate) mode: ReloadMode,
|
||||
pub(crate) failure_policy: ReloadFailurePolicy,
|
||||
}
|
||||
|
||||
pub(crate) struct ReloadCommand {
|
||||
pub(crate) reload_id: u64,
|
||||
pub(crate) target_generation: u64,
|
||||
pub(crate) config: Arc<ProxyConfig>,
|
||||
pub(crate) config_revision: String,
|
||||
pub(crate) request: ReloadRequest,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum ReloadSubmitError {
|
||||
InProgress(u64),
|
||||
MaestroUnavailable,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct ReloadControl {
|
||||
command_tx: mpsc::Sender<ReloadCommand>,
|
||||
status_store: Arc<ReloadStatusStore>,
|
||||
active_generation: Arc<AtomicU64>,
|
||||
}
|
||||
|
||||
pub(crate) struct ReloadCommandReceiver {
|
||||
command_rx: mpsc::Receiver<ReloadCommand>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ReloadStatusState {
|
||||
next_reload_id: u64,
|
||||
active_reload_id: Option<u64>,
|
||||
statuses: VecDeque<ReloadStatus>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ReloadStatusStore {
|
||||
state: Mutex<ReloadStatusState>,
|
||||
}
|
||||
|
||||
impl ReloadControl {
|
||||
pub(crate) fn channel(initial_generation: u64) -> (Self, ReloadCommandReceiver) {
|
||||
let (command_tx, command_rx) = mpsc::channel(RELOAD_COMMAND_CAPACITY);
|
||||
(
|
||||
Self {
|
||||
command_tx,
|
||||
status_store: Arc::new(ReloadStatusStore::default()),
|
||||
active_generation: Arc::new(AtomicU64::new(initial_generation)),
|
||||
},
|
||||
ReloadCommandReceiver { command_rx },
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) async fn submit(
|
||||
&self,
|
||||
config: Arc<ProxyConfig>,
|
||||
config_revision: String,
|
||||
request: ReloadRequest,
|
||||
) -> Result<ReloadAccepted, ReloadSubmitError> {
|
||||
let target_generation = self
|
||||
.active_generation
|
||||
.load(Ordering::Acquire)
|
||||
.saturating_add(1);
|
||||
let status = self
|
||||
.status_store
|
||||
.reserve(target_generation, config_revision, request.clone())
|
||||
.await?;
|
||||
let command = ReloadCommand {
|
||||
reload_id: status.reload_id,
|
||||
target_generation,
|
||||
config,
|
||||
config_revision: status.config_revision.clone(),
|
||||
request,
|
||||
};
|
||||
if self.command_tx.send(command).await.is_err() {
|
||||
self.status_store
|
||||
.finish(
|
||||
status.reload_id,
|
||||
ReloadPhase::Failed,
|
||||
Some("maestro command channel is closed".to_string()),
|
||||
)
|
||||
.await;
|
||||
return Err(ReloadSubmitError::MaestroUnavailable);
|
||||
}
|
||||
Ok(ReloadAccepted {
|
||||
reload_id: status.reload_id,
|
||||
target_generation,
|
||||
config_revision: status.config_revision,
|
||||
state: ReloadPhase::Accepted,
|
||||
mode: status.mode,
|
||||
failure_policy: status.failure_policy,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn status(&self, reload_id: u64) -> Option<ReloadStatus> {
|
||||
self.status_store.get(reload_id).await
|
||||
}
|
||||
|
||||
pub(crate) async fn in_progress(&self) -> Option<u64> {
|
||||
self.status_store.state.lock().await.active_reload_id
|
||||
}
|
||||
|
||||
pub(crate) async fn mark_phase(&self, reload_id: u64, phase: ReloadPhase) {
|
||||
self.status_store.mark_phase(reload_id, phase).await;
|
||||
}
|
||||
|
||||
pub(crate) async fn set_deferred_fields(&self, reload_id: u64, fields: Vec<String>) {
|
||||
self.status_store
|
||||
.update(reload_id, |status| status.deferred_fields = fields)
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(crate) async fn succeed(&self, reload_id: u64, generation: u64) {
|
||||
self.active_generation.store(generation, Ordering::Release);
|
||||
self.status_store
|
||||
.finish(reload_id, ReloadPhase::Succeeded, None)
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(crate) async fn fail(&self, reload_id: u64, error: impl Into<String>) {
|
||||
self.status_store
|
||||
.finish(reload_id, ReloadPhase::Failed, Some(error.into()))
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(crate) async fn rolled_back(&self, reload_id: u64, error: impl Into<String>) {
|
||||
self.status_store
|
||||
.finish(reload_id, ReloadPhase::RolledBack, Some(error.into()))
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(crate) async fn add_warning(&self, reload_id: u64, warning: impl Into<String>) {
|
||||
let warning = warning.into();
|
||||
self.status_store
|
||||
.update(reload_id, |status| status.warnings.push(warning))
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
impl ReloadCommandReceiver {
|
||||
pub(crate) async fn recv(&mut self) -> Option<ReloadCommand> {
|
||||
self.command_rx.recv().await
|
||||
}
|
||||
}
|
||||
|
||||
impl ReloadStatusStore {
|
||||
async fn reserve(
|
||||
&self,
|
||||
target_generation: u64,
|
||||
config_revision: String,
|
||||
request: ReloadRequest,
|
||||
) -> Result<ReloadStatus, ReloadSubmitError> {
|
||||
let mut state = self.state.lock().await;
|
||||
if let Some(reload_id) = state.active_reload_id {
|
||||
return Err(ReloadSubmitError::InProgress(reload_id));
|
||||
}
|
||||
state.next_reload_id = state.next_reload_id.saturating_add(1).max(1);
|
||||
let reload_id = state.next_reload_id;
|
||||
let status = ReloadStatus {
|
||||
reload_id,
|
||||
target_generation,
|
||||
config_revision,
|
||||
state: ReloadPhase::Accepted,
|
||||
mode: request.mode,
|
||||
failure_policy: request.failure_policy,
|
||||
requested_at_epoch_secs: now_epoch_secs(),
|
||||
started_at_epoch_secs: None,
|
||||
finished_at_epoch_secs: None,
|
||||
deferred_fields: Vec::new(),
|
||||
warnings: Vec::new(),
|
||||
error: None,
|
||||
};
|
||||
state.active_reload_id = Some(reload_id);
|
||||
state.statuses.push_back(status.clone());
|
||||
while state.statuses.len() > RELOAD_HISTORY_CAPACITY {
|
||||
state.statuses.pop_front();
|
||||
}
|
||||
Ok(status)
|
||||
}
|
||||
|
||||
async fn get(&self, reload_id: u64) -> Option<ReloadStatus> {
|
||||
self.state
|
||||
.lock()
|
||||
.await
|
||||
.statuses
|
||||
.iter()
|
||||
.find(|status| status.reload_id == reload_id)
|
||||
.cloned()
|
||||
}
|
||||
|
||||
async fn mark_phase(&self, reload_id: u64, phase: ReloadPhase) {
|
||||
self.update(reload_id, |status| {
|
||||
status.state = phase;
|
||||
if status.started_at_epoch_secs.is_none() && phase != ReloadPhase::Accepted {
|
||||
status.started_at_epoch_secs = Some(now_epoch_secs());
|
||||
}
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn finish(&self, reload_id: u64, phase: ReloadPhase, error: Option<String>) {
|
||||
debug_assert!(phase.is_terminal());
|
||||
let mut state = self.state.lock().await;
|
||||
if let Some(status) = state
|
||||
.statuses
|
||||
.iter_mut()
|
||||
.find(|status| status.reload_id == reload_id)
|
||||
{
|
||||
status.state = phase;
|
||||
status.error = error;
|
||||
status.finished_at_epoch_secs = Some(now_epoch_secs());
|
||||
}
|
||||
if state.active_reload_id == Some(reload_id) {
|
||||
state.active_reload_id = None;
|
||||
}
|
||||
}
|
||||
|
||||
async fn update(&self, reload_id: u64, update: impl FnOnce(&mut ReloadStatus)) {
|
||||
let mut state = self.state.lock().await;
|
||||
if let Some(status) = state
|
||||
.statuses
|
||||
.iter_mut()
|
||||
.find(|status| status.reload_id == reload_id)
|
||||
{
|
||||
update(status);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn now_epoch_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn request_defaults_to_instant_keep_new() {
|
||||
let request: ReloadRequest = serde_json::from_str("{}").unwrap();
|
||||
assert_eq!(request, ReloadRequest::default());
|
||||
assert_eq!(request.validate(), Ok(()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drain_requires_bounded_timeout() {
|
||||
let missing = ReloadRequest {
|
||||
mode: ReloadMode::Drain,
|
||||
..ReloadRequest::default()
|
||||
};
|
||||
assert!(missing.validate().is_err());
|
||||
let valid = ReloadRequest {
|
||||
mode: ReloadMode::Drain,
|
||||
timeout_secs: Some(30),
|
||||
..ReloadRequest::default()
|
||||
};
|
||||
assert_eq!(valid.validate(), Ok(()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn patch_query_parses_reload_policy() {
|
||||
let request =
|
||||
ReloadRequest::from_query(Some("reload=drain&timeout_secs=30&failure_policy=rollback"))
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(request.mode, ReloadMode::Drain);
|
||||
assert_eq!(request.timeout_secs, Some(30));
|
||||
assert_eq!(request.failure_policy, ReloadFailurePolicy::Rollback);
|
||||
assert!(ReloadRequest::from_query(Some("timeout_secs=30")).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn status_uses_documented_deferred_process_fields_key() {
|
||||
let status = ReloadStatus {
|
||||
reload_id: 1,
|
||||
target_generation: 2,
|
||||
config_revision: "revision".to_string(),
|
||||
state: ReloadPhase::Succeeded,
|
||||
mode: ReloadMode::Instant,
|
||||
failure_policy: ReloadFailurePolicy::KeepNew,
|
||||
requested_at_epoch_secs: 10,
|
||||
started_at_epoch_secs: Some(11),
|
||||
finished_at_epoch_secs: Some(12),
|
||||
deferred_fields: vec!["server.listeners".to_string()],
|
||||
warnings: Vec::new(),
|
||||
error: None,
|
||||
};
|
||||
let value = serde_json::to_value(status).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
value["deferred_process_fields"],
|
||||
serde_json::json!(["server.listeners"])
|
||||
);
|
||||
assert!(value.get("deferred_fields").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn coordinator_rejects_concurrent_reload_and_releases_terminal_slot() {
|
||||
let (control, mut receiver) = ReloadControl::channel(1);
|
||||
let first = control
|
||||
.submit(
|
||||
Arc::new(ProxyConfig::default()),
|
||||
"rev-1".to_string(),
|
||||
ReloadRequest::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let _command = receiver.recv().await.unwrap();
|
||||
let second = control
|
||||
.submit(
|
||||
Arc::new(ProxyConfig::default()),
|
||||
"rev-2".to_string(),
|
||||
ReloadRequest::default(),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(second, Err(ReloadSubmitError::InProgress(first.reload_id)));
|
||||
control
|
||||
.succeed(first.reload_id, first.target_generation)
|
||||
.await;
|
||||
let third = control
|
||||
.submit(
|
||||
Arc::new(ProxyConfig::default()),
|
||||
"rev-3".to_string(),
|
||||
ReloadRequest::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(third.reload_id, first.reload_id + 1);
|
||||
}
|
||||
}
|
||||
256
src/maestro/reload_supervisor.rs
Normal file
256
src/maestro/reload_supervisor.rs
Normal file
@@ -0,0 +1,256 @@
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use arc_swap::ArcSwap;
|
||||
use tokio::sync::watch;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::stats::QuotaStore;
|
||||
|
||||
use super::generation::{RuntimeGeneration, RuntimeWatchState};
|
||||
use super::reload::{
|
||||
ReloadCommand, ReloadCommandReceiver, ReloadControl, ReloadFailurePolicy, ReloadMode,
|
||||
ReloadPhase,
|
||||
};
|
||||
use super::runtime_build::{deferred_process_fields, prepare_runtime};
|
||||
use super::runtime_tasks::RuntimeLogFilter;
|
||||
|
||||
pub(crate) struct ReloadSupervisor {
|
||||
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
||||
control: ReloadControl,
|
||||
commands: ReloadCommandReceiver,
|
||||
config_path: PathBuf,
|
||||
quota_store: Arc<QuotaStore>,
|
||||
detected_ips_tx: watch::Sender<(Option<std::net::IpAddr>, Option<std::net::IpAddr>)>,
|
||||
runtime_log_filter: RuntimeLogFilter,
|
||||
runtime_watch_tx: watch::Sender<Option<RuntimeWatchState>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
enum RevisionGateAction {
|
||||
Proceed,
|
||||
Warn(String),
|
||||
Rollback(String),
|
||||
}
|
||||
|
||||
fn revision_gate_action(
|
||||
accepted_revision: &str,
|
||||
current_revision: Result<String, String>,
|
||||
failure_policy: ReloadFailurePolicy,
|
||||
) -> RevisionGateAction {
|
||||
let warning = match current_revision {
|
||||
Ok(current) if current == accepted_revision => return RevisionGateAction::Proceed,
|
||||
Ok(current) => format!(
|
||||
"config revision changed during preparation: accepted={} current={}",
|
||||
accepted_revision, current
|
||||
),
|
||||
Err(error) => format!("config revision verification failed: {}", error),
|
||||
};
|
||||
match failure_policy {
|
||||
ReloadFailurePolicy::KeepNew => RevisionGateAction::Warn(warning),
|
||||
ReloadFailurePolicy::Rollback => RevisionGateAction::Rollback(warning),
|
||||
}
|
||||
}
|
||||
|
||||
async fn stop_background_and_middle_end(generation: &RuntimeGeneration) -> bool {
|
||||
generation.stop_background_tasks().await;
|
||||
let Some(pool) = generation.current_me_pool().await else {
|
||||
return false;
|
||||
};
|
||||
tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all())
|
||||
.await
|
||||
.is_err()
|
||||
}
|
||||
|
||||
async fn cleanup_candidate(generation: &RuntimeGeneration) -> bool {
|
||||
generation.stop_sessions().await;
|
||||
stop_background_and_middle_end(generation).await
|
||||
}
|
||||
|
||||
impl ReloadSupervisor {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) fn spawn(
|
||||
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
||||
control: ReloadControl,
|
||||
commands: ReloadCommandReceiver,
|
||||
config_path: PathBuf,
|
||||
quota_store: Arc<QuotaStore>,
|
||||
detected_ips_tx: watch::Sender<(Option<std::net::IpAddr>, Option<std::net::IpAddr>)>,
|
||||
runtime_log_filter: RuntimeLogFilter,
|
||||
runtime_watch_tx: watch::Sender<Option<RuntimeWatchState>>,
|
||||
) {
|
||||
let supervisor = Self {
|
||||
active_runtime,
|
||||
control,
|
||||
commands,
|
||||
config_path,
|
||||
quota_store,
|
||||
detected_ips_tx,
|
||||
runtime_log_filter,
|
||||
runtime_watch_tx,
|
||||
};
|
||||
tokio::spawn(supervisor.run());
|
||||
}
|
||||
|
||||
async fn run(mut self) {
|
||||
while let Some(command) = self.commands.recv().await {
|
||||
self.reload(command).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn reload(&self, command: ReloadCommand) {
|
||||
self.control
|
||||
.mark_phase(command.reload_id, ReloadPhase::Preparing)
|
||||
.await;
|
||||
let old_runtime = self.active_runtime.load_full();
|
||||
let deferred = deferred_process_fields(&old_runtime.config(), &command.config);
|
||||
self.control
|
||||
.set_deferred_fields(command.reload_id, deferred)
|
||||
.await;
|
||||
|
||||
let prepared = match prepare_runtime(
|
||||
command.target_generation,
|
||||
command.config.as_ref().clone(),
|
||||
&self.config_path,
|
||||
self.quota_store.clone(),
|
||||
self.runtime_log_filter.clone(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(prepared) => prepared,
|
||||
Err(error) => {
|
||||
self.control.fail(command.reload_id, error).await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let revision_action = revision_gate_action(
|
||||
&command.config_revision,
|
||||
crate::api::config_store::current_revision_for_maestro(&self.config_path).await,
|
||||
command.request.failure_policy,
|
||||
);
|
||||
match revision_action {
|
||||
RevisionGateAction::Proceed => {}
|
||||
RevisionGateAction::Warn(warning) => {
|
||||
self.control.add_warning(command.reload_id, warning).await;
|
||||
}
|
||||
RevisionGateAction::Rollback(warning) => {
|
||||
let _ = cleanup_candidate(&prepared.generation).await;
|
||||
self.runtime_log_filter
|
||||
.apply_reload(&old_runtime.config().general.log_level);
|
||||
self.control.rolled_back(command.reload_id, warning).await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
self.control
|
||||
.mark_phase(command.reload_id, ReloadPhase::Activating)
|
||||
.await;
|
||||
let new_runtime = prepared.generation;
|
||||
old_runtime.stop_accepting_sessions();
|
||||
if let Err(error) = crate::network::dns_overrides::install_entries(
|
||||
&new_runtime.config().network.dns_overrides,
|
||||
) {
|
||||
let message = format!("runtime DNS activation failed: {}", error);
|
||||
if command.request.failure_policy == ReloadFailurePolicy::Rollback {
|
||||
old_runtime.resume_accepting_sessions();
|
||||
let _ = cleanup_candidate(&new_runtime).await;
|
||||
self.runtime_log_filter
|
||||
.apply_reload(&old_runtime.config().general.log_level);
|
||||
self.control.rolled_back(command.reload_id, message).await;
|
||||
return;
|
||||
}
|
||||
self.control.add_warning(command.reload_id, message).await;
|
||||
}
|
||||
let replaced = self.active_runtime.swap(new_runtime.clone());
|
||||
self.detected_ips_tx.send_replace(prepared.detected_ips);
|
||||
self.runtime_log_filter
|
||||
.apply_reload(&new_runtime.config().general.log_level);
|
||||
self.runtime_watch_tx
|
||||
.send_replace(Some(new_runtime.watch_state()));
|
||||
|
||||
info!(
|
||||
reload_id = command.reload_id,
|
||||
old_generation = replaced.id,
|
||||
new_generation = new_runtime.id,
|
||||
config_revision = %command.config_revision,
|
||||
"Runtime generation activated"
|
||||
);
|
||||
|
||||
match command.request.mode {
|
||||
ReloadMode::Instant => {
|
||||
replaced.stop_sessions().await;
|
||||
}
|
||||
ReloadMode::Drain => {
|
||||
self.control
|
||||
.mark_phase(command.reload_id, ReloadPhase::Draining)
|
||||
.await;
|
||||
let timeout = Duration::from_secs(
|
||||
command
|
||||
.request
|
||||
.timeout_secs
|
||||
.expect("validated drain request must carry timeout_secs"),
|
||||
);
|
||||
if !replaced.drain_sessions(timeout).await {
|
||||
let warning = format!(
|
||||
"generation {} exceeded drain timeout; remaining sessions were cancelled",
|
||||
replaced.id
|
||||
);
|
||||
warn!(reload_id = command.reload_id, warning = %warning);
|
||||
self.control.add_warning(command.reload_id, warning).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if stop_background_and_middle_end(&replaced).await {
|
||||
let warning = format!(
|
||||
"generation {} Middle-End close broadcast timed out",
|
||||
replaced.id
|
||||
);
|
||||
warn!(reload_id = command.reload_id, warning = %warning);
|
||||
self.control.add_warning(command.reload_id, warning).await;
|
||||
}
|
||||
self.control
|
||||
.succeed(command.reload_id, new_runtime.id)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn revision_gate_proceeds_only_on_verified_match() {
|
||||
assert_eq!(
|
||||
revision_gate_action(
|
||||
"accepted",
|
||||
Ok("accepted".to_string()),
|
||||
ReloadFailurePolicy::Rollback,
|
||||
),
|
||||
RevisionGateAction::Proceed
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn revision_gate_applies_failure_policy_to_mismatch_and_read_error() {
|
||||
for result in [
|
||||
Ok("changed".to_string()),
|
||||
Err("read failed".to_string()),
|
||||
] {
|
||||
assert!(matches!(
|
||||
revision_gate_action(
|
||||
"accepted",
|
||||
result.clone(),
|
||||
ReloadFailurePolicy::KeepNew,
|
||||
),
|
||||
RevisionGateAction::Warn(_)
|
||||
));
|
||||
assert!(matches!(
|
||||
revision_gate_action("accepted", result, ReloadFailurePolicy::Rollback),
|
||||
RevisionGateAction::Rollback(_)
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
383
src/maestro/runtime_build.rs
Normal file
383
src/maestro/runtime_build.rs
Normal file
@@ -0,0 +1,383 @@
|
||||
use std::net::IpAddr;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use tokio::sync::{RwLock, Semaphore, watch};
|
||||
|
||||
use crate::config::ProxyConfig;
|
||||
use crate::crypto::SecureRandom;
|
||||
use crate::ip_tracker::UserIpTracker;
|
||||
use crate::network::probe::{decide_network_capabilities, run_probe};
|
||||
use crate::proxy::direct_buffer_budget::{
|
||||
DirectBufferBudget, resolve_direct_buffer_hard_limit, run_direct_buffer_budget_controller,
|
||||
};
|
||||
use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController};
|
||||
use crate::proxy::shared_state::ProxySharedState;
|
||||
use crate::startup::StartupTracker;
|
||||
use crate::stats::beobachten::BeobachtenStore;
|
||||
use crate::stats::telemetry::TelemetryPolicy;
|
||||
use crate::stats::{QuotaStore, ReplayChecker, Stats};
|
||||
use crate::stream::BufferPool;
|
||||
use crate::transport::UpstreamManager;
|
||||
use crate::transport::middle_proxy::MePool;
|
||||
|
||||
use super::admission;
|
||||
use super::generation::{RuntimeGeneration, RuntimeTaskScope};
|
||||
use super::runtime_tasks::RuntimeLogFilter;
|
||||
use super::{me_startup, runtime_tasks, tls_bootstrap};
|
||||
|
||||
pub(crate) struct PreparedRuntime {
|
||||
pub(crate) generation: Arc<RuntimeGeneration>,
|
||||
pub(crate) detected_ips: (Option<IpAddr>, Option<IpAddr>),
|
||||
}
|
||||
|
||||
pub(crate) async fn prepare_runtime(
|
||||
generation_id: u64,
|
||||
config: ProxyConfig,
|
||||
config_path: &Path,
|
||||
quota_store: Arc<QuotaStore>,
|
||||
runtime_log_filter: RuntimeLogFilter,
|
||||
) -> Result<PreparedRuntime, String> {
|
||||
let started_at_epoch_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
let startup_tracker = Arc::new(StartupTracker::new(started_at_epoch_secs));
|
||||
let task_scope = RuntimeTaskScope::new();
|
||||
let stats = Arc::new(Stats::with_quota_store(quota_store));
|
||||
stats.apply_telemetry_policy(TelemetryPolicy::from_config(&config.general.telemetry));
|
||||
|
||||
let upstream_manager = Arc::new(
|
||||
UpstreamManager::new(
|
||||
config.upstreams.clone(),
|
||||
config.general.upstream_connect_retry_attempts,
|
||||
config.general.upstream_connect_retry_backoff_ms,
|
||||
config.general.upstream_connect_budget_ms,
|
||||
config.general.tg_connect,
|
||||
config.general.upstream_unhealthy_fail_threshold,
|
||||
config.general.upstream_connect_failfast_hard_errors,
|
||||
stats.clone(),
|
||||
)
|
||||
.with_dns_overrides(&config.network.dns_overrides)
|
||||
.map_err(|error| format!("DNS override preparation failed: {}", error))?,
|
||||
);
|
||||
let ip_tracker = Arc::new(UserIpTracker::new());
|
||||
ip_tracker
|
||||
.load_limits(
|
||||
config.access.user_max_unique_ips_global_each,
|
||||
&config.access.user_max_unique_ips,
|
||||
)
|
||||
.await;
|
||||
ip_tracker
|
||||
.set_limit_policy(
|
||||
config.access.user_max_unique_ips_mode,
|
||||
config.access.user_max_unique_ips_window_secs,
|
||||
)
|
||||
.await;
|
||||
|
||||
let hard_limit =
|
||||
resolve_direct_buffer_hard_limit(config.general.direct_relay_buffer_budget_max_bytes).await;
|
||||
let direct_buffer_budget = DirectBufferBudget::new(hard_limit);
|
||||
let proxy_shared =
|
||||
ProxySharedState::new_with_direct_buffer_budget(direct_buffer_budget.clone());
|
||||
proxy_shared.apply_user_enabled_config(&config.access.user_enabled);
|
||||
proxy_shared.traffic_limiter.apply_policy(
|
||||
config.access.user_rate_limits.clone(),
|
||||
config.access.cidr_rate_limits.clone(),
|
||||
);
|
||||
|
||||
let probe = run_probe(
|
||||
&config.network,
|
||||
&config.upstreams,
|
||||
config.general.middle_proxy_nat_probe,
|
||||
config.general.stun_nat_probe_concurrency,
|
||||
)
|
||||
.await
|
||||
.map_err(|error| format!("network probe failed: {}", error))?;
|
||||
let decision =
|
||||
decide_network_capabilities(&config.network, &probe, config.general.middle_proxy_nat_ip);
|
||||
let prefer_ipv6 = decision.prefer_ipv6();
|
||||
|
||||
let mut tls_domains = Vec::with_capacity(1 + config.censorship.tls_domains.len());
|
||||
tls_domains.push(config.censorship.tls_domain.clone());
|
||||
for domain in &config.censorship.tls_domains {
|
||||
if !tls_domains.contains(domain) {
|
||||
tls_domains.push(domain.clone());
|
||||
}
|
||||
}
|
||||
let tls_cache = tls_bootstrap::bootstrap_tls_front(
|
||||
&config,
|
||||
&tls_domains,
|
||||
upstream_manager.clone(),
|
||||
&startup_tracker,
|
||||
task_scope.clone(),
|
||||
tls_bootstrap::TlsBootstrapPolicy::RequireReady,
|
||||
)
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
|
||||
let beobachten = Arc::new(BeobachtenStore::new());
|
||||
let rng = Arc::new(SecureRandom::new());
|
||||
let route_mode = if !config.general.use_middle_proxy || config.general.me2dc_fallback {
|
||||
RelayRouteMode::Direct
|
||||
} else {
|
||||
RelayRouteMode::Middle
|
||||
};
|
||||
let route_runtime = Arc::new(RouteRuntimeController::new(route_mode));
|
||||
let me_pool_runtime = Arc::new(RwLock::new(None::<Arc<MePool>>));
|
||||
let (me_ready_tx, me_ready_rx) = watch::channel(0_u64);
|
||||
let direct_first_startup = config.general.use_middle_proxy && config.general.me2dc_fallback;
|
||||
let me_pool = if direct_first_startup {
|
||||
None
|
||||
} else {
|
||||
me_startup::initialize_me_pool(
|
||||
config.general.use_middle_proxy,
|
||||
&config,
|
||||
&decision,
|
||||
&probe,
|
||||
&startup_tracker,
|
||||
upstream_manager.clone(),
|
||||
rng.clone(),
|
||||
stats.clone(),
|
||||
me_pool_runtime.clone(),
|
||||
me_ready_tx.clone(),
|
||||
task_scope.clone(),
|
||||
)
|
||||
.await
|
||||
};
|
||||
if strict_middle_proxy_unavailable(
|
||||
config.general.use_middle_proxy,
|
||||
direct_first_startup,
|
||||
me_pool.is_some(),
|
||||
) {
|
||||
task_scope.stop().await;
|
||||
return Err(
|
||||
"Middle-End pool is required but did not become ready during reload preparation"
|
||||
.to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
let config = Arc::new(config);
|
||||
let replay_checker = Arc::new(ReplayChecker::new(
|
||||
config.access.replay_check_len,
|
||||
Duration::from_secs(config.access.replay_window_secs),
|
||||
));
|
||||
let buffer_pool = Arc::new(BufferPool::with_config(64 * 1024, 4096));
|
||||
let max_connections_limit = if config.server.max_connections == 0 {
|
||||
Semaphore::MAX_PERMITS
|
||||
} else {
|
||||
config.server.max_connections as usize
|
||||
};
|
||||
let max_connections = Arc::new(Semaphore::new(max_connections_limit));
|
||||
let watches = runtime_tasks::spawn_runtime_tasks(
|
||||
&config,
|
||||
config_path,
|
||||
&probe,
|
||||
prefer_ipv6,
|
||||
decision.ipv4_dc,
|
||||
decision.ipv6_dc,
|
||||
&startup_tracker,
|
||||
stats.clone(),
|
||||
upstream_manager.clone(),
|
||||
replay_checker.clone(),
|
||||
me_pool.clone(),
|
||||
rng.clone(),
|
||||
ip_tracker.clone(),
|
||||
beobachten.clone(),
|
||||
me_pool.clone(),
|
||||
proxy_shared.clone(),
|
||||
me_ready_tx.clone(),
|
||||
task_scope.clone(),
|
||||
)
|
||||
.await;
|
||||
let config_rx = watches.config_rx;
|
||||
runtime_log_filter.spawn_watcher(watches.log_level_rx, task_scope.clone());
|
||||
let initial_admission_open = !config.general.use_middle_proxy || me_pool.is_some();
|
||||
let (admission_tx, admission_rx) = watch::channel(initial_admission_open);
|
||||
admission::configure_admission_gate(
|
||||
&config,
|
||||
me_pool.clone(),
|
||||
me_pool_runtime.clone(),
|
||||
route_runtime.clone(),
|
||||
&admission_tx,
|
||||
config_rx.clone(),
|
||||
me_ready_rx,
|
||||
task_scope.clone(),
|
||||
)
|
||||
.await;
|
||||
|
||||
if direct_first_startup {
|
||||
let config_bg = config.clone();
|
||||
let decision_bg = decision.clone();
|
||||
let probe_bg = probe.clone();
|
||||
let startup_tracker_bg = startup_tracker.clone();
|
||||
let upstream_manager_bg = upstream_manager.clone();
|
||||
let rng_bg = rng.clone();
|
||||
let stats_bg = stats.clone();
|
||||
let me_pool_runtime_bg = me_pool_runtime.clone();
|
||||
let me_ready_tx_bg = me_ready_tx.clone();
|
||||
let config_rx_bg = config_rx.clone();
|
||||
let task_scope_bg = task_scope.clone();
|
||||
let retry_limit = config.general.me_init_retry_attempts;
|
||||
task_scope.spawn(async move {
|
||||
let mut attempt = 0_u32;
|
||||
loop {
|
||||
attempt = attempt.saturating_add(1);
|
||||
let pool = me_startup::initialize_me_pool(
|
||||
true,
|
||||
config_bg.as_ref(),
|
||||
&decision_bg,
|
||||
&probe_bg,
|
||||
&startup_tracker_bg,
|
||||
upstream_manager_bg.clone(),
|
||||
rng_bg.clone(),
|
||||
stats_bg.clone(),
|
||||
me_pool_runtime_bg.clone(),
|
||||
me_ready_tx_bg.clone(),
|
||||
task_scope_bg.clone(),
|
||||
)
|
||||
.await;
|
||||
if let Some(pool) = pool {
|
||||
runtime_tasks::spawn_middle_proxy_runtime_tasks(
|
||||
config_bg.as_ref(),
|
||||
config_rx_bg,
|
||||
pool,
|
||||
rng_bg,
|
||||
me_ready_tx_bg,
|
||||
task_scope_bg,
|
||||
);
|
||||
break;
|
||||
}
|
||||
if retry_limit > 0 && attempt >= retry_limit {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_secs(2)).await;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
let conntrack_scope = task_scope.clone();
|
||||
task_scope.spawn(crate::conntrack_control::run_conntrack_controller(
|
||||
config_rx.clone(),
|
||||
stats.clone(),
|
||||
proxy_shared.clone(),
|
||||
conntrack_scope.cancellation_token(),
|
||||
));
|
||||
task_scope.spawn(run_direct_buffer_budget_controller(
|
||||
direct_buffer_budget,
|
||||
buffer_pool.clone(),
|
||||
stats.clone(),
|
||||
proxy_shared.clone(),
|
||||
config.server.max_connections,
|
||||
));
|
||||
let generation = RuntimeGeneration::new(
|
||||
generation_id,
|
||||
config_rx,
|
||||
admission_rx,
|
||||
stats,
|
||||
upstream_manager,
|
||||
replay_checker,
|
||||
buffer_pool,
|
||||
rng,
|
||||
me_pool,
|
||||
me_pool_runtime,
|
||||
route_runtime,
|
||||
tls_cache,
|
||||
ip_tracker,
|
||||
beobachten,
|
||||
proxy_shared,
|
||||
max_connections,
|
||||
task_scope,
|
||||
);
|
||||
drop(admission_tx);
|
||||
|
||||
Ok(PreparedRuntime {
|
||||
generation,
|
||||
detected_ips: (
|
||||
probe.detected_ipv4.map(IpAddr::V4),
|
||||
probe.detected_ipv6.map(IpAddr::V6),
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
fn strict_middle_proxy_unavailable(
|
||||
use_middle_proxy: bool,
|
||||
direct_first_startup: bool,
|
||||
pool_available: bool,
|
||||
) -> bool {
|
||||
use_middle_proxy && !direct_first_startup && !pool_available
|
||||
}
|
||||
|
||||
pub(crate) fn deferred_process_fields(old: &ProxyConfig, new: &ProxyConfig) -> Vec<String> {
|
||||
let mut fields = Vec::new();
|
||||
if old.server.port != new.server.port
|
||||
|| old.server.proxy_protocol != new.server.proxy_protocol
|
||||
|| old.server.listen_backlog != new.server.listen_backlog
|
||||
|| serde_json::to_value(&old.server.listeners).ok()
|
||||
!= serde_json::to_value(&new.server.listeners).ok()
|
||||
{
|
||||
fields.push("server.listeners".to_string());
|
||||
}
|
||||
if old.server.listen_unix_sock != new.server.listen_unix_sock
|
||||
|| old.server.listen_unix_sock_perm != new.server.listen_unix_sock_perm
|
||||
{
|
||||
fields.push("server.listen_unix_sock".to_string());
|
||||
}
|
||||
if old.server.api.listen != new.server.api.listen
|
||||
|| old.server.api.enabled != new.server.api.enabled
|
||||
{
|
||||
fields.push("server.api.listen".to_string());
|
||||
}
|
||||
if old.server.metrics_listen != new.server.metrics_listen
|
||||
|| old.server.metrics_port != new.server.metrics_port
|
||||
{
|
||||
fields.push("server.metrics_listen".to_string());
|
||||
}
|
||||
if old.general.quota_state_path != new.general.quota_state_path {
|
||||
fields.push("general.quota_state_path".to_string());
|
||||
}
|
||||
if old.general.disable_colors != new.general.disable_colors {
|
||||
fields.push("general.disable_colors".to_string());
|
||||
}
|
||||
if old.general.data_path != new.general.data_path {
|
||||
fields.push("general.data_path".to_string());
|
||||
}
|
||||
if serde_json::to_value(&old.logging).ok() != serde_json::to_value(&new.logging).ok() {
|
||||
fields.push("logging".to_string());
|
||||
}
|
||||
fields
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn process_socket_and_logging_changes_are_deferred() {
|
||||
let old = ProxyConfig::default();
|
||||
let mut new = old.clone();
|
||||
new.server.listen_backlog = new.server.listen_backlog.saturating_add(1);
|
||||
new.general.disable_colors = !new.general.disable_colors;
|
||||
|
||||
let fields = deferred_process_fields(&old, &new);
|
||||
assert!(fields.contains(&"server.listeners".to_string()));
|
||||
assert!(fields.contains(&"general.disable_colors".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runtime_only_change_does_not_require_process_rebind() {
|
||||
let old = ProxyConfig::default();
|
||||
let mut new = old.clone();
|
||||
new.censorship.tls_domain = "reload.example".to_string();
|
||||
assert!(deferred_process_fields(&old, &new).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strict_middle_proxy_requires_a_prepared_pool() {
|
||||
assert!(strict_middle_proxy_unavailable(true, false, false));
|
||||
assert!(!strict_middle_proxy_unavailable(true, false, true));
|
||||
assert!(!strict_middle_proxy_unavailable(true, true, false));
|
||||
assert!(!strict_middle_proxy_unavailable(false, false, false));
|
||||
}
|
||||
}
|
||||
@@ -2,9 +2,11 @@ use std::net::IpAddr;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use arc_swap::ArcSwap;
|
||||
use tokio::sync::{mpsc, watch};
|
||||
use tracing::{debug, info, warn};
|
||||
use tracing_subscriber::EnvFilter;
|
||||
use tracing_subscriber::Registry;
|
||||
use tracing_subscriber::reload;
|
||||
|
||||
use crate::config::hot_reload::spawn_config_watcher;
|
||||
@@ -21,10 +23,11 @@ use crate::startup::{
|
||||
use crate::stats::beobachten::BeobachtenStore;
|
||||
use crate::stats::telemetry::TelemetryPolicy;
|
||||
use crate::stats::{ReplayChecker, Stats};
|
||||
use crate::tls_front::TlsFrontCache;
|
||||
use crate::transport::UpstreamManager;
|
||||
use crate::transport::middle_proxy::{MePool, MeReinitTrigger};
|
||||
|
||||
use super::generation::RuntimeGeneration;
|
||||
use super::generation::RuntimeTaskScope;
|
||||
use super::helpers::write_beobachten_snapshot;
|
||||
|
||||
pub(crate) struct RuntimeWatches {
|
||||
@@ -34,6 +37,56 @@ pub(crate) struct RuntimeWatches {
|
||||
pub(crate) detected_ip_v6: Option<IpAddr>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct RuntimeLogFilter {
|
||||
handle: reload::Handle<EnvFilter, Registry>,
|
||||
}
|
||||
|
||||
impl RuntimeLogFilter {
|
||||
pub(crate) fn new(handle: reload::Handle<EnvFilter, Registry>) -> Self {
|
||||
Self { handle }
|
||||
}
|
||||
|
||||
pub(crate) fn start(
|
||||
&self,
|
||||
has_rust_log: bool,
|
||||
effective_log_level: &LogLevel,
|
||||
log_level_rx: watch::Receiver<LogLevel>,
|
||||
task_scope: RuntimeTaskScope,
|
||||
) {
|
||||
self.apply(effective_log_level, has_rust_log);
|
||||
self.spawn_watcher(log_level_rx, task_scope);
|
||||
}
|
||||
|
||||
pub(crate) fn apply_reload(&self, level: &LogLevel) {
|
||||
self.apply(level, false);
|
||||
}
|
||||
|
||||
pub(crate) fn spawn_watcher(
|
||||
&self,
|
||||
mut log_level_rx: watch::Receiver<LogLevel>,
|
||||
task_scope: RuntimeTaskScope,
|
||||
) {
|
||||
let filter = self.clone();
|
||||
task_scope.spawn(async move {
|
||||
loop {
|
||||
if log_level_rx.changed().await.is_err() {
|
||||
break;
|
||||
}
|
||||
let level = log_level_rx.borrow_and_update().clone();
|
||||
filter.apply_reload(&level);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn apply(&self, level: &LogLevel, has_rust_log: bool) {
|
||||
let runtime_filter = EnvFilter::new(log_filter_spec(has_rust_log, level));
|
||||
if let Err(error) = self.handle.reload(runtime_filter) {
|
||||
tracing::error!(error = %error, "Failed to update runtime log filter");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn spawn_runtime_tasks(
|
||||
config: &Arc<ProxyConfig>,
|
||||
@@ -50,14 +103,14 @@ pub(crate) async fn spawn_runtime_tasks(
|
||||
rng: Arc<SecureRandom>,
|
||||
ip_tracker: Arc<UserIpTracker>,
|
||||
beobachten: Arc<BeobachtenStore>,
|
||||
api_config_tx: watch::Sender<Arc<ProxyConfig>>,
|
||||
me_pool_for_policy: Option<Arc<MePool>>,
|
||||
shared_state: Arc<ProxySharedState>,
|
||||
me_ready_tx: watch::Sender<u64>,
|
||||
task_scope: RuntimeTaskScope,
|
||||
) -> RuntimeWatches {
|
||||
let um_clone = upstream_manager.clone();
|
||||
let dc_overrides_for_health = config.dc_overrides.clone();
|
||||
tokio::spawn(async move {
|
||||
task_scope.spawn(async move {
|
||||
um_clone
|
||||
.run_health_checks(
|
||||
prefer_ipv6,
|
||||
@@ -69,19 +122,19 @@ pub(crate) async fn spawn_runtime_tasks(
|
||||
});
|
||||
|
||||
let rc_clone = replay_checker.clone();
|
||||
tokio::spawn(async move {
|
||||
task_scope.spawn(async move {
|
||||
rc_clone.run_periodic_cleanup().await;
|
||||
});
|
||||
|
||||
let stats_maintenance = stats.clone();
|
||||
tokio::spawn(async move {
|
||||
task_scope.spawn(async move {
|
||||
stats_maintenance
|
||||
.run_periodic_user_stats_maintenance()
|
||||
.await;
|
||||
});
|
||||
|
||||
let ip_tracker_maintenance = ip_tracker.clone();
|
||||
tokio::spawn(async move {
|
||||
task_scope.spawn(async move {
|
||||
ip_tracker_maintenance.run_periodic_maintenance().await;
|
||||
});
|
||||
|
||||
@@ -104,6 +157,7 @@ pub(crate) async fn spawn_runtime_tasks(
|
||||
config.clone(),
|
||||
detected_ip_v4,
|
||||
detected_ip_v6,
|
||||
task_scope.cancellation_token(),
|
||||
);
|
||||
startup_tracker
|
||||
.complete_component(
|
||||
@@ -111,21 +165,10 @@ pub(crate) async fn spawn_runtime_tasks(
|
||||
Some("config hot-reload watcher started".to_string()),
|
||||
)
|
||||
.await;
|
||||
let mut config_rx_api_bridge = config_rx.clone();
|
||||
let api_config_tx_bridge = api_config_tx.clone();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
if config_rx_api_bridge.changed().await.is_err() {
|
||||
break;
|
||||
}
|
||||
let cfg = config_rx_api_bridge.borrow_and_update().clone();
|
||||
api_config_tx_bridge.send_replace(cfg);
|
||||
}
|
||||
});
|
||||
|
||||
let stats_policy = stats.clone();
|
||||
let upstream_policy = upstream_manager.clone();
|
||||
let mut config_rx_policy = config_rx.clone();
|
||||
tokio::spawn(async move {
|
||||
task_scope.spawn(async move {
|
||||
loop {
|
||||
if config_rx_policy.changed().await.is_err() {
|
||||
break;
|
||||
@@ -133,6 +176,9 @@ pub(crate) async fn spawn_runtime_tasks(
|
||||
let cfg = config_rx_policy.borrow_and_update().clone();
|
||||
stats_policy
|
||||
.apply_telemetry_policy(TelemetryPolicy::from_config(&cfg.general.telemetry));
|
||||
if let Err(error) = upstream_policy.update_dns_overrides(&cfg.network.dns_overrides) {
|
||||
warn!(error = %error, "Failed to update generation DNS overrides");
|
||||
}
|
||||
if let Some(pool) = &me_pool_for_policy {
|
||||
pool.update_runtime_transport_policy(
|
||||
cfg.general.me_socks_kdf_policy,
|
||||
@@ -149,7 +195,7 @@ pub(crate) async fn spawn_runtime_tasks(
|
||||
|
||||
let ip_tracker_policy = ip_tracker.clone();
|
||||
let mut config_rx_ip_limits = config_rx.clone();
|
||||
tokio::spawn(async move {
|
||||
task_scope.spawn(async move {
|
||||
let mut prev_limits = config_rx_ip_limits
|
||||
.borrow()
|
||||
.access
|
||||
@@ -205,7 +251,7 @@ pub(crate) async fn spawn_runtime_tasks(
|
||||
config.access.cidr_rate_limits.clone(),
|
||||
);
|
||||
let mut config_rx_rate_limits = config_rx.clone();
|
||||
tokio::spawn(async move {
|
||||
task_scope.spawn(async move {
|
||||
let mut prev_user_limits = config_rx_rate_limits
|
||||
.borrow()
|
||||
.access
|
||||
@@ -236,7 +282,7 @@ pub(crate) async fn spawn_runtime_tasks(
|
||||
|
||||
let shared_user_enabled = shared_state.clone();
|
||||
let mut config_rx_user_enabled = config_rx.clone();
|
||||
tokio::spawn(async move {
|
||||
task_scope.spawn(async move {
|
||||
loop {
|
||||
if config_rx_user_enabled.changed().await.is_err() {
|
||||
break;
|
||||
@@ -257,7 +303,7 @@ pub(crate) async fn spawn_runtime_tasks(
|
||||
|
||||
let beobachten_writer = beobachten.clone();
|
||||
let config_rx_beobachten = config_rx.clone();
|
||||
tokio::spawn(async move {
|
||||
task_scope.spawn(async move {
|
||||
loop {
|
||||
let cfg = config_rx_beobachten.borrow().clone();
|
||||
let sleep_secs = cfg.general.beobachten_flush_secs.max(1);
|
||||
@@ -278,7 +324,14 @@ pub(crate) async fn spawn_runtime_tasks(
|
||||
});
|
||||
|
||||
if let Some(pool) = me_pool {
|
||||
spawn_middle_proxy_runtime_tasks(config, config_rx.clone(), pool, rng, me_ready_tx);
|
||||
spawn_middle_proxy_runtime_tasks(
|
||||
config,
|
||||
config_rx.clone(),
|
||||
pool,
|
||||
rng,
|
||||
me_ready_tx,
|
||||
task_scope,
|
||||
);
|
||||
}
|
||||
|
||||
RuntimeWatches {
|
||||
@@ -295,6 +348,7 @@ pub(crate) fn spawn_middle_proxy_runtime_tasks(
|
||||
pool: Arc<MePool>,
|
||||
rng: Arc<SecureRandom>,
|
||||
me_ready_tx: watch::Sender<u64>,
|
||||
task_scope: RuntimeTaskScope,
|
||||
) {
|
||||
let reinit_trigger_capacity = config.general.me_reinit_trigger_channel.max(1);
|
||||
let (reinit_tx, reinit_rx) = mpsc::channel::<MeReinitTrigger>(reinit_trigger_capacity);
|
||||
@@ -303,7 +357,7 @@ pub(crate) fn spawn_middle_proxy_runtime_tasks(
|
||||
let rng_clone_sched = rng.clone();
|
||||
let config_rx_clone_sched = config_rx.clone();
|
||||
let me_ready_tx_sched = me_ready_tx.clone();
|
||||
tokio::spawn(async move {
|
||||
task_scope.spawn(async move {
|
||||
crate::transport::middle_proxy::me_reinit_scheduler(
|
||||
pool_clone_sched,
|
||||
rng_clone_sched,
|
||||
@@ -317,7 +371,7 @@ pub(crate) fn spawn_middle_proxy_runtime_tasks(
|
||||
let pool_clone = pool.clone();
|
||||
let config_rx_clone = config_rx.clone();
|
||||
let reinit_tx_updater = reinit_tx.clone();
|
||||
tokio::spawn(async move {
|
||||
task_scope.spawn(async move {
|
||||
crate::transport::middle_proxy::me_config_updater(
|
||||
pool_clone,
|
||||
config_rx_clone,
|
||||
@@ -328,37 +382,12 @@ pub(crate) fn spawn_middle_proxy_runtime_tasks(
|
||||
|
||||
let config_rx_clone_rot = config_rx.clone();
|
||||
let reinit_tx_rotation = reinit_tx.clone();
|
||||
tokio::spawn(async move {
|
||||
task_scope.spawn(async move {
|
||||
crate::transport::middle_proxy::me_rotation_task(config_rx_clone_rot, reinit_tx_rotation)
|
||||
.await;
|
||||
});
|
||||
}
|
||||
|
||||
pub(crate) async fn apply_runtime_log_filter(
|
||||
has_rust_log: bool,
|
||||
effective_log_level: &LogLevel,
|
||||
filter_handle: reload::Handle<EnvFilter, tracing_subscriber::Registry>,
|
||||
mut log_level_rx: watch::Receiver<LogLevel>,
|
||||
) {
|
||||
let runtime_filter = EnvFilter::new(log_filter_spec(has_rust_log, effective_log_level));
|
||||
filter_handle
|
||||
.reload(runtime_filter)
|
||||
.expect("Failed to switch log filter");
|
||||
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
if log_level_rx.changed().await.is_err() {
|
||||
break;
|
||||
}
|
||||
let level = log_level_rx.borrow_and_update().clone();
|
||||
let new_filter = tracing_subscriber::EnvFilter::new(log_filter_spec(false, &level));
|
||||
if let Err(e) = filter_handle.reload(new_filter) {
|
||||
tracing::error!("config reload: failed to update log filter: {}", e);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
pub(crate) fn log_filter_spec(has_rust_log: bool, effective_log_level: &LogLevel) -> String {
|
||||
if has_rust_log {
|
||||
std::env::var("RUST_LOG")
|
||||
@@ -373,12 +402,7 @@ pub(crate) fn log_filter_spec(has_rust_log: bool, effective_log_level: &LogLevel
|
||||
pub(crate) async fn spawn_metrics_if_configured(
|
||||
config: &Arc<ProxyConfig>,
|
||||
startup_tracker: &Arc<StartupTracker>,
|
||||
stats: Arc<Stats>,
|
||||
beobachten: Arc<BeobachtenStore>,
|
||||
shared_state: Arc<ProxySharedState>,
|
||||
ip_tracker: Arc<UserIpTracker>,
|
||||
tls_cache: Option<Arc<TlsFrontCache>>,
|
||||
config_rx: watch::Receiver<Arc<ProxyConfig>>,
|
||||
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
||||
) {
|
||||
// metrics_listen takes precedence; fall back to metrics_port for backward compat.
|
||||
let metrics_target: Option<(u16, Option<String>)> =
|
||||
@@ -408,28 +432,10 @@ pub(crate) async fn spawn_metrics_if_configured(
|
||||
Some(format!("spawn metrics endpoint on {}", label)),
|
||||
)
|
||||
.await;
|
||||
let stats = stats.clone();
|
||||
let beobachten = beobachten.clone();
|
||||
let shared_state = shared_state.clone();
|
||||
let config_rx_metrics = config_rx.clone();
|
||||
let ip_tracker_metrics = ip_tracker.clone();
|
||||
let tls_cache_metrics = tls_cache.clone();
|
||||
let whitelist = config.server.metrics_whitelist.clone();
|
||||
let active_runtime = active_runtime.clone();
|
||||
let listen_backlog = config.server.listen_backlog;
|
||||
tokio::spawn(async move {
|
||||
metrics::serve(
|
||||
port,
|
||||
listen,
|
||||
listen_backlog,
|
||||
stats,
|
||||
beobachten,
|
||||
shared_state,
|
||||
ip_tracker_metrics,
|
||||
tls_cache_metrics,
|
||||
config_rx_metrics,
|
||||
whitelist,
|
||||
)
|
||||
.await;
|
||||
metrics::serve(port, listen, listen_backlog, active_runtime).await;
|
||||
});
|
||||
startup_tracker
|
||||
.complete_component(
|
||||
|
||||
@@ -12,17 +12,17 @@ use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use arc_swap::ArcSwap;
|
||||
#[cfg(not(unix))]
|
||||
use tokio::signal;
|
||||
#[cfg(unix)]
|
||||
use tokio::signal::unix::{SignalKind, signal};
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::generation::RuntimeGeneration;
|
||||
use super::helpers::{format_uptime, unit_label};
|
||||
use crate::stats::Stats;
|
||||
use crate::synlimit_control;
|
||||
use crate::transport::middle_proxy::MePool;
|
||||
|
||||
use super::helpers::{format_uptime, unit_label};
|
||||
|
||||
/// Signal that triggered shutdown.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
@@ -48,17 +48,17 @@ impl std::fmt::Display for ShutdownSignal {
|
||||
/// Waits for a shutdown signal and performs graceful shutdown.
|
||||
pub(crate) async fn wait_for_shutdown(
|
||||
process_started_at: Instant,
|
||||
me_pool: Option<Arc<MePool>>,
|
||||
stats: Arc<Stats>,
|
||||
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
||||
quota_state_path: PathBuf,
|
||||
synlimit_controller: tokio::task::JoinHandle<()>,
|
||||
) {
|
||||
let signal = wait_for_shutdown_signal().await;
|
||||
perform_shutdown(
|
||||
signal,
|
||||
process_started_at,
|
||||
me_pool,
|
||||
&stats,
|
||||
active_runtime,
|
||||
quota_state_path,
|
||||
synlimit_controller,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -87,10 +87,12 @@ async fn wait_for_shutdown_signal() -> ShutdownSignal {
|
||||
async fn perform_shutdown(
|
||||
signal: ShutdownSignal,
|
||||
process_started_at: Instant,
|
||||
me_pool: Option<Arc<MePool>>,
|
||||
stats: &Stats,
|
||||
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
||||
quota_state_path: PathBuf,
|
||||
synlimit_controller: tokio::task::JoinHandle<()>,
|
||||
) {
|
||||
let runtime = active_runtime.load_full();
|
||||
let stats = runtime.stats.as_ref();
|
||||
let shutdown_started_at = Instant::now();
|
||||
info!(signal = %signal, "Received shutdown signal");
|
||||
|
||||
@@ -103,12 +105,10 @@ async fn perform_shutdown(
|
||||
let uptime_secs = process_started_at.elapsed().as_secs();
|
||||
info!("Uptime: {}", format_uptime(uptime_secs));
|
||||
|
||||
if let Err(error) = synlimit_control::clear_synlimit_rules_all_backends().await {
|
||||
warn!(error = %error, "Failed to clear SYN limiter rules during shutdown");
|
||||
}
|
||||
|
||||
// Graceful ME pool shutdown
|
||||
if let Some(pool) = &me_pool {
|
||||
runtime.stop_sessions().await;
|
||||
runtime.stop_background_tasks().await;
|
||||
if let Some(pool) = runtime.current_me_pool().await {
|
||||
match tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all())
|
||||
.await
|
||||
{
|
||||
@@ -124,6 +124,12 @@ async fn perform_shutdown(
|
||||
}
|
||||
}
|
||||
|
||||
synlimit_controller.abort();
|
||||
let _ = synlimit_controller.await;
|
||||
if let Err(error) = synlimit_control::clear_synlimit_rules_all_backends().await {
|
||||
warn!(error = %error, "Failed to clear SYN limiter rules during shutdown");
|
||||
}
|
||||
|
||||
match crate::quota_state::save_quota_state("a_state_path, stats).await {
|
||||
Ok(()) => {
|
||||
info!(
|
||||
@@ -191,7 +197,10 @@ fn dump_stats(stats: &Stats, process_started_at: Instant) {
|
||||
/// - SIGUSR1: Log rotation acknowledgment (for external log rotation tools)
|
||||
/// - SIGUSR2: Dump runtime status to log
|
||||
#[cfg(unix)]
|
||||
pub(crate) fn spawn_signal_handlers(stats: Arc<Stats>, process_started_at: Instant) {
|
||||
pub(crate) fn spawn_signal_handlers(
|
||||
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
||||
process_started_at: Instant,
|
||||
) {
|
||||
tokio::spawn(async move {
|
||||
let mut sigusr1 =
|
||||
signal(SignalKind::user_defined1()).expect("Failed to register SIGUSR1 handler");
|
||||
@@ -204,7 +213,8 @@ pub(crate) fn spawn_signal_handlers(stats: Arc<Stats>, process_started_at: Insta
|
||||
handle_sigusr1();
|
||||
}
|
||||
_ = sigusr2.recv() => {
|
||||
handle_sigusr2(&stats, process_started_at);
|
||||
let runtime = active_runtime.load_full();
|
||||
handle_sigusr2(runtime.stats.as_ref(), process_started_at);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -213,7 +223,10 @@ pub(crate) fn spawn_signal_handlers(stats: Arc<Stats>, process_started_at: Insta
|
||||
|
||||
/// No-op on non-Unix platforms.
|
||||
#[cfg(not(unix))]
|
||||
pub(crate) fn spawn_signal_handlers(_stats: Arc<Stats>, _process_started_at: Instant) {
|
||||
pub(crate) fn spawn_signal_handlers(
|
||||
_active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
||||
_process_started_at: Instant,
|
||||
) {
|
||||
// No SIGUSR1/SIGUSR2 on non-Unix
|
||||
}
|
||||
|
||||
|
||||
@@ -5,11 +5,90 @@ use rand::RngExt;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::config::ProxyConfig;
|
||||
use crate::error::{ProxyError, Result};
|
||||
use crate::startup::{COMPONENT_TLS_FRONT_BOOTSTRAP, StartupTracker};
|
||||
use crate::tls_front::TlsFrontCache;
|
||||
use crate::tls_front::fetcher::TlsFetchStrategy;
|
||||
use crate::transport::UpstreamManager;
|
||||
|
||||
use super::generation::RuntimeTaskScope;
|
||||
|
||||
/// Readiness requirement for TLS-front cache initialization.
|
||||
#[derive(Clone, Copy)]
|
||||
pub(crate) enum TlsBootstrapPolicy {
|
||||
BestEffort,
|
||||
RequireReady,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct TlsFetchContext {
|
||||
cache: Arc<TlsFrontCache>,
|
||||
domains: Vec<String>,
|
||||
mask_host: String,
|
||||
primary_domain: String,
|
||||
mask_unix_sock: Option<String>,
|
||||
tls_fetch_scope: Option<String>,
|
||||
upstream_manager: Arc<UpstreamManager>,
|
||||
strategy: TlsFetchStrategy,
|
||||
port: u16,
|
||||
proxy_protocol: u8,
|
||||
}
|
||||
|
||||
impl TlsFetchContext {
|
||||
async fn fetch_all(&self, failure_message: &'static str) {
|
||||
let mut join = tokio::task::JoinSet::new();
|
||||
for domain in self.domains.clone() {
|
||||
let cache = self.cache.clone();
|
||||
let host = tls_fetch_host_for_domain(
|
||||
&self.mask_host,
|
||||
&self.primary_domain,
|
||||
&domain,
|
||||
);
|
||||
let unix_sock = self.mask_unix_sock.clone();
|
||||
let scope = self.tls_fetch_scope.clone();
|
||||
let upstream = self.upstream_manager.clone();
|
||||
let strategy = self.strategy.clone();
|
||||
let port = self.port;
|
||||
let proxy_protocol = self.proxy_protocol;
|
||||
join.spawn(async move {
|
||||
match crate::tls_front::fetcher::fetch_real_tls_with_strategy(
|
||||
&host,
|
||||
port,
|
||||
&domain,
|
||||
&strategy,
|
||||
Some(upstream),
|
||||
scope.as_deref(),
|
||||
proxy_protocol,
|
||||
unix_sock.as_deref(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(result) => cache.update_from_fetch(&domain, result).await,
|
||||
Err(error) => warn!(domain = %domain, error = %error, failure_message),
|
||||
}
|
||||
});
|
||||
}
|
||||
while let Some(result) = join.join_next().await {
|
||||
if let Err(error) = result {
|
||||
warn!(error = %error, "TLS emulation fetch task join failed");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn fetch_all_with_budget(&self, phase: &'static str) {
|
||||
if tokio::time::timeout(self.strategy.total_budget, self.fetch_all(phase))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
warn!(
|
||||
phase,
|
||||
timeout_ms = self.strategy.total_budget.as_millis(),
|
||||
"TLS emulation fetch budget exhausted"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn tls_fetch_host_for_domain(mask_host: &str, primary_tls_domain: &str, domain: &str) -> String {
|
||||
if mask_host.eq_ignore_ascii_case(primary_tls_domain) {
|
||||
domain.to_string()
|
||||
@@ -18,12 +97,24 @@ fn tls_fetch_host_for_domain(mask_host: &str, primary_tls_domain: &str, domain:
|
||||
}
|
||||
}
|
||||
|
||||
fn readiness_error(default_domains: &[String]) -> Option<String> {
|
||||
(!default_domains.is_empty()).then(|| {
|
||||
format!(
|
||||
"TLS-front profiles are not ready for domains: {}",
|
||||
default_domains.join(", ")
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
/// Initializes the TLS-front cache and generation-owned refresh tasks.
|
||||
pub(crate) async fn bootstrap_tls_front(
|
||||
config: &ProxyConfig,
|
||||
tls_domains: &[String],
|
||||
upstream_manager: Arc<UpstreamManager>,
|
||||
startup_tracker: &Arc<StartupTracker>,
|
||||
) -> Option<Arc<TlsFrontCache>> {
|
||||
task_scope: RuntimeTaskScope,
|
||||
policy: TlsBootstrapPolicy,
|
||||
) -> Result<Option<Arc<TlsFrontCache>>> {
|
||||
startup_tracker
|
||||
.start_component(
|
||||
COMPONENT_TLS_FRONT_BOOTSTRAP,
|
||||
@@ -31,26 +122,38 @@ pub(crate) async fn bootstrap_tls_front(
|
||||
)
|
||||
.await;
|
||||
|
||||
let tls_cache: Option<Arc<TlsFrontCache>> = if config.censorship.tls_emulation {
|
||||
let cache = Arc::new(TlsFrontCache::new(
|
||||
tls_domains,
|
||||
config.censorship.fake_cert_len,
|
||||
&config.censorship.tls_front_dir,
|
||||
));
|
||||
cache.load_from_disk().await;
|
||||
if !config.censorship.tls_emulation {
|
||||
startup_tracker
|
||||
.skip_component(
|
||||
COMPONENT_TLS_FRONT_BOOTSTRAP,
|
||||
Some("censorship.tls_emulation is false".to_string()),
|
||||
)
|
||||
.await;
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let port = config.censorship.mask_port;
|
||||
let proxy_protocol = config.censorship.mask_proxy_protocol;
|
||||
let mask_host = config
|
||||
let cache = Arc::new(TlsFrontCache::new(
|
||||
tls_domains,
|
||||
config.censorship.fake_cert_len,
|
||||
&config.censorship.tls_front_dir,
|
||||
));
|
||||
cache.load_from_disk().await;
|
||||
|
||||
let tls_fetch = config.censorship.tls_fetch.clone();
|
||||
let fetch_context = TlsFetchContext {
|
||||
cache: cache.clone(),
|
||||
domains: tls_domains.to_vec(),
|
||||
mask_host: config
|
||||
.censorship
|
||||
.mask_host
|
||||
.clone()
|
||||
.unwrap_or_else(|| config.censorship.tls_domain.clone());
|
||||
let mask_unix_sock = config.censorship.mask_unix_sock.clone();
|
||||
let tls_fetch_scope = (!config.censorship.tls_fetch_scope.is_empty())
|
||||
.then(|| config.censorship.tls_fetch_scope.clone());
|
||||
let tls_fetch = config.censorship.tls_fetch.clone();
|
||||
let fetch_strategy = TlsFetchStrategy {
|
||||
.unwrap_or_else(|| config.censorship.tls_domain.clone()),
|
||||
primary_domain: config.censorship.tls_domain.clone(),
|
||||
mask_unix_sock: config.censorship.mask_unix_sock.clone(),
|
||||
tls_fetch_scope: (!config.censorship.tls_fetch_scope.is_empty())
|
||||
.then(|| config.censorship.tls_fetch_scope.clone()),
|
||||
upstream_manager,
|
||||
strategy: TlsFetchStrategy {
|
||||
profiles: tls_fetch.profiles,
|
||||
strict_route: tls_fetch.strict_route,
|
||||
attempt_timeout: Duration::from_millis(tls_fetch.attempt_timeout_ms.max(1)),
|
||||
@@ -58,150 +161,71 @@ pub(crate) async fn bootstrap_tls_front(
|
||||
grease_enabled: tls_fetch.grease_enabled,
|
||||
deterministic: tls_fetch.deterministic,
|
||||
profile_cache_ttl: Duration::from_secs(tls_fetch.profile_cache_ttl_secs),
|
||||
};
|
||||
let fetch_timeout = fetch_strategy.total_budget;
|
||||
},
|
||||
port: config.censorship.mask_port,
|
||||
proxy_protocol: config.censorship.mask_proxy_protocol,
|
||||
};
|
||||
|
||||
let cache_initial = cache.clone();
|
||||
let domains_initial = tls_domains.to_vec();
|
||||
let host_initial = mask_host.clone();
|
||||
let primary_initial = config.censorship.tls_domain.clone();
|
||||
let unix_sock_initial = mask_unix_sock.clone();
|
||||
let scope_initial = tls_fetch_scope.clone();
|
||||
let upstream_initial = upstream_manager.clone();
|
||||
let strategy_initial = fetch_strategy.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut join = tokio::task::JoinSet::new();
|
||||
for domain in domains_initial {
|
||||
let cache_domain = cache_initial.clone();
|
||||
let host_domain =
|
||||
tls_fetch_host_for_domain(&host_initial, &primary_initial, &domain);
|
||||
let unix_sock_domain = unix_sock_initial.clone();
|
||||
let scope_domain = scope_initial.clone();
|
||||
let upstream_domain = upstream_initial.clone();
|
||||
let strategy_domain = strategy_initial.clone();
|
||||
join.spawn(async move {
|
||||
match crate::tls_front::fetcher::fetch_real_tls_with_strategy(
|
||||
&host_domain,
|
||||
port,
|
||||
&domain,
|
||||
&strategy_domain,
|
||||
Some(upstream_domain),
|
||||
scope_domain.as_deref(),
|
||||
proxy_protocol,
|
||||
unix_sock_domain.as_deref(),
|
||||
)
|
||||
match policy {
|
||||
TlsBootstrapPolicy::BestEffort => {
|
||||
let initial_fetch = fetch_context.clone();
|
||||
let fake_cert_len = config.censorship.fake_cert_len;
|
||||
task_scope.spawn(async move {
|
||||
initial_fetch
|
||||
.fetch_all_with_budget("TLS emulation initial fetch failed")
|
||||
.await;
|
||||
for domain in initial_fetch
|
||||
.cache
|
||||
.default_profile_domains(&initial_fetch.domains)
|
||||
.await
|
||||
{
|
||||
Ok(res) => cache_domain.update_from_fetch(&domain, res).await,
|
||||
Err(e) => {
|
||||
warn!(domain = %domain, error = %e, "TLS emulation initial fetch failed")
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
while let Some(res) = join.join_next().await {
|
||||
if let Err(e) = res {
|
||||
warn!(error = %e, "TLS emulation initial fetch task join failed");
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let cache_timeout = cache.clone();
|
||||
let domains_timeout = tls_domains.to_vec();
|
||||
let fake_cert_len = config.censorship.fake_cert_len;
|
||||
tokio::spawn(async move {
|
||||
tokio::time::sleep(fetch_timeout).await;
|
||||
for domain in domains_timeout {
|
||||
let cached = cache_timeout.get(&domain).await;
|
||||
if cached.domain == "default" {
|
||||
{
|
||||
warn!(
|
||||
domain = %domain,
|
||||
timeout_secs = fetch_timeout.as_secs(),
|
||||
timeout_ms = initial_fetch.strategy.total_budget.as_millis(),
|
||||
fake_cert_len,
|
||||
"TLS-front fetch not ready within timeout; using cache/default fake cert fallback"
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
TlsBootstrapPolicy::RequireReady => {
|
||||
fetch_context
|
||||
.fetch_all_with_budget("TLS emulation initial fetch failed")
|
||||
.await;
|
||||
let default_domains = cache.default_profile_domains(tls_domains).await;
|
||||
if let Some(error) = readiness_error(&default_domains) {
|
||||
startup_tracker
|
||||
.fail_component(COMPONENT_TLS_FRONT_BOOTSTRAP, Some(error.clone()))
|
||||
.await;
|
||||
return Err(ProxyError::Proxy(error));
|
||||
}
|
||||
});
|
||||
|
||||
let cache_refresh = cache.clone();
|
||||
let domains_refresh = tls_domains.to_vec();
|
||||
let host_refresh = mask_host.clone();
|
||||
let primary_refresh = config.censorship.tls_domain.clone();
|
||||
let unix_sock_refresh = mask_unix_sock.clone();
|
||||
let scope_refresh = tls_fetch_scope.clone();
|
||||
let upstream_refresh = upstream_manager.clone();
|
||||
let strategy_refresh = fetch_strategy.clone();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let base_secs = rand::rng().random_range(4 * 3600..=6 * 3600);
|
||||
let jitter_secs = rand::rng().random_range(0..=7200);
|
||||
tokio::time::sleep(Duration::from_secs(base_secs + jitter_secs)).await;
|
||||
|
||||
let mut join = tokio::task::JoinSet::new();
|
||||
for domain in domains_refresh.clone() {
|
||||
let cache_domain = cache_refresh.clone();
|
||||
let host_domain =
|
||||
tls_fetch_host_for_domain(&host_refresh, &primary_refresh, &domain);
|
||||
let unix_sock_domain = unix_sock_refresh.clone();
|
||||
let scope_domain = scope_refresh.clone();
|
||||
let upstream_domain = upstream_refresh.clone();
|
||||
let strategy_domain = strategy_refresh.clone();
|
||||
join.spawn(async move {
|
||||
match crate::tls_front::fetcher::fetch_real_tls_with_strategy(
|
||||
&host_domain,
|
||||
port,
|
||||
&domain,
|
||||
&strategy_domain,
|
||||
Some(upstream_domain),
|
||||
scope_domain.as_deref(),
|
||||
proxy_protocol,
|
||||
unix_sock_domain.as_deref(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(res) => cache_domain.update_from_fetch(&domain, res).await,
|
||||
Err(e) => {
|
||||
warn!(domain = %domain, error = %e, "TLS emulation refresh failed")
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
while let Some(res) = join.join_next().await {
|
||||
if let Err(e) = res {
|
||||
warn!(error = %e, "TLS emulation refresh task join failed");
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Some(cache)
|
||||
} else {
|
||||
startup_tracker
|
||||
.skip_component(
|
||||
COMPONENT_TLS_FRONT_BOOTSTRAP,
|
||||
Some("censorship.tls_emulation is false".to_string()),
|
||||
)
|
||||
.await;
|
||||
None
|
||||
};
|
||||
|
||||
if tls_cache.is_some() {
|
||||
startup_tracker
|
||||
.complete_component(
|
||||
COMPONENT_TLS_FRONT_BOOTSTRAP,
|
||||
Some("tls front cache is initialized".to_string()),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
tls_cache
|
||||
let refresh_context = fetch_context;
|
||||
task_scope.spawn(async move {
|
||||
loop {
|
||||
let base_secs = rand::rng().random_range(4 * 3600..=6 * 3600);
|
||||
let jitter_secs = rand::rng().random_range(0..=7200);
|
||||
tokio::time::sleep(Duration::from_secs(base_secs + jitter_secs)).await;
|
||||
refresh_context
|
||||
.fetch_all_with_budget("TLS emulation refresh failed")
|
||||
.await;
|
||||
}
|
||||
});
|
||||
|
||||
startup_tracker
|
||||
.complete_component(
|
||||
COMPONENT_TLS_FRONT_BOOTSTRAP,
|
||||
Some("tls front cache is initialized".to_string()),
|
||||
)
|
||||
.await;
|
||||
Ok(Some(cache))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::tls_fetch_host_for_domain;
|
||||
use super::{readiness_error, tls_fetch_host_for_domain};
|
||||
|
||||
#[test]
|
||||
fn tls_fetch_host_uses_each_domain_when_mask_host_is_primary_default() {
|
||||
@@ -218,4 +242,13 @@ mod tests {
|
||||
"origin.example"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn readiness_rejects_only_default_profiles() {
|
||||
assert!(readiness_error(&[]).is_none());
|
||||
assert_eq!(
|
||||
readiness_error(&["front.example".to_string()]),
|
||||
Some("TLS-front profiles are not ready for domains: front.example".to_string())
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
261
src/metrics.rs
261
src/metrics.rs
@@ -4,12 +4,12 @@ use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use arc_swap::ArcSwap;
|
||||
use http_body_util::Full;
|
||||
use hyper::body::Bytes;
|
||||
use hyper::server::conn::http1;
|
||||
use hyper::service::service_fn;
|
||||
use hyper::{Request, Response, StatusCode};
|
||||
use ipnetwork::IpNetwork;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::Semaphore;
|
||||
use tokio::time::timeout;
|
||||
@@ -17,6 +17,7 @@ use tracing::{debug, info, warn};
|
||||
|
||||
use crate::config::ProxyConfig;
|
||||
use crate::ip_tracker::UserIpTracker;
|
||||
use crate::maestro::generation::RuntimeGeneration;
|
||||
use crate::proxy::shared_state::ProxySharedState;
|
||||
use crate::stats::Stats;
|
||||
use crate::stats::beobachten::BeobachtenStore;
|
||||
@@ -36,16 +37,8 @@ pub async fn serve(
|
||||
port: u16,
|
||||
listen: Option<String>,
|
||||
listen_backlog: u32,
|
||||
stats: Arc<Stats>,
|
||||
beobachten: Arc<BeobachtenStore>,
|
||||
shared_state: Arc<ProxySharedState>,
|
||||
ip_tracker: Arc<UserIpTracker>,
|
||||
tls_cache: Option<Arc<TlsFrontCache>>,
|
||||
config_rx: tokio::sync::watch::Receiver<Arc<ProxyConfig>>,
|
||||
whitelist: Vec<IpNetwork>,
|
||||
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
|
||||
) {
|
||||
let whitelist = Arc::new(whitelist);
|
||||
|
||||
// If `metrics_listen` is set, bind on that single address only.
|
||||
if let Some(ref listen_addr) = listen {
|
||||
let addr: SocketAddr = match listen_addr.parse() {
|
||||
@@ -61,17 +54,7 @@ pub async fn serve(
|
||||
match bind_metrics_listener(addr, ipv6_only, listen_backlog) {
|
||||
Ok(listener) => {
|
||||
info!("Metrics endpoint: http://{}/metrics and /beobachten", addr);
|
||||
serve_listener(
|
||||
listener,
|
||||
stats,
|
||||
beobachten,
|
||||
shared_state,
|
||||
ip_tracker,
|
||||
tls_cache,
|
||||
config_rx,
|
||||
whitelist,
|
||||
)
|
||||
.await;
|
||||
serve_listener(listener, active_runtime).await;
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(error = %e, "Failed to bind metrics on {}", addr);
|
||||
@@ -117,50 +100,14 @@ pub async fn serve(
|
||||
warn!("Metrics listener is unavailable on both IPv4 and IPv6");
|
||||
}
|
||||
(Some(listener), None) | (None, Some(listener)) => {
|
||||
serve_listener(
|
||||
listener,
|
||||
stats,
|
||||
beobachten,
|
||||
shared_state,
|
||||
ip_tracker,
|
||||
tls_cache,
|
||||
config_rx,
|
||||
whitelist,
|
||||
)
|
||||
.await;
|
||||
serve_listener(listener, active_runtime).await;
|
||||
}
|
||||
(Some(listener4), Some(listener6)) => {
|
||||
let stats_v6 = stats.clone();
|
||||
let beobachten_v6 = beobachten.clone();
|
||||
let shared_state_v6 = shared_state.clone();
|
||||
let ip_tracker_v6 = ip_tracker.clone();
|
||||
let tls_cache_v6 = tls_cache.clone();
|
||||
let config_rx_v6 = config_rx.clone();
|
||||
let whitelist_v6 = whitelist.clone();
|
||||
let active_runtime_v6 = active_runtime.clone();
|
||||
tokio::spawn(async move {
|
||||
serve_listener(
|
||||
listener6,
|
||||
stats_v6,
|
||||
beobachten_v6,
|
||||
shared_state_v6,
|
||||
ip_tracker_v6,
|
||||
tls_cache_v6,
|
||||
config_rx_v6,
|
||||
whitelist_v6,
|
||||
)
|
||||
.await;
|
||||
serve_listener(listener6, active_runtime_v6).await;
|
||||
});
|
||||
serve_listener(
|
||||
listener4,
|
||||
stats,
|
||||
beobachten,
|
||||
shared_state,
|
||||
ip_tracker,
|
||||
tls_cache,
|
||||
config_rx,
|
||||
whitelist,
|
||||
)
|
||||
.await;
|
||||
serve_listener(listener4, active_runtime).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -180,16 +127,7 @@ fn bind_metrics_listener(
|
||||
TcpListener::from_std(socket.into())
|
||||
}
|
||||
|
||||
async fn serve_listener(
|
||||
listener: TcpListener,
|
||||
stats: Arc<Stats>,
|
||||
beobachten: Arc<BeobachtenStore>,
|
||||
shared_state: Arc<ProxySharedState>,
|
||||
ip_tracker: Arc<UserIpTracker>,
|
||||
tls_cache: Option<Arc<TlsFrontCache>>,
|
||||
config_rx: tokio::sync::watch::Receiver<Arc<ProxyConfig>>,
|
||||
whitelist: Arc<Vec<IpNetwork>>,
|
||||
) {
|
||||
async fn serve_listener(listener: TcpListener, active_runtime: Arc<ArcSwap<RuntimeGeneration>>) {
|
||||
let connection_permits = Arc::new(Semaphore::new(METRICS_MAX_CONTROL_CONNECTIONS));
|
||||
|
||||
loop {
|
||||
@@ -201,7 +139,15 @@ async fn serve_listener(
|
||||
}
|
||||
};
|
||||
|
||||
if !whitelist.is_empty() && !whitelist.iter().any(|net| net.contains(peer.ip())) {
|
||||
let runtime = active_runtime.load_full();
|
||||
let config = runtime.config();
|
||||
if !config.server.metrics_whitelist.is_empty()
|
||||
&& !config
|
||||
.server
|
||||
.metrics_whitelist
|
||||
.iter()
|
||||
.any(|net| net.contains(peer.ip()))
|
||||
{
|
||||
debug!(peer = %peer, "Metrics request denied by whitelist");
|
||||
continue;
|
||||
}
|
||||
@@ -218,21 +164,17 @@ async fn serve_listener(
|
||||
}
|
||||
};
|
||||
|
||||
let stats = stats.clone();
|
||||
let beobachten = beobachten.clone();
|
||||
let shared_state = shared_state.clone();
|
||||
let ip_tracker = ip_tracker.clone();
|
||||
let tls_cache = tls_cache.clone();
|
||||
let config_rx_conn = config_rx.clone();
|
||||
let active_runtime = active_runtime.clone();
|
||||
tokio::spawn(async move {
|
||||
let _connection_permit = connection_permit;
|
||||
let svc = service_fn(move |req| {
|
||||
let stats = stats.clone();
|
||||
let beobachten = beobachten.clone();
|
||||
let shared_state = shared_state.clone();
|
||||
let ip_tracker = ip_tracker.clone();
|
||||
let tls_cache = tls_cache.clone();
|
||||
let config = config_rx_conn.borrow().clone();
|
||||
let runtime = active_runtime.load_full();
|
||||
let stats = runtime.stats.clone();
|
||||
let beobachten = runtime.beobachten.clone();
|
||||
let shared_state = runtime.proxy_shared.clone();
|
||||
let ip_tracker = runtime.ip_tracker.clone();
|
||||
let tls_cache = runtime.tls_cache.clone();
|
||||
let config = runtime.config();
|
||||
async move {
|
||||
handle(
|
||||
req,
|
||||
@@ -595,6 +537,81 @@ async fn render_metrics(
|
||||
"telemt_buffer_pool_buffers_total{{kind=\"in_use\"}} {}",
|
||||
stats.get_buffer_pool_in_use_gauge()
|
||||
);
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"# HELP telemt_buffer_pool_events_total Buffer-pool allocation lifecycle events"
|
||||
);
|
||||
let _ = writeln!(out, "# TYPE telemt_buffer_pool_events_total counter");
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_buffer_pool_events_total{{event=\"replaced_nonstandard\"}} {}",
|
||||
stats.get_buffer_pool_replaced_nonstandard_total()
|
||||
);
|
||||
|
||||
let direct_budget = shared_state.direct_buffer_budget.snapshot();
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"# HELP telemt_direct_relay_buffer_budget_bytes Direct relay copy-buffer budget and memory inputs"
|
||||
);
|
||||
let _ = writeln!(out, "# TYPE telemt_direct_relay_buffer_budget_bytes gauge");
|
||||
for (kind, value) in [
|
||||
("hard_limit", direct_budget.hard_limit_bytes),
|
||||
("target", direct_budget.target_bytes),
|
||||
("reserved", direct_budget.reserved_bytes),
|
||||
("memory_total", direct_budget.memory_total_bytes),
|
||||
("memory_available", direct_budget.memory_available_bytes),
|
||||
("process_rss", direct_budget.process_rss_bytes),
|
||||
] {
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_direct_relay_buffer_budget_bytes{{kind=\"{}\"}} {}",
|
||||
kind, value
|
||||
);
|
||||
}
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"# HELP telemt_direct_relay_buffer_budget_events_total Direct relay buffer-budget lifecycle events"
|
||||
);
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"# TYPE telemt_direct_relay_buffer_budget_events_total counter"
|
||||
);
|
||||
for (result, value) in [
|
||||
("promotion", direct_budget.promotion_total),
|
||||
("promotion_denied", direct_budget.promotion_denied_total),
|
||||
("minimum_fallback", direct_budget.minimum_fallback_total),
|
||||
("admission_rejected", direct_budget.admission_rejected_total),
|
||||
("quiet_demotion", direct_budget.quiet_demotion_total),
|
||||
(
|
||||
"write_pressure_demotion",
|
||||
direct_budget.write_pressure_demotion_total,
|
||||
),
|
||||
(
|
||||
"global_pressure_demotion",
|
||||
direct_budget.global_pressure_demotion_total,
|
||||
),
|
||||
] {
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_direct_relay_buffer_budget_events_total{{result=\"{}\"}} {}",
|
||||
result, value
|
||||
);
|
||||
}
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"# HELP telemt_direct_relay_buffer_sessions Current Direct relay sessions by adaptive tier"
|
||||
);
|
||||
let _ = writeln!(out, "# TYPE telemt_direct_relay_buffer_sessions gauge");
|
||||
for (tier, value) in ["base", "tier1", "tier2", "tier3"]
|
||||
.into_iter()
|
||||
.zip(direct_budget.tier_sessions)
|
||||
{
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_direct_relay_buffer_sessions{{tier=\"{}\"}} {}",
|
||||
tier, value
|
||||
);
|
||||
}
|
||||
|
||||
let _ = writeln!(
|
||||
out,
|
||||
@@ -2435,6 +2452,82 @@ async fn render_metrics(
|
||||
}
|
||||
);
|
||||
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"# HELP telemt_me_writer_byte_budget_limit_bytes Configured resident-memory budget per ME writer"
|
||||
);
|
||||
let _ = writeln!(out, "# TYPE telemt_me_writer_byte_budget_limit_bytes gauge");
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_me_writer_byte_budget_limit_bytes {}",
|
||||
if me_allows_normal {
|
||||
stats.get_me_writer_byte_budget_limit_bytes_gauge()
|
||||
} else {
|
||||
0
|
||||
}
|
||||
);
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"# HELP telemt_me_writer_byte_budget_reserved_bytes Aggregate ME writer memory reservations by lifecycle state"
|
||||
);
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"# TYPE telemt_me_writer_byte_budget_reserved_bytes gauge"
|
||||
);
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_me_writer_byte_budget_reserved_bytes{{state=\"queued\"}} {}",
|
||||
if me_allows_normal {
|
||||
stats.get_me_writer_byte_budget_queued_bytes_gauge()
|
||||
} else {
|
||||
0
|
||||
}
|
||||
);
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_me_writer_byte_budget_reserved_bytes{{state=\"inflight\"}} {}",
|
||||
if me_allows_normal {
|
||||
stats.get_me_writer_byte_budget_inflight_bytes_gauge()
|
||||
} else {
|
||||
0
|
||||
}
|
||||
);
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"# HELP telemt_me_writer_byte_budget_events_total ME writer byte-budget outcomes"
|
||||
);
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"# TYPE telemt_me_writer_byte_budget_events_total counter"
|
||||
);
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_me_writer_byte_budget_events_total{{result=\"wait\"}} {}",
|
||||
if me_allows_normal {
|
||||
stats.get_me_writer_byte_budget_wait_total()
|
||||
} else {
|
||||
0
|
||||
}
|
||||
);
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_me_writer_byte_budget_events_total{{result=\"timeout\"}} {}",
|
||||
if me_allows_normal {
|
||||
stats.get_me_writer_byte_budget_timeout_total()
|
||||
} else {
|
||||
0
|
||||
}
|
||||
);
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_me_writer_byte_budget_events_total{{result=\"oversize\"}} {}",
|
||||
if me_allows_normal {
|
||||
stats.get_me_writer_byte_budget_oversize_total()
|
||||
} else {
|
||||
0
|
||||
}
|
||||
);
|
||||
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"# HELP telemt_me_writer_pick_total ME writer-pick outcomes by mode and result"
|
||||
|
||||
@@ -8,6 +8,33 @@ use crate::error::{ProxyError, Result};
|
||||
|
||||
type OverrideMap = HashMap<(String, u16), IpAddr>;
|
||||
|
||||
/// Immutable DNS override snapshot owned by one runtime generation.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct DnsOverrides {
|
||||
entries: std::sync::Arc<OverrideMap>,
|
||||
}
|
||||
|
||||
impl DnsOverrides {
|
||||
/// Parses a validated generation-local override snapshot.
|
||||
pub fn from_entries(entries: &[String]) -> Result<Self> {
|
||||
Ok(Self {
|
||||
entries: std::sync::Arc::new(parse_entries(entries)?),
|
||||
})
|
||||
}
|
||||
|
||||
/// Resolves a generation-local hostname override.
|
||||
pub fn resolve(&self, host: &str, port: u16) -> Option<IpAddr> {
|
||||
self.entries
|
||||
.get(&(host.to_ascii_lowercase(), port))
|
||||
.copied()
|
||||
}
|
||||
|
||||
/// Resolves a generation-local override as a socket address.
|
||||
pub fn resolve_socket_addr(&self, host: &str, port: u16) -> Option<SocketAddr> {
|
||||
self.resolve(host, port).map(|ip| SocketAddr::new(ip, port))
|
||||
}
|
||||
}
|
||||
|
||||
static DNS_OVERRIDES: OnceLock<RwLock<OverrideMap>> = OnceLock::new();
|
||||
|
||||
fn overrides_store() -> &'static RwLock<OverrideMap> {
|
||||
@@ -180,6 +207,22 @@ mod tests {
|
||||
assert_eq!(resolved, Some("127.0.0.1".parse().unwrap()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn generation_snapshots_do_not_observe_each_other() {
|
||||
let first = DnsOverrides::from_entries(&["example.com:443:127.0.0.1".to_string()]).unwrap();
|
||||
let second =
|
||||
DnsOverrides::from_entries(&["example.com:443:127.0.0.2".to_string()]).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
first.resolve("example.com", 443),
|
||||
Some("127.0.0.1".parse().unwrap())
|
||||
);
|
||||
assert_eq!(
|
||||
second.resolve("example.com", 443),
|
||||
Some("127.0.0.2".parse().unwrap())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_host_port_parses_supported_shapes() {
|
||||
assert_eq!(
|
||||
|
||||
@@ -1,7 +1,4 @@
|
||||
#![allow(dead_code)]
|
||||
|
||||
// Adaptive buffer policy is staged and retained for deterministic rollout.
|
||||
// Keep definitions compiled for compatibility and security test scaffolding.
|
||||
// Adaptive buffer policy shared by active Direct relay sessions.
|
||||
|
||||
use dashmap::DashMap;
|
||||
use std::cmp::max;
|
||||
@@ -13,29 +10,42 @@ const PROFILE_TTL: Duration = Duration::from_secs(300);
|
||||
const THROUGHPUT_UP_BPS: f64 = 8_000_000.0;
|
||||
const THROUGHPUT_DOWN_BPS: f64 = 2_000_000.0;
|
||||
const RATIO_CONFIRM_THRESHOLD: f64 = 1.12;
|
||||
const TIER1_HOLD_TICKS: u32 = 8;
|
||||
const TIER2_HOLD_TICKS: u32 = 4;
|
||||
const QUIET_DEMOTE_TICKS: u32 = 480;
|
||||
const HARD_COOLDOWN_TICKS: u32 = 20;
|
||||
const TIER1_HOLD: Duration = Duration::from_secs(2);
|
||||
const TIER2_HOLD: Duration = Duration::from_secs(1);
|
||||
const QUIET_DEMOTE: Duration = Duration::from_secs(120);
|
||||
const HARD_COOLDOWN: Duration = Duration::from_secs(5);
|
||||
const SUSTAINED_PRESSURE_DEMOTE: Duration = Duration::from_secs(30);
|
||||
const PRESSURE_DEMOTE_COOLDOWN: Duration = Duration::from_secs(60);
|
||||
const HARD_PENDING_THRESHOLD: u32 = 3;
|
||||
const HARD_PARTIAL_RATIO_THRESHOLD: f64 = 0.25;
|
||||
#[cfg(test)]
|
||||
const DIRECT_C2S_CAP_BYTES: usize = 128 * 1024;
|
||||
#[cfg(test)]
|
||||
const DIRECT_S2C_CAP_BYTES: usize = 512 * 1024;
|
||||
#[cfg(test)]
|
||||
const ME_FRAMES_CAP: usize = 96;
|
||||
#[cfg(test)]
|
||||
const ME_BYTES_CAP: usize = 384 * 1024;
|
||||
#[cfg(test)]
|
||||
const ME_DELAY_MIN_US: u64 = 150;
|
||||
const MAX_USER_PROFILES_ENTRIES: usize = 50_000;
|
||||
const MAX_USER_KEY_BYTES: usize = 512;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
||||
/// Per-session Direct copy-buffer capacity tier.
|
||||
pub enum AdaptiveTier {
|
||||
/// Conservative baseline capacity.
|
||||
Base = 0,
|
||||
/// First throughput promotion.
|
||||
Tier1 = 1,
|
||||
/// Sustained bidirectional pressure promotion.
|
||||
Tier2 = 2,
|
||||
/// Configured per-direction ceilings.
|
||||
Tier3 = 3,
|
||||
}
|
||||
|
||||
impl AdaptiveTier {
|
||||
/// Returns the next larger tier, saturating at `Tier3`.
|
||||
pub fn promote(self) -> Self {
|
||||
match self {
|
||||
Self::Base => Self::Tier1,
|
||||
@@ -45,6 +55,7 @@ impl AdaptiveTier {
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the next smaller tier, saturating at `Base`.
|
||||
pub fn demote(self) -> Self {
|
||||
match self {
|
||||
Self::Base => Self::Base,
|
||||
@@ -54,6 +65,7 @@ impl AdaptiveTier {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn ratio(self) -> (usize, usize) {
|
||||
match self {
|
||||
Self::Base => (1, 1),
|
||||
@@ -63,49 +75,70 @@ impl AdaptiveTier {
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the stable numeric tier used by bounded metrics.
|
||||
pub fn as_u8(self) -> u8 {
|
||||
self as u8
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
/// Signal that caused an accepted adaptive tier transition.
|
||||
pub enum TierTransitionReason {
|
||||
/// Sustained throughput and directional ratio confirmation.
|
||||
SoftConfirmed,
|
||||
/// Short pending or partial-write pressure burst.
|
||||
HardPressure,
|
||||
/// Sustained low-throughput period.
|
||||
QuietDemotion,
|
||||
/// Sustained pending or partial-write pressure.
|
||||
SustainedWritePressure,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
/// Proposed transition emitted by the per-session controller.
|
||||
pub struct TierTransition {
|
||||
/// Tier active before the observation.
|
||||
pub from: AdaptiveTier,
|
||||
/// Tier requested after the observation.
|
||||
pub to: AdaptiveTier,
|
||||
/// Pressure or throughput condition that requested the transition.
|
||||
pub reason: TierTransitionReason,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
/// Directional byte and write-pressure deltas for one observation period.
|
||||
pub struct RelaySignalSample {
|
||||
/// Client-to-DC bytes copied during the period.
|
||||
pub c2s_bytes: u64,
|
||||
/// Bytes offered to DC-to-client writes during the period.
|
||||
pub s2c_requested_bytes: u64,
|
||||
/// Bytes accepted by DC-to-client writes during the period.
|
||||
pub s2c_written_bytes: u64,
|
||||
/// Successful DC-to-client write operations during the period.
|
||||
pub s2c_write_ops: u64,
|
||||
/// Partial DC-to-client write operations during the period.
|
||||
pub s2c_partial_writes: u64,
|
||||
/// Consecutive pending DC-to-client writes at sample time.
|
||||
pub s2c_consecutive_pending_writes: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
/// Stateful hysteresis controller for one active Direct session.
|
||||
pub struct SessionAdaptiveController {
|
||||
tier: AdaptiveTier,
|
||||
max_tier_seen: AdaptiveTier,
|
||||
throughput_ema_bps: f64,
|
||||
incoming_ema_bps: f64,
|
||||
outgoing_ema_bps: f64,
|
||||
tier1_hold_ticks: u32,
|
||||
tier2_hold_ticks: u32,
|
||||
quiet_ticks: u32,
|
||||
hard_cooldown_ticks: u32,
|
||||
tier1_hold: Duration,
|
||||
tier2_hold: Duration,
|
||||
quiet: Duration,
|
||||
hard_cooldown: Duration,
|
||||
sustained_pressure: Duration,
|
||||
}
|
||||
|
||||
impl SessionAdaptiveController {
|
||||
/// Creates a controller at the tier whose memory reservation was accepted.
|
||||
pub fn new(initial_tier: AdaptiveTier) -> Self {
|
||||
Self {
|
||||
tier: initial_tier,
|
||||
@@ -113,25 +146,33 @@ impl SessionAdaptiveController {
|
||||
throughput_ema_bps: 0.0,
|
||||
incoming_ema_bps: 0.0,
|
||||
outgoing_ema_bps: 0.0,
|
||||
tier1_hold_ticks: 0,
|
||||
tier2_hold_ticks: 0,
|
||||
quiet_ticks: 0,
|
||||
hard_cooldown_ticks: 0,
|
||||
tier1_hold: Duration::ZERO,
|
||||
tier2_hold: Duration::ZERO,
|
||||
quiet: Duration::ZERO,
|
||||
hard_cooldown: Duration::ZERO,
|
||||
sustained_pressure: Duration::ZERO,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the highest tier proposed by controller observations.
|
||||
#[allow(dead_code)]
|
||||
pub fn max_tier_seen(&self) -> AdaptiveTier {
|
||||
self.max_tier_seen
|
||||
}
|
||||
|
||||
/// Returns the controller's current logical tier.
|
||||
pub fn tier(&self) -> AdaptiveTier {
|
||||
self.tier
|
||||
}
|
||||
|
||||
/// Observes one period and returns at most one hysteresis-controlled transition.
|
||||
pub fn observe(&mut self, sample: RelaySignalSample, tick_secs: f64) -> Option<TierTransition> {
|
||||
if tick_secs <= f64::EPSILON {
|
||||
return None;
|
||||
}
|
||||
|
||||
if self.hard_cooldown_ticks > 0 {
|
||||
self.hard_cooldown_ticks -= 1;
|
||||
}
|
||||
let tick = Duration::from_secs_f64(tick_secs);
|
||||
self.hard_cooldown = self.hard_cooldown.saturating_sub(tick);
|
||||
|
||||
let c2s_bps = (sample.c2s_bytes as f64 * 8.0) / tick_secs;
|
||||
let incoming_bps = (sample.s2c_requested_bytes as f64 * 8.0) / tick_secs;
|
||||
@@ -144,9 +185,9 @@ impl SessionAdaptiveController {
|
||||
|
||||
let tier1_now = self.throughput_ema_bps >= THROUGHPUT_UP_BPS;
|
||||
if tier1_now {
|
||||
self.tier1_hold_ticks = self.tier1_hold_ticks.saturating_add(1);
|
||||
self.tier1_hold = self.tier1_hold.saturating_add(tick);
|
||||
} else {
|
||||
self.tier1_hold_ticks = 0;
|
||||
self.tier1_hold = Duration::ZERO;
|
||||
}
|
||||
|
||||
let ratio = if self.outgoing_ema_bps <= f64::EPSILON {
|
||||
@@ -156,9 +197,9 @@ impl SessionAdaptiveController {
|
||||
};
|
||||
let tier2_now = ratio >= RATIO_CONFIRM_THRESHOLD;
|
||||
if tier2_now {
|
||||
self.tier2_hold_ticks = self.tier2_hold_ticks.saturating_add(1);
|
||||
self.tier2_hold = self.tier2_hold.saturating_add(tick);
|
||||
} else {
|
||||
self.tier2_hold_ticks = 0;
|
||||
self.tier2_hold = Duration::ZERO;
|
||||
}
|
||||
|
||||
let partial_ratio = if sample.s2c_write_ops == 0 {
|
||||
@@ -169,24 +210,37 @@ impl SessionAdaptiveController {
|
||||
let hard_now = sample.s2c_consecutive_pending_writes >= HARD_PENDING_THRESHOLD
|
||||
|| partial_ratio >= HARD_PARTIAL_RATIO_THRESHOLD;
|
||||
|
||||
if hard_now && self.hard_cooldown_ticks == 0 {
|
||||
return self.promote(TierTransitionReason::HardPressure, HARD_COOLDOWN_TICKS);
|
||||
if hard_now {
|
||||
self.sustained_pressure = self.sustained_pressure.saturating_add(tick);
|
||||
if self.sustained_pressure >= SUSTAINED_PRESSURE_DEMOTE {
|
||||
self.sustained_pressure = Duration::ZERO;
|
||||
return self.demote(
|
||||
TierTransitionReason::SustainedWritePressure,
|
||||
PRESSURE_DEMOTE_COOLDOWN,
|
||||
);
|
||||
}
|
||||
} else {
|
||||
self.sustained_pressure = Duration::ZERO;
|
||||
}
|
||||
|
||||
if self.tier1_hold_ticks >= TIER1_HOLD_TICKS && self.tier2_hold_ticks >= TIER2_HOLD_TICKS {
|
||||
return self.promote(TierTransitionReason::SoftConfirmed, 0);
|
||||
if hard_now && self.hard_cooldown.is_zero() {
|
||||
return self.promote(TierTransitionReason::HardPressure, HARD_COOLDOWN);
|
||||
}
|
||||
|
||||
if self.tier1_hold >= TIER1_HOLD && self.tier2_hold >= TIER2_HOLD {
|
||||
return self.promote(TierTransitionReason::SoftConfirmed, Duration::ZERO);
|
||||
}
|
||||
|
||||
let demote_candidate =
|
||||
self.throughput_ema_bps < THROUGHPUT_DOWN_BPS && !tier2_now && !hard_now;
|
||||
if demote_candidate {
|
||||
self.quiet_ticks = self.quiet_ticks.saturating_add(1);
|
||||
if self.quiet_ticks >= QUIET_DEMOTE_TICKS {
|
||||
self.quiet_ticks = 0;
|
||||
return self.demote(TierTransitionReason::QuietDemotion);
|
||||
self.quiet = self.quiet.saturating_add(tick);
|
||||
if self.quiet >= QUIET_DEMOTE {
|
||||
self.quiet = Duration::ZERO;
|
||||
return self.demote(TierTransitionReason::QuietDemotion, Duration::ZERO);
|
||||
}
|
||||
} else {
|
||||
self.quiet_ticks = 0;
|
||||
self.quiet = Duration::ZERO;
|
||||
}
|
||||
|
||||
None
|
||||
@@ -195,7 +249,7 @@ impl SessionAdaptiveController {
|
||||
fn promote(
|
||||
&mut self,
|
||||
reason: TierTransitionReason,
|
||||
hard_cooldown_ticks: u32,
|
||||
hard_cooldown: Duration,
|
||||
) -> Option<TierTransition> {
|
||||
let from = self.tier;
|
||||
let to = from.promote();
|
||||
@@ -204,22 +258,27 @@ impl SessionAdaptiveController {
|
||||
}
|
||||
self.tier = to;
|
||||
self.max_tier_seen = max(self.max_tier_seen, to);
|
||||
self.hard_cooldown_ticks = hard_cooldown_ticks;
|
||||
self.tier1_hold_ticks = 0;
|
||||
self.tier2_hold_ticks = 0;
|
||||
self.quiet_ticks = 0;
|
||||
self.hard_cooldown = hard_cooldown;
|
||||
self.tier1_hold = Duration::ZERO;
|
||||
self.tier2_hold = Duration::ZERO;
|
||||
self.quiet = Duration::ZERO;
|
||||
Some(TierTransition { from, to, reason })
|
||||
}
|
||||
|
||||
fn demote(&mut self, reason: TierTransitionReason) -> Option<TierTransition> {
|
||||
fn demote(
|
||||
&mut self,
|
||||
reason: TierTransitionReason,
|
||||
hard_cooldown: Duration,
|
||||
) -> Option<TierTransition> {
|
||||
let from = self.tier;
|
||||
let to = from.demote();
|
||||
if from == to {
|
||||
return None;
|
||||
}
|
||||
self.tier = to;
|
||||
self.tier1_hold_ticks = 0;
|
||||
self.tier2_hold_ticks = 0;
|
||||
self.hard_cooldown = hard_cooldown;
|
||||
self.tier1_hold = Duration::ZERO;
|
||||
self.tier2_hold = Duration::ZERO;
|
||||
Some(TierTransition { from, to, reason })
|
||||
}
|
||||
}
|
||||
@@ -235,6 +294,8 @@ fn profiles() -> &'static DashMap<String, UserAdaptiveProfile> {
|
||||
USER_PROFILES.get_or_init(DashMap::new)
|
||||
}
|
||||
|
||||
/// Returns a fresh user's recent successful Direct tier, or `Base` when stale.
|
||||
#[allow(dead_code)]
|
||||
pub fn seed_tier_for_user(user: &str) -> AdaptiveTier {
|
||||
if user.len() > MAX_USER_KEY_BYTES {
|
||||
return AdaptiveTier::Base;
|
||||
@@ -253,6 +314,8 @@ pub fn seed_tier_for_user(user: &str) -> AdaptiveTier {
|
||||
AdaptiveTier::Base
|
||||
}
|
||||
|
||||
/// Records the highest successfully allocated tier for bounded session seeding.
|
||||
#[allow(dead_code)]
|
||||
pub fn record_user_tier(user: &str, tier: AdaptiveTier) {
|
||||
if user.len() > MAX_USER_KEY_BYTES {
|
||||
return;
|
||||
@@ -282,6 +345,8 @@ pub fn record_user_tier(user: &str, tier: AdaptiveTier) {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
/// Returns the legacy staged scaling policy retained by security fixtures.
|
||||
pub fn direct_copy_buffers_for_tier(
|
||||
tier: AdaptiveTier,
|
||||
base_c2s: usize,
|
||||
@@ -294,6 +359,32 @@ pub fn direct_copy_buffers_for_tier(
|
||||
)
|
||||
}
|
||||
|
||||
/// Maps an adaptive tier to independent capacities within configured ceilings.
|
||||
pub(crate) fn direct_copy_buffers_for_tier_with_ceilings(
|
||||
tier: AdaptiveTier,
|
||||
base_c2s: usize,
|
||||
base_s2c: usize,
|
||||
ceiling_c2s: usize,
|
||||
ceiling_s2c: usize,
|
||||
) -> (usize, usize) {
|
||||
(
|
||||
direct_direction_size(tier, base_c2s, ceiling_c2s),
|
||||
direct_direction_size(tier, base_s2c, ceiling_s2c),
|
||||
)
|
||||
}
|
||||
|
||||
fn direct_direction_size(tier: AdaptiveTier, base: usize, ceiling: usize) -> usize {
|
||||
let target = match tier {
|
||||
AdaptiveTier::Base => base,
|
||||
AdaptiveTier::Tier1 => ceiling / 4,
|
||||
AdaptiveTier::Tier2 => ceiling / 2,
|
||||
AdaptiveTier::Tier3 => ceiling,
|
||||
};
|
||||
target.max(base).min(ceiling.max(base)).max(1)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
/// Returns the staged Middle-End flush policy retained by security fixtures.
|
||||
pub fn me_flush_policy_for_tier(
|
||||
tier: AdaptiveTier,
|
||||
base_frames: usize,
|
||||
@@ -323,6 +414,7 @@ fn ema(prev: f64, value: f64) -> f64 {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn scale(base: usize, numerator: usize, denominator: usize, cap: usize) -> usize {
|
||||
let scaled = base
|
||||
.saturating_mul(numerator)
|
||||
@@ -338,6 +430,10 @@ mod adaptive_buffers_security_tests;
|
||||
#[path = "tests/adaptive_buffers_record_race_security_tests.rs"]
|
||||
mod adaptive_buffers_record_race_security_tests;
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "tests/adaptive_direct_budget_policy_tests.rs"]
|
||||
mod adaptive_direct_budget_policy_tests;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -396,7 +492,7 @@ mod tests {
|
||||
fn test_quiet_demotion_is_slow_and_stepwise() {
|
||||
let mut ctrl = SessionAdaptiveController::new(AdaptiveTier::Tier2);
|
||||
let mut demotion = None;
|
||||
for _ in 0..QUIET_DEMOTE_TICKS {
|
||||
for _ in 0..480 {
|
||||
demotion = ctrl.observe(sample(1, 1, 1, 1, 0, 0), 0.25);
|
||||
}
|
||||
|
||||
|
||||
568
src/proxy/direct_buffer_budget.rs
Normal file
568
src/proxy/direct_buffer_budget.rs
Normal file
@@ -0,0 +1,568 @@
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::sync::watch;
|
||||
|
||||
use crate::stats::Stats;
|
||||
use crate::stream::BufferPool;
|
||||
|
||||
use super::shared_state::ProxySharedState;
|
||||
|
||||
/// Accounting granularity for process-wide Direct copy-buffer reservations.
|
||||
pub(crate) const DIRECT_BUFFER_UNIT_BYTES: usize = 4 * 1024;
|
||||
/// Minimum client-to-DC copy-buffer capacity for one Direct session.
|
||||
pub(crate) const DIRECT_BASE_C2S_BYTES: usize = 4 * 1024;
|
||||
/// Minimum DC-to-client copy-buffer capacity for one Direct session.
|
||||
pub(crate) const DIRECT_BASE_S2C_BYTES: usize = 8 * 1024;
|
||||
|
||||
const AUTO_HARD_MIN_BYTES: usize = 64 * 1024 * 1024;
|
||||
const AUTO_HARD_MAX_BYTES: usize = 2 * 1024 * 1024 * 1024;
|
||||
const AUTO_HARD_FALLBACK_BYTES: usize = 512 * 1024 * 1024;
|
||||
const TARGET_FLOOR_MIN_BYTES: usize = 16 * 1024 * 1024;
|
||||
const CONTROL_INTERVAL: Duration = Duration::from_secs(1);
|
||||
const HEALTHY_RECOVERY_SAMPLES: u8 = 30;
|
||||
const BUFFER_POOL_TRIM_LOW_WATERMARK: usize = 64;
|
||||
const BUFFER_POOL_TRIM_HIGH_WATERMARK: usize = 128;
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
/// Lock-free observability snapshot of the Direct copy-buffer envelope.
|
||||
pub(crate) struct DirectBufferBudgetSnapshot {
|
||||
/// Absolute process-wide copy-buffer ceiling.
|
||||
pub(crate) hard_limit_bytes: u64,
|
||||
/// Current pressure-adjusted promotion target.
|
||||
pub(crate) target_bytes: u64,
|
||||
/// Bytes currently covered by active session leases.
|
||||
pub(crate) reserved_bytes: u64,
|
||||
/// Effective host or cgroup memory limit.
|
||||
pub(crate) memory_total_bytes: u64,
|
||||
/// Effective host or cgroup memory headroom.
|
||||
pub(crate) memory_available_bytes: u64,
|
||||
/// Current process resident set size.
|
||||
pub(crate) process_rss_bytes: u64,
|
||||
/// Successful tier growth reservations.
|
||||
pub(crate) promotion_total: u64,
|
||||
/// Tier growth attempts rejected by the adaptive target.
|
||||
pub(crate) promotion_denied_total: u64,
|
||||
/// Sessions admitted at minimum size above the adaptive target.
|
||||
pub(crate) minimum_fallback_total: u64,
|
||||
/// Sessions rejected by the absolute ceiling.
|
||||
pub(crate) admission_rejected_total: u64,
|
||||
/// Quiet-period tier reductions.
|
||||
pub(crate) quiet_demotion_total: u64,
|
||||
/// Sustained write-pressure tier reductions.
|
||||
pub(crate) write_pressure_demotion_total: u64,
|
||||
/// Process-wide pressure tier reductions.
|
||||
pub(crate) global_pressure_demotion_total: u64,
|
||||
/// Current sessions for Base through Tier3.
|
||||
pub(crate) tier_sessions: [u64; 4],
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
struct SystemMemorySample {
|
||||
total_bytes: u64,
|
||||
available_bytes: u64,
|
||||
process_rss_bytes: u64,
|
||||
}
|
||||
|
||||
/// Process-wide hard envelope and adaptive target for Direct copy buffers.
|
||||
pub(crate) struct DirectBufferBudget {
|
||||
hard_limit_bytes: u64,
|
||||
target_bytes: AtomicU64,
|
||||
reserved_bytes: AtomicU64,
|
||||
pressure_generation: AtomicU64,
|
||||
pressure_tx: watch::Sender<u64>,
|
||||
memory_total_bytes: AtomicU64,
|
||||
memory_available_bytes: AtomicU64,
|
||||
process_rss_bytes: AtomicU64,
|
||||
promotion_total: AtomicU64,
|
||||
promotion_denied_total: AtomicU64,
|
||||
minimum_fallback_total: AtomicU64,
|
||||
admission_rejected_total: AtomicU64,
|
||||
quiet_demotion_total: AtomicU64,
|
||||
write_pressure_demotion_total: AtomicU64,
|
||||
global_pressure_demotion_total: AtomicU64,
|
||||
tier_sessions: [AtomicU64; 4],
|
||||
}
|
||||
|
||||
impl DirectBufferBudget {
|
||||
/// Creates an envelope with a fixed absolute ceiling.
|
||||
pub(crate) fn new(hard_limit_bytes: usize) -> Arc<Self> {
|
||||
let hard_limit_bytes = align_down(hard_limit_bytes.max(DIRECT_BUFFER_UNIT_BYTES)) as u64;
|
||||
let (pressure_tx, _) = watch::channel(0);
|
||||
Arc::new(Self {
|
||||
hard_limit_bytes,
|
||||
target_bytes: AtomicU64::new(hard_limit_bytes),
|
||||
reserved_bytes: AtomicU64::new(0),
|
||||
pressure_generation: AtomicU64::new(0),
|
||||
pressure_tx,
|
||||
memory_total_bytes: AtomicU64::new(0),
|
||||
memory_available_bytes: AtomicU64::new(0),
|
||||
process_rss_bytes: AtomicU64::new(0),
|
||||
promotion_total: AtomicU64::new(0),
|
||||
promotion_denied_total: AtomicU64::new(0),
|
||||
minimum_fallback_total: AtomicU64::new(0),
|
||||
admission_rejected_total: AtomicU64::new(0),
|
||||
quiet_demotion_total: AtomicU64::new(0),
|
||||
write_pressure_demotion_total: AtomicU64::new(0),
|
||||
global_pressure_demotion_total: AtomicU64::new(0),
|
||||
tier_sessions: std::array::from_fn(|_| AtomicU64::new(0)),
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns the current pressure-adjusted reservation target.
|
||||
pub(crate) fn target_bytes(&self) -> usize {
|
||||
self.target_bytes.load(Ordering::Relaxed) as usize
|
||||
}
|
||||
|
||||
/// Subscribes to target reductions that require prompt session demotion.
|
||||
pub(crate) fn subscribe_pressure(&self) -> watch::Receiver<u64> {
|
||||
self.pressure_tx.subscribe()
|
||||
}
|
||||
|
||||
/// Reserves bytes against either the adaptive target or the absolute ceiling.
|
||||
pub(crate) fn try_reserve(
|
||||
self: &Arc<Self>,
|
||||
bytes: usize,
|
||||
allow_above_target: bool,
|
||||
) -> Option<DirectBufferLease> {
|
||||
let bytes = align_up(bytes) as u64;
|
||||
let limit = if allow_above_target {
|
||||
self.hard_limit_bytes
|
||||
} else {
|
||||
self.target_bytes
|
||||
.load(Ordering::Relaxed)
|
||||
.min(self.hard_limit_bytes)
|
||||
};
|
||||
if !self.try_add_reserved(bytes, limit) {
|
||||
return None;
|
||||
}
|
||||
self.tier_sessions[0].fetch_add(1, Ordering::Relaxed);
|
||||
Some(DirectBufferLease {
|
||||
budget: Arc::clone(self),
|
||||
reserved_bytes: bytes,
|
||||
tier: 0,
|
||||
})
|
||||
}
|
||||
|
||||
fn try_add_reserved(&self, bytes: u64, limit: u64) -> bool {
|
||||
let mut current = self.reserved_bytes.load(Ordering::Acquire);
|
||||
loop {
|
||||
if bytes > limit.saturating_sub(current) {
|
||||
return false;
|
||||
}
|
||||
match self.reserved_bytes.compare_exchange_weak(
|
||||
current,
|
||||
current + bytes,
|
||||
Ordering::AcqRel,
|
||||
Ordering::Acquire,
|
||||
) {
|
||||
Ok(_) => return true,
|
||||
Err(observed) => current = observed,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn target_floor_bytes(&self) -> u64 {
|
||||
(self.hard_limit_bytes / 8)
|
||||
.max(TARGET_FLOOR_MIN_BYTES as u64)
|
||||
.min(self.hard_limit_bytes)
|
||||
}
|
||||
|
||||
fn set_target_bytes(&self, target: u64) {
|
||||
let target =
|
||||
align_down(target.clamp(self.target_floor_bytes(), self.hard_limit_bytes) as usize)
|
||||
as u64;
|
||||
let previous = self.target_bytes.swap(target, Ordering::AcqRel);
|
||||
if target < previous {
|
||||
let generation = self
|
||||
.pressure_generation
|
||||
.fetch_add(1, Ordering::AcqRel)
|
||||
.wrapping_add(1);
|
||||
self.pressure_tx.send_replace(generation);
|
||||
}
|
||||
}
|
||||
|
||||
fn update_system_sample(&self, sample: SystemMemorySample) {
|
||||
self.memory_total_bytes
|
||||
.store(sample.total_bytes, Ordering::Relaxed);
|
||||
self.memory_available_bytes
|
||||
.store(sample.available_bytes, Ordering::Relaxed);
|
||||
self.process_rss_bytes
|
||||
.store(sample.process_rss_bytes, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Records a session that had to bypass the adaptive target at minimum size.
|
||||
pub(crate) fn increment_minimum_fallback(&self) {
|
||||
self.minimum_fallback_total.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Records a session rejected because the absolute ceiling was exhausted.
|
||||
pub(crate) fn increment_admission_rejected(&self) {
|
||||
self.admission_rejected_total
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Records a tier reduction after sustained low throughput.
|
||||
pub(crate) fn increment_quiet_demotion(&self) {
|
||||
self.quiet_demotion_total.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Records a tier reduction after sustained partial or pending writes.
|
||||
pub(crate) fn increment_write_pressure_demotion(&self) {
|
||||
self.write_pressure_demotion_total
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Records a tier reduction requested by the process-wide controller.
|
||||
pub(crate) fn increment_global_pressure_demotion(&self) {
|
||||
self.global_pressure_demotion_total
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Captures all bounded metrics without allocating or locking.
|
||||
pub(crate) fn snapshot(&self) -> DirectBufferBudgetSnapshot {
|
||||
DirectBufferBudgetSnapshot {
|
||||
hard_limit_bytes: self.hard_limit_bytes,
|
||||
target_bytes: self.target_bytes.load(Ordering::Relaxed),
|
||||
reserved_bytes: self.reserved_bytes.load(Ordering::Relaxed),
|
||||
memory_total_bytes: self.memory_total_bytes.load(Ordering::Relaxed),
|
||||
memory_available_bytes: self.memory_available_bytes.load(Ordering::Relaxed),
|
||||
process_rss_bytes: self.process_rss_bytes.load(Ordering::Relaxed),
|
||||
promotion_total: self.promotion_total.load(Ordering::Relaxed),
|
||||
promotion_denied_total: self.promotion_denied_total.load(Ordering::Relaxed),
|
||||
minimum_fallback_total: self.minimum_fallback_total.load(Ordering::Relaxed),
|
||||
admission_rejected_total: self.admission_rejected_total.load(Ordering::Relaxed),
|
||||
quiet_demotion_total: self.quiet_demotion_total.load(Ordering::Relaxed),
|
||||
write_pressure_demotion_total: self
|
||||
.write_pressure_demotion_total
|
||||
.load(Ordering::Relaxed),
|
||||
global_pressure_demotion_total: self
|
||||
.global_pressure_demotion_total
|
||||
.load(Ordering::Relaxed),
|
||||
tier_sessions: std::array::from_fn(|index| {
|
||||
self.tier_sessions[index].load(Ordering::Relaxed)
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the conservative ceiling used when memory discovery is unavailable.
|
||||
pub(crate) fn fallback_direct_buffer_hard_limit() -> usize {
|
||||
AUTO_HARD_FALLBACK_BYTES
|
||||
}
|
||||
|
||||
/// RAII ownership of all copy-buffer bytes retained by one Direct session.
|
||||
pub(crate) struct DirectBufferLease {
|
||||
budget: Arc<DirectBufferBudget>,
|
||||
reserved_bytes: u64,
|
||||
tier: usize,
|
||||
}
|
||||
|
||||
impl DirectBufferLease {
|
||||
/// Returns the currently covered allocation rounded to accounting units.
|
||||
pub(crate) fn reserved_bytes(&self) -> usize {
|
||||
self.reserved_bytes as usize
|
||||
}
|
||||
|
||||
/// Attempts to cover a larger tier before its buffers are resized.
|
||||
pub(crate) fn try_grow_to(&mut self, bytes: usize) -> bool {
|
||||
let bytes = align_up(bytes) as u64;
|
||||
if bytes <= self.reserved_bytes {
|
||||
return true;
|
||||
}
|
||||
let delta = bytes - self.reserved_bytes;
|
||||
let limit = self
|
||||
.budget
|
||||
.target_bytes
|
||||
.load(Ordering::Relaxed)
|
||||
.min(self.budget.hard_limit_bytes);
|
||||
if !self.budget.try_add_reserved(delta, limit) {
|
||||
self.budget
|
||||
.promotion_denied_total
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
return false;
|
||||
}
|
||||
self.reserved_bytes = bytes;
|
||||
self.budget.promotion_total.fetch_add(1, Ordering::Relaxed);
|
||||
true
|
||||
}
|
||||
|
||||
/// Releases bytes only after both directional buffers report smaller coverage.
|
||||
pub(crate) fn shrink_to(&mut self, bytes: usize) {
|
||||
let bytes = align_up(bytes) as u64;
|
||||
if bytes >= self.reserved_bytes {
|
||||
return;
|
||||
}
|
||||
let released = self.reserved_bytes - bytes;
|
||||
self.reserved_bytes = bytes;
|
||||
self.budget
|
||||
.reserved_bytes
|
||||
.fetch_sub(released, Ordering::AcqRel);
|
||||
}
|
||||
|
||||
/// Updates bounded per-tier session gauges for an accepted transition.
|
||||
pub(crate) fn set_tier(&mut self, tier: usize) {
|
||||
let tier = tier.min(self.budget.tier_sessions.len() - 1);
|
||||
if tier == self.tier {
|
||||
return;
|
||||
}
|
||||
decrement_saturating(&self.budget.tier_sessions[self.tier]);
|
||||
self.budget.tier_sessions[tier].fetch_add(1, Ordering::Relaxed);
|
||||
self.tier = tier;
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for DirectBufferLease {
|
||||
fn drop(&mut self) {
|
||||
self.budget
|
||||
.reserved_bytes
|
||||
.fetch_sub(self.reserved_bytes, Ordering::AcqRel);
|
||||
decrement_saturating(&self.budget.tier_sessions[self.tier]);
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolves the startup hard ceiling from config, cgroup, and host memory.
|
||||
pub(crate) async fn resolve_direct_buffer_hard_limit(configured: usize) -> usize {
|
||||
if configured != 0 {
|
||||
return align_down(configured);
|
||||
}
|
||||
let sample = read_system_memory_sample().await;
|
||||
if sample.total_bytes == 0 {
|
||||
return AUTO_HARD_FALLBACK_BYTES;
|
||||
}
|
||||
let derived = (sample.total_bytes / 4)
|
||||
.clamp(AUTO_HARD_MIN_BYTES as u64, AUTO_HARD_MAX_BYTES as u64)
|
||||
.min(sample.total_bytes);
|
||||
align_down(derived as usize).max(DIRECT_BUFFER_UNIT_BYTES)
|
||||
}
|
||||
|
||||
/// Runs the control-plane loop for Direct budget and shared pool pressure.
|
||||
pub(crate) async fn run_direct_buffer_budget_controller(
|
||||
budget: Arc<DirectBufferBudget>,
|
||||
buffer_pool: Arc<BufferPool>,
|
||||
stats: Arc<Stats>,
|
||||
shared: Arc<ProxySharedState>,
|
||||
max_connections: u32,
|
||||
) {
|
||||
let mut interval = tokio::time::interval(CONTROL_INTERVAL);
|
||||
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
let mut healthy_streak = 0u8;
|
||||
let mut previous_denied = 0u64;
|
||||
let mut previous_fallback = 0u64;
|
||||
let mut previous_rejected = 0u64;
|
||||
let pool_trim_low = buffer_pool
|
||||
.max_buffers()
|
||||
.min(BUFFER_POOL_TRIM_LOW_WATERMARK);
|
||||
let pool_trim_high = buffer_pool
|
||||
.max_buffers()
|
||||
.min(BUFFER_POOL_TRIM_HIGH_WATERMARK);
|
||||
let mut pool_trim_armed = true;
|
||||
|
||||
loop {
|
||||
interval.tick().await;
|
||||
let sample = read_system_memory_sample().await;
|
||||
budget.update_system_sample(sample);
|
||||
|
||||
let snapshot = budget.snapshot();
|
||||
let denied_delta = snapshot
|
||||
.promotion_denied_total
|
||||
.saturating_sub(previous_denied);
|
||||
previous_denied = snapshot.promotion_denied_total;
|
||||
let fallback_delta = snapshot
|
||||
.minimum_fallback_total
|
||||
.saturating_sub(previous_fallback);
|
||||
previous_fallback = snapshot.minimum_fallback_total;
|
||||
let rejected_delta = snapshot
|
||||
.admission_rejected_total
|
||||
.saturating_sub(previous_rejected);
|
||||
previous_rejected = snapshot.admission_rejected_total;
|
||||
|
||||
let connection_pct = connection_fill_pct(stats.as_ref(), max_connections);
|
||||
let memory_available_pct = percentage(sample.available_bytes, sample.total_bytes);
|
||||
let target_utilization_pct = percentage(snapshot.reserved_bytes, snapshot.target_bytes);
|
||||
let pressure = shared.conntrack_pressure_active()
|
||||
|| connection_pct.is_some_and(|value| value >= 85)
|
||||
|| memory_available_pct.is_some_and(|value| value <= 15)
|
||||
|| target_utilization_pct.is_some_and(|value| value >= 90)
|
||||
|| denied_delta > 0
|
||||
|| fallback_delta > 0
|
||||
|| rejected_delta > 0;
|
||||
|
||||
if !pressure {
|
||||
pool_trim_armed = true;
|
||||
} else if pool_trim_armed && buffer_pool.pooled() > pool_trim_high {
|
||||
buffer_pool.trim_to(pool_trim_low);
|
||||
pool_trim_armed = false;
|
||||
}
|
||||
|
||||
let pool_snapshot = buffer_pool.stats();
|
||||
stats.set_buffer_pool_gauges(
|
||||
pool_snapshot.pooled,
|
||||
pool_snapshot.allocated,
|
||||
pool_snapshot.allocated.saturating_sub(pool_snapshot.pooled),
|
||||
);
|
||||
stats.set_buffer_pool_replaced_nonstandard_total(pool_snapshot.replaced_nonstandard);
|
||||
|
||||
let headroom_target = if sample.total_bytes == 0 {
|
||||
snapshot.hard_limit_bytes
|
||||
} else {
|
||||
snapshot
|
||||
.reserved_bytes
|
||||
.saturating_add(sample.available_bytes / 4)
|
||||
.min(snapshot.hard_limit_bytes)
|
||||
};
|
||||
|
||||
if pressure {
|
||||
healthy_streak = 0;
|
||||
let reduced = snapshot.target_bytes.saturating_mul(3) / 4;
|
||||
budget.set_target_bytes(reduced.min(headroom_target));
|
||||
continue;
|
||||
}
|
||||
|
||||
let healthy = memory_available_pct.is_none_or(|value| value >= 30)
|
||||
&& connection_pct.is_none_or(|value| value <= 70);
|
||||
if !healthy {
|
||||
healthy_streak = 0;
|
||||
if headroom_target < snapshot.target_bytes {
|
||||
budget.set_target_bytes(headroom_target);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
healthy_streak = healthy_streak.saturating_add(1);
|
||||
if healthy_streak >= HEALTHY_RECOVERY_SAMPLES {
|
||||
healthy_streak = 0;
|
||||
let increment = (snapshot.target_bytes / 16).max(4 * 1024 * 1024);
|
||||
budget.set_target_bytes(
|
||||
snapshot
|
||||
.target_bytes
|
||||
.saturating_add(increment)
|
||||
.min(headroom_target),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn connection_fill_pct(stats: &Stats, max_connections: u32) -> Option<u8> {
|
||||
if max_connections == 0 {
|
||||
return None;
|
||||
}
|
||||
Some(
|
||||
((stats.get_current_connections_total().saturating_mul(100)) / u64::from(max_connections))
|
||||
.min(100) as u8,
|
||||
)
|
||||
}
|
||||
|
||||
fn percentage(value: u64, total: u64) -> Option<u8> {
|
||||
if total == 0 {
|
||||
return None;
|
||||
}
|
||||
Some(((value.saturating_mul(100)) / total).min(100) as u8)
|
||||
}
|
||||
|
||||
async fn read_system_memory_sample() -> SystemMemorySample {
|
||||
#[cfg(target_os = "linux")]
|
||||
{
|
||||
let meminfo = tokio::fs::read_to_string("/proc/meminfo")
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
let status = tokio::fs::read_to_string("/proc/self/status")
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
let host_total = parse_kib_field(&meminfo, "MemTotal:");
|
||||
let host_available = parse_kib_field(&meminfo, "MemAvailable:");
|
||||
let process_rss = parse_kib_field(&status, "VmRSS:");
|
||||
|
||||
let cgroup_v2_max = read_cgroup_limit("/sys/fs/cgroup/memory.max").await;
|
||||
let cgroup_v2_current = read_u64_file("/sys/fs/cgroup/memory.current").await;
|
||||
let cgroup_v1_max = read_cgroup_limit("/sys/fs/cgroup/memory/memory.limit_in_bytes").await;
|
||||
let cgroup_v1_current = read_u64_file("/sys/fs/cgroup/memory/memory.usage_in_bytes").await;
|
||||
let cgroup_max = cgroup_v2_max.or(cgroup_v1_max);
|
||||
let cgroup_current = cgroup_v2_current.or(cgroup_v1_current);
|
||||
|
||||
let total = match (host_total, cgroup_max) {
|
||||
(0, Some(limit)) => limit,
|
||||
(host, Some(limit)) => host.min(limit),
|
||||
(host, None) => host,
|
||||
};
|
||||
let cgroup_available = cgroup_max
|
||||
.zip(cgroup_current)
|
||||
.map(|(limit, current)| limit.saturating_sub(current));
|
||||
let available = match (host_available, cgroup_available) {
|
||||
(0, Some(value)) => value,
|
||||
(host, Some(value)) => host.min(value),
|
||||
(host, None) => host,
|
||||
};
|
||||
return SystemMemorySample {
|
||||
total_bytes: total,
|
||||
available_bytes: available,
|
||||
process_rss_bytes: process_rss,
|
||||
};
|
||||
}
|
||||
#[cfg(not(target_os = "linux"))]
|
||||
{
|
||||
SystemMemorySample::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
async fn read_cgroup_limit(path: &str) -> Option<u64> {
|
||||
let raw = tokio::fs::read_to_string(path).await.ok()?;
|
||||
let raw = raw.trim();
|
||||
if raw == "max" {
|
||||
return None;
|
||||
}
|
||||
let value = raw.parse::<u64>().ok()?;
|
||||
(value < (1u64 << 60)).then_some(value)
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
async fn read_u64_file(path: &str) -> Option<u64> {
|
||||
tokio::fs::read_to_string(path)
|
||||
.await
|
||||
.ok()?
|
||||
.trim()
|
||||
.parse()
|
||||
.ok()
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
fn parse_kib_field(raw: &str, key: &str) -> u64 {
|
||||
raw.lines()
|
||||
.find_map(|line| {
|
||||
let value = line.strip_prefix(key)?.split_whitespace().next()?;
|
||||
value.parse::<u64>().ok()
|
||||
})
|
||||
.unwrap_or(0)
|
||||
.saturating_mul(1024)
|
||||
}
|
||||
|
||||
fn align_up(bytes: usize) -> usize {
|
||||
bytes
|
||||
.div_ceil(DIRECT_BUFFER_UNIT_BYTES)
|
||||
.saturating_mul(DIRECT_BUFFER_UNIT_BYTES)
|
||||
}
|
||||
|
||||
fn align_down(bytes: usize) -> usize {
|
||||
bytes / DIRECT_BUFFER_UNIT_BYTES * DIRECT_BUFFER_UNIT_BYTES
|
||||
}
|
||||
|
||||
fn decrement_saturating(value: &AtomicU64) {
|
||||
let mut current = value.load(Ordering::Relaxed);
|
||||
while current != 0 {
|
||||
match value.compare_exchange_weak(
|
||||
current,
|
||||
current - 1,
|
||||
Ordering::Relaxed,
|
||||
Ordering::Relaxed,
|
||||
) {
|
||||
Ok(_) => return,
|
||||
Err(observed) => current = observed,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "tests/direct_buffer_budget_tests.rs"]
|
||||
mod tests;
|
||||
@@ -345,22 +345,22 @@ where
|
||||
} else {
|
||||
Duration::from_secs(1800)
|
||||
};
|
||||
let relay_result =
|
||||
crate::proxy::relay::relay_bidirectional_with_activity_timeout_lease_and_cancel(
|
||||
client_reader,
|
||||
client_writer,
|
||||
tg_reader,
|
||||
tg_writer,
|
||||
config.general.direct_relay_copy_buf_c2s_bytes,
|
||||
config.general.direct_relay_copy_buf_s2c_bytes,
|
||||
user,
|
||||
Arc::clone(&stats),
|
||||
config.access.user_data_quota.get(user).copied(),
|
||||
buffer_pool,
|
||||
traffic_lease,
|
||||
relay_activity_timeout,
|
||||
session_cancel.clone(),
|
||||
);
|
||||
let relay_result = crate::proxy::relay::relay_direct_adaptive(
|
||||
client_reader,
|
||||
client_writer,
|
||||
tg_reader,
|
||||
tg_writer,
|
||||
config.general.direct_relay_copy_buf_c2s_bytes,
|
||||
config.general.direct_relay_copy_buf_s2c_bytes,
|
||||
config.server.max_connections,
|
||||
user,
|
||||
Arc::clone(&stats),
|
||||
config.access.user_data_quota.get(user).copied(),
|
||||
traffic_lease,
|
||||
relay_activity_timeout,
|
||||
session_cancel.clone(),
|
||||
Arc::clone(&shared.direct_buffer_budget),
|
||||
);
|
||||
tokio::pin!(relay_result);
|
||||
let relay_result = loop {
|
||||
if let Some(cutover) =
|
||||
@@ -400,7 +400,6 @@ where
|
||||
Err(e) => debug!(user = %user, error = %e, "Direct relay ended with error"),
|
||||
}
|
||||
|
||||
buffer_pool_trim.trim_to(buffer_pool_trim.max_buffers().min(64));
|
||||
let pool_snapshot = buffer_pool_trim.stats();
|
||||
stats.set_buffer_pool_gauges(
|
||||
pool_snapshot.pooled,
|
||||
|
||||
@@ -959,6 +959,7 @@ fn configure_mask_backend_socket(stream: &TcpStream) {
|
||||
}
|
||||
|
||||
/// Handles a bad client by forwarding it to the configured mask target.
|
||||
#[cfg(test)]
|
||||
pub async fn handle_bad_client<R, W>(
|
||||
reader: R,
|
||||
writer: W,
|
||||
|
||||
@@ -149,6 +149,7 @@ where
|
||||
peer,
|
||||
translated_local_addr,
|
||||
payload,
|
||||
_permit,
|
||||
flags,
|
||||
effective_tag_array,
|
||||
)
|
||||
@@ -822,7 +823,6 @@ where
|
||||
|
||||
clear_relay_idle_candidate_in(shared.as_ref(), conn_id);
|
||||
me_pool.registry().unregister(conn_id).await;
|
||||
buffer_pool.trim_to(buffer_pool.max_buffers().min(64));
|
||||
let pool_snapshot = buffer_pool.stats();
|
||||
stats.set_buffer_pool_gauges(
|
||||
pool_snapshot.pooled,
|
||||
|
||||
@@ -60,6 +60,8 @@
|
||||
|
||||
pub mod adaptive_buffers;
|
||||
pub mod client;
|
||||
// Process-wide Direct relay copy-buffer ownership and pressure policy.
|
||||
pub(crate) mod direct_buffer_budget;
|
||||
pub mod direct_relay;
|
||||
pub mod handshake;
|
||||
pub mod masking;
|
||||
|
||||
@@ -84,8 +84,11 @@ fn watchdog_delta(current: u64, previous: u64) -> u64 {
|
||||
current.saturating_sub(previous)
|
||||
}
|
||||
|
||||
mod adaptive_copy;
|
||||
mod io;
|
||||
|
||||
pub(crate) use self::adaptive_copy::relay_direct_adaptive;
|
||||
|
||||
use self::io::{CombinedStream, SharedCounters, StatsIo, is_quota_io_error};
|
||||
#[cfg(test)]
|
||||
use self::io::{quota_adaptive_interval_bytes, should_immediate_quota_check};
|
||||
@@ -217,6 +220,7 @@ where
|
||||
.await
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub async fn relay_bidirectional_with_activity_timeout_lease_and_cancel<CR, CW, SR, SW>(
|
||||
client_reader: CR,
|
||||
client_writer: CW,
|
||||
|
||||
527
src/proxy/relay/adaptive_copy.rs
Normal file
527
src/proxy/relay/adaptive_copy.rs
Normal file
@@ -0,0 +1,527 @@
|
||||
use std::io;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
||||
use std::task::{Context, Poll};
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf, copy_buf};
|
||||
use tokio::time::Instant;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{debug, warn};
|
||||
|
||||
use crate::error::{ProxyError, Result};
|
||||
use crate::proxy::adaptive_buffers::{
|
||||
AdaptiveTier, RelaySignalSample, SessionAdaptiveController, TierTransitionReason,
|
||||
direct_copy_buffers_for_tier_with_ceilings,
|
||||
};
|
||||
use crate::proxy::direct_buffer_budget::{
|
||||
DIRECT_BASE_C2S_BYTES, DIRECT_BASE_S2C_BYTES, DirectBufferBudget, DirectBufferLease,
|
||||
};
|
||||
use crate::proxy::traffic_limiter::TrafficLease;
|
||||
use crate::stats::Stats;
|
||||
|
||||
use super::WATCHDOG_INTERVAL;
|
||||
use super::io::{SharedCounters, StatsIo, is_quota_io_error};
|
||||
use super::watchdog_delta;
|
||||
|
||||
mod write_pressure;
|
||||
|
||||
use self::write_pressure::WritePressureIo;
|
||||
|
||||
struct AdaptiveBufferState {
|
||||
desired_bytes: AtomicUsize,
|
||||
actual_bytes: AtomicUsize,
|
||||
}
|
||||
|
||||
impl AdaptiveBufferState {
|
||||
fn new(bytes: usize) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
desired_bytes: AtomicUsize::new(bytes.max(1)),
|
||||
actual_bytes: AtomicUsize::new(bytes.max(1)),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct AdaptiveBufReader<R> {
|
||||
inner: R,
|
||||
buffer: Box<[u8]>,
|
||||
pos: usize,
|
||||
cap: usize,
|
||||
state: Arc<AdaptiveBufferState>,
|
||||
}
|
||||
|
||||
impl<R> AdaptiveBufReader<R> {
|
||||
fn new(inner: R, state: Arc<AdaptiveBufferState>) -> Self {
|
||||
let bytes = state.actual_bytes.load(Ordering::Relaxed).max(1);
|
||||
Self {
|
||||
inner,
|
||||
buffer: vec![0; bytes].into_boxed_slice(),
|
||||
pos: 0,
|
||||
cap: 0,
|
||||
state,
|
||||
}
|
||||
}
|
||||
|
||||
fn resize_if_drained(&mut self) {
|
||||
if self.pos != self.cap {
|
||||
return;
|
||||
}
|
||||
let desired = self.state.desired_bytes.load(Ordering::Acquire).max(1);
|
||||
if desired == self.buffer.len() {
|
||||
return;
|
||||
}
|
||||
self.buffer = vec![0; desired].into_boxed_slice();
|
||||
self.pos = 0;
|
||||
self.cap = 0;
|
||||
self.state.actual_bytes.store(desired, Ordering::Release);
|
||||
}
|
||||
}
|
||||
|
||||
impl<R: AsyncRead + Unpin> AsyncRead for AdaptiveBufReader<R> {
|
||||
fn poll_read(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
output: &mut ReadBuf<'_>,
|
||||
) -> Poll<io::Result<()>> {
|
||||
let this = self.get_mut();
|
||||
if this.pos < this.cap {
|
||||
let available = &this.buffer[this.pos..this.cap];
|
||||
let copied = available.len().min(output.remaining());
|
||||
output.put_slice(&available[..copied]);
|
||||
this.pos += copied;
|
||||
return Poll::Ready(Ok(()));
|
||||
}
|
||||
this.resize_if_drained();
|
||||
Pin::new(&mut this.inner).poll_read(cx, output)
|
||||
}
|
||||
}
|
||||
|
||||
impl<R: AsyncRead + Unpin> AsyncBufRead for AdaptiveBufReader<R> {
|
||||
fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
|
||||
let this = self.get_mut();
|
||||
if this.pos < this.cap {
|
||||
return Poll::Ready(Ok(&this.buffer[this.pos..this.cap]));
|
||||
}
|
||||
|
||||
this.resize_if_drained();
|
||||
let mut read_buf = ReadBuf::new(&mut this.buffer);
|
||||
match Pin::new(&mut this.inner).poll_read(cx, &mut read_buf) {
|
||||
Poll::Ready(Ok(())) => {
|
||||
this.pos = 0;
|
||||
this.cap = read_buf.filled().len();
|
||||
Poll::Ready(Ok(&this.buffer[..this.cap]))
|
||||
}
|
||||
Poll::Ready(Err(error)) => Poll::Ready(Err(error)),
|
||||
Poll::Pending => Poll::Pending,
|
||||
}
|
||||
}
|
||||
|
||||
fn consume(self: Pin<&mut Self>, amount: usize) {
|
||||
let this = self.get_mut();
|
||||
this.pos = this.pos.saturating_add(amount).min(this.cap);
|
||||
}
|
||||
}
|
||||
|
||||
enum AdaptiveRelayOutcome {
|
||||
Copy(io::Result<(u64, u64)>),
|
||||
ActivityTimeout,
|
||||
UserDisabled,
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
/// Relays one Direct session with independently resizable directional buffers.
|
||||
pub(crate) async fn relay_direct_adaptive<CR, CW, SR, SW>(
|
||||
client_reader: CR,
|
||||
client_writer: CW,
|
||||
server_reader: SR,
|
||||
server_writer: SW,
|
||||
ceiling_c2s_bytes: usize,
|
||||
ceiling_s2c_bytes: usize,
|
||||
max_connections: u32,
|
||||
user: &str,
|
||||
stats: Arc<Stats>,
|
||||
quota_limit: Option<u64>,
|
||||
traffic_lease: Option<Arc<TrafficLease>>,
|
||||
activity_timeout: Duration,
|
||||
session_cancel: CancellationToken,
|
||||
budget: Arc<DirectBufferBudget>,
|
||||
) -> Result<()>
|
||||
where
|
||||
CR: AsyncRead + Unpin + Send + 'static,
|
||||
CW: AsyncWrite + Unpin + Send + 'static,
|
||||
SR: AsyncRead + Unpin + Send + 'static,
|
||||
SW: AsyncWrite + Unpin + Send + 'static,
|
||||
{
|
||||
let activity_timeout = activity_timeout.max(Duration::from_secs(1));
|
||||
let epoch = Instant::now();
|
||||
let counters = Arc::new(SharedCounters::new());
|
||||
let quota_exceeded = Arc::new(AtomicBool::new(false));
|
||||
let user_owned = user.to_string();
|
||||
|
||||
let (base_c2s, base_s2c) = initial_base_sizes(
|
||||
ceiling_c2s_bytes,
|
||||
ceiling_s2c_bytes,
|
||||
max_connections,
|
||||
budget.target_bytes(),
|
||||
);
|
||||
let base_total = base_c2s.saturating_add(base_s2c);
|
||||
let mut lease = match budget.try_reserve(base_total, false) {
|
||||
Some(lease) => lease,
|
||||
None => {
|
||||
let minimum_total = DIRECT_BASE_C2S_BYTES + DIRECT_BASE_S2C_BYTES;
|
||||
match budget.try_reserve(minimum_total, true) {
|
||||
Some(lease) => {
|
||||
budget.increment_minimum_fallback();
|
||||
lease
|
||||
}
|
||||
None => {
|
||||
budget.increment_admission_rejected();
|
||||
return Err(ProxyError::Proxy(
|
||||
"Direct relay buffer pressure: budget exhausted".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let effective_base = if lease.reserved_bytes() < base_total {
|
||||
(DIRECT_BASE_C2S_BYTES, DIRECT_BASE_S2C_BYTES)
|
||||
} else {
|
||||
(base_c2s, base_s2c)
|
||||
};
|
||||
let c2s_state = AdaptiveBufferState::new(effective_base.0);
|
||||
let s2c_state = AdaptiveBufferState::new(effective_base.1);
|
||||
|
||||
let mut controller = SessionAdaptiveController::new(AdaptiveTier::Base);
|
||||
|
||||
let c2s_client = StatsIo::new_with_traffic_lease(
|
||||
client_reader,
|
||||
Arc::clone(&counters),
|
||||
Arc::clone(&stats),
|
||||
user_owned.clone(),
|
||||
traffic_lease.clone(),
|
||||
quota_limit,
|
||||
Arc::clone("a_exceeded),
|
||||
epoch,
|
||||
);
|
||||
let client_writer = StatsIo::new_with_traffic_lease(
|
||||
client_writer,
|
||||
Arc::clone(&counters),
|
||||
Arc::clone(&stats),
|
||||
user_owned.clone(),
|
||||
traffic_lease,
|
||||
quota_limit,
|
||||
Arc::clone("a_exceeded),
|
||||
epoch,
|
||||
);
|
||||
let mut client_writer = WritePressureIo::new(client_writer, Arc::clone(&counters));
|
||||
let mut c2s_reader = AdaptiveBufReader::new(c2s_client, Arc::clone(&c2s_state));
|
||||
let mut s2c_reader = AdaptiveBufReader::new(server_reader, Arc::clone(&s2c_state));
|
||||
let mut server_writer = server_writer;
|
||||
let mut pressure_rx = budget.subscribe_pressure();
|
||||
|
||||
let relay_outcome = {
|
||||
let copy = async {
|
||||
let c2s = async {
|
||||
let copied = copy_buf(&mut c2s_reader, &mut server_writer).await?;
|
||||
server_writer.shutdown().await?;
|
||||
Ok::<u64, io::Error>(copied)
|
||||
};
|
||||
let s2c = async {
|
||||
let copied = copy_buf(&mut s2c_reader, &mut client_writer).await?;
|
||||
client_writer.shutdown().await?;
|
||||
Ok::<u64, io::Error>(copied)
|
||||
};
|
||||
tokio::try_join!(c2s, s2c)
|
||||
};
|
||||
tokio::pin!(copy);
|
||||
|
||||
let mut interval = tokio::time::interval(WATCHDOG_INTERVAL);
|
||||
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
interval.tick().await;
|
||||
let mut previous = RelaySignalSample::default();
|
||||
let mut previous_log_c2s = 0u64;
|
||||
let mut previous_log_s2c = 0u64;
|
||||
let mut previous_sample_at = epoch;
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
result = &mut copy => break AdaptiveRelayOutcome::Copy(result),
|
||||
_ = session_cancel.cancelled() => break AdaptiveRelayOutcome::UserDisabled,
|
||||
changed = pressure_rx.changed() => {
|
||||
if changed.is_ok() {
|
||||
apply_global_pressure_demotion(
|
||||
&mut controller,
|
||||
&mut lease,
|
||||
&c2s_state,
|
||||
&s2c_state,
|
||||
effective_base,
|
||||
(ceiling_c2s_bytes, ceiling_s2c_bytes),
|
||||
budget.as_ref(),
|
||||
);
|
||||
reconcile_reservation(&mut lease, &c2s_state, &s2c_state);
|
||||
}
|
||||
}
|
||||
_ = interval.tick() => {
|
||||
let now = Instant::now();
|
||||
let idle = counters.idle_duration(now, epoch);
|
||||
if quota_exceeded.load(Ordering::Acquire) {
|
||||
warn!(user = %user_owned, "User data quota reached, closing relay");
|
||||
break AdaptiveRelayOutcome::ActivityTimeout;
|
||||
}
|
||||
if idle >= activity_timeout {
|
||||
warn!(
|
||||
user = %user_owned,
|
||||
c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed),
|
||||
s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed),
|
||||
idle_secs = idle.as_secs(),
|
||||
"Activity timeout"
|
||||
);
|
||||
break AdaptiveRelayOutcome::ActivityTimeout;
|
||||
}
|
||||
|
||||
let sample = current_sample(counters.as_ref());
|
||||
let c2s_delta = watchdog_delta(sample.c2s_bytes, previous_log_c2s);
|
||||
let s2c_delta = watchdog_delta(sample.s2c_written_bytes, previous_log_s2c);
|
||||
if c2s_delta > 0 || s2c_delta > 0 {
|
||||
let secs = now.saturating_duration_since(previous_sample_at).as_secs_f64();
|
||||
debug!(
|
||||
user = %user_owned,
|
||||
c2s_kbps = (c2s_delta as f64 / secs / 1024.0) as u64,
|
||||
s2c_kbps = (s2c_delta as f64 / secs / 1024.0) as u64,
|
||||
c2s_total = sample.c2s_bytes,
|
||||
s2c_total = sample.s2c_written_bytes,
|
||||
"Relay active"
|
||||
);
|
||||
}
|
||||
|
||||
let delta = sample_delta(sample, previous);
|
||||
let tick_secs = now.saturating_duration_since(previous_sample_at).as_secs_f64();
|
||||
if let Some(transition) = controller.observe(delta, tick_secs) {
|
||||
apply_controller_transition(
|
||||
transition,
|
||||
&mut controller,
|
||||
&mut lease,
|
||||
&c2s_state,
|
||||
&s2c_state,
|
||||
effective_base,
|
||||
(ceiling_c2s_bytes, ceiling_s2c_bytes),
|
||||
budget.as_ref(),
|
||||
);
|
||||
}
|
||||
reconcile_reservation(&mut lease, &c2s_state, &s2c_state);
|
||||
previous = sample;
|
||||
previous_log_c2s = sample.c2s_bytes;
|
||||
previous_log_s2c = sample.s2c_written_bytes;
|
||||
previous_sample_at = now;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let _ = client_writer.shutdown().await;
|
||||
let _ = server_writer.shutdown().await;
|
||||
let c2s_ops = counters.c2s_ops.load(Ordering::Relaxed);
|
||||
let s2c_ops = counters.s2c_ops.load(Ordering::Relaxed);
|
||||
let duration = epoch.elapsed();
|
||||
match relay_outcome {
|
||||
AdaptiveRelayOutcome::Copy(Ok((c2s, s2c))) => {
|
||||
debug!(
|
||||
user = %user_owned,
|
||||
c2s_bytes = c2s,
|
||||
s2c_bytes = s2c,
|
||||
c2s_msgs = c2s_ops,
|
||||
s2c_msgs = s2c_ops,
|
||||
duration_secs = duration.as_secs(),
|
||||
"Relay finished"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
AdaptiveRelayOutcome::Copy(Err(error)) if is_quota_io_error(&error) => {
|
||||
warn!(
|
||||
user = %user_owned,
|
||||
c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed),
|
||||
s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed),
|
||||
c2s_msgs = c2s_ops,
|
||||
s2c_msgs = s2c_ops,
|
||||
duration_secs = duration.as_secs(),
|
||||
"Data quota reached, closing relay"
|
||||
);
|
||||
Err(ProxyError::DataQuotaExceeded { user: user_owned })
|
||||
}
|
||||
AdaptiveRelayOutcome::Copy(Err(error)) => {
|
||||
debug!(
|
||||
user = %user_owned,
|
||||
c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed),
|
||||
s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed),
|
||||
c2s_msgs = c2s_ops,
|
||||
s2c_msgs = s2c_ops,
|
||||
duration_secs = duration.as_secs(),
|
||||
error = %error,
|
||||
"Relay error"
|
||||
);
|
||||
Err(error.into())
|
||||
}
|
||||
AdaptiveRelayOutcome::ActivityTimeout => {
|
||||
debug!(
|
||||
user = %user_owned,
|
||||
c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed),
|
||||
s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed),
|
||||
c2s_msgs = c2s_ops,
|
||||
s2c_msgs = s2c_ops,
|
||||
duration_secs = duration.as_secs(),
|
||||
"Relay finished (activity timeout)"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
AdaptiveRelayOutcome::UserDisabled => {
|
||||
debug!(
|
||||
user = %user_owned,
|
||||
c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed),
|
||||
s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed),
|
||||
c2s_msgs = c2s_ops,
|
||||
s2c_msgs = s2c_ops,
|
||||
duration_secs = duration.as_secs(),
|
||||
"Relay finished (user disabled)"
|
||||
);
|
||||
Err(ProxyError::UserDisabled { user: user_owned })
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn initial_base_sizes(
|
||||
ceiling_c2s: usize,
|
||||
ceiling_s2c: usize,
|
||||
max_connections: u32,
|
||||
target_bytes: usize,
|
||||
) -> (usize, usize) {
|
||||
let configured_total = ceiling_c2s.saturating_add(ceiling_s2c);
|
||||
let configured_worst_case = configured_total.saturating_mul(max_connections as usize);
|
||||
if max_connections != 0 && configured_worst_case <= target_bytes {
|
||||
return (ceiling_c2s, ceiling_s2c);
|
||||
}
|
||||
(
|
||||
DIRECT_BASE_C2S_BYTES.min(ceiling_c2s),
|
||||
DIRECT_BASE_S2C_BYTES.min(ceiling_s2c),
|
||||
)
|
||||
}
|
||||
|
||||
fn current_sample(counters: &SharedCounters) -> RelaySignalSample {
|
||||
RelaySignalSample {
|
||||
c2s_bytes: counters.c2s_bytes.load(Ordering::Relaxed),
|
||||
s2c_requested_bytes: counters.s2c_requested_bytes.load(Ordering::Relaxed),
|
||||
s2c_written_bytes: counters.s2c_bytes.load(Ordering::Relaxed),
|
||||
s2c_write_ops: counters.s2c_ops.load(Ordering::Relaxed),
|
||||
s2c_partial_writes: counters.s2c_partial_writes.load(Ordering::Relaxed),
|
||||
s2c_consecutive_pending_writes: counters
|
||||
.s2c_consecutive_pending_writes
|
||||
.load(Ordering::Relaxed),
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_delta(current: RelaySignalSample, previous: RelaySignalSample) -> RelaySignalSample {
|
||||
RelaySignalSample {
|
||||
c2s_bytes: current.c2s_bytes.saturating_sub(previous.c2s_bytes),
|
||||
s2c_requested_bytes: current
|
||||
.s2c_requested_bytes
|
||||
.saturating_sub(previous.s2c_requested_bytes),
|
||||
s2c_written_bytes: current
|
||||
.s2c_written_bytes
|
||||
.saturating_sub(previous.s2c_written_bytes),
|
||||
s2c_write_ops: current.s2c_write_ops.saturating_sub(previous.s2c_write_ops),
|
||||
s2c_partial_writes: current
|
||||
.s2c_partial_writes
|
||||
.saturating_sub(previous.s2c_partial_writes),
|
||||
s2c_consecutive_pending_writes: current.s2c_consecutive_pending_writes,
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn apply_controller_transition(
|
||||
transition: crate::proxy::adaptive_buffers::TierTransition,
|
||||
controller: &mut SessionAdaptiveController,
|
||||
lease: &mut DirectBufferLease,
|
||||
c2s_state: &AdaptiveBufferState,
|
||||
s2c_state: &AdaptiveBufferState,
|
||||
base: (usize, usize),
|
||||
ceilings: (usize, usize),
|
||||
budget: &DirectBufferBudget,
|
||||
) {
|
||||
let sizes = direct_copy_buffers_for_tier_with_ceilings(
|
||||
transition.to,
|
||||
base.0,
|
||||
base.1,
|
||||
ceilings.0,
|
||||
ceilings.1,
|
||||
);
|
||||
if transition.to > transition.from {
|
||||
if !lease.try_grow_to(sizes.0.saturating_add(sizes.1)) {
|
||||
*controller = SessionAdaptiveController::new(transition.from);
|
||||
return;
|
||||
}
|
||||
} else {
|
||||
match transition.reason {
|
||||
TierTransitionReason::QuietDemotion => budget.increment_quiet_demotion(),
|
||||
TierTransitionReason::SustainedWritePressure => {
|
||||
budget.increment_write_pressure_demotion();
|
||||
}
|
||||
TierTransitionReason::SoftConfirmed | TierTransitionReason::HardPressure => {}
|
||||
}
|
||||
}
|
||||
set_desired_sizes(c2s_state, s2c_state, sizes);
|
||||
lease.set_tier(transition.to.as_u8() as usize);
|
||||
}
|
||||
|
||||
fn apply_global_pressure_demotion(
|
||||
controller: &mut SessionAdaptiveController,
|
||||
lease: &mut DirectBufferLease,
|
||||
c2s_state: &AdaptiveBufferState,
|
||||
s2c_state: &AdaptiveBufferState,
|
||||
base: (usize, usize),
|
||||
ceilings: (usize, usize),
|
||||
budget: &DirectBufferBudget,
|
||||
) {
|
||||
let current = controller.tier();
|
||||
let target = current.demote();
|
||||
if target == current {
|
||||
return;
|
||||
}
|
||||
*controller = SessionAdaptiveController::new(target);
|
||||
let sizes =
|
||||
direct_copy_buffers_for_tier_with_ceilings(target, base.0, base.1, ceilings.0, ceilings.1);
|
||||
set_desired_sizes(c2s_state, s2c_state, sizes);
|
||||
lease.set_tier(target.as_u8() as usize);
|
||||
budget.increment_global_pressure_demotion();
|
||||
}
|
||||
|
||||
fn set_desired_sizes(
|
||||
c2s_state: &AdaptiveBufferState,
|
||||
s2c_state: &AdaptiveBufferState,
|
||||
sizes: (usize, usize),
|
||||
) {
|
||||
c2s_state
|
||||
.desired_bytes
|
||||
.store(sizes.0.max(1), Ordering::Release);
|
||||
s2c_state
|
||||
.desired_bytes
|
||||
.store(sizes.1.max(1), Ordering::Release);
|
||||
}
|
||||
|
||||
fn reconcile_reservation(
|
||||
lease: &mut DirectBufferLease,
|
||||
c2s_state: &AdaptiveBufferState,
|
||||
s2c_state: &AdaptiveBufferState,
|
||||
) {
|
||||
// Promotion reserves the desired allocation before either reader grows.
|
||||
// Demotion keeps the actual allocation covered until its buffered bytes drain.
|
||||
let covered_c2s = c2s_state
|
||||
.actual_bytes
|
||||
.load(Ordering::Acquire)
|
||||
.max(c2s_state.desired_bytes.load(Ordering::Acquire));
|
||||
let covered_s2c = s2c_state
|
||||
.actual_bytes
|
||||
.load(Ordering::Acquire)
|
||||
.max(s2c_state.desired_bytes.load(Ordering::Acquire));
|
||||
lease.shrink_to(covered_c2s.saturating_add(covered_s2c));
|
||||
}
|
||||
72
src/proxy/relay/adaptive_copy/write_pressure.rs
Normal file
72
src/proxy/relay/adaptive_copy/write_pressure.rs
Normal file
@@ -0,0 +1,72 @@
|
||||
use std::io;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
use tokio::io::AsyncWrite;
|
||||
|
||||
use super::super::io::SharedCounters;
|
||||
|
||||
/// Direct-only writer wrapper that exposes bounded backpressure signals.
|
||||
pub(super) struct WritePressureIo<W> {
|
||||
inner: W,
|
||||
counters: Arc<SharedCounters>,
|
||||
}
|
||||
|
||||
impl<W> WritePressureIo<W> {
|
||||
/// Wraps the client writer without changing its I/O or error contract.
|
||||
pub(super) fn new(inner: W, counters: Arc<SharedCounters>) -> Self {
|
||||
Self { inner, counters }
|
||||
}
|
||||
}
|
||||
|
||||
impl<W: AsyncWrite + Unpin> AsyncWrite for WritePressureIo<W> {
|
||||
fn poll_write(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
buffer: &[u8],
|
||||
) -> Poll<io::Result<usize>> {
|
||||
let this = self.get_mut();
|
||||
if !buffer.is_empty() {
|
||||
this.counters
|
||||
.s2c_requested_bytes
|
||||
.fetch_add(buffer.len() as u64, Ordering::Relaxed);
|
||||
}
|
||||
match Pin::new(&mut this.inner).poll_write(cx, buffer) {
|
||||
Poll::Ready(Ok(written)) => {
|
||||
this.counters
|
||||
.s2c_consecutive_pending_writes
|
||||
.store(0, Ordering::Relaxed);
|
||||
if written < buffer.len() {
|
||||
this.counters
|
||||
.s2c_partial_writes
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
Poll::Ready(Ok(written))
|
||||
}
|
||||
Poll::Ready(Err(error)) => {
|
||||
this.counters
|
||||
.s2c_consecutive_pending_writes
|
||||
.store(0, Ordering::Relaxed);
|
||||
Poll::Ready(Err(error))
|
||||
}
|
||||
Poll::Pending => {
|
||||
let _ = this.counters.s2c_consecutive_pending_writes.fetch_update(
|
||||
Ordering::Relaxed,
|
||||
Ordering::Relaxed,
|
||||
|current| Some(current.saturating_add(1)),
|
||||
);
|
||||
Poll::Pending
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||
Pin::new(&mut self.get_mut().inner).poll_flush(cx)
|
||||
}
|
||||
|
||||
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||
Pin::new(&mut self.get_mut().inner).poll_shutdown(cx)
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,4 @@
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::time::Instant;
|
||||
@@ -20,6 +20,12 @@ pub(in crate::proxy::relay) struct SharedCounters {
|
||||
pub(in crate::proxy::relay) c2s_ops: AtomicU64,
|
||||
/// Number of poll_write completions (≈ S→C chunks)
|
||||
pub(in crate::proxy::relay) s2c_ops: AtomicU64,
|
||||
/// Bytes presented to client writes, including retried pending writes.
|
||||
pub(in crate::proxy::relay) s2c_requested_bytes: AtomicU64,
|
||||
/// Successful client writes that consumed only part of the offered slice.
|
||||
pub(in crate::proxy::relay) s2c_partial_writes: AtomicU64,
|
||||
/// Consecutive pending client writes observed by the active copy loop.
|
||||
pub(in crate::proxy::relay) s2c_consecutive_pending_writes: AtomicU32,
|
||||
/// Milliseconds since relay epoch of last I/O activity
|
||||
last_activity_ms: AtomicU64,
|
||||
}
|
||||
@@ -31,6 +37,9 @@ impl SharedCounters {
|
||||
s2c_bytes: AtomicU64::new(0),
|
||||
c2s_ops: AtomicU64::new(0),
|
||||
s2c_ops: AtomicU64::new(0),
|
||||
s2c_requested_bytes: AtomicU64::new(0),
|
||||
s2c_partial_writes: AtomicU64::new(0),
|
||||
s2c_consecutive_pending_writes: AtomicU32::new(0),
|
||||
last_activity_ms: AtomicU64::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
use crate::stats::UserStats;
|
||||
use std::io;
|
||||
use std::sync::atomic::Ordering;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct QuotaIoSentinel;
|
||||
@@ -52,17 +51,5 @@ pub(super) fn refund_reserved_quota_bytes(user_stats: &UserStats, reserved_bytes
|
||||
if reserved_bytes == 0 {
|
||||
return;
|
||||
}
|
||||
let mut current = user_stats.quota_used.load(Ordering::Relaxed);
|
||||
loop {
|
||||
let next = current.saturating_sub(reserved_bytes);
|
||||
match user_stats.quota_used.compare_exchange_weak(
|
||||
current,
|
||||
next,
|
||||
Ordering::Relaxed,
|
||||
Ordering::Relaxed,
|
||||
) {
|
||||
Ok(_) => return,
|
||||
Err(observed) => current = observed,
|
||||
}
|
||||
}
|
||||
user_stats.refund_quota(reserved_bytes);
|
||||
}
|
||||
|
||||
@@ -9,6 +9,7 @@ use dashmap::DashMap;
|
||||
use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::proxy::direct_buffer_budget::{DirectBufferBudget, fallback_direct_buffer_hard_limit};
|
||||
use crate::proxy::handshake::{AuthProbeSaturationState, AuthProbeState};
|
||||
use crate::proxy::middle_relay::{DesyncDedupRotationState, RelayIdleCandidateRegistry};
|
||||
use crate::proxy::traffic_limiter::TrafficLimiter;
|
||||
@@ -69,6 +70,7 @@ pub(crate) struct ProxySharedState {
|
||||
pub(crate) handshake: HandshakeSharedState,
|
||||
pub(crate) middle_relay: MiddleRelaySharedState,
|
||||
pub(crate) traffic_limiter: Arc<TrafficLimiter>,
|
||||
pub(crate) direct_buffer_budget: Arc<DirectBufferBudget>,
|
||||
disabled_users: DashMap<String, ()>,
|
||||
active_user_sessions: DashMap<(String, u64), CancellationToken>,
|
||||
pub(crate) conntrack_pressure_active: AtomicBool,
|
||||
@@ -101,6 +103,15 @@ impl Drop for UserSessionGuard {
|
||||
|
||||
impl ProxySharedState {
|
||||
pub(crate) fn new() -> Arc<Self> {
|
||||
Self::new_with_direct_buffer_budget(DirectBufferBudget::new(
|
||||
fallback_direct_buffer_hard_limit(),
|
||||
))
|
||||
}
|
||||
|
||||
/// Creates process state with the startup-resolved Direct buffer envelope.
|
||||
pub(crate) fn new_with_direct_buffer_budget(
|
||||
direct_buffer_budget: Arc<DirectBufferBudget>,
|
||||
) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
handshake: HandshakeSharedState {
|
||||
auth_probe: DashMap::new(),
|
||||
@@ -129,6 +140,7 @@ impl ProxySharedState {
|
||||
relay_idle_mark_seq: AtomicU64::new(0),
|
||||
},
|
||||
traffic_limiter: TrafficLimiter::new(),
|
||||
direct_buffer_budget,
|
||||
disabled_users: DashMap::new(),
|
||||
active_user_sessions: DashMap::new(),
|
||||
conntrack_pressure_active: AtomicBool::new(false),
|
||||
|
||||
58
src/proxy/tests/adaptive_direct_budget_policy_tests.rs
Normal file
58
src/proxy/tests/adaptive_direct_budget_policy_tests.rs
Normal file
@@ -0,0 +1,58 @@
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn configured_direct_sizes_are_strict_tier_ceilings() {
|
||||
let base = (4 * 1024, 8 * 1024);
|
||||
let ceilings = (64 * 1024, 256 * 1024);
|
||||
assert_eq!(
|
||||
direct_copy_buffers_for_tier_with_ceilings(
|
||||
AdaptiveTier::Base,
|
||||
base.0,
|
||||
base.1,
|
||||
ceilings.0,
|
||||
ceilings.1,
|
||||
),
|
||||
base
|
||||
);
|
||||
assert_eq!(
|
||||
direct_copy_buffers_for_tier_with_ceilings(
|
||||
AdaptiveTier::Tier3,
|
||||
base.0,
|
||||
base.1,
|
||||
ceilings.0,
|
||||
ceilings.1,
|
||||
),
|
||||
ceilings
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sustained_pending_pressure_demotes_after_transient_promotion() {
|
||||
let mut controller = SessionAdaptiveController::new(AdaptiveTier::Tier1);
|
||||
let pressure = RelaySignalSample {
|
||||
c2s_bytes: 0,
|
||||
s2c_requested_bytes: 1024,
|
||||
s2c_written_bytes: 0,
|
||||
s2c_write_ops: 0,
|
||||
s2c_partial_writes: 0,
|
||||
s2c_consecutive_pending_writes: 3,
|
||||
};
|
||||
|
||||
let first = controller
|
||||
.observe(pressure, 10.0)
|
||||
.expect("transient pressure must retain the staged promotion");
|
||||
assert_eq!(first.reason, TierTransitionReason::HardPressure);
|
||||
|
||||
let second = controller
|
||||
.observe(pressure, 10.0)
|
||||
.expect("bounded transient pressure may promote one additional tier");
|
||||
assert_eq!(second.reason, TierTransitionReason::HardPressure);
|
||||
let sustained = controller
|
||||
.observe(pressure, 10.0)
|
||||
.expect("sustained pressure must release one tier");
|
||||
assert_eq!(
|
||||
sustained.reason,
|
||||
TierTransitionReason::SustainedWritePressure
|
||||
);
|
||||
assert_eq!(sustained.to, AdaptiveTier::Tier2);
|
||||
}
|
||||
38
src/proxy/tests/direct_buffer_budget_tests.rs
Normal file
38
src/proxy/tests/direct_buffer_budget_tests.rs
Normal file
@@ -0,0 +1,38 @@
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn lease_drop_releases_the_complete_reservation() {
|
||||
let budget = DirectBufferBudget::new(16 * 1024);
|
||||
{
|
||||
let lease = budget
|
||||
.try_reserve(12 * 1024, false)
|
||||
.expect("minimum reservation must fit");
|
||||
assert_eq!(lease.reserved_bytes(), 12 * 1024);
|
||||
assert_eq!(budget.snapshot().reserved_bytes, 12 * 1024);
|
||||
}
|
||||
assert_eq!(budget.snapshot().reserved_bytes, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn absolute_ceiling_rejects_excess_minimum_reservations() {
|
||||
let budget = DirectBufferBudget::new(16 * 1024);
|
||||
let _lease = budget
|
||||
.try_reserve(12 * 1024, true)
|
||||
.expect("first minimum reservation must fit");
|
||||
assert!(budget.try_reserve(8 * 1024, true).is_none());
|
||||
assert_eq!(budget.snapshot().reserved_bytes, 12 * 1024);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn growth_and_shrink_keep_accounting_balanced() {
|
||||
let budget = DirectBufferBudget::new(32 * 1024);
|
||||
let mut lease = budget
|
||||
.try_reserve(12 * 1024, false)
|
||||
.expect("base reservation must fit");
|
||||
assert!(lease.try_grow_to(24 * 1024));
|
||||
assert_eq!(budget.snapshot().reserved_bytes, 24 * 1024);
|
||||
lease.shrink_to(16 * 1024);
|
||||
assert_eq!(budget.snapshot().reserved_bytes, 16 * 1024);
|
||||
drop(lease);
|
||||
assert_eq!(budget.snapshot().reserved_bytes, 0);
|
||||
}
|
||||
@@ -80,7 +80,11 @@ impl Stats {
|
||||
return handle;
|
||||
}
|
||||
|
||||
let entry = self.user_stats.entry(user.to_string()).or_default();
|
||||
let quota = self.quota_store.user(user);
|
||||
let entry = self
|
||||
.user_stats
|
||||
.entry(user.to_string())
|
||||
.or_insert_with(|| Arc::new(UserStats::with_quota(quota)));
|
||||
if entry.last_seen_epoch_secs.load(Ordering::Relaxed) == 0 {
|
||||
self.touch_user_stats(entry.value().as_ref());
|
||||
}
|
||||
@@ -166,10 +170,7 @@ impl Stats {
|
||||
#[inline]
|
||||
pub(crate) fn quota_charge_post_write(&self, user_stats: &UserStats, bytes: u64) -> u64 {
|
||||
self.touch_user_stats(user_stats);
|
||||
user_stats
|
||||
.quota_used
|
||||
.fetch_add(bytes, Ordering::Relaxed)
|
||||
.saturating_add(bytes)
|
||||
user_stats.quota.charge(bytes)
|
||||
}
|
||||
|
||||
pub(super) fn maybe_cleanup_user_stats(&self) {
|
||||
|
||||
@@ -204,6 +204,12 @@ impl Stats {
|
||||
self.buffer_pool_in_use_gauge.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Returns the count of non-standard buffers replaced before pooling.
|
||||
pub fn get_buffer_pool_replaced_nonstandard_total(&self) -> u64 {
|
||||
self.buffer_pool_replaced_nonstandard_total
|
||||
.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
pub fn get_me_c2me_send_full_total(&self) -> u64 {
|
||||
self.me_c2me_send_full_total.load(Ordering::Relaxed)
|
||||
}
|
||||
@@ -269,6 +275,36 @@ impl Stats {
|
||||
self.me_writer_pick_mode_switch_total
|
||||
.load(Ordering::Relaxed)
|
||||
}
|
||||
/// Returns the configured resident-memory limit per ME writer.
|
||||
pub fn get_me_writer_byte_budget_limit_bytes_gauge(&self) -> u64 {
|
||||
self.me_writer_byte_budget_limit_bytes_gauge
|
||||
.load(Ordering::Relaxed)
|
||||
}
|
||||
/// Returns aggregate queued or enqueueing writer memory reservations.
|
||||
pub fn get_me_writer_byte_budget_queued_bytes_gauge(&self) -> u64 {
|
||||
self.me_writer_byte_budget_queued_bytes_gauge
|
||||
.load(Ordering::Relaxed)
|
||||
}
|
||||
/// Returns aggregate writer reservations currently owned by socket writes.
|
||||
pub fn get_me_writer_byte_budget_inflight_bytes_gauge(&self) -> u64 {
|
||||
self.me_writer_byte_budget_inflight_bytes_gauge
|
||||
.load(Ordering::Relaxed)
|
||||
}
|
||||
/// Returns the count of blocking writer byte-budget waits.
|
||||
pub fn get_me_writer_byte_budget_wait_total(&self) -> u64 {
|
||||
self.me_writer_byte_budget_wait_total
|
||||
.load(Ordering::Relaxed)
|
||||
}
|
||||
/// Returns the count of writer byte-budget wait timeouts.
|
||||
pub fn get_me_writer_byte_budget_timeout_total(&self) -> u64 {
|
||||
self.me_writer_byte_budget_timeout_total
|
||||
.load(Ordering::Relaxed)
|
||||
}
|
||||
/// Returns the count of payloads that cannot fit the configured writer budget.
|
||||
pub fn get_me_writer_byte_budget_oversize_total(&self) -> u64 {
|
||||
self.me_writer_byte_budget_oversize_total
|
||||
.load(Ordering::Relaxed)
|
||||
}
|
||||
pub fn get_me_socks_kdf_strict_reject(&self) -> u64 {
|
||||
self.me_socks_kdf_strict_reject.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
@@ -8,6 +8,7 @@ mod core_getters;
|
||||
mod helpers;
|
||||
mod me_counters;
|
||||
mod me_getters;
|
||||
mod quota_store;
|
||||
mod replay;
|
||||
pub mod telemetry;
|
||||
pub mod tls_fingerprints;
|
||||
@@ -20,6 +21,7 @@ use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering};
|
||||
use std::time::Instant;
|
||||
|
||||
pub(crate) use self::quota_store::QuotaStore;
|
||||
#[allow(unused_imports)]
|
||||
pub use self::replay::{ReplayChecker, ReplayStats};
|
||||
use self::telemetry::TelemetryPolicy;
|
||||
@@ -274,6 +276,7 @@ pub struct Stats {
|
||||
buffer_pool_pooled_gauge: AtomicU64,
|
||||
buffer_pool_allocated_gauge: AtomicU64,
|
||||
buffer_pool_in_use_gauge: AtomicU64,
|
||||
buffer_pool_replaced_nonstandard_total: AtomicU64,
|
||||
// C2ME enqueue observability
|
||||
me_c2me_send_full_total: AtomicU64,
|
||||
me_c2me_send_high_water_total: AtomicU64,
|
||||
@@ -292,6 +295,12 @@ pub struct Stats {
|
||||
me_writer_pick_p2c_no_candidate_total: AtomicU64,
|
||||
me_writer_pick_blocking_fallback_total: AtomicU64,
|
||||
me_writer_pick_mode_switch_total: AtomicU64,
|
||||
me_writer_byte_budget_limit_bytes_gauge: AtomicU64,
|
||||
me_writer_byte_budget_queued_bytes_gauge: AtomicU64,
|
||||
me_writer_byte_budget_inflight_bytes_gauge: AtomicU64,
|
||||
me_writer_byte_budget_wait_total: AtomicU64,
|
||||
me_writer_byte_budget_timeout_total: AtomicU64,
|
||||
me_writer_byte_budget_oversize_total: AtomicU64,
|
||||
me_socks_kdf_strict_reject: AtomicU64,
|
||||
me_socks_kdf_compat_fallback: AtomicU64,
|
||||
secure_padding_invalid: AtomicU64,
|
||||
@@ -337,6 +346,7 @@ pub struct Stats {
|
||||
cached_epoch_secs: AtomicU64,
|
||||
tls_fingerprints: tls_fingerprints::TlsFingerprintCollector,
|
||||
user_stats: DashMap<String, Arc<UserStats>>,
|
||||
quota_store: Arc<QuotaStore>,
|
||||
user_stats_last_cleanup_epoch_secs: AtomicU64,
|
||||
start_time: parking_lot::RwLock<Option<Instant>>,
|
||||
}
|
||||
@@ -349,12 +359,7 @@ pub struct UserStats {
|
||||
pub octets_to_client: AtomicU64,
|
||||
pub msgs_from_client: AtomicU64,
|
||||
pub msgs_to_client: AtomicU64,
|
||||
/// Total bytes charged against per-user quota admission.
|
||||
///
|
||||
/// This counter is the single source of truth for quota enforcement and
|
||||
/// intentionally tracks attempted traffic, not guaranteed delivery.
|
||||
pub quota_used: AtomicU64,
|
||||
pub quota_last_reset_epoch_secs: AtomicU64,
|
||||
quota: Arc<quota_store::UserQuotaCounters>,
|
||||
pub last_seen_epoch_secs: AtomicU64,
|
||||
}
|
||||
|
||||
@@ -371,9 +376,21 @@ pub enum QuotaReserveError {
|
||||
}
|
||||
|
||||
impl UserStats {
|
||||
fn with_quota(quota: Arc<quota_store::UserQuotaCounters>) -> Self {
|
||||
Self {
|
||||
quota,
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn quota_used(&self) -> u64 {
|
||||
self.quota_used.load(Ordering::Relaxed)
|
||||
self.quota.used()
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub(crate) fn refund_quota(&self, bytes: u64) {
|
||||
self.quota.refund(bytes);
|
||||
}
|
||||
|
||||
/// Attempts one CAS reservation step against the quota counter.
|
||||
@@ -383,27 +400,20 @@ impl UserStats {
|
||||
/// with their own contention strategy.
|
||||
#[inline]
|
||||
pub fn quota_try_reserve(&self, bytes: u64, limit: u64) -> Result<u64, QuotaReserveError> {
|
||||
let current = self.quota_used.load(Ordering::Relaxed);
|
||||
if bytes > limit.saturating_sub(current) {
|
||||
return Err(QuotaReserveError::LimitExceeded);
|
||||
}
|
||||
|
||||
let next = current.saturating_add(bytes);
|
||||
match self.quota_used.compare_exchange_weak(
|
||||
current,
|
||||
next,
|
||||
Ordering::Relaxed,
|
||||
Ordering::Relaxed,
|
||||
) {
|
||||
Ok(_) => Ok(next),
|
||||
Err(_) => Err(QuotaReserveError::Contended),
|
||||
}
|
||||
self.quota.try_reserve(bytes, limit)
|
||||
}
|
||||
}
|
||||
|
||||
impl Stats {
|
||||
pub fn new() -> Self {
|
||||
let stats = Self::default();
|
||||
Self::with_quota_store(Arc::new(QuotaStore::default()))
|
||||
}
|
||||
|
||||
pub(crate) fn with_quota_store(quota_store: Arc<QuotaStore>) -> Self {
|
||||
let stats = Self {
|
||||
quota_store,
|
||||
..Self::default()
|
||||
};
|
||||
stats.apply_telemetry_policy(TelemetryPolicy::default());
|
||||
stats.refresh_cached_epoch_secs();
|
||||
*stats.start_time.write() = Some(Instant::now());
|
||||
|
||||
147
src/stats/quota_store.rs
Normal file
147
src/stats/quota_store.rs
Normal file
@@ -0,0 +1,147 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
use dashmap::DashMap;
|
||||
|
||||
use super::{QuotaReserveError, UserQuotaSnapshot};
|
||||
|
||||
/// Process-scoped per-user quota accounting shared by runtime generations.
|
||||
#[derive(Default)]
|
||||
pub struct QuotaStore {
|
||||
users: DashMap<String, Arc<UserQuotaCounters>>,
|
||||
}
|
||||
|
||||
/// Atomic quota state for one configured user.
|
||||
#[derive(Default)]
|
||||
pub(crate) struct UserQuotaCounters {
|
||||
used_bytes: AtomicU64,
|
||||
last_reset_epoch_secs: AtomicU64,
|
||||
}
|
||||
|
||||
impl QuotaStore {
|
||||
pub(crate) fn user(&self, user: &str) -> Arc<UserQuotaCounters> {
|
||||
if let Some(existing) = self.users.get(user) {
|
||||
return Arc::clone(existing.value());
|
||||
}
|
||||
Arc::clone(
|
||||
self.users
|
||||
.entry(user.to_string())
|
||||
.or_insert_with(|| Arc::new(UserQuotaCounters::default()))
|
||||
.value(),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn used(&self, user: &str) -> u64 {
|
||||
self.users.get(user).map(|state| state.used()).unwrap_or(0)
|
||||
}
|
||||
|
||||
pub(crate) fn load(&self, user: &str, used_bytes: u64, last_reset_epoch_secs: u64) {
|
||||
let state = self.user(user);
|
||||
state.used_bytes.store(used_bytes, Ordering::Relaxed);
|
||||
state
|
||||
.last_reset_epoch_secs
|
||||
.store(last_reset_epoch_secs, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
pub(crate) fn reset(&self, user: &str, now_epoch_secs: u64) -> UserQuotaSnapshot {
|
||||
let state = self.user(user);
|
||||
state.used_bytes.store(0, Ordering::Relaxed);
|
||||
state
|
||||
.last_reset_epoch_secs
|
||||
.store(now_epoch_secs, Ordering::Relaxed);
|
||||
UserQuotaSnapshot {
|
||||
used_bytes: 0,
|
||||
last_reset_epoch_secs: now_epoch_secs,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn snapshot(&self) -> HashMap<String, UserQuotaSnapshot> {
|
||||
let mut out = HashMap::new();
|
||||
for entry in self.users.iter() {
|
||||
let state = entry.value();
|
||||
let used_bytes = state.used();
|
||||
let last_reset_epoch_secs = state.last_reset_epoch_secs.load(Ordering::Relaxed);
|
||||
if used_bytes == 0 && last_reset_epoch_secs == 0 {
|
||||
continue;
|
||||
}
|
||||
out.insert(
|
||||
entry.key().clone(),
|
||||
UserQuotaSnapshot {
|
||||
used_bytes,
|
||||
last_reset_epoch_secs,
|
||||
},
|
||||
);
|
||||
}
|
||||
out
|
||||
}
|
||||
}
|
||||
|
||||
impl UserQuotaCounters {
|
||||
#[inline]
|
||||
pub(crate) fn used(&self) -> u64 {
|
||||
self.used_bytes.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub(crate) fn charge(&self, bytes: u64) -> u64 {
|
||||
self.used_bytes
|
||||
.fetch_add(bytes, Ordering::Relaxed)
|
||||
.saturating_add(bytes)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub(crate) fn refund(&self, bytes: u64) {
|
||||
let mut current = self.used_bytes.load(Ordering::Relaxed);
|
||||
loop {
|
||||
let next = current.saturating_sub(bytes);
|
||||
match self.used_bytes.compare_exchange_weak(
|
||||
current,
|
||||
next,
|
||||
Ordering::Relaxed,
|
||||
Ordering::Relaxed,
|
||||
) {
|
||||
Ok(_) => return,
|
||||
Err(observed) => current = observed,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub(crate) fn try_reserve(&self, bytes: u64, limit: u64) -> Result<u64, QuotaReserveError> {
|
||||
let current = self.used_bytes.load(Ordering::Relaxed);
|
||||
if bytes > limit.saturating_sub(current) {
|
||||
return Err(QuotaReserveError::LimitExceeded);
|
||||
}
|
||||
|
||||
let next = current.saturating_add(bytes);
|
||||
match self.used_bytes.compare_exchange_weak(
|
||||
current,
|
||||
next,
|
||||
Ordering::Relaxed,
|
||||
Ordering::Relaxed,
|
||||
) {
|
||||
Ok(_) => Ok(next),
|
||||
Err(_) => Err(QuotaReserveError::Contended),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::stats::Stats;
|
||||
|
||||
#[test]
|
||||
fn quota_counters_are_shared_across_stats_generations() {
|
||||
let store = Arc::new(QuotaStore::default());
|
||||
let first = Stats::with_quota_store(store.clone());
|
||||
store.user("alice").charge(512);
|
||||
assert_eq!(first.get_user_quota_used("alice"), 512);
|
||||
|
||||
let second = Stats::with_quota_store(store);
|
||||
assert_eq!(second.get_user_quota_used("alice"), 512);
|
||||
second.reset_user_quota("alice");
|
||||
assert_eq!(first.get_user_quota_used("alice"), 0);
|
||||
}
|
||||
}
|
||||
@@ -119,51 +119,21 @@ impl Stats {
|
||||
}
|
||||
|
||||
pub fn get_user_quota_used(&self, user: &str) -> u64 {
|
||||
self.user_stats
|
||||
.get(user)
|
||||
.map(|s| s.quota_used.load(Ordering::Relaxed))
|
||||
.unwrap_or(0)
|
||||
self.quota_store.used(user)
|
||||
}
|
||||
|
||||
pub fn load_user_quota_state(&self, user: &str, used_bytes: u64, last_reset_epoch_secs: u64) {
|
||||
let stats = self.get_or_create_user_stats_handle(user);
|
||||
stats.quota_used.store(used_bytes, Ordering::Relaxed);
|
||||
stats
|
||||
.quota_last_reset_epoch_secs
|
||||
.store(last_reset_epoch_secs, Ordering::Relaxed);
|
||||
self.quota_store
|
||||
.load(user, used_bytes, last_reset_epoch_secs);
|
||||
}
|
||||
|
||||
pub fn reset_user_quota(&self, user: &str) -> UserQuotaSnapshot {
|
||||
let stats = self.get_or_create_user_stats_handle(user);
|
||||
let last_reset_epoch_secs = Self::now_epoch_secs();
|
||||
stats.quota_used.store(0, Ordering::Relaxed);
|
||||
stats
|
||||
.quota_last_reset_epoch_secs
|
||||
.store(last_reset_epoch_secs, Ordering::Relaxed);
|
||||
UserQuotaSnapshot {
|
||||
used_bytes: 0,
|
||||
last_reset_epoch_secs,
|
||||
}
|
||||
self.quota_store.reset(user, last_reset_epoch_secs)
|
||||
}
|
||||
|
||||
pub fn user_quota_snapshot(&self) -> HashMap<String, UserQuotaSnapshot> {
|
||||
let mut out = HashMap::new();
|
||||
for entry in self.user_stats.iter() {
|
||||
let stats = entry.value();
|
||||
let used_bytes = stats.quota_used.load(Ordering::Relaxed);
|
||||
let last_reset_epoch_secs = stats.quota_last_reset_epoch_secs.load(Ordering::Relaxed);
|
||||
if used_bytes == 0 && last_reset_epoch_secs == 0 {
|
||||
continue;
|
||||
}
|
||||
out.insert(
|
||||
entry.key().clone(),
|
||||
UserQuotaSnapshot {
|
||||
used_bytes,
|
||||
last_reset_epoch_secs,
|
||||
},
|
||||
);
|
||||
}
|
||||
out
|
||||
self.quota_store.snapshot()
|
||||
}
|
||||
|
||||
pub fn get_handshake_timeouts(&self) -> u64 {
|
||||
|
||||
@@ -88,6 +88,56 @@ impl Stats {
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
/// Publishes the configured resident-memory limit applied to each ME writer.
|
||||
pub fn set_me_writer_byte_budget_limit_bytes(&self, bytes: usize) {
|
||||
self.me_writer_byte_budget_limit_bytes_gauge
|
||||
.store(bytes as u64, Ordering::Relaxed);
|
||||
}
|
||||
pub(crate) fn add_me_writer_byte_budget_queued_bytes(&self, bytes: u64) {
|
||||
self.me_writer_byte_budget_queued_bytes_gauge
|
||||
.fetch_add(bytes, Ordering::Relaxed);
|
||||
}
|
||||
pub(crate) fn move_me_writer_byte_budget_to_inflight(&self, bytes: u64) {
|
||||
let _ = self.me_writer_byte_budget_queued_bytes_gauge.fetch_update(
|
||||
Ordering::Relaxed,
|
||||
Ordering::Relaxed,
|
||||
|current| Some(current.saturating_sub(bytes)),
|
||||
);
|
||||
self.me_writer_byte_budget_inflight_bytes_gauge
|
||||
.fetch_add(bytes, Ordering::Relaxed);
|
||||
}
|
||||
pub(crate) fn release_me_writer_byte_budget_queued_bytes(&self, bytes: u64) {
|
||||
let _ = self.me_writer_byte_budget_queued_bytes_gauge.fetch_update(
|
||||
Ordering::Relaxed,
|
||||
Ordering::Relaxed,
|
||||
|current| Some(current.saturating_sub(bytes)),
|
||||
);
|
||||
}
|
||||
pub(crate) fn release_me_writer_byte_budget_inflight_bytes(&self, bytes: u64) {
|
||||
let _ = self
|
||||
.me_writer_byte_budget_inflight_bytes_gauge
|
||||
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
|
||||
Some(current.saturating_sub(bytes))
|
||||
});
|
||||
}
|
||||
pub(crate) fn increment_me_writer_byte_budget_wait_total(&self) {
|
||||
if self.telemetry_me_allows_normal() {
|
||||
self.me_writer_byte_budget_wait_total
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
pub(crate) fn increment_me_writer_byte_budget_timeout_total(&self) {
|
||||
if self.telemetry_me_allows_normal() {
|
||||
self.me_writer_byte_budget_timeout_total
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
pub(crate) fn increment_me_writer_byte_budget_oversize_total(&self) {
|
||||
if self.telemetry_me_allows_normal() {
|
||||
self.me_writer_byte_budget_oversize_total
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
pub fn increment_me_socks_kdf_strict_reject(&self) {
|
||||
if self.telemetry_me_allows_normal() {
|
||||
self.me_socks_kdf_strict_reject
|
||||
@@ -502,6 +552,14 @@ impl Stats {
|
||||
}
|
||||
}
|
||||
|
||||
/// Publishes the cumulative count of non-standard pool buffer replacements.
|
||||
pub fn set_buffer_pool_replaced_nonstandard_total(&self, value: usize) {
|
||||
if self.telemetry_me_allows_normal() {
|
||||
self.buffer_pool_replaced_nonstandard_total
|
||||
.store(value as u64, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn increment_me_c2me_send_full_total(&self) {
|
||||
if self.telemetry_me_allows_normal() {
|
||||
self.me_c2me_send_full_total.fetch_add(1, Ordering::Relaxed);
|
||||
|
||||
@@ -68,8 +68,8 @@ use crate::crypto::AesCtr;
|
||||
/// Actual limit is supplied at runtime from configuration.
|
||||
const DEFAULT_MAX_PENDING_WRITE: usize = 64 * 1024;
|
||||
|
||||
/// Default read buffer capacity (reader mostly decrypts in-place into caller buffer).
|
||||
const DEFAULT_READ_CAPACITY: usize = 16 * 1024;
|
||||
/// Maximum scratch capacity retained after a completed write.
|
||||
const MAX_RETAINED_SCRATCH_CAPACITY: usize = 32 * 1024;
|
||||
|
||||
// ============= CryptoReader State =============
|
||||
|
||||
@@ -110,10 +110,6 @@ pub struct CryptoReader<R> {
|
||||
upstream: R,
|
||||
decryptor: AesCtr,
|
||||
state: CryptoReaderState,
|
||||
|
||||
/// Reserved for future coalescing optimizations.
|
||||
#[allow(dead_code)]
|
||||
read_buf: BytesMut,
|
||||
}
|
||||
|
||||
impl<R> CryptoReader<R> {
|
||||
@@ -122,7 +118,6 @@ impl<R> CryptoReader<R> {
|
||||
upstream,
|
||||
decryptor,
|
||||
state: CryptoReaderState::Idle,
|
||||
read_buf: BytesMut::with_capacity(DEFAULT_READ_CAPACITY),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -321,7 +316,7 @@ struct PendingCiphertext {
|
||||
impl PendingCiphertext {
|
||||
fn new(max_len: usize) -> Self {
|
||||
Self {
|
||||
buf: BytesMut::with_capacity(16 * 1024),
|
||||
buf: BytesMut::new(),
|
||||
pos: 0,
|
||||
max_len,
|
||||
}
|
||||
@@ -372,15 +367,12 @@ impl PendingCiphertext {
|
||||
}
|
||||
}
|
||||
|
||||
/// Replace the entire pending ciphertext by moving `src` in (swap, no copy).
|
||||
fn replace_with(&mut self, mut src: BytesMut) {
|
||||
/// Replace the entire pending ciphertext by moving `src` in without copying.
|
||||
fn replace_with(&mut self, src: BytesMut) {
|
||||
debug_assert!(src.len() <= self.max_len);
|
||||
|
||||
self.buf.clear();
|
||||
self.buf = src;
|
||||
self.pos = 0;
|
||||
|
||||
// Swap: keep allocations hot and avoid copying bytes.
|
||||
std::mem::swap(&mut self.buf, &mut src);
|
||||
}
|
||||
|
||||
/// Append plaintext and encrypt appended range in-place.
|
||||
@@ -465,7 +457,7 @@ impl<W> CryptoWriter<W> {
|
||||
upstream,
|
||||
encryptor,
|
||||
state: CryptoWriterState::Idle,
|
||||
scratch: BytesMut::with_capacity(16 * 1024),
|
||||
scratch: BytesMut::new(),
|
||||
max_pending_write: max_pending.max(4 * 1024),
|
||||
}
|
||||
}
|
||||
@@ -502,6 +494,7 @@ impl<W> CryptoWriter<W> {
|
||||
}
|
||||
|
||||
fn poison(&mut self, error: io::Error) {
|
||||
self.scratch = BytesMut::new();
|
||||
self.state = CryptoWriterState::Poisoned { error: Some(error) };
|
||||
}
|
||||
|
||||
@@ -552,6 +545,15 @@ impl<W> CryptoWriter<W> {
|
||||
scratch.extend_from_slice(plaintext);
|
||||
encryptor.apply(&mut scratch[..]);
|
||||
}
|
||||
|
||||
/// Clear reusable scratch while releasing allocations inflated by large writes.
|
||||
fn recycle_scratch(scratch: &mut BytesMut) {
|
||||
if scratch.capacity() > MAX_RETAINED_SCRATCH_CAPACITY {
|
||||
*scratch = BytesMut::new();
|
||||
} else {
|
||||
scratch.clear();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<W: AsyncWrite + Unpin> CryptoWriter<W> {
|
||||
@@ -698,13 +700,13 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for CryptoWriter<W> {
|
||||
|
||||
Poll::Ready(Ok(n)) => {
|
||||
if n == this.scratch.len() {
|
||||
this.scratch.clear();
|
||||
Self::recycle_scratch(&mut this.scratch);
|
||||
return Poll::Ready(Ok(to_accept));
|
||||
}
|
||||
|
||||
// Partial upstream write of ciphertext
|
||||
let remainder = this.scratch.split_off(n);
|
||||
this.scratch.clear();
|
||||
let mut remainder = std::mem::take(&mut this.scratch);
|
||||
let _ = remainder.split_to(n);
|
||||
|
||||
let pending = Self::ensure_pending(&mut this.state, this.max_pending_write);
|
||||
pending.replace_with(remainder);
|
||||
|
||||
@@ -40,7 +40,7 @@ use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf};
|
||||
|
||||
use super::state::{HeaderBuffer, StreamState, WriteBuffer, YieldBuffer};
|
||||
use super::state::{HeaderBuffer, StreamState, YieldBuffer};
|
||||
use crate::protocol::constants::{
|
||||
MAX_TLS_CIPHERTEXT_SIZE, MAX_TLS_PLAINTEXT_SIZE, TLS_RECORD_ALERT, TLS_RECORD_APPLICATION,
|
||||
TLS_RECORD_CHANGE_CIPHER, TLS_RECORD_HANDSHAKE, TLS_VERSION,
|
||||
@@ -59,6 +59,9 @@ const MAX_TLS_PAYLOAD: usize = MAX_TLS_CIPHERTEXT_SIZE;
|
||||
/// Note: we never queue unlimited amount of data here; state holds at most one record.
|
||||
const MAX_PENDING_WRITE: usize = 64 * 1024;
|
||||
|
||||
/// Maximum record buffer capacity retained between writes.
|
||||
const MAX_RETAINED_RECORD_CAPACITY: usize = 2 * (TLS_HEADER_SIZE + MAX_TLS_PAYLOAD);
|
||||
|
||||
// ============= TLS Record Types =============
|
||||
|
||||
/// Parsed TLS record header (5 bytes)
|
||||
@@ -639,7 +642,7 @@ enum TlsWriterState {
|
||||
|
||||
/// Writing a complete TLS record (header + body), possibly partially
|
||||
WritingRecord {
|
||||
record: WriteBuffer,
|
||||
position: usize,
|
||||
payload_size: usize,
|
||||
},
|
||||
|
||||
@@ -676,6 +679,7 @@ impl StreamState for TlsWriterState {
|
||||
pub struct FakeTlsWriter<W> {
|
||||
upstream: W,
|
||||
state: TlsWriterState,
|
||||
record_buffer: BytesMut,
|
||||
}
|
||||
|
||||
impl<W> FakeTlsWriter<W> {
|
||||
@@ -683,6 +687,7 @@ impl<W> FakeTlsWriter<W> {
|
||||
Self {
|
||||
upstream,
|
||||
state: TlsWriterState::Idle,
|
||||
record_buffer: BytesMut::new(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -704,10 +709,15 @@ impl<W> FakeTlsWriter<W> {
|
||||
}
|
||||
|
||||
pub fn has_pending(&self) -> bool {
|
||||
matches!(&self.state, TlsWriterState::WritingRecord { record, .. } if !record.is_empty())
|
||||
matches!(
|
||||
&self.state,
|
||||
TlsWriterState::WritingRecord { position, .. }
|
||||
if *position < self.record_buffer.len()
|
||||
)
|
||||
}
|
||||
|
||||
fn poison(&mut self, error: io::Error) {
|
||||
self.record_buffer = BytesMut::new();
|
||||
self.state = TlsWriterState::Poisoned { error: Some(error) };
|
||||
}
|
||||
|
||||
@@ -720,17 +730,26 @@ impl<W> FakeTlsWriter<W> {
|
||||
}
|
||||
}
|
||||
|
||||
fn build_record(data: &[u8]) -> BytesMut {
|
||||
fn build_record(&mut self, data: &[u8]) {
|
||||
let header = TlsRecordHeader {
|
||||
record_type: TLS_RECORD_APPLICATION,
|
||||
version: TLS_VERSION,
|
||||
length: data.len() as u16,
|
||||
};
|
||||
|
||||
let mut record = BytesMut::with_capacity(TLS_HEADER_SIZE + data.len());
|
||||
record.extend_from_slice(&header.to_bytes());
|
||||
record.extend_from_slice(data);
|
||||
record
|
||||
self.record_buffer.clear();
|
||||
self.record_buffer.reserve(TLS_HEADER_SIZE + data.len());
|
||||
self.record_buffer.extend_from_slice(&header.to_bytes());
|
||||
self.record_buffer.extend_from_slice(data);
|
||||
debug_assert!(self.record_buffer.len() <= MAX_PENDING_WRITE);
|
||||
}
|
||||
|
||||
fn recycle_record_buffer(&mut self) {
|
||||
if self.record_buffer.capacity() > MAX_RETAINED_RECORD_CAPACITY {
|
||||
self.record_buffer = BytesMut::new();
|
||||
} else {
|
||||
self.record_buffer.clear();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -744,10 +763,11 @@ impl<W: AsyncWrite + Unpin> FakeTlsWriter<W> {
|
||||
fn poll_flush_record_inner(
|
||||
upstream: &mut W,
|
||||
cx: &mut Context<'_>,
|
||||
record: &mut WriteBuffer,
|
||||
record: &[u8],
|
||||
position: &mut usize,
|
||||
) -> FlushResult {
|
||||
while !record.is_empty() {
|
||||
let data = record.pending();
|
||||
while *position < record.len() {
|
||||
let data = &record[*position..];
|
||||
match Pin::new(&mut *upstream).poll_write(cx, data) {
|
||||
Poll::Pending => return FlushResult::Pending,
|
||||
Poll::Ready(Err(e)) => return FlushResult::Error(e),
|
||||
@@ -757,7 +777,7 @@ impl<W: AsyncWrite + Unpin> FakeTlsWriter<W> {
|
||||
"upstream returned 0 bytes written",
|
||||
));
|
||||
}
|
||||
Poll::Ready(Ok(n)) => record.advance(n),
|
||||
Poll::Ready(Ok(n)) => *position += n,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -780,14 +800,19 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
|
||||
}
|
||||
|
||||
TlsWriterState::WritingRecord {
|
||||
mut record,
|
||||
mut position,
|
||||
payload_size,
|
||||
} => {
|
||||
// Finish writing previous record before accepting new bytes.
|
||||
match Self::poll_flush_record_inner(&mut this.upstream, cx, &mut record) {
|
||||
match Self::poll_flush_record_inner(
|
||||
&mut this.upstream,
|
||||
cx,
|
||||
&this.record_buffer,
|
||||
&mut position,
|
||||
) {
|
||||
FlushResult::Pending => {
|
||||
this.state = TlsWriterState::WritingRecord {
|
||||
record,
|
||||
position,
|
||||
payload_size,
|
||||
};
|
||||
return Poll::Pending;
|
||||
@@ -797,6 +822,7 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
|
||||
return Poll::Ready(Err(e));
|
||||
}
|
||||
FlushResult::Complete(_) => {
|
||||
this.recycle_record_buffer();
|
||||
this.state = TlsWriterState::Idle;
|
||||
// continue to accept new buf below
|
||||
}
|
||||
@@ -818,19 +844,17 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
|
||||
let chunk = &buf[..chunk_size];
|
||||
|
||||
// Build the complete record (header + payload)
|
||||
let record_data = Self::build_record(chunk);
|
||||
this.build_record(chunk);
|
||||
|
||||
match Pin::new(&mut this.upstream).poll_write(cx, &record_data) {
|
||||
Poll::Ready(Ok(n)) if n == record_data.len() => Poll::Ready(Ok(chunk_size)),
|
||||
match Pin::new(&mut this.upstream).poll_write(cx, &this.record_buffer) {
|
||||
Poll::Ready(Ok(n)) if n == this.record_buffer.len() => {
|
||||
this.recycle_record_buffer();
|
||||
Poll::Ready(Ok(chunk_size))
|
||||
}
|
||||
|
||||
Poll::Ready(Ok(n)) => {
|
||||
// Partial write of the record: store remainder.
|
||||
let mut write_buffer = WriteBuffer::with_max_size(MAX_PENDING_WRITE);
|
||||
// record_data length is <= 16389, fits MAX_PENDING_WRITE
|
||||
let _ = write_buffer.extend(&record_data[n..]);
|
||||
|
||||
this.state = TlsWriterState::WritingRecord {
|
||||
record: write_buffer,
|
||||
position: n,
|
||||
payload_size: chunk_size,
|
||||
};
|
||||
|
||||
@@ -844,12 +868,8 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
|
||||
}
|
||||
|
||||
Poll::Pending => {
|
||||
// Buffer entire record and report success for this chunk.
|
||||
let mut write_buffer = WriteBuffer::with_max_size(MAX_PENDING_WRITE);
|
||||
let _ = write_buffer.extend(&record_data);
|
||||
|
||||
this.state = TlsWriterState::WritingRecord {
|
||||
record: write_buffer,
|
||||
position: 0,
|
||||
payload_size: chunk_size,
|
||||
};
|
||||
|
||||
@@ -871,12 +891,17 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
|
||||
}
|
||||
|
||||
TlsWriterState::WritingRecord {
|
||||
mut record,
|
||||
mut position,
|
||||
payload_size,
|
||||
} => match Self::poll_flush_record_inner(&mut this.upstream, cx, &mut record) {
|
||||
} => match Self::poll_flush_record_inner(
|
||||
&mut this.upstream,
|
||||
cx,
|
||||
&this.record_buffer,
|
||||
&mut position,
|
||||
) {
|
||||
FlushResult::Pending => {
|
||||
this.state = TlsWriterState::WritingRecord {
|
||||
record,
|
||||
position,
|
||||
payload_size,
|
||||
};
|
||||
return Poll::Pending;
|
||||
@@ -886,6 +911,7 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
|
||||
return Poll::Ready(Err(e));
|
||||
}
|
||||
FlushResult::Complete(_) => {
|
||||
this.recycle_record_buffer();
|
||||
this.state = TlsWriterState::Idle;
|
||||
}
|
||||
},
|
||||
@@ -905,11 +931,17 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
|
||||
|
||||
match state {
|
||||
TlsWriterState::WritingRecord {
|
||||
mut record,
|
||||
mut position,
|
||||
payload_size: _,
|
||||
} => {
|
||||
// Best-effort flush (do not block shutdown forever).
|
||||
let _ = Self::poll_flush_record_inner(&mut this.upstream, cx, &mut record);
|
||||
let _ = Self::poll_flush_record_inner(
|
||||
&mut this.upstream,
|
||||
cx,
|
||||
&this.record_buffer,
|
||||
&mut position,
|
||||
);
|
||||
this.recycle_record_buffer();
|
||||
this.state = TlsWriterState::Idle;
|
||||
}
|
||||
_ => {
|
||||
|
||||
@@ -4,6 +4,7 @@ use tokio::sync::watch;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::config::{ProxyConfig, SynLimitMode};
|
||||
use crate::maestro::generation::RuntimeWatchState;
|
||||
|
||||
mod command;
|
||||
mod iptables;
|
||||
@@ -15,28 +16,105 @@ use self::model::{SynLimitNamespace, synlimit_namespace, synlimit_targets};
|
||||
|
||||
static ACTIVE_SYNLIMIT_NAMESPACE: Mutex<Option<SynLimitNamespace>> = Mutex::new(None);
|
||||
|
||||
pub(crate) fn spawn_synlimit_controller(config_rx: watch::Receiver<Arc<ProxyConfig>>) {
|
||||
/// Spawns the process-scoped SYN limiter reconciler for active generations.
|
||||
pub(crate) fn spawn_synlimit_controller(
|
||||
runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
||||
) -> tokio::task::JoinHandle<()> {
|
||||
if !cfg!(target_os = "linux") {
|
||||
if has_synlimit_config(&config_rx.borrow()) {
|
||||
warn!("SYN limiter is configured but unsupported on this OS; skipping netfilter rules");
|
||||
}
|
||||
return;
|
||||
return tokio::spawn(watch_active_runtime_configs(
|
||||
runtime_watch_rx,
|
||||
|_generation_id, cfg| async move {
|
||||
if has_synlimit_config(&cfg) {
|
||||
warn!(
|
||||
"SYN limiter is configured but unsupported on this OS; skipping netfilter rules"
|
||||
);
|
||||
}
|
||||
},
|
||||
));
|
||||
}
|
||||
|
||||
tokio::spawn(async move {
|
||||
wait_for_config_channel_close_and_reconcile(config_rx).await;
|
||||
if let Err(error) = clear_synlimit_rules_all_backends().await {
|
||||
warn!(error = %error, "Failed to clear SYN limiter rules after config channel close");
|
||||
}
|
||||
});
|
||||
tokio::spawn(watch_active_runtime_configs(
|
||||
runtime_watch_rx,
|
||||
|_generation_id, cfg| async move {
|
||||
reconcile_synlimit_rules(&cfg).await;
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
async fn wait_for_config_channel_close_and_reconcile(
|
||||
mut config_rx: watch::Receiver<Arc<ProxyConfig>>,
|
||||
) {
|
||||
while config_rx.changed().await.is_ok() {
|
||||
let cfg = config_rx.borrow_and_update().clone();
|
||||
reconcile_synlimit_rules(&cfg).await;
|
||||
async fn watch_active_runtime_configs<F, Fut>(
|
||||
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
||||
mut on_config: F,
|
||||
)
|
||||
where
|
||||
F: FnMut(u64, Arc<ProxyConfig>) -> Fut,
|
||||
Fut: std::future::Future<Output = ()>,
|
||||
{
|
||||
let mut current = loop {
|
||||
if let Some(state) = runtime_watch_rx.borrow().clone() {
|
||||
break state;
|
||||
}
|
||||
if runtime_watch_rx.changed().await.is_err() {
|
||||
return;
|
||||
}
|
||||
};
|
||||
let initial_config = current.config_rx.borrow().clone();
|
||||
on_config(current.generation_id, initial_config).await;
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
biased;
|
||||
changed = runtime_watch_rx.changed() => {
|
||||
if changed.is_err() {
|
||||
break;
|
||||
}
|
||||
let Some(next) = runtime_watch_rx.borrow().clone() else {
|
||||
continue;
|
||||
};
|
||||
if next.generation_id != current.generation_id {
|
||||
current = next;
|
||||
let config = current.config_rx.borrow().clone();
|
||||
on_config(current.generation_id, config).await;
|
||||
}
|
||||
}
|
||||
changed = current.config_rx.changed() => {
|
||||
if changed.is_err() {
|
||||
let Some(next) = wait_for_new_runtime(
|
||||
&mut runtime_watch_rx,
|
||||
current.generation_id,
|
||||
).await else {
|
||||
break;
|
||||
};
|
||||
current = next;
|
||||
let config = current.config_rx.borrow().clone();
|
||||
on_config(current.generation_id, config).await;
|
||||
continue;
|
||||
}
|
||||
let active_generation_id = runtime_watch_rx
|
||||
.borrow()
|
||||
.as_ref()
|
||||
.map(|state| state.generation_id);
|
||||
if active_generation_id == Some(current.generation_id) {
|
||||
let cfg = current.config_rx.borrow_and_update().clone();
|
||||
on_config(current.generation_id, cfg).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn wait_for_new_runtime(
|
||||
runtime_watch_rx: &mut watch::Receiver<Option<RuntimeWatchState>>,
|
||||
previous_generation_id: u64,
|
||||
) -> Option<RuntimeWatchState> {
|
||||
loop {
|
||||
if let Some(state) = runtime_watch_rx.borrow().clone()
|
||||
&& state.generation_id != previous_generation_id
|
||||
{
|
||||
return Some(state);
|
||||
}
|
||||
if runtime_watch_rx.changed().await.is_err() {
|
||||
return None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -168,3 +246,70 @@ fn has_synlimit_config(cfg: &ProxyConfig) -> bool {
|
||||
.iter()
|
||||
.any(|listener| !matches!(listener.synlimit, SynLimitMode::Off))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
fn runtime_state(
|
||||
generation_id: u64,
|
||||
max_connections: u32,
|
||||
) -> (
|
||||
RuntimeWatchState,
|
||||
watch::Sender<Arc<ProxyConfig>>,
|
||||
watch::Sender<bool>,
|
||||
) {
|
||||
let mut config = ProxyConfig::default();
|
||||
config.server.max_connections = max_connections;
|
||||
let (config_tx, config_rx) = watch::channel(Arc::new(config));
|
||||
let (admission_tx, admission_rx) = watch::channel(true);
|
||||
(
|
||||
RuntimeWatchState {
|
||||
generation_id,
|
||||
config_rx,
|
||||
admission_rx,
|
||||
},
|
||||
config_tx,
|
||||
admission_tx,
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn config_watcher_ignores_retired_generation_updates() {
|
||||
let (initial, initial_config_tx, _initial_admission_tx) = runtime_state(1, 10);
|
||||
let (runtime_tx, runtime_rx) = watch::channel(Some(initial));
|
||||
let (observed_tx, mut observed_rx) = mpsc::unbounded_channel();
|
||||
let watcher = tokio::spawn(watch_active_runtime_configs(
|
||||
runtime_rx,
|
||||
move |generation_id, cfg| {
|
||||
let observed_tx = observed_tx.clone();
|
||||
async move {
|
||||
let _ = observed_tx.send((generation_id, cfg.server.max_connections));
|
||||
}
|
||||
},
|
||||
));
|
||||
|
||||
assert_eq!(observed_rx.recv().await, Some((1, 10)));
|
||||
let (next, next_config_tx, _next_admission_tx) = runtime_state(2, 20);
|
||||
runtime_tx.send_replace(Some(next));
|
||||
assert_eq!(observed_rx.recv().await, Some((2, 20)));
|
||||
|
||||
let mut stale = ProxyConfig::default();
|
||||
stale.server.max_connections = 30;
|
||||
initial_config_tx.send_replace(Arc::new(stale));
|
||||
assert!(
|
||||
tokio::time::timeout(Duration::from_millis(50), observed_rx.recv())
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
|
||||
let mut active = ProxyConfig::default();
|
||||
active.server.max_connections = 40;
|
||||
next_config_tx.send_replace(Arc::new(active));
|
||||
assert_eq!(observed_rx.recv().await, Some((2, 40)));
|
||||
|
||||
watcher.abort();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -190,6 +190,22 @@ impl TlsFrontCache {
|
||||
(snapshot, suppressed)
|
||||
}
|
||||
|
||||
/// Returns configured domains that still resolve to the synthetic default profile.
|
||||
pub(crate) async fn default_profile_domains(&self, domains: &[String]) -> Vec<String> {
|
||||
let guard = self.memory.read().await;
|
||||
domains
|
||||
.iter()
|
||||
.filter(|domain| {
|
||||
guard
|
||||
.get(domain.as_str())
|
||||
.unwrap_or(&self.default)
|
||||
.domain
|
||||
== "default"
|
||||
})
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn full_cert_sent_shard_index(client_ip: IpAddr) -> usize {
|
||||
let mut hasher = DefaultHasher::new();
|
||||
client_ip.hash(&mut hasher);
|
||||
@@ -546,6 +562,23 @@ mod tests {
|
||||
assert!(!cert_info_matches_domain(&cached));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn default_profile_domains_reports_only_unprepared_entries() {
|
||||
let domains = vec!["ready.example".to_string(), "pending.example".to_string()];
|
||||
let cache = TlsFrontCache::new(&domains, 1024, "tlsfront-test-cache");
|
||||
cache
|
||||
.set(
|
||||
"ready.example",
|
||||
cached_with_cert_info("ready.example", None, Vec::new()),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
cache.default_profile_domains(&domains).await,
|
||||
vec!["pending.example".to_string()]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_take_full_cert_budget_for_ip_uses_ttl() {
|
||||
let cache = TlsFrontCache::new(&["example.com".to_string()], 1024, "tlsfront-test-cache");
|
||||
|
||||
@@ -916,26 +916,22 @@ async fn connect_tcp_with_upstream(
|
||||
strict_route: bool,
|
||||
) -> Result<UpstreamStream> {
|
||||
if let Some(manager) = upstream {
|
||||
let resolved = if let Some(addr) = resolve_socket_addr(host, port) {
|
||||
Some(addr)
|
||||
} else {
|
||||
match tokio::net::lookup_host((host, port)).await {
|
||||
Ok(mut addrs) => addrs.find(|a| a.is_ipv4()),
|
||||
Err(e) => {
|
||||
if strict_route {
|
||||
return Err(anyhow!(
|
||||
"upstream route DNS resolution failed for {host}:{port}: {e}"
|
||||
));
|
||||
}
|
||||
warn!(
|
||||
host = %host,
|
||||
port = port,
|
||||
scope = ?scope,
|
||||
error = %e,
|
||||
"Upstream DNS resolution failed, using direct connect"
|
||||
);
|
||||
None
|
||||
let resolved = match manager.resolve_hostname(host, port).await {
|
||||
Ok(addr) => Some(addr),
|
||||
Err(e) => {
|
||||
if strict_route {
|
||||
return Err(anyhow!(
|
||||
"upstream route DNS resolution failed for {host}:{port}: {e}"
|
||||
));
|
||||
}
|
||||
warn!(
|
||||
host = %host,
|
||||
port = port,
|
||||
scope = ?scope,
|
||||
error = %e,
|
||||
"Upstream DNS resolution failed, using direct connect"
|
||||
);
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
@@ -955,6 +951,9 @@ async fn connect_tcp_with_upstream(
|
||||
error = %e,
|
||||
"Upstream connect failed, using direct connect"
|
||||
);
|
||||
return Ok(UpstreamStream::Tcp(
|
||||
timeout(connect_timeout, TcpStream::connect(addr)).await??,
|
||||
));
|
||||
}
|
||||
}
|
||||
} else if strict_route {
|
||||
|
||||
@@ -1,9 +1,12 @@
|
||||
use bytes::Bytes;
|
||||
use std::sync::Arc;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::sync::OwnedSemaphorePermit;
|
||||
|
||||
use crate::crypto::{AesCbc, crc32, crc32c};
|
||||
use crate::error::{ProxyError, Result};
|
||||
use crate::protocol::constants::*;
|
||||
use crate::stats::Stats;
|
||||
use crate::stream::PooledBuffer;
|
||||
|
||||
use super::wire::{append_proxy_req_payload_into, proxy_req_payload_len};
|
||||
@@ -11,9 +14,66 @@ use super::wire::{append_proxy_req_payload_into, proxy_req_payload_len};
|
||||
const RPC_WRITER_FRAME_BUF_SHRINK_THRESHOLD: usize = 256 * 1024;
|
||||
const RPC_WRITER_FRAME_BUF_RETAIN: usize = 64 * 1024;
|
||||
|
||||
enum WriterBytePermitState {
|
||||
Queued,
|
||||
Inflight,
|
||||
}
|
||||
|
||||
/// Holds one writer memory reservation through queueing and socket write completion.
|
||||
pub(crate) struct WriterBytePermit {
|
||||
_permit: OwnedSemaphorePermit,
|
||||
reserved_bytes: u64,
|
||||
stats: Arc<Stats>,
|
||||
state: WriterBytePermitState,
|
||||
}
|
||||
|
||||
impl WriterBytePermit {
|
||||
/// Creates a queued reservation and publishes its rounded resident-byte cost.
|
||||
pub(crate) fn new(
|
||||
permit: OwnedSemaphorePermit,
|
||||
reserved_bytes: usize,
|
||||
stats: Arc<Stats>,
|
||||
) -> Self {
|
||||
let reserved_bytes = reserved_bytes as u64;
|
||||
stats.add_me_writer_byte_budget_queued_bytes(reserved_bytes);
|
||||
Self {
|
||||
_permit: permit,
|
||||
reserved_bytes,
|
||||
stats,
|
||||
state: WriterBytePermitState::Queued,
|
||||
}
|
||||
}
|
||||
|
||||
/// Moves the reservation from queued to in-flight accounting exactly once.
|
||||
pub(crate) fn mark_inflight(&mut self) {
|
||||
if matches!(self.state, WriterBytePermitState::Queued) {
|
||||
self.stats
|
||||
.move_me_writer_byte_budget_to_inflight(self.reserved_bytes);
|
||||
self.state = WriterBytePermitState::Inflight;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for WriterBytePermit {
|
||||
fn drop(&mut self) {
|
||||
match self.state {
|
||||
WriterBytePermitState::Queued => self
|
||||
.stats
|
||||
.release_me_writer_byte_budget_queued_bytes(self.reserved_bytes),
|
||||
WriterBytePermitState::Inflight => self
|
||||
.stats
|
||||
.release_me_writer_byte_budget_inflight_bytes(self.reserved_bytes),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Commands sent to dedicated writer tasks to avoid mutex contention on TCP writes.
|
||||
pub(crate) enum WriterCommand {
|
||||
Data(Bytes),
|
||||
Data {
|
||||
payload: Bytes,
|
||||
_permit: Option<OwnedSemaphorePermit>,
|
||||
writer_permit: WriterBytePermit,
|
||||
},
|
||||
DataAndFlush(Bytes),
|
||||
ProxyReq(ProxyReqCommand),
|
||||
ControlAndFlush([u8; 12]),
|
||||
@@ -28,6 +88,8 @@ pub(crate) struct ProxyReqCommand {
|
||||
pub(crate) proto_flags: u32,
|
||||
pub(crate) proxy_tag: Option<[u8; 16]>,
|
||||
pub(crate) payload: PooledBuffer,
|
||||
pub(crate) _permit: OwnedSemaphorePermit,
|
||||
pub(crate) writer_permit: WriterBytePermit,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
@@ -82,7 +144,7 @@ fn build_rpc_frame_into(
|
||||
) {
|
||||
let total_len = 4 + 4 + payload.len() + 4;
|
||||
frame.clear();
|
||||
frame.reserve(total_len + 15);
|
||||
frame.reserve_exact(total_len + 15);
|
||||
let total_len = total_len as u32;
|
||||
frame.extend_from_slice(&total_len.to_le_bytes());
|
||||
frame.extend_from_slice(&seq_no.to_le_bytes());
|
||||
@@ -274,7 +336,7 @@ impl RpcWriter {
|
||||
);
|
||||
let total_len = 4 + 4 + payload_len + 4;
|
||||
self.frame_buf.clear();
|
||||
self.frame_buf.reserve(total_len + 15);
|
||||
self.frame_buf.reserve_exact(total_len + 15);
|
||||
self.frame_buf
|
||||
.extend_from_slice(&(total_len as u32).to_le_bytes());
|
||||
self.frame_buf.extend_from_slice(&self.seq_no.to_le_bytes());
|
||||
@@ -322,3 +384,38 @@ impl RpcWriter {
|
||||
self.writer.flush().await.map_err(ProxyError::Io)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
#[test]
|
||||
fn writer_byte_permit_tracks_queued_and_inflight_lifecycle() {
|
||||
let stats = Arc::new(Stats::default());
|
||||
let semaphore = Arc::new(Semaphore::new(2));
|
||||
let permit = semaphore
|
||||
.clone()
|
||||
.try_acquire_many_owned(2)
|
||||
.expect("writer byte permits must be available");
|
||||
let mut writer_permit = WriterBytePermit::new(permit, 32 * 1024, stats.clone());
|
||||
|
||||
assert_eq!(
|
||||
stats.get_me_writer_byte_budget_queued_bytes_gauge(),
|
||||
32 * 1024
|
||||
);
|
||||
assert_eq!(stats.get_me_writer_byte_budget_inflight_bytes_gauge(), 0);
|
||||
|
||||
writer_permit.mark_inflight();
|
||||
assert_eq!(stats.get_me_writer_byte_budget_queued_bytes_gauge(), 0);
|
||||
assert_eq!(
|
||||
stats.get_me_writer_byte_budget_inflight_bytes_gauge(),
|
||||
32 * 1024
|
||||
);
|
||||
|
||||
drop(writer_permit);
|
||||
assert_eq!(stats.get_me_writer_byte_budget_queued_bytes_gauge(), 0);
|
||||
assert_eq!(stats.get_me_writer_byte_budget_inflight_bytes_gauge(), 0);
|
||||
assert_eq!(semaphore.available_permits(), 2);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1795,6 +1795,7 @@ mod tests {
|
||||
general.me_writer_pick_sample_size,
|
||||
MeSocksKdfPolicy::default(),
|
||||
general.me_writer_cmd_channel_capacity,
|
||||
general.me_writer_byte_budget_bytes,
|
||||
general.me_route_channel_capacity,
|
||||
general.me_route_backpressure_enabled,
|
||||
general.me_route_fairshare_enabled,
|
||||
@@ -1821,6 +1822,7 @@ mod tests {
|
||||
) -> u64 {
|
||||
let (conn_id, _rx) = pool.registry.register().await;
|
||||
let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8);
|
||||
let byte_budget = pool.new_writer_byte_budget();
|
||||
let writer = MeWriter {
|
||||
id: writer_id,
|
||||
addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 4000 + writer_id as u16),
|
||||
@@ -1830,6 +1832,7 @@ mod tests {
|
||||
contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())),
|
||||
created_at: Instant::now() - Duration::from_secs(writer_id),
|
||||
tx: tx.clone(),
|
||||
byte_budget: byte_budget.clone(),
|
||||
cancel: CancellationToken::new(),
|
||||
degraded: Arc::new(AtomicBool::new(false)),
|
||||
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
|
||||
@@ -1839,7 +1842,9 @@ mod tests {
|
||||
allow_drain_fallback: Arc::new(AtomicBool::new(false)),
|
||||
};
|
||||
pool.writers.write().await.push(writer);
|
||||
pool.registry.register_writer(writer_id, tx).await;
|
||||
pool.registry
|
||||
.register_writer(writer_id, tx, byte_budget)
|
||||
.await;
|
||||
pool.conn_count.fetch_add(1, Ordering::Relaxed);
|
||||
assert!(
|
||||
pool.registry
|
||||
@@ -1860,6 +1865,7 @@ mod tests {
|
||||
|
||||
async fn insert_live_writer(pool: &Arc<MePool>, writer_id: u64, writer_dc: i32) {
|
||||
let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8);
|
||||
let byte_budget = pool.new_writer_byte_budget();
|
||||
let writer = MeWriter {
|
||||
id: writer_id,
|
||||
addr: SocketAddr::new(
|
||||
@@ -1877,6 +1883,7 @@ mod tests {
|
||||
contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())),
|
||||
created_at: Instant::now(),
|
||||
tx: tx.clone(),
|
||||
byte_budget: byte_budget.clone(),
|
||||
cancel: CancellationToken::new(),
|
||||
degraded: Arc::new(AtomicBool::new(false)),
|
||||
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
|
||||
@@ -1886,7 +1893,9 @@ mod tests {
|
||||
allow_drain_fallback: Arc::new(AtomicBool::new(false)),
|
||||
};
|
||||
pool.writers.write().await.push(writer);
|
||||
pool.registry.register_writer(writer_id, tx).await;
|
||||
pool.registry
|
||||
.register_writer(writer_id, tx, byte_budget)
|
||||
.await;
|
||||
pool.conn_count.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
|
||||
@@ -66,33 +66,13 @@ fn extract_host_port_path(url: &str) -> Result<(String, u16, String)> {
|
||||
Ok((host, port, path))
|
||||
}
|
||||
|
||||
async fn resolve_target_addr(host: &str, port: u16) -> Result<std::net::SocketAddr> {
|
||||
if let Some(addr) = resolve_socket_addr(host, port) {
|
||||
return Ok(addr);
|
||||
}
|
||||
|
||||
let addrs: Vec<std::net::SocketAddr> = tokio::net::lookup_host((host, port))
|
||||
.await
|
||||
.map_err(|e| ProxyError::Proxy(format!("DNS resolve failed for {host}:{port}: {e}")))?
|
||||
.collect();
|
||||
|
||||
if let Some(addr) = addrs.iter().copied().find(|addr| addr.is_ipv4()) {
|
||||
return Ok(addr);
|
||||
}
|
||||
|
||||
addrs
|
||||
.first()
|
||||
.copied()
|
||||
.ok_or_else(|| ProxyError::Proxy(format!("DNS returned no addresses for {host}:{port}")))
|
||||
}
|
||||
|
||||
async fn connect_https_transport(
|
||||
host: &str,
|
||||
port: u16,
|
||||
upstream: Option<Arc<UpstreamManager>>,
|
||||
) -> Result<UpstreamStream> {
|
||||
if let Some(manager) = upstream {
|
||||
let target = resolve_target_addr(host, port).await?;
|
||||
let target = manager.resolve_hostname(host, port).await?;
|
||||
return timeout(HTTP_CONNECT_TIMEOUT, manager.connect(target, None, None))
|
||||
.await
|
||||
.map_err(|_| ProxyError::Proxy(format!("upstream connect timeout for {host}:{port}")))?
|
||||
|
||||
@@ -10,7 +10,7 @@ use std::sync::atomic::{
|
||||
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use arc_swap::ArcSwap;
|
||||
use tokio::sync::{Mutex, RwLock, mpsc, watch};
|
||||
use tokio::sync::{Mutex, RwLock, Semaphore, mpsc, watch};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::config::{
|
||||
@@ -48,6 +48,8 @@ pub struct MeWriter {
|
||||
pub contour: Arc<AtomicU8>,
|
||||
pub created_at: Instant,
|
||||
pub tx: mpsc::Sender<WriterCommand>,
|
||||
/// Aggregate resident-memory budget shared by all data commands for this writer.
|
||||
pub byte_budget: Arc<Semaphore>,
|
||||
pub cancel: CancellationToken,
|
||||
pub degraded: Arc<AtomicBool>,
|
||||
pub rtt_ema_ms_x10: Arc<AtomicU32>,
|
||||
@@ -277,6 +279,7 @@ pub(super) struct WriterLifecycleCore {
|
||||
pub(super) me_keepalive_payload_random: bool,
|
||||
pub(super) rpc_proxy_req_every_secs: AtomicU64,
|
||||
pub(super) writer_cmd_channel_capacity: usize,
|
||||
pub(super) writer_byte_budget_permits: usize,
|
||||
}
|
||||
|
||||
pub(super) struct RouteRuntimeCore {
|
||||
@@ -553,6 +556,7 @@ impl MePool {
|
||||
me_writer_pick_sample_size: u8,
|
||||
me_socks_kdf_policy: MeSocksKdfPolicy,
|
||||
me_writer_cmd_channel_capacity: usize,
|
||||
me_writer_byte_budget_bytes: usize,
|
||||
me_route_channel_capacity: usize,
|
||||
me_route_backpressure_enabled: bool,
|
||||
me_route_fairshare_enabled: bool,
|
||||
@@ -583,6 +587,7 @@ impl MePool {
|
||||
);
|
||||
let (writer_epoch, _) = watch::channel(0u64);
|
||||
let now_epoch_secs = Self::now_epoch_secs();
|
||||
stats.set_me_writer_byte_budget_limit_bytes(me_writer_byte_budget_bytes);
|
||||
Arc::new(Self {
|
||||
routing: Arc::new(RoutingCore {
|
||||
registry,
|
||||
@@ -615,6 +620,9 @@ impl MePool {
|
||||
me_keepalive_payload_random,
|
||||
rpc_proxy_req_every_secs: AtomicU64::new(rpc_proxy_req_every_secs),
|
||||
writer_cmd_channel_capacity: me_writer_cmd_channel_capacity.max(1),
|
||||
writer_byte_budget_permits: me_writer_byte_budget_bytes
|
||||
.div_ceil(crate::config::defaults::ME_WRITER_BYTE_PERMIT_UNIT_BYTES)
|
||||
.max(1),
|
||||
}),
|
||||
route_runtime: Arc::new(RouteRuntimeCore {
|
||||
me_route_no_writer_mode: AtomicU8::new(me_route_no_writer_mode.as_u8()),
|
||||
@@ -833,6 +841,13 @@ impl MePool {
|
||||
})
|
||||
}
|
||||
|
||||
/// Creates the immutable byte semaphore assigned to one ME writer generation.
|
||||
pub(crate) fn new_writer_byte_budget(&self) -> Arc<Semaphore> {
|
||||
Arc::new(Semaphore::new(
|
||||
self.writer_lifecycle.writer_byte_budget_permits,
|
||||
))
|
||||
}
|
||||
|
||||
pub fn current_generation(&self) -> u64 {
|
||||
self.reinit.active_generation.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
@@ -63,13 +63,19 @@ async fn writer_command_loop(
|
||||
_ = cancel.cancelled() => return Ok(()),
|
||||
cmd = rx.recv() => {
|
||||
match cmd {
|
||||
Some(WriterCommand::Data(payload)) => {
|
||||
Some(WriterCommand::Data {
|
||||
payload,
|
||||
_permit,
|
||||
mut writer_permit,
|
||||
}) => {
|
||||
writer_permit.mark_inflight();
|
||||
rpc_writer.send(&payload).await?;
|
||||
}
|
||||
Some(WriterCommand::DataAndFlush(payload)) => {
|
||||
rpc_writer.send_and_flush(&payload).await?;
|
||||
}
|
||||
Some(WriterCommand::ProxyReq(command)) => {
|
||||
Some(WriterCommand::ProxyReq(mut command)) => {
|
||||
command.writer_permit.mark_inflight();
|
||||
rpc_writer.send_proxy_req(&command).await?;
|
||||
}
|
||||
Some(WriterCommand::ControlAndFlush(payload)) => {
|
||||
@@ -419,6 +425,7 @@ impl MePool {
|
||||
let draining_started_at_epoch_secs = Arc::new(AtomicU64::new(0));
|
||||
let drain_deadline_epoch_secs = Arc::new(AtomicU64::new(0));
|
||||
let allow_drain_fallback = Arc::new(AtomicBool::new(false));
|
||||
let byte_budget = self.new_writer_byte_budget();
|
||||
let (tx, rx) =
|
||||
mpsc::channel::<WriterCommand>(self.writer_lifecycle.writer_cmd_channel_capacity);
|
||||
let rpc_writer = RpcWriter {
|
||||
@@ -438,6 +445,7 @@ impl MePool {
|
||||
contour: contour.clone(),
|
||||
created_at: Instant::now(),
|
||||
tx: tx.clone(),
|
||||
byte_budget: byte_budget.clone(),
|
||||
cancel: cancel.clone(),
|
||||
degraded: degraded.clone(),
|
||||
rtt_ema_ms_x10: rtt_ema_ms_x10.clone(),
|
||||
@@ -449,7 +457,9 @@ impl MePool {
|
||||
self.writers
|
||||
.update(|writers| writers.push(writer.clone()))
|
||||
.await;
|
||||
self.registry.register_writer(writer_id, tx.clone()).await;
|
||||
self.registry
|
||||
.register_writer(writer_id, tx.clone(), byte_budget)
|
||||
.await;
|
||||
self.registry.mark_writer_idle(writer_id).await;
|
||||
self.conn_count.fetch_add(1, Ordering::Relaxed);
|
||||
self.notify_writer_epoch();
|
||||
|
||||
@@ -48,6 +48,8 @@ pub struct BoundConn {
|
||||
pub struct ConnWriter {
|
||||
pub writer_id: u64,
|
||||
pub tx: mpsc::Sender<WriterCommand>,
|
||||
/// Writer-local memory budget used by the hot bound-client route.
|
||||
pub byte_budget: Arc<Semaphore>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
@@ -62,7 +64,13 @@ struct RoutingTable {
|
||||
}
|
||||
|
||||
struct WriterTable {
|
||||
map: DashMap<u64, mpsc::Sender<WriterCommand>>,
|
||||
map: DashMap<u64, WriterRoute>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct WriterRoute {
|
||||
tx: mpsc::Sender<WriterCommand>,
|
||||
byte_budget: Arc<Semaphore>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
|
||||
@@ -1,10 +1,16 @@
|
||||
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
|
||||
use std::sync::Arc;
|
||||
|
||||
use bytes::Bytes;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use super::{ConnMeta, ConnRegistry, RouteResult};
|
||||
use crate::transport::middle_proxy::MeResponse;
|
||||
|
||||
fn writer_byte_budget() -> Arc<Semaphore> {
|
||||
Arc::new(Semaphore::new(2049))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn writer_activity_snapshot_tracks_writer_and_dc_load() {
|
||||
let registry = ConnRegistry::new();
|
||||
@@ -14,8 +20,12 @@ async fn writer_activity_snapshot_tracks_writer_and_dc_load() {
|
||||
let (conn_c, _rx_c) = registry.register().await;
|
||||
let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8);
|
||||
let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8);
|
||||
registry.register_writer(10, writer_tx_a.clone()).await;
|
||||
registry.register_writer(20, writer_tx_b.clone()).await;
|
||||
registry
|
||||
.register_writer(10, writer_tx_a.clone(), writer_byte_budget())
|
||||
.await;
|
||||
registry
|
||||
.register_writer(20, writer_tx_b.clone(), writer_byte_budget())
|
||||
.await;
|
||||
|
||||
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
|
||||
assert!(
|
||||
@@ -124,8 +134,12 @@ async fn bind_writer_rebinds_conn_atomically() {
|
||||
let (conn_id, _rx) = registry.register().await;
|
||||
let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8);
|
||||
let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8);
|
||||
registry.register_writer(10, writer_tx_a).await;
|
||||
registry.register_writer(20, writer_tx_b).await;
|
||||
registry
|
||||
.register_writer(10, writer_tx_a, writer_byte_budget())
|
||||
.await;
|
||||
registry
|
||||
.register_writer(20, writer_tx_b, writer_byte_budget())
|
||||
.await;
|
||||
|
||||
let client_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
|
||||
let first_our_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)), 443);
|
||||
@@ -184,8 +198,12 @@ async fn writer_lost_does_not_drop_rebound_conn() {
|
||||
let (conn_id, _rx) = registry.register().await;
|
||||
let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8);
|
||||
let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8);
|
||||
registry.register_writer(10, writer_tx_a).await;
|
||||
registry.register_writer(20, writer_tx_b).await;
|
||||
registry
|
||||
.register_writer(10, writer_tx_a, writer_byte_budget())
|
||||
.await;
|
||||
registry
|
||||
.register_writer(20, writer_tx_b, writer_byte_budget())
|
||||
.await;
|
||||
|
||||
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
|
||||
assert!(
|
||||
@@ -262,8 +280,12 @@ async fn non_empty_writer_ids_returns_only_writers_with_bound_clients() {
|
||||
let (conn_id, _rx) = registry.register().await;
|
||||
let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8);
|
||||
let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8);
|
||||
registry.register_writer(10, writer_tx_a).await;
|
||||
registry.register_writer(20, writer_tx_b).await;
|
||||
registry
|
||||
.register_writer(10, writer_tx_a, writer_byte_budget())
|
||||
.await;
|
||||
registry
|
||||
.register_writer(20, writer_tx_b, writer_byte_budget())
|
||||
.await;
|
||||
|
||||
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
|
||||
assert!(
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::time::Duration;
|
||||
|
||||
@@ -56,7 +57,13 @@ impl ConnRegistry {
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn register_writer(&self, writer_id: u64, tx: mpsc::Sender<WriterCommand>) {
|
||||
/// Registers one writer command route and its matching memory budget atomically.
|
||||
pub async fn register_writer(
|
||||
&self,
|
||||
writer_id: u64,
|
||||
tx: mpsc::Sender<WriterCommand>,
|
||||
byte_budget: Arc<tokio::sync::Semaphore>,
|
||||
) {
|
||||
let mut binding = self.binding.inner.lock().await;
|
||||
binding
|
||||
.conns_for_writer
|
||||
@@ -70,7 +77,9 @@ impl ConnRegistry {
|
||||
.writer_idle_since_epoch_secs
|
||||
.entry(writer_id)
|
||||
.or_insert_with(Self::now_epoch_secs);
|
||||
self.writers.map.insert(writer_id, tx);
|
||||
self.writers
|
||||
.map
|
||||
.insert(writer_id, super::WriterRoute { tx, byte_budget });
|
||||
}
|
||||
|
||||
/// Unregister connection, returning associated writer_id if any.
|
||||
@@ -417,7 +426,8 @@ impl ConnRegistry {
|
||||
.map(|entry| entry.value().clone())?;
|
||||
Some(ConnWriter {
|
||||
writer_id,
|
||||
tx: writer,
|
||||
tx: writer.tx,
|
||||
byte_budget: writer.byte_budget,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -438,7 +448,8 @@ impl ConnRegistry {
|
||||
Some((
|
||||
ConnWriter {
|
||||
writer_id,
|
||||
tx: writer,
|
||||
tx: writer.tx,
|
||||
byte_budget: writer.byte_budget,
|
||||
},
|
||||
meta,
|
||||
))
|
||||
|
||||
@@ -6,16 +6,18 @@ use std::sync::Arc;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::mpsc::error::TrySendError;
|
||||
use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError, mpsc};
|
||||
use tracing::{debug, warn};
|
||||
|
||||
use super::MePool;
|
||||
use super::codec::{ProxyReqCommand, WriterCommand};
|
||||
use super::codec::{ProxyReqCommand, WriterBytePermit, WriterCommand};
|
||||
use super::registry::ConnMeta;
|
||||
use super::wire::build_proxy_req_payload;
|
||||
use super::wire::{build_proxy_req_payload, proxy_req_payload_len};
|
||||
use crate::config::defaults::ME_WRITER_BYTE_PERMIT_UNIT_BYTES;
|
||||
use crate::config::{MeRouteNoWriterMode, MeWriterPickMode};
|
||||
use crate::error::{ProxyError, Result};
|
||||
use crate::stats::Stats;
|
||||
use crate::stream::PooledBuffer;
|
||||
use rand::seq::SliceRandom;
|
||||
|
||||
@@ -29,6 +31,8 @@ const PICK_PENALTY_WARM: u64 = 200;
|
||||
const PICK_PENALTY_DRAINING: u64 = 600;
|
||||
const PICK_PENALTY_STALE: u64 = 300;
|
||||
const PICK_PENALTY_DEGRADED: u64 = 250;
|
||||
const RPC_WRITER_FRAME_CAPACITY_OVERHEAD_BYTES: usize = 27;
|
||||
const LEGACY_PROXY_REQ_SOURCE_CAPACITY_OVERHEAD_BYTES: usize = 128;
|
||||
|
||||
mod close;
|
||||
mod recovery;
|
||||
@@ -39,34 +43,133 @@ enum WriterCommandReserveError {
|
||||
TimedOut,
|
||||
}
|
||||
|
||||
enum WriterByteReserveError {
|
||||
Closed,
|
||||
TimedOut,
|
||||
}
|
||||
|
||||
fn proxy_tag_array(tag: Option<&[u8]>) -> Option<[u8; 16]> {
|
||||
tag.and_then(|tag| <[u8; 16]>::try_from(tag).ok())
|
||||
}
|
||||
|
||||
fn proxy_req_payload_from_command(cmd: WriterCommand) -> Option<PooledBuffer> {
|
||||
fn proxy_req_payload_from_command(
|
||||
cmd: WriterCommand,
|
||||
) -> Option<(PooledBuffer, OwnedSemaphorePermit)> {
|
||||
match cmd {
|
||||
WriterCommand::ProxyReq(command) => Some(command.payload),
|
||||
WriterCommand::ProxyReq(command) => Some((command.payload, command._permit)),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn payload_permit_from_data_command(cmd: WriterCommand) -> Option<OwnedSemaphorePermit> {
|
||||
match cmd {
|
||||
WriterCommand::Data { _permit, .. } => _permit,
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
async fn reserve_writer_command_slot(
|
||||
tx: &mpsc::Sender<WriterCommand>,
|
||||
wait: Option<Duration>,
|
||||
deadline: Option<Instant>,
|
||||
) -> std::result::Result<mpsc::OwnedPermit<WriterCommand>, WriterCommandReserveError> {
|
||||
let reserve = tx.clone().reserve_owned();
|
||||
match wait {
|
||||
Some(wait) => match tokio::time::timeout(wait, reserve).await {
|
||||
Ok(Ok(permit)) => Ok(permit),
|
||||
Ok(Err(_)) => Err(WriterCommandReserveError::Closed),
|
||||
Err(_) => Err(WriterCommandReserveError::TimedOut),
|
||||
},
|
||||
match deadline {
|
||||
Some(deadline) => {
|
||||
match tokio::time::timeout(deadline.saturating_duration_since(Instant::now()), reserve)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(permit)) => Ok(permit),
|
||||
Ok(Err(_)) => Err(WriterCommandReserveError::Closed),
|
||||
Err(_) => Err(WriterCommandReserveError::TimedOut),
|
||||
}
|
||||
}
|
||||
None => reserve.await.map_err(|_| WriterCommandReserveError::Closed),
|
||||
}
|
||||
}
|
||||
|
||||
fn writer_send_deadline(wait: Option<Duration>) -> Option<Instant> {
|
||||
wait.map(|wait| Instant::now() + wait)
|
||||
}
|
||||
|
||||
fn writer_resident_permits(
|
||||
source_capacity: usize,
|
||||
encoded_payload_len: usize,
|
||||
) -> Option<(u32, usize)> {
|
||||
let resident_bytes = source_capacity
|
||||
.checked_add(encoded_payload_len)?
|
||||
.checked_add(RPC_WRITER_FRAME_CAPACITY_OVERHEAD_BYTES)?;
|
||||
let permits = resident_bytes.div_ceil(ME_WRITER_BYTE_PERMIT_UNIT_BYTES);
|
||||
let permits = u32::try_from(permits).ok()?;
|
||||
let reserved_bytes = (permits as usize).checked_mul(ME_WRITER_BYTE_PERMIT_UNIT_BYTES)?;
|
||||
Some((
|
||||
permits.max(1),
|
||||
reserved_bytes.max(ME_WRITER_BYTE_PERMIT_UNIT_BYTES),
|
||||
))
|
||||
}
|
||||
|
||||
fn proxy_req_resident_permits(
|
||||
source_capacity: usize,
|
||||
data_len: usize,
|
||||
proxy_tag: Option<&[u8]>,
|
||||
proto_flags: u32,
|
||||
) -> Option<(u32, usize)> {
|
||||
writer_resident_permits(
|
||||
source_capacity,
|
||||
proxy_req_payload_len(data_len, proxy_tag, proto_flags),
|
||||
)
|
||||
}
|
||||
|
||||
fn try_reserve_writer_bytes(
|
||||
byte_budget: &Arc<Semaphore>,
|
||||
permits: u32,
|
||||
reserved_bytes: usize,
|
||||
stats: &Arc<Stats>,
|
||||
) -> std::result::Result<WriterBytePermit, TryAcquireError> {
|
||||
byte_budget
|
||||
.clone()
|
||||
.try_acquire_many_owned(permits)
|
||||
.map(|permit| WriterBytePermit::new(permit, reserved_bytes, stats.clone()))
|
||||
}
|
||||
|
||||
async fn reserve_writer_bytes(
|
||||
byte_budget: &Arc<Semaphore>,
|
||||
permits: u32,
|
||||
reserved_bytes: usize,
|
||||
deadline: Option<Instant>,
|
||||
stats: &Arc<Stats>,
|
||||
) -> std::result::Result<WriterBytePermit, WriterByteReserveError> {
|
||||
match try_reserve_writer_bytes(byte_budget, permits, reserved_bytes, stats) {
|
||||
Ok(permit) => return Ok(permit),
|
||||
Err(TryAcquireError::Closed) => return Err(WriterByteReserveError::Closed),
|
||||
Err(TryAcquireError::NoPermits) => {
|
||||
stats.increment_me_writer_byte_budget_wait_total();
|
||||
}
|
||||
}
|
||||
|
||||
let acquire = byte_budget.clone().acquire_many_owned(permits);
|
||||
match deadline {
|
||||
Some(deadline) => {
|
||||
match tokio::time::timeout(deadline.saturating_duration_since(Instant::now()), acquire)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(permit)) => Ok(WriterBytePermit::new(permit, reserved_bytes, stats.clone())),
|
||||
Ok(Err(_)) => Err(WriterByteReserveError::Closed),
|
||||
Err(_) => {
|
||||
stats.increment_me_writer_byte_budget_timeout_total();
|
||||
Err(WriterByteReserveError::TimedOut)
|
||||
}
|
||||
}
|
||||
}
|
||||
None => acquire
|
||||
.await
|
||||
.map(|permit| WriterBytePermit::new(permit, reserved_bytes, stats.clone()))
|
||||
.map_err(|_| WriterByteReserveError::Closed),
|
||||
}
|
||||
}
|
||||
|
||||
impl MePool {
|
||||
/// Send RPC_PROXY_REQ. `tag_override`: per-user ad_tag (from access.user_ad_tags); if None, uses pool default.
|
||||
/// `payload_permit` keeps optional client byte accounting alive until the writer consumes the command.
|
||||
pub async fn send_proxy_req(
|
||||
self: &Arc<Self>,
|
||||
conn_id: u64,
|
||||
@@ -76,8 +179,32 @@ impl MePool {
|
||||
data: &[u8],
|
||||
proto_flags: u32,
|
||||
tag_override: Option<&[u8]>,
|
||||
mut payload_permit: Option<OwnedSemaphorePermit>,
|
||||
) -> Result<()> {
|
||||
let tag = tag_override.or(self.proxy_tag.as_deref());
|
||||
let Some(source_capacity) = data
|
||||
.len()
|
||||
.checked_add(LEGACY_PROXY_REQ_SOURCE_CAPACITY_OVERHEAD_BYTES)
|
||||
else {
|
||||
self.stats.increment_me_writer_byte_budget_oversize_total();
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer payload residency calculation overflow".into(),
|
||||
));
|
||||
};
|
||||
let Some((writer_byte_permits, writer_reserved_bytes)) =
|
||||
proxy_req_resident_permits(source_capacity, data.len(), tag, proto_flags)
|
||||
else {
|
||||
self.stats.increment_me_writer_byte_budget_oversize_total();
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer payload residency calculation overflow".into(),
|
||||
));
|
||||
};
|
||||
if writer_byte_permits as usize > self.writer_lifecycle.writer_byte_budget_permits {
|
||||
self.stats.increment_me_writer_byte_budget_oversize_total();
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer payload exceeds configured byte budget".into(),
|
||||
));
|
||||
}
|
||||
let build_routed_payload = |effective_our_addr: SocketAddr| {
|
||||
(
|
||||
build_proxy_req_payload(
|
||||
@@ -118,19 +245,48 @@ impl MePool {
|
||||
loop {
|
||||
if let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await
|
||||
{
|
||||
let deadline =
|
||||
writer_send_deadline(self.route_runtime.me_route_blocking_send_timeout);
|
||||
let writer_permit = match reserve_writer_bytes(
|
||||
¤t.byte_budget,
|
||||
writer_byte_permits,
|
||||
writer_reserved_bytes,
|
||||
deadline,
|
||||
&self.stats,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(permit) => permit,
|
||||
Err(WriterByteReserveError::TimedOut) => {
|
||||
self.stats
|
||||
.increment_me_writer_pick_full_total(self.writer_pick_mode());
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer byte budget full within blocking send timeout".into(),
|
||||
));
|
||||
}
|
||||
Err(WriterByteReserveError::Closed) => {
|
||||
warn!(
|
||||
writer_id = current.writer_id,
|
||||
"ME writer byte budget closed"
|
||||
);
|
||||
self.remove_writer_and_close_clients(current.writer_id)
|
||||
.await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let (current_payload, _) = build_routed_payload(current_meta.our_addr);
|
||||
match current.tx.try_send(WriterCommand::Data(current_payload)) {
|
||||
let command = WriterCommand::Data {
|
||||
payload: current_payload,
|
||||
_permit: payload_permit.take(),
|
||||
writer_permit,
|
||||
};
|
||||
match current.tx.try_send(command) {
|
||||
Ok(()) => {
|
||||
self.note_hybrid_route_success();
|
||||
return Ok(());
|
||||
}
|
||||
Err(TrySendError::Full(cmd)) => {
|
||||
match reserve_writer_command_slot(
|
||||
¤t.tx,
|
||||
self.route_runtime.me_route_blocking_send_timeout,
|
||||
)
|
||||
.await
|
||||
{
|
||||
match reserve_writer_command_slot(¤t.tx, deadline).await {
|
||||
Ok(permit) => {
|
||||
permit.send(cmd);
|
||||
self.note_hybrid_route_success();
|
||||
@@ -143,14 +299,17 @@ impl MePool {
|
||||
"ME writer channel full within blocking send timeout".into(),
|
||||
));
|
||||
}
|
||||
Err(WriterCommandReserveError::Closed) => {}
|
||||
Err(WriterCommandReserveError::Closed) => {
|
||||
payload_permit = payload_permit_from_data_command(cmd);
|
||||
}
|
||||
}
|
||||
warn!(writer_id = current.writer_id, "ME writer channel closed");
|
||||
self.remove_writer_and_close_clients(current.writer_id)
|
||||
.await;
|
||||
continue;
|
||||
}
|
||||
Err(TrySendError::Closed(_)) => {
|
||||
Err(TrySendError::Closed(cmd)) => {
|
||||
payload_permit = payload_permit_from_data_command(cmd);
|
||||
warn!(writer_id = current.writer_id, "ME writer channel closed");
|
||||
self.remove_writer_and_close_clients(current.writer_id)
|
||||
.await;
|
||||
@@ -464,11 +623,31 @@ impl MePool {
|
||||
if !self.writer_accepts_new_binding(w) {
|
||||
continue;
|
||||
}
|
||||
// Keep the advertised proxy IP aligned with the selected ME writer source.
|
||||
let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port());
|
||||
let (payload, meta) = build_routed_payload(effective_our_addr);
|
||||
let writer_permit = match try_reserve_writer_bytes(
|
||||
&w.byte_budget,
|
||||
writer_byte_permits,
|
||||
writer_reserved_bytes,
|
||||
&self.stats,
|
||||
) {
|
||||
Ok(permit) => permit,
|
||||
Err(TryAcquireError::NoPermits) => {
|
||||
if fallback_blocking_idx.is_none() {
|
||||
fallback_blocking_idx = Some(idx);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
Err(TryAcquireError::Closed) => {
|
||||
self.stats.increment_me_writer_pick_closed_total(pick_mode);
|
||||
warn!(writer_id = w.id, "ME writer byte budget closed");
|
||||
self.remove_writer_and_close_clients(w.id).await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
match w.tx.clone().try_reserve_owned() {
|
||||
Ok(permit) => {
|
||||
// Keep the advertised proxy IP aligned with the selected ME writer source.
|
||||
let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port());
|
||||
let (payload, meta) = build_routed_payload(effective_our_addr);
|
||||
if !self.registry.bind_writer(conn_id, w.id, meta).await {
|
||||
debug!(
|
||||
conn_id,
|
||||
@@ -479,7 +658,11 @@ impl MePool {
|
||||
self.remove_writer_and_close_clients(w.id).await;
|
||||
continue;
|
||||
}
|
||||
permit.send(WriterCommand::Data(payload));
|
||||
permit.send(WriterCommand::Data {
|
||||
payload,
|
||||
_permit: payload_permit.take(),
|
||||
writer_permit,
|
||||
});
|
||||
self.stats
|
||||
.increment_me_writer_pick_success_try_total(pick_mode);
|
||||
if w.generation < self.current_generation() {
|
||||
@@ -521,52 +704,71 @@ impl MePool {
|
||||
}
|
||||
self.stats
|
||||
.increment_me_writer_pick_blocking_fallback_total();
|
||||
// Keep the advertised proxy IP aligned with the selected ME writer source.
|
||||
let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port());
|
||||
let (payload, meta) = build_routed_payload(effective_our_addr);
|
||||
let reserve_result =
|
||||
if let Some(timeout) = self.route_runtime.me_route_blocking_send_timeout {
|
||||
match tokio::time::timeout(timeout, w.tx.clone().reserve_owned()).await {
|
||||
Ok(result) => result,
|
||||
Err(_) => {
|
||||
self.stats.increment_me_writer_pick_full_total(pick_mode);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
w.tx.clone().reserve_owned().await
|
||||
};
|
||||
match reserve_result {
|
||||
Ok(permit) => {
|
||||
if !self.registry.bind_writer(conn_id, w.id, meta).await {
|
||||
debug!(
|
||||
conn_id,
|
||||
writer_id = w.id,
|
||||
"ME writer disappeared before fallback bind commit, pruning stale writer"
|
||||
);
|
||||
drop(permit);
|
||||
self.remove_writer_and_close_clients(w.id).await;
|
||||
continue;
|
||||
}
|
||||
permit.send(WriterCommand::Data(payload));
|
||||
self.stats
|
||||
.increment_me_writer_pick_success_fallback_total(pick_mode);
|
||||
if w.generation < self.current_generation() {
|
||||
self.stats.increment_pool_stale_pick_total();
|
||||
}
|
||||
self.note_hybrid_route_success();
|
||||
return Ok(());
|
||||
let deadline = writer_send_deadline(self.route_runtime.me_route_blocking_send_timeout);
|
||||
let writer_permit = match reserve_writer_bytes(
|
||||
&w.byte_budget,
|
||||
writer_byte_permits,
|
||||
writer_reserved_bytes,
|
||||
deadline,
|
||||
&self.stats,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(permit) => permit,
|
||||
Err(WriterByteReserveError::TimedOut) => {
|
||||
self.stats.increment_me_writer_pick_full_total(pick_mode);
|
||||
continue;
|
||||
}
|
||||
Err(_) => {
|
||||
Err(WriterByteReserveError::Closed) => {
|
||||
self.stats.increment_me_writer_pick_closed_total(pick_mode);
|
||||
warn!(writer_id = w.id, "ME writer byte budget closed (blocking)");
|
||||
self.remove_writer_and_close_clients(w.id).await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let permit = match reserve_writer_command_slot(&w.tx, deadline).await {
|
||||
Ok(permit) => permit,
|
||||
Err(WriterCommandReserveError::TimedOut) => {
|
||||
self.stats.increment_me_writer_pick_full_total(pick_mode);
|
||||
continue;
|
||||
}
|
||||
Err(WriterCommandReserveError::Closed) => {
|
||||
self.stats.increment_me_writer_pick_closed_total(pick_mode);
|
||||
warn!(writer_id = w.id, "ME writer channel closed (blocking)");
|
||||
self.remove_writer_and_close_clients(w.id).await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
// Keep the advertised proxy IP aligned with the selected ME writer source.
|
||||
let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port());
|
||||
let (payload, meta) = build_routed_payload(effective_our_addr);
|
||||
if !self.registry.bind_writer(conn_id, w.id, meta).await {
|
||||
debug!(
|
||||
conn_id,
|
||||
writer_id = w.id,
|
||||
"ME writer disappeared before fallback bind commit, pruning stale writer"
|
||||
);
|
||||
drop(permit);
|
||||
self.remove_writer_and_close_clients(w.id).await;
|
||||
continue;
|
||||
}
|
||||
permit.send(WriterCommand::Data {
|
||||
payload,
|
||||
_permit: payload_permit.take(),
|
||||
writer_permit,
|
||||
});
|
||||
self.stats
|
||||
.increment_me_writer_pick_success_fallback_total(pick_mode);
|
||||
if w.generation < self.current_generation() {
|
||||
self.stats.increment_pool_stale_pick_total();
|
||||
}
|
||||
self.note_hybrid_route_success();
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
/// Send RPC_PROXY_REQ while keeping the first bound-writer path allocation-light.
|
||||
/// The client byte permit follows the payload until writer completion or command drop.
|
||||
pub async fn send_proxy_req_pooled(
|
||||
self: &Arc<Self>,
|
||||
conn_id: u64,
|
||||
@@ -574,12 +776,69 @@ impl MePool {
|
||||
client_addr: SocketAddr,
|
||||
our_addr: SocketAddr,
|
||||
payload: PooledBuffer,
|
||||
_permit: OwnedSemaphorePermit,
|
||||
proto_flags: u32,
|
||||
tag_override: Option<[u8; 16]>,
|
||||
) -> Result<()> {
|
||||
let tag = tag_override.or_else(|| proxy_tag_array(self.proxy_tag.as_deref()));
|
||||
let Some((writer_byte_permits, writer_reserved_bytes)) = proxy_req_resident_permits(
|
||||
payload.capacity(),
|
||||
payload.len(),
|
||||
tag.as_ref().map(|tag| tag.as_slice()),
|
||||
proto_flags,
|
||||
) else {
|
||||
self.stats.increment_me_writer_byte_budget_oversize_total();
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer payload residency calculation overflow".into(),
|
||||
));
|
||||
};
|
||||
if writer_byte_permits as usize > self.writer_lifecycle.writer_byte_budget_permits {
|
||||
self.stats.increment_me_writer_byte_budget_oversize_total();
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer payload exceeds configured byte budget".into(),
|
||||
));
|
||||
}
|
||||
|
||||
if let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await {
|
||||
let deadline = writer_send_deadline(self.route_runtime.me_route_blocking_send_timeout);
|
||||
let writer_permit = match reserve_writer_bytes(
|
||||
¤t.byte_budget,
|
||||
writer_byte_permits,
|
||||
writer_reserved_bytes,
|
||||
deadline,
|
||||
&self.stats,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(permit) => permit,
|
||||
Err(WriterByteReserveError::TimedOut) => {
|
||||
self.stats
|
||||
.increment_me_writer_pick_full_total(self.writer_pick_mode());
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer byte budget full within blocking send timeout".into(),
|
||||
));
|
||||
}
|
||||
Err(WriterByteReserveError::Closed) => {
|
||||
warn!(
|
||||
writer_id = current.writer_id,
|
||||
"ME writer byte budget closed"
|
||||
);
|
||||
self.remove_writer_and_close_clients(current.writer_id)
|
||||
.await;
|
||||
return self
|
||||
.send_proxy_req(
|
||||
conn_id,
|
||||
target_dc,
|
||||
client_addr,
|
||||
our_addr,
|
||||
payload.as_ref(),
|
||||
proto_flags,
|
||||
tag.as_ref().map(|tag| tag.as_slice()),
|
||||
Some(_permit),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
};
|
||||
let command = WriterCommand::ProxyReq(ProxyReqCommand {
|
||||
conn_id,
|
||||
client_addr,
|
||||
@@ -587,6 +846,8 @@ impl MePool {
|
||||
proto_flags,
|
||||
proxy_tag: tag,
|
||||
payload,
|
||||
_permit,
|
||||
writer_permit,
|
||||
});
|
||||
match current.tx.try_send(command) {
|
||||
Ok(()) => {
|
||||
@@ -594,12 +855,7 @@ impl MePool {
|
||||
return Ok(());
|
||||
}
|
||||
Err(TrySendError::Full(cmd)) => {
|
||||
match reserve_writer_command_slot(
|
||||
¤t.tx,
|
||||
self.route_runtime.me_route_blocking_send_timeout,
|
||||
)
|
||||
.await
|
||||
{
|
||||
match reserve_writer_command_slot(¤t.tx, deadline).await {
|
||||
Ok(permit) => {
|
||||
permit.send(cmd);
|
||||
self.note_hybrid_route_success();
|
||||
@@ -613,7 +869,8 @@ impl MePool {
|
||||
));
|
||||
}
|
||||
Err(WriterCommandReserveError::Closed) => {
|
||||
let Some(payload) = proxy_req_payload_from_command(cmd) else {
|
||||
let Some((payload, _permit)) = proxy_req_payload_from_command(cmd)
|
||||
else {
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer rejected unexpected command type".into(),
|
||||
));
|
||||
@@ -630,13 +887,14 @@ impl MePool {
|
||||
payload.as_ref(),
|
||||
proto_flags,
|
||||
tag.as_ref().map(|tag| tag.as_slice()),
|
||||
Some(_permit),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(TrySendError::Closed(cmd)) => {
|
||||
let Some(payload) = proxy_req_payload_from_command(cmd) else {
|
||||
let Some((payload, _permit)) = proxy_req_payload_from_command(cmd) else {
|
||||
return Err(ProxyError::Proxy(
|
||||
"ME writer rejected unexpected command type".into(),
|
||||
));
|
||||
@@ -653,6 +911,7 @@ impl MePool {
|
||||
payload.as_ref(),
|
||||
proto_flags,
|
||||
tag.as_ref().map(|tag| tag.as_slice()),
|
||||
Some(_permit),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -667,6 +926,7 @@ impl MePool {
|
||||
payload.as_ref(),
|
||||
proto_flags,
|
||||
tag.as_ref().map(|tag| tag.as_slice()),
|
||||
Some(_permit),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
@@ -10,7 +10,7 @@ use crate::protocol::constants::{RPC_CLOSE_CONN_U32, RPC_CLOSE_EXT_U32};
|
||||
|
||||
use super::super::MePool;
|
||||
use super::super::codec::{WriterCommand, build_control_payload};
|
||||
use super::{WriterCommandReserveError, reserve_writer_command_slot};
|
||||
use super::{WriterCommandReserveError, reserve_writer_command_slot, writer_send_deadline};
|
||||
|
||||
const ME_CLOSE_SIGNAL_SEND_TIMEOUT: Duration = Duration::from_millis(50);
|
||||
|
||||
@@ -22,8 +22,11 @@ impl MePool {
|
||||
match w.tx.try_send(WriterCommand::ControlAndFlush(payload)) {
|
||||
Ok(()) => {}
|
||||
Err(TrySendError::Full(cmd)) => {
|
||||
match reserve_writer_command_slot(&w.tx, Some(ME_CLOSE_SIGNAL_SEND_TIMEOUT))
|
||||
.await
|
||||
match reserve_writer_command_slot(
|
||||
&w.tx,
|
||||
writer_send_deadline(Some(ME_CLOSE_SIGNAL_SEND_TIMEOUT)),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(permit) => {
|
||||
permit.send(cmd);
|
||||
@@ -63,9 +66,12 @@ impl MePool {
|
||||
match w.tx.try_send(WriterCommand::ControlAndFlush(payload)) {
|
||||
Ok(()) => {}
|
||||
Err(TrySendError::Full(cmd)) => {
|
||||
let _ = reserve_writer_command_slot(&w.tx, Some(ME_CLOSE_SIGNAL_SEND_TIMEOUT))
|
||||
.await
|
||||
.map(|permit| permit.send(cmd));
|
||||
let _ = reserve_writer_command_slot(
|
||||
&w.tx,
|
||||
writer_send_deadline(Some(ME_CLOSE_SIGNAL_SEND_TIMEOUT)),
|
||||
)
|
||||
.await
|
||||
.map(|permit| permit.send(cmd));
|
||||
}
|
||||
Err(TrySendError::Closed(_)) => {
|
||||
debug!(conn_id, "ME close_conn skipped: writer channel closed");
|
||||
|
||||
@@ -105,6 +105,7 @@ async fn make_pool(
|
||||
general.me_writer_pick_sample_size,
|
||||
MeSocksKdfPolicy::default(),
|
||||
general.me_writer_cmd_channel_capacity,
|
||||
general.me_writer_byte_budget_bytes,
|
||||
general.me_route_channel_capacity,
|
||||
general.me_route_backpressure_enabled,
|
||||
general.me_route_fairshare_enabled,
|
||||
@@ -134,6 +135,7 @@ async fn insert_draining_writer(
|
||||
drain_deadline_epoch_secs: u64,
|
||||
) {
|
||||
let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8);
|
||||
let byte_budget = pool.new_writer_byte_budget();
|
||||
let writer = MeWriter {
|
||||
id: writer_id,
|
||||
addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 6000 + writer_id as u16),
|
||||
@@ -143,6 +145,7 @@ async fn insert_draining_writer(
|
||||
contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())),
|
||||
created_at: Instant::now() - Duration::from_secs(writer_id),
|
||||
tx: tx.clone(),
|
||||
byte_budget: byte_budget.clone(),
|
||||
cancel: CancellationToken::new(),
|
||||
degraded: Arc::new(AtomicBool::new(false)),
|
||||
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
|
||||
@@ -153,7 +156,9 @@ async fn insert_draining_writer(
|
||||
};
|
||||
|
||||
pool.writers.write().await.push(writer);
|
||||
pool.registry.register_writer(writer_id, tx).await;
|
||||
pool.registry
|
||||
.register_writer(writer_id, tx, byte_budget)
|
||||
.await;
|
||||
pool.conn_count.fetch_add(1, Ordering::Relaxed);
|
||||
|
||||
for idx in 0..bound_clients {
|
||||
|
||||
@@ -103,6 +103,7 @@ async fn make_pool(
|
||||
general.me_writer_pick_sample_size,
|
||||
MeSocksKdfPolicy::default(),
|
||||
general.me_writer_cmd_channel_capacity,
|
||||
general.me_writer_byte_budget_bytes,
|
||||
general.me_route_channel_capacity,
|
||||
general.me_route_backpressure_enabled,
|
||||
general.me_route_fairshare_enabled,
|
||||
@@ -131,6 +132,7 @@ async fn insert_draining_writer(
|
||||
drain_deadline_epoch_secs: u64,
|
||||
) {
|
||||
let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8);
|
||||
let byte_budget = pool.new_writer_byte_budget();
|
||||
let writer = MeWriter {
|
||||
id: writer_id,
|
||||
addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 5500 + writer_id as u16),
|
||||
@@ -140,6 +142,7 @@ async fn insert_draining_writer(
|
||||
contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())),
|
||||
created_at: Instant::now() - Duration::from_secs(writer_id),
|
||||
tx: tx.clone(),
|
||||
byte_budget: byte_budget.clone(),
|
||||
cancel: CancellationToken::new(),
|
||||
degraded: Arc::new(AtomicBool::new(false)),
|
||||
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
|
||||
@@ -149,7 +152,9 @@ async fn insert_draining_writer(
|
||||
allow_drain_fallback: Arc::new(AtomicBool::new(false)),
|
||||
};
|
||||
pool.writers.write().await.push(writer);
|
||||
pool.registry.register_writer(writer_id, tx).await;
|
||||
pool.registry
|
||||
.register_writer(writer_id, tx, byte_budget)
|
||||
.await;
|
||||
pool.conn_count.fetch_add(1, Ordering::Relaxed);
|
||||
for idx in 0..bound_clients {
|
||||
let (conn_id, _rx) = pool.registry.register().await;
|
||||
|
||||
@@ -98,6 +98,7 @@ async fn make_pool(me_pool_drain_threshold: u64) -> Arc<MePool> {
|
||||
general.me_writer_pick_sample_size,
|
||||
MeSocksKdfPolicy::default(),
|
||||
general.me_writer_cmd_channel_capacity,
|
||||
general.me_writer_byte_budget_bytes,
|
||||
general.me_route_channel_capacity,
|
||||
general.me_route_backpressure_enabled,
|
||||
general.me_route_fairshare_enabled,
|
||||
@@ -126,6 +127,7 @@ async fn insert_draining_writer(
|
||||
) -> Vec<u64> {
|
||||
let mut conn_ids = Vec::with_capacity(bound_clients);
|
||||
let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8);
|
||||
let byte_budget = pool.new_writer_byte_budget();
|
||||
let writer = MeWriter {
|
||||
id: writer_id,
|
||||
addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 4500 + writer_id as u16),
|
||||
@@ -135,6 +137,7 @@ async fn insert_draining_writer(
|
||||
contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())),
|
||||
created_at: Instant::now() - Duration::from_secs(writer_id),
|
||||
tx: tx.clone(),
|
||||
byte_budget: byte_budget.clone(),
|
||||
cancel: CancellationToken::new(),
|
||||
degraded: Arc::new(AtomicBool::new(false)),
|
||||
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
|
||||
@@ -144,7 +147,9 @@ async fn insert_draining_writer(
|
||||
allow_drain_fallback: Arc::new(AtomicBool::new(false)),
|
||||
};
|
||||
pool.writers.write().await.push(writer);
|
||||
pool.registry.register_writer(writer_id, tx).await;
|
||||
pool.registry
|
||||
.register_writer(writer_id, tx, byte_budget)
|
||||
.await;
|
||||
pool.conn_count.fetch_add(1, Ordering::Relaxed);
|
||||
for idx in 0..bound_clients {
|
||||
let (conn_id, _rx) = pool.registry.register().await;
|
||||
|
||||
@@ -87,6 +87,7 @@ async fn make_pool() -> Arc<MePool> {
|
||||
general.me_writer_pick_sample_size,
|
||||
MeSocksKdfPolicy::default(),
|
||||
general.me_writer_cmd_channel_capacity,
|
||||
general.me_writer_byte_budget_bytes,
|
||||
general.me_route_channel_capacity,
|
||||
general.me_route_backpressure_enabled,
|
||||
general.me_route_fairshare_enabled,
|
||||
|
||||
@@ -92,6 +92,7 @@ async fn make_pool() -> Arc<MePool> {
|
||||
general.me_writer_pick_sample_size,
|
||||
MeSocksKdfPolicy::default(),
|
||||
general.me_writer_cmd_channel_capacity,
|
||||
general.me_writer_byte_budget_bytes,
|
||||
general.me_route_channel_capacity,
|
||||
general.me_route_backpressure_enabled,
|
||||
general.me_route_fairshare_enabled,
|
||||
@@ -120,6 +121,7 @@ async fn insert_writer(
|
||||
created_at: Instant,
|
||||
) {
|
||||
let (tx, _rx) = mpsc::channel::<WriterCommand>(8);
|
||||
let byte_budget = pool.new_writer_byte_budget();
|
||||
let contour = if draining {
|
||||
WriterContour::Draining
|
||||
} else {
|
||||
@@ -134,6 +136,7 @@ async fn insert_writer(
|
||||
contour: Arc::new(AtomicU8::new(contour.as_u8())),
|
||||
created_at,
|
||||
tx: tx.clone(),
|
||||
byte_budget: byte_budget.clone(),
|
||||
cancel: CancellationToken::new(),
|
||||
degraded: Arc::new(AtomicBool::new(false)),
|
||||
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
|
||||
@@ -144,7 +147,9 @@ async fn insert_writer(
|
||||
};
|
||||
|
||||
pool.writers.write().await.push(writer);
|
||||
pool.registry.register_writer(writer_id, tx).await;
|
||||
pool.registry
|
||||
.register_writer(writer_id, tx, byte_budget)
|
||||
.await;
|
||||
pool.conn_count.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
|
||||
@@ -98,6 +98,7 @@ async fn make_pool() -> (Arc<MePool>, Arc<SecureRandom>) {
|
||||
general.me_writer_pick_sample_size,
|
||||
MeSocksKdfPolicy::default(),
|
||||
general.me_writer_cmd_channel_capacity,
|
||||
general.me_writer_byte_budget_bytes,
|
||||
general.me_route_channel_capacity,
|
||||
general.me_route_backpressure_enabled,
|
||||
general.me_route_fairshare_enabled,
|
||||
@@ -127,6 +128,7 @@ async fn insert_writer(
|
||||
register_in_registry: bool,
|
||||
) -> mpsc::Receiver<WriterCommand> {
|
||||
let (tx, rx) = mpsc::channel::<WriterCommand>(8);
|
||||
let byte_budget = pool.new_writer_byte_budget();
|
||||
let writer = MeWriter {
|
||||
id: writer_id,
|
||||
addr,
|
||||
@@ -136,6 +138,7 @@ async fn insert_writer(
|
||||
contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())),
|
||||
created_at: Instant::now(),
|
||||
tx: tx.clone(),
|
||||
byte_budget: byte_budget.clone(),
|
||||
cancel: CancellationToken::new(),
|
||||
degraded: Arc::new(AtomicBool::new(false)),
|
||||
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
|
||||
@@ -154,7 +157,9 @@ async fn insert_writer(
|
||||
}
|
||||
pool.rebuild_endpoint_dc_map().await;
|
||||
if register_in_registry {
|
||||
pool.registry.register_writer(writer_id, tx).await;
|
||||
pool.registry
|
||||
.register_writer(writer_id, tx, byte_budget)
|
||||
.await;
|
||||
}
|
||||
rx
|
||||
}
|
||||
@@ -165,7 +170,7 @@ async fn recv_data_count(rx: &mut mpsc::Receiver<WriterCommand>, budget: Duratio
|
||||
while Instant::now().duration_since(start) < budget {
|
||||
let remaining = budget.saturating_sub(Instant::now().duration_since(start));
|
||||
match tokio::time::timeout(remaining.min(Duration::from_millis(10)), rx.recv()).await {
|
||||
Ok(Some(WriterCommand::Data(_))) => data_count += 1,
|
||||
Ok(Some(WriterCommand::Data { .. })) => data_count += 1,
|
||||
Ok(Some(WriterCommand::DataAndFlush(_))) => data_count += 1,
|
||||
Ok(Some(WriterCommand::ProxyReq(_))) => data_count += 1,
|
||||
Ok(Some(WriterCommand::ControlAndFlush(_))) => data_count += 1,
|
||||
@@ -185,7 +190,7 @@ async fn recv_first_data_payload(
|
||||
while Instant::now().duration_since(start) < budget {
|
||||
let remaining = budget.saturating_sub(Instant::now().duration_since(start));
|
||||
match tokio::time::timeout(remaining.min(Duration::from_millis(10)), rx.recv()).await {
|
||||
Ok(Some(WriterCommand::Data(payload))) => return Some(payload.to_vec()),
|
||||
Ok(Some(WriterCommand::Data { payload, .. })) => return Some(payload.to_vec()),
|
||||
Ok(Some(WriterCommand::DataAndFlush(payload))) => return Some(payload.to_vec()),
|
||||
Ok(Some(_)) => {}
|
||||
Ok(None) => break,
|
||||
@@ -240,6 +245,7 @@ async fn send_proxy_req_does_not_replay_when_first_bind_commit_fails() {
|
||||
b"hello",
|
||||
0,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -299,6 +305,7 @@ async fn send_proxy_req_prunes_iterative_stale_bind_failures_without_data_replay
|
||||
b"storm",
|
||||
0,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -356,6 +363,7 @@ async fn send_proxy_req_uses_writer_source_ip_when_advertised_our_addr_differs()
|
||||
b"route",
|
||||
0,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -411,6 +419,7 @@ async fn send_proxy_req_blocking_fallback_uses_writer_source_ip() {
|
||||
b"blocking",
|
||||
0,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
});
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
|
||||
#![allow(deprecated)]
|
||||
|
||||
use arc_swap::ArcSwap;
|
||||
use rand::RngExt;
|
||||
use std::collections::{BTreeSet, HashMap};
|
||||
use std::net::{IpAddr, SocketAddr};
|
||||
@@ -20,7 +21,7 @@ use tracing::{debug, info, trace, warn};
|
||||
|
||||
use crate::config::{UpstreamConfig, UpstreamType};
|
||||
use crate::error::{ProxyError, Result};
|
||||
use crate::network::dns_overrides::{resolve_socket_addr, split_host_port};
|
||||
use crate::network::dns_overrides::{DnsOverrides, split_host_port};
|
||||
use crate::protocol::constants::{TG_DATACENTER_PORT, TG_DATACENTERS_V4, TG_DATACENTERS_V6};
|
||||
use crate::stats::Stats;
|
||||
use crate::transport::shadowsocks::{
|
||||
@@ -333,6 +334,7 @@ pub struct UpstreamManager {
|
||||
no_upstreams_warn_epoch_ms: Arc<AtomicU64>,
|
||||
no_healthy_warn_epoch_ms: Arc<AtomicU64>,
|
||||
stats: Arc<Stats>,
|
||||
dns_overrides: Arc<ArcSwap<DnsOverrides>>,
|
||||
}
|
||||
|
||||
impl UpstreamManager {
|
||||
@@ -374,9 +376,37 @@ impl UpstreamManager {
|
||||
no_upstreams_warn_epoch_ms: Arc::new(AtomicU64::new(0)),
|
||||
no_healthy_warn_epoch_ms: Arc::new(AtomicU64::new(0)),
|
||||
stats,
|
||||
dns_overrides: Arc::new(ArcSwap::from_pointee(DnsOverrides::default())),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn with_dns_overrides(mut self, entries: &[String]) -> Result<Self> {
|
||||
self.dns_overrides = Arc::new(ArcSwap::from_pointee(DnsOverrides::from_entries(entries)?));
|
||||
Ok(self)
|
||||
}
|
||||
|
||||
pub(crate) fn update_dns_overrides(&self, entries: &[String]) -> Result<()> {
|
||||
let snapshot = DnsOverrides::from_entries(entries)?;
|
||||
self.dns_overrides.store(Arc::new(snapshot));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_hostname(&self, host: &str, port: u16) -> Result<SocketAddr> {
|
||||
if let Some(addr) = self.dns_overrides.load().resolve_socket_addr(host, port) {
|
||||
return Ok(addr);
|
||||
}
|
||||
let addrs: Vec<SocketAddr> = tokio::net::lookup_host((host, port))
|
||||
.await
|
||||
.map_err(ProxyError::Io)?
|
||||
.collect();
|
||||
if let Some(addr) = addrs.iter().copied().find(SocketAddr::is_ipv4) {
|
||||
return Ok(addr);
|
||||
}
|
||||
addrs.first().copied().ok_or_else(|| {
|
||||
ProxyError::Proxy(format!("DNS returned no addresses for {host}:{port}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn now_epoch_ms() -> u64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
@@ -709,11 +739,12 @@ impl UpstreamManager {
|
||||
}
|
||||
|
||||
async fn connect_hostname_with_dns_override(
|
||||
&self,
|
||||
address: &str,
|
||||
connect_timeout: Duration,
|
||||
) -> Result<TcpStream> {
|
||||
if let Some((host, port)) = split_host_port(address)
|
||||
&& let Some(addr) = resolve_socket_addr(&host, port)
|
||||
&& let Some(addr) = self.dns_overrides.load().resolve_socket_addr(&host, port)
|
||||
{
|
||||
return match tokio::time::timeout(connect_timeout, TcpStream::connect(addr)).await {
|
||||
Ok(Ok(stream)) => Ok(stream),
|
||||
@@ -1203,7 +1234,8 @@ impl UpstreamManager {
|
||||
"SOCKS4 interface binding is not supported for hostname addresses, ignoring"
|
||||
);
|
||||
}
|
||||
Self::connect_hostname_with_dns_override(address, connect_timeout).await?
|
||||
self.connect_hostname_with_dns_override(address, connect_timeout)
|
||||
.await?
|
||||
};
|
||||
|
||||
// replace socks user_id with config.selected_scope, if set
|
||||
@@ -1291,7 +1323,8 @@ impl UpstreamManager {
|
||||
"SOCKS5 interface binding is not supported for hostname addresses, ignoring"
|
||||
);
|
||||
}
|
||||
Self::connect_hostname_with_dns_override(address, connect_timeout).await?
|
||||
self.connect_hostname_with_dns_override(address, connect_timeout)
|
||||
.await?
|
||||
};
|
||||
|
||||
debug!(config = ?config, "Socks5 connection");
|
||||
|
||||
Reference in New Issue
Block a user