mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 09:57:47 +08:00
elky
579f2c7cc1
Consolidate subscription usage policy enforcement, privacy-safe persistence, and gateway security hardening into one reviewable change. Includes bounded HTTP and execution envelopes, header and protocol guards, DNS and relay validation, authentication and secret projection hardening, secure backup/install paths, and regression coverage.
1977 lines
67 KiB
Rust
1977 lines
67 KiB
Rust
use async_trait::async_trait;
|
|
use futures_util::TryStreamExt;
|
|
use sha2::{Digest, Sha256};
|
|
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
|
|
|
|
use aether_data_contracts::repository::proxy_nodes::{
|
|
bucket_start_unix_secs, build_tunnel_error_event_detail, build_tunnel_metrics_sample,
|
|
merge_proxy_metadata_for_registration, normalize_heartbeat_proxy_metadata,
|
|
normalize_proxy_metadata, reconcile_remote_config_after_heartbeat, ProxyNodeEventQuery,
|
|
ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation,
|
|
ProxyNodeMetricsCleanupSummary, ProxyNodeMetricsStep, ProxyNodeReadRepository,
|
|
ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation,
|
|
ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyFleetMetricsBucket,
|
|
StoredProxyNode, StoredProxyNodeEvent, StoredProxyNodeMetricsBucket, TunnelErrorEventRecord,
|
|
TunnelMetricsSample, PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR,
|
|
};
|
|
use aether_data_contracts::DataLayerError;
|
|
use aether_data_query::{push_eq, push_limit, WhereClause};
|
|
|
|
use crate::error::{postgres_error, SqlxResultExt};
|
|
|
|
fn log_reported_tunnel_error_event(
|
|
node_id: &str,
|
|
event: &TunnelErrorEventRecord,
|
|
received_at_unix_secs: u64,
|
|
) {
|
|
tracing::warn!(
|
|
event_name = "proxy_tunnel_error_reported",
|
|
source = "heartbeat",
|
|
node_id = %node_id,
|
|
category = %event.category,
|
|
message = %event.message,
|
|
severity = ?event.severity,
|
|
component = ?event.component,
|
|
summary = ?event.summary,
|
|
operator_action = ?event.operator_action,
|
|
error_reported_at_unix_secs = event.timestamp_unix_secs,
|
|
error_reported_at_unix_ms = ?event.timestamp_unix_ms,
|
|
report_received_at_unix_secs = received_at_unix_secs,
|
|
"proxy reported tunnel error via heartbeat"
|
|
);
|
|
}
|
|
|
|
const FIND_PROXY_NODE_SQL: &str = r#"
|
|
SELECT
|
|
id,
|
|
tunnel_generation,
|
|
name,
|
|
ip,
|
|
port,
|
|
region,
|
|
is_manual,
|
|
proxy_url,
|
|
proxy_username,
|
|
proxy_password,
|
|
CAST(status AS TEXT) AS status,
|
|
registered_by,
|
|
EXTRACT(EPOCH FROM last_heartbeat_at)::bigint AS last_heartbeat_at_unix_secs,
|
|
heartbeat_interval,
|
|
active_connections,
|
|
total_requests,
|
|
CAST(avg_latency_ms AS DOUBLE PRECISION) AS avg_latency_ms,
|
|
failed_requests,
|
|
dns_failures,
|
|
stream_errors,
|
|
proxy_metadata,
|
|
hardware_info,
|
|
estimated_max_concurrency,
|
|
tunnel_mode,
|
|
tunnel_connected,
|
|
EXTRACT(EPOCH FROM tunnel_connected_at)::bigint AS tunnel_connected_at_unix_secs,
|
|
remote_config,
|
|
config_version,
|
|
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms,
|
|
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
|
|
FROM proxy_nodes
|
|
WHERE id = $1
|
|
LIMIT 1
|
|
"#;
|
|
|
|
const LIST_PROXY_NODE_EVENTS_SQL: &str = r#"
|
|
SELECT
|
|
id,
|
|
node_id,
|
|
CAST(event_type AS TEXT) AS event_type,
|
|
detail,
|
|
event_metadata,
|
|
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms
|
|
FROM proxy_node_events
|
|
WHERE node_id = $1
|
|
ORDER BY created_at DESC, id DESC
|
|
LIMIT $2
|
|
"#;
|
|
|
|
const APPLY_HEARTBEAT_SQL: &str = r#"
|
|
UPDATE proxy_nodes
|
|
SET
|
|
last_heartbeat_at = NOW(),
|
|
status = CASE
|
|
WHEN status <> 'online'::proxynodestatus OR tunnel_connected = FALSE
|
|
THEN 'online'::proxynodestatus
|
|
ELSE status
|
|
END,
|
|
tunnel_connected = CASE
|
|
WHEN status <> 'online'::proxynodestatus OR tunnel_connected = FALSE
|
|
THEN TRUE
|
|
ELSE tunnel_connected
|
|
END,
|
|
tunnel_connected_at = CASE
|
|
WHEN status <> 'online'::proxynodestatus OR tunnel_connected = FALSE
|
|
THEN NOW()
|
|
ELSE tunnel_connected_at
|
|
END,
|
|
updated_at = CASE
|
|
WHEN status <> 'online'::proxynodestatus OR tunnel_connected = FALSE
|
|
THEN NOW()
|
|
ELSE updated_at
|
|
END,
|
|
heartbeat_interval = COALESCE($2, heartbeat_interval),
|
|
active_connections = COALESCE($3, active_connections),
|
|
avg_latency_ms = COALESCE($4, avg_latency_ms),
|
|
total_requests = total_requests + GREATEST(COALESCE($5, 0), 0),
|
|
failed_requests = failed_requests + GREATEST(COALESCE($6, 0), 0),
|
|
dns_failures = dns_failures + GREATEST(COALESCE($7, 0), 0),
|
|
stream_errors = stream_errors + GREATEST(COALESCE($8, 0), 0)
|
|
WHERE id = $1
|
|
AND tunnel_mode = TRUE
|
|
AND tunnel_generation = $9
|
|
"#;
|
|
|
|
const CAS_HEARTBEAT_PROXY_METADATA_SQL: &str = r#"
|
|
UPDATE proxy_nodes
|
|
SET proxy_metadata = $2::json, updated_at = NOW()
|
|
WHERE id = $1
|
|
AND tunnel_generation = $3
|
|
AND proxy_metadata::jsonb IS NOT DISTINCT FROM $4::jsonb
|
|
"#;
|
|
|
|
const UPDATE_TUNNEL_STATUS_SQL: &str = r#"
|
|
UPDATE proxy_nodes
|
|
SET
|
|
tunnel_connected = $2,
|
|
active_connections = CASE
|
|
WHEN $2 THEN active_connections
|
|
ELSE 0
|
|
END,
|
|
tunnel_connected_at = CASE
|
|
WHEN $3::double precision IS NULL THEN NOW()
|
|
ELSE TO_TIMESTAMP($3::double precision)
|
|
END,
|
|
status = CASE
|
|
WHEN $2 THEN 'online'::proxynodestatus
|
|
ELSE 'offline'::proxynodestatus
|
|
END,
|
|
updated_at = CASE
|
|
WHEN $3::double precision IS NULL THEN NOW()
|
|
ELSE TO_TIMESTAMP($3::double precision)
|
|
END
|
|
WHERE id = $1
|
|
AND tunnel_generation = $4
|
|
AND (
|
|
tunnel_connected_at IS NULL
|
|
OR tunnel_connected_at <= CASE
|
|
WHEN $3::double precision IS NULL THEN NOW()
|
|
ELSE TO_TIMESTAMP($3::double precision)
|
|
END
|
|
)
|
|
"#;
|
|
|
|
const FIND_EXISTING_TUNNEL_NODE_SQL: &str = r#"
|
|
SELECT
|
|
id,
|
|
tunnel_generation,
|
|
name,
|
|
ip,
|
|
port,
|
|
region,
|
|
is_manual,
|
|
proxy_url,
|
|
proxy_username,
|
|
proxy_password,
|
|
CAST(status AS TEXT) AS status,
|
|
registered_by,
|
|
EXTRACT(EPOCH FROM last_heartbeat_at)::bigint AS last_heartbeat_at_unix_secs,
|
|
heartbeat_interval,
|
|
active_connections,
|
|
total_requests,
|
|
CAST(avg_latency_ms AS DOUBLE PRECISION) AS avg_latency_ms,
|
|
failed_requests,
|
|
dns_failures,
|
|
stream_errors,
|
|
proxy_metadata,
|
|
hardware_info,
|
|
estimated_max_concurrency,
|
|
tunnel_mode,
|
|
tunnel_connected,
|
|
EXTRACT(EPOCH FROM tunnel_connected_at)::bigint AS tunnel_connected_at_unix_secs,
|
|
remote_config,
|
|
config_version,
|
|
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms,
|
|
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
|
|
FROM proxy_nodes
|
|
WHERE ip = $1
|
|
AND port = $2
|
|
AND is_manual = FALSE
|
|
ORDER BY created_at ASC, id ASC
|
|
LIMIT 1
|
|
FOR UPDATE
|
|
"#;
|
|
|
|
const INSERT_PROXY_NODE_SQL: &str = r#"
|
|
INSERT INTO proxy_nodes (
|
|
id,
|
|
name,
|
|
ip,
|
|
port,
|
|
region,
|
|
status,
|
|
registered_by,
|
|
last_heartbeat_at,
|
|
heartbeat_interval,
|
|
active_connections,
|
|
total_requests,
|
|
avg_latency_ms,
|
|
hardware_info,
|
|
estimated_max_concurrency,
|
|
tunnel_mode,
|
|
tunnel_connected,
|
|
proxy_metadata,
|
|
tunnel_generation
|
|
)
|
|
VALUES (
|
|
$1,
|
|
$2,
|
|
$3,
|
|
$4,
|
|
$5,
|
|
'offline'::proxynodestatus,
|
|
$6,
|
|
NOW(),
|
|
$7,
|
|
COALESCE($8, 0),
|
|
COALESCE($9, 0),
|
|
$10,
|
|
$11::json,
|
|
$12,
|
|
$13,
|
|
FALSE,
|
|
$14::json,
|
|
$15
|
|
)
|
|
"#;
|
|
|
|
const FIND_DUPLICATE_PROXY_NODE_SQL: &str = r#"
|
|
SELECT
|
|
id,
|
|
name,
|
|
ip,
|
|
port
|
|
FROM proxy_nodes
|
|
WHERE ip = $1
|
|
AND port = $2
|
|
LIMIT 1
|
|
FOR UPDATE
|
|
"#;
|
|
|
|
const FIND_DUPLICATE_PROXY_NODE_EXCLUDING_ID_SQL: &str = r#"
|
|
SELECT
|
|
id,
|
|
name,
|
|
ip,
|
|
port
|
|
FROM proxy_nodes
|
|
WHERE ip = $1
|
|
AND port = $2
|
|
AND id <> $3
|
|
LIMIT 1
|
|
FOR UPDATE
|
|
"#;
|
|
|
|
const INSERT_MANUAL_PROXY_NODE_SQL: &str = r#"
|
|
INSERT INTO proxy_nodes (
|
|
id,
|
|
name,
|
|
ip,
|
|
port,
|
|
region,
|
|
is_manual,
|
|
proxy_url,
|
|
proxy_username,
|
|
proxy_password,
|
|
status,
|
|
registered_by,
|
|
last_heartbeat_at,
|
|
heartbeat_interval,
|
|
active_connections,
|
|
total_requests,
|
|
tunnel_mode,
|
|
tunnel_connected,
|
|
config_version,
|
|
tunnel_generation
|
|
)
|
|
VALUES (
|
|
$1,
|
|
$2,
|
|
$3,
|
|
$4,
|
|
$5,
|
|
TRUE,
|
|
$6,
|
|
$7,
|
|
$8,
|
|
'online'::proxynodestatus,
|
|
$9,
|
|
NULL,
|
|
0,
|
|
0,
|
|
0,
|
|
FALSE,
|
|
FALSE,
|
|
0,
|
|
$10
|
|
)
|
|
"#;
|
|
|
|
const UPDATE_PROXY_NODE_REGISTRATION_SQL: &str = r#"
|
|
UPDATE proxy_nodes
|
|
SET
|
|
name = $2,
|
|
ip = $3,
|
|
port = $4,
|
|
region = $5,
|
|
registered_by = $6,
|
|
last_heartbeat_at = NOW(),
|
|
heartbeat_interval = $7,
|
|
active_connections = COALESCE($8, active_connections),
|
|
total_requests = COALESCE($9, total_requests),
|
|
avg_latency_ms = COALESCE($10, avg_latency_ms),
|
|
hardware_info = COALESCE($11::json, hardware_info),
|
|
estimated_max_concurrency = COALESCE($12, estimated_max_concurrency),
|
|
tunnel_mode = $13,
|
|
proxy_metadata = COALESCE($14::json, proxy_metadata),
|
|
updated_at = NOW()
|
|
WHERE id = $1
|
|
"#;
|
|
|
|
const UPDATE_MANUAL_PROXY_NODE_SQL: &str = r#"
|
|
UPDATE proxy_nodes
|
|
SET
|
|
name = COALESCE($2, name),
|
|
ip = COALESCE($3, ip),
|
|
port = COALESCE($4, port),
|
|
region = COALESCE($5, region),
|
|
proxy_url = COALESCE($6, proxy_url),
|
|
proxy_username = COALESCE($7, proxy_username),
|
|
proxy_password = COALESCE($8, proxy_password),
|
|
updated_at = NOW()
|
|
WHERE id = $1
|
|
AND is_manual = TRUE
|
|
AND tunnel_generation = $9
|
|
"#;
|
|
|
|
const RECORD_PROXY_NODE_TRAFFIC_SQL: &str = r#"
|
|
UPDATE proxy_nodes
|
|
SET
|
|
total_requests = total_requests + GREATEST($2, 0),
|
|
failed_requests = failed_requests + GREATEST($3, 0),
|
|
dns_failures = dns_failures + GREATEST($4, 0),
|
|
stream_errors = stream_errors + GREATEST($5, 0),
|
|
updated_at = NOW()
|
|
WHERE id = $1
|
|
AND is_manual = TRUE
|
|
AND tunnel_generation = $6
|
|
"#;
|
|
|
|
const UNREGISTER_PROXY_NODE_SQL: &str = r#"
|
|
UPDATE proxy_nodes
|
|
SET
|
|
status = 'offline'::proxynodestatus,
|
|
tunnel_connected = FALSE,
|
|
tunnel_connected_at = NOW(),
|
|
updated_at = NOW()
|
|
WHERE id = $1
|
|
AND tunnel_generation = $2
|
|
"#;
|
|
|
|
const DELETE_PROXY_NODE_SQL: &str = r#"
|
|
DELETE FROM proxy_nodes
|
|
WHERE id = $1
|
|
AND tunnel_generation = $2
|
|
"#;
|
|
|
|
// Run after the parent delete commits so delete never waits on an outbox row
|
|
// already claimed by the flusher (which acquires locks in the opposite order).
|
|
const RETIRE_PROXY_NODE_PENDING_COUNTERS_SQL: &str = r#"
|
|
DELETE FROM usage_counter_deltas
|
|
WHERE kind = 'proxy_node'
|
|
AND target_id = $1
|
|
AND target_tunnel_generation = $2
|
|
AND processed_at IS NULL
|
|
"#;
|
|
|
|
const DELETE_PROXY_NODE_EVENTS_SQL: &str = r#"
|
|
DELETE FROM proxy_node_events
|
|
WHERE node_id = $1
|
|
AND EXISTS (
|
|
SELECT 1
|
|
FROM proxy_nodes
|
|
WHERE id = $1
|
|
AND tunnel_generation = $2
|
|
)
|
|
"#;
|
|
|
|
const DELETE_PROXY_NODE_METRICS_1M_SQL: &str = r#"
|
|
DELETE FROM proxy_node_metrics_1m
|
|
WHERE node_id = $1
|
|
AND EXISTS (
|
|
SELECT 1
|
|
FROM proxy_nodes
|
|
WHERE id = $1
|
|
AND tunnel_generation = $2
|
|
)
|
|
"#;
|
|
|
|
const DELETE_PROXY_NODE_METRICS_1H_SQL: &str = r#"
|
|
DELETE FROM proxy_node_metrics_1h
|
|
WHERE node_id = $1
|
|
AND EXISTS (
|
|
SELECT 1
|
|
FROM proxy_nodes
|
|
WHERE id = $1
|
|
AND tunnel_generation = $2
|
|
)
|
|
"#;
|
|
|
|
const UPDATE_PROXY_NODE_REMOTE_CONFIG_SQL: &str = r#"
|
|
UPDATE proxy_nodes
|
|
SET
|
|
name = COALESCE($2, name),
|
|
remote_config = $3::json,
|
|
config_version = config_version + 1,
|
|
updated_at = NOW()
|
|
WHERE id = $1
|
|
AND tunnel_generation = $4
|
|
AND config_version = $5
|
|
AND is_manual = FALSE
|
|
"#;
|
|
|
|
const RESET_STALE_TUNNEL_STATUSES_SQL: &str = r#"
|
|
UPDATE proxy_nodes
|
|
SET
|
|
tunnel_connected = FALSE,
|
|
status = 'offline'::proxynodestatus,
|
|
active_connections = 0,
|
|
tunnel_connected_at = NOW(),
|
|
updated_at = NOW()
|
|
WHERE is_manual = FALSE
|
|
AND tunnel_connected = TRUE
|
|
"#;
|
|
|
|
const INCREMENT_MANUAL_PROXY_NODE_REQUESTS_SQL: &str = r#"
|
|
UPDATE proxy_nodes
|
|
SET
|
|
total_requests = total_requests + GREATEST($1::bigint, 0),
|
|
failed_requests = failed_requests + GREATEST($2::bigint, 0),
|
|
avg_latency_ms = COALESCE($3, avg_latency_ms),
|
|
last_heartbeat_at = NOW(),
|
|
updated_at = NOW()
|
|
WHERE id = $4
|
|
AND is_manual = TRUE
|
|
AND tunnel_generation = $5
|
|
"#;
|
|
|
|
const INSERT_PROXY_NODE_EVENT_SQL: &str = r#"
|
|
INSERT INTO proxy_node_events (node_id, event_type, detail, event_metadata, created_at)
|
|
SELECT
|
|
$1,
|
|
$2,
|
|
$3,
|
|
$4::json,
|
|
CASE
|
|
WHEN $5::double precision IS NULL THEN NOW()
|
|
ELSE TO_TIMESTAMP($5::double precision)
|
|
END
|
|
FROM proxy_nodes
|
|
WHERE id = $1 AND ($6::text IS NULL OR tunnel_generation = $6)
|
|
"#;
|
|
|
|
const UPSERT_PROXY_NODE_METRICS_1M_SQL: &str = r#"
|
|
INSERT INTO proxy_node_metrics_1m (
|
|
node_id,
|
|
bucket_start_unix_secs,
|
|
samples,
|
|
uptime_samples,
|
|
active_connections_sum,
|
|
active_connections_max,
|
|
heartbeat_rtt_ms_sum,
|
|
heartbeat_rtt_ms_max,
|
|
connect_errors_delta,
|
|
disconnects_delta,
|
|
error_events_delta,
|
|
ws_in_bytes_delta,
|
|
ws_out_bytes_delta,
|
|
ws_in_frames_delta,
|
|
ws_out_frames_delta
|
|
)
|
|
SELECT $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15
|
|
FROM proxy_nodes
|
|
WHERE id = $1 AND ($16::text IS NULL OR tunnel_generation = $16)
|
|
ON CONFLICT (node_id, bucket_start_unix_secs) DO UPDATE SET
|
|
samples = proxy_node_metrics_1m.samples + EXCLUDED.samples,
|
|
uptime_samples = proxy_node_metrics_1m.uptime_samples + EXCLUDED.uptime_samples,
|
|
active_connections_sum = proxy_node_metrics_1m.active_connections_sum + EXCLUDED.active_connections_sum,
|
|
active_connections_max = GREATEST(proxy_node_metrics_1m.active_connections_max, EXCLUDED.active_connections_max),
|
|
heartbeat_rtt_ms_sum = proxy_node_metrics_1m.heartbeat_rtt_ms_sum + EXCLUDED.heartbeat_rtt_ms_sum,
|
|
heartbeat_rtt_ms_max = GREATEST(proxy_node_metrics_1m.heartbeat_rtt_ms_max, EXCLUDED.heartbeat_rtt_ms_max),
|
|
connect_errors_delta = proxy_node_metrics_1m.connect_errors_delta + EXCLUDED.connect_errors_delta,
|
|
disconnects_delta = proxy_node_metrics_1m.disconnects_delta + EXCLUDED.disconnects_delta,
|
|
error_events_delta = proxy_node_metrics_1m.error_events_delta + EXCLUDED.error_events_delta,
|
|
ws_in_bytes_delta = proxy_node_metrics_1m.ws_in_bytes_delta + EXCLUDED.ws_in_bytes_delta,
|
|
ws_out_bytes_delta = proxy_node_metrics_1m.ws_out_bytes_delta + EXCLUDED.ws_out_bytes_delta,
|
|
ws_in_frames_delta = proxy_node_metrics_1m.ws_in_frames_delta + EXCLUDED.ws_in_frames_delta,
|
|
ws_out_frames_delta = proxy_node_metrics_1m.ws_out_frames_delta + EXCLUDED.ws_out_frames_delta
|
|
"#;
|
|
|
|
const UPSERT_PROXY_NODE_METRICS_1H_SQL: &str = r#"
|
|
INSERT INTO proxy_node_metrics_1h (
|
|
node_id,
|
|
bucket_start_unix_secs,
|
|
samples,
|
|
uptime_samples,
|
|
active_connections_sum,
|
|
active_connections_max,
|
|
heartbeat_rtt_ms_sum,
|
|
heartbeat_rtt_ms_max,
|
|
connect_errors_delta,
|
|
disconnects_delta,
|
|
error_events_delta,
|
|
ws_in_bytes_delta,
|
|
ws_out_bytes_delta,
|
|
ws_in_frames_delta,
|
|
ws_out_frames_delta
|
|
)
|
|
SELECT $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15
|
|
FROM proxy_nodes
|
|
WHERE id = $1 AND ($16::text IS NULL OR tunnel_generation = $16)
|
|
ON CONFLICT (node_id, bucket_start_unix_secs) DO UPDATE SET
|
|
samples = proxy_node_metrics_1h.samples + EXCLUDED.samples,
|
|
uptime_samples = proxy_node_metrics_1h.uptime_samples + EXCLUDED.uptime_samples,
|
|
active_connections_sum = proxy_node_metrics_1h.active_connections_sum + EXCLUDED.active_connections_sum,
|
|
active_connections_max = GREATEST(proxy_node_metrics_1h.active_connections_max, EXCLUDED.active_connections_max),
|
|
heartbeat_rtt_ms_sum = proxy_node_metrics_1h.heartbeat_rtt_ms_sum + EXCLUDED.heartbeat_rtt_ms_sum,
|
|
heartbeat_rtt_ms_max = GREATEST(proxy_node_metrics_1h.heartbeat_rtt_ms_max, EXCLUDED.heartbeat_rtt_ms_max),
|
|
connect_errors_delta = proxy_node_metrics_1h.connect_errors_delta + EXCLUDED.connect_errors_delta,
|
|
disconnects_delta = proxy_node_metrics_1h.disconnects_delta + EXCLUDED.disconnects_delta,
|
|
error_events_delta = proxy_node_metrics_1h.error_events_delta + EXCLUDED.error_events_delta,
|
|
ws_in_bytes_delta = proxy_node_metrics_1h.ws_in_bytes_delta + EXCLUDED.ws_in_bytes_delta,
|
|
ws_out_bytes_delta = proxy_node_metrics_1h.ws_out_bytes_delta + EXCLUDED.ws_out_bytes_delta,
|
|
ws_in_frames_delta = proxy_node_metrics_1h.ws_in_frames_delta + EXCLUDED.ws_in_frames_delta,
|
|
ws_out_frames_delta = proxy_node_metrics_1h.ws_out_frames_delta + EXCLUDED.ws_out_frames_delta
|
|
"#;
|
|
|
|
const LIST_PROXY_NODE_METRICS_1M_SQL: &str = r#"
|
|
SELECT
|
|
node_id,
|
|
bucket_start_unix_secs,
|
|
samples,
|
|
uptime_samples,
|
|
active_connections_sum,
|
|
active_connections_max,
|
|
heartbeat_rtt_ms_sum,
|
|
heartbeat_rtt_ms_max,
|
|
connect_errors_delta,
|
|
disconnects_delta,
|
|
error_events_delta,
|
|
ws_in_bytes_delta,
|
|
ws_out_bytes_delta,
|
|
ws_in_frames_delta,
|
|
ws_out_frames_delta
|
|
FROM proxy_node_metrics_1m
|
|
WHERE node_id = $1
|
|
AND bucket_start_unix_secs >= $2
|
|
AND bucket_start_unix_secs <= $3
|
|
ORDER BY bucket_start_unix_secs ASC
|
|
LIMIT $4
|
|
"#;
|
|
|
|
const LIST_PROXY_NODE_METRICS_1H_SQL: &str = r#"
|
|
SELECT
|
|
node_id,
|
|
bucket_start_unix_secs,
|
|
samples,
|
|
uptime_samples,
|
|
active_connections_sum,
|
|
active_connections_max,
|
|
heartbeat_rtt_ms_sum,
|
|
heartbeat_rtt_ms_max,
|
|
connect_errors_delta,
|
|
disconnects_delta,
|
|
error_events_delta,
|
|
ws_in_bytes_delta,
|
|
ws_out_bytes_delta,
|
|
ws_in_frames_delta,
|
|
ws_out_frames_delta
|
|
FROM proxy_node_metrics_1h
|
|
WHERE node_id = $1
|
|
AND bucket_start_unix_secs >= $2
|
|
AND bucket_start_unix_secs <= $3
|
|
ORDER BY bucket_start_unix_secs ASC
|
|
LIMIT $4
|
|
"#;
|
|
|
|
const LIST_PROXY_FLEET_METRICS_1M_SQL: &str = r#"
|
|
SELECT
|
|
bucket_start_unix_secs,
|
|
SUM(samples) AS samples,
|
|
SUM(uptime_samples) AS uptime_samples,
|
|
SUM(active_connections_sum) AS active_connections_sum,
|
|
MAX(active_connections_max) AS active_connections_max,
|
|
SUM(heartbeat_rtt_ms_sum) AS heartbeat_rtt_ms_sum,
|
|
MAX(heartbeat_rtt_ms_max) AS heartbeat_rtt_ms_max,
|
|
SUM(connect_errors_delta) AS connect_errors_delta,
|
|
SUM(disconnects_delta) AS disconnects_delta,
|
|
SUM(error_events_delta) AS error_events_delta,
|
|
SUM(ws_in_bytes_delta) AS ws_in_bytes_delta,
|
|
SUM(ws_out_bytes_delta) AS ws_out_bytes_delta,
|
|
SUM(ws_in_frames_delta) AS ws_in_frames_delta,
|
|
SUM(ws_out_frames_delta) AS ws_out_frames_delta
|
|
FROM proxy_node_metrics_1m
|
|
WHERE bucket_start_unix_secs >= $1
|
|
AND bucket_start_unix_secs <= $2
|
|
GROUP BY bucket_start_unix_secs
|
|
ORDER BY bucket_start_unix_secs ASC
|
|
LIMIT $3
|
|
"#;
|
|
|
|
const LIST_PROXY_FLEET_METRICS_1H_SQL: &str = r#"
|
|
SELECT
|
|
bucket_start_unix_secs,
|
|
SUM(samples) AS samples,
|
|
SUM(uptime_samples) AS uptime_samples,
|
|
SUM(active_connections_sum) AS active_connections_sum,
|
|
MAX(active_connections_max) AS active_connections_max,
|
|
SUM(heartbeat_rtt_ms_sum) AS heartbeat_rtt_ms_sum,
|
|
MAX(heartbeat_rtt_ms_max) AS heartbeat_rtt_ms_max,
|
|
SUM(connect_errors_delta) AS connect_errors_delta,
|
|
SUM(disconnects_delta) AS disconnects_delta,
|
|
SUM(error_events_delta) AS error_events_delta,
|
|
SUM(ws_in_bytes_delta) AS ws_in_bytes_delta,
|
|
SUM(ws_out_bytes_delta) AS ws_out_bytes_delta,
|
|
SUM(ws_in_frames_delta) AS ws_in_frames_delta,
|
|
SUM(ws_out_frames_delta) AS ws_out_frames_delta
|
|
FROM proxy_node_metrics_1h
|
|
WHERE bucket_start_unix_secs >= $1
|
|
AND bucket_start_unix_secs <= $2
|
|
GROUP BY bucket_start_unix_secs
|
|
ORDER BY bucket_start_unix_secs ASC
|
|
LIMIT $3
|
|
"#;
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct SqlxProxyNodeRepository {
|
|
pool: PgPool,
|
|
}
|
|
|
|
impl SqlxProxyNodeRepository {
|
|
pub fn new(pool: PgPool) -> Self {
|
|
Self { pool }
|
|
}
|
|
|
|
fn optional_unix_secs(value: Option<i64>) -> Option<u64> {
|
|
value.and_then(|value| u64::try_from(value).ok())
|
|
}
|
|
|
|
fn row_to_stored(row: &PgRow) -> Result<StoredProxyNode, DataLayerError> {
|
|
let tunnel_generation: String = row.try_get("tunnel_generation").map_postgres_err()?;
|
|
if tunnel_generation.trim().is_empty() {
|
|
return Err(DataLayerError::UnexpectedValue(
|
|
"proxy_nodes.tunnel_generation must not be empty".to_string(),
|
|
));
|
|
}
|
|
Ok(StoredProxyNode::new(
|
|
row.try_get("id").map_postgres_err()?,
|
|
row.try_get("name").map_postgres_err()?,
|
|
row.try_get("ip").map_postgres_err()?,
|
|
row.try_get("port").map_postgres_err()?,
|
|
row.try_get("is_manual").map_postgres_err()?,
|
|
row.try_get("status").map_postgres_err()?,
|
|
row.try_get("heartbeat_interval").map_postgres_err()?,
|
|
row.try_get("active_connections").map_postgres_err()?,
|
|
row.try_get("total_requests").map_postgres_err()?,
|
|
row.try_get("failed_requests").map_postgres_err()?,
|
|
row.try_get("dns_failures").map_postgres_err()?,
|
|
row.try_get("stream_errors").map_postgres_err()?,
|
|
row.try_get("tunnel_mode").map_postgres_err()?,
|
|
row.try_get("tunnel_connected").map_postgres_err()?,
|
|
row.try_get("config_version").map_postgres_err()?,
|
|
)?
|
|
.with_tunnel_generation(tunnel_generation)
|
|
.with_manual_proxy_fields(
|
|
row.try_get("proxy_url").map_postgres_err()?,
|
|
row.try_get("proxy_username").map_postgres_err()?,
|
|
row.try_get("proxy_password").map_postgres_err()?,
|
|
)
|
|
.with_runtime_fields(
|
|
row.try_get("region").map_postgres_err()?,
|
|
row.try_get("registered_by").map_postgres_err()?,
|
|
Self::optional_unix_secs(
|
|
row.try_get("last_heartbeat_at_unix_secs")
|
|
.map_postgres_err()?,
|
|
),
|
|
row.try_get("avg_latency_ms").map_postgres_err()?,
|
|
row.try_get("proxy_metadata").map_postgres_err()?,
|
|
row.try_get("hardware_info").map_postgres_err()?,
|
|
row.try_get("estimated_max_concurrency")
|
|
.map_postgres_err()?,
|
|
Self::optional_unix_secs(
|
|
row.try_get("tunnel_connected_at_unix_secs")
|
|
.map_postgres_err()?,
|
|
),
|
|
row.try_get("remote_config").map_postgres_err()?,
|
|
Self::optional_unix_secs(row.try_get("created_at_unix_ms").map_postgres_err()?),
|
|
Self::optional_unix_secs(row.try_get("updated_at_unix_secs").map_postgres_err()?),
|
|
))
|
|
}
|
|
|
|
fn row_to_event(row: &PgRow) -> Result<StoredProxyNodeEvent, DataLayerError> {
|
|
Ok(StoredProxyNodeEvent {
|
|
id: row.try_get("id").map_postgres_err()?,
|
|
node_id: row.try_get("node_id").map_postgres_err()?,
|
|
event_type: row.try_get("event_type").map_postgres_err()?,
|
|
detail: row.try_get("detail").map_postgres_err()?,
|
|
event_metadata: row.try_get("event_metadata").map_postgres_err()?,
|
|
created_at_unix_ms: Self::optional_unix_secs(
|
|
row.try_get("created_at_unix_ms").map_postgres_err()?,
|
|
),
|
|
})
|
|
}
|
|
|
|
fn row_to_node_metric(row: &PgRow) -> Result<StoredProxyNodeMetricsBucket, DataLayerError> {
|
|
Ok(StoredProxyNodeMetricsBucket {
|
|
node_id: row.try_get("node_id").map_postgres_err()?,
|
|
bucket_start_unix_secs: Self::optional_unix_secs(
|
|
row.try_get("bucket_start_unix_secs").map_postgres_err()?,
|
|
)
|
|
.unwrap_or_default(),
|
|
samples: row.try_get("samples").map_postgres_err()?,
|
|
uptime_samples: row.try_get("uptime_samples").map_postgres_err()?,
|
|
active_connections_sum: row.try_get("active_connections_sum").map_postgres_err()?,
|
|
active_connections_max: row.try_get("active_connections_max").map_postgres_err()?,
|
|
heartbeat_rtt_ms_sum: row.try_get("heartbeat_rtt_ms_sum").map_postgres_err()?,
|
|
heartbeat_rtt_ms_max: row.try_get("heartbeat_rtt_ms_max").map_postgres_err()?,
|
|
connect_errors_delta: row.try_get("connect_errors_delta").map_postgres_err()?,
|
|
disconnects_delta: row.try_get("disconnects_delta").map_postgres_err()?,
|
|
error_events_delta: row.try_get("error_events_delta").map_postgres_err()?,
|
|
ws_in_bytes_delta: row.try_get("ws_in_bytes_delta").map_postgres_err()?,
|
|
ws_out_bytes_delta: row.try_get("ws_out_bytes_delta").map_postgres_err()?,
|
|
ws_in_frames_delta: row.try_get("ws_in_frames_delta").map_postgres_err()?,
|
|
ws_out_frames_delta: row.try_get("ws_out_frames_delta").map_postgres_err()?,
|
|
})
|
|
}
|
|
|
|
fn row_to_fleet_metric(row: &PgRow) -> Result<StoredProxyFleetMetricsBucket, DataLayerError> {
|
|
Ok(StoredProxyFleetMetricsBucket {
|
|
bucket_start_unix_secs: Self::optional_unix_secs(
|
|
row.try_get("bucket_start_unix_secs").map_postgres_err()?,
|
|
)
|
|
.unwrap_or_default(),
|
|
samples: row.try_get("samples").map_postgres_err()?,
|
|
uptime_samples: row.try_get("uptime_samples").map_postgres_err()?,
|
|
active_connections_sum: row.try_get("active_connections_sum").map_postgres_err()?,
|
|
active_connections_max: row.try_get("active_connections_max").map_postgres_err()?,
|
|
heartbeat_rtt_ms_sum: row.try_get("heartbeat_rtt_ms_sum").map_postgres_err()?,
|
|
heartbeat_rtt_ms_max: row.try_get("heartbeat_rtt_ms_max").map_postgres_err()?,
|
|
connect_errors_delta: row.try_get("connect_errors_delta").map_postgres_err()?,
|
|
disconnects_delta: row.try_get("disconnects_delta").map_postgres_err()?,
|
|
error_events_delta: row.try_get("error_events_delta").map_postgres_err()?,
|
|
ws_in_bytes_delta: row.try_get("ws_in_bytes_delta").map_postgres_err()?,
|
|
ws_out_bytes_delta: row.try_get("ws_out_bytes_delta").map_postgres_err()?,
|
|
ws_in_frames_delta: row.try_get("ws_in_frames_delta").map_postgres_err()?,
|
|
ws_out_frames_delta: row.try_get("ws_out_frames_delta").map_postgres_err()?,
|
|
})
|
|
}
|
|
|
|
async fn insert_event(
|
|
&self,
|
|
node_id: &str,
|
|
expected_tunnel_generation: Option<&str>,
|
|
event_type: &str,
|
|
detail: Option<&str>,
|
|
event_metadata: Option<&serde_json::Value>,
|
|
created_at_unix_secs: Option<u64>,
|
|
) -> Result<(), DataLayerError> {
|
|
sqlx::query(INSERT_PROXY_NODE_EVENT_SQL)
|
|
.bind(node_id)
|
|
.bind(event_type)
|
|
.bind(detail)
|
|
.bind(event_metadata)
|
|
.bind(created_at_unix_secs.map(|value| value as f64))
|
|
.bind(expected_tunnel_generation)
|
|
.execute(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
Ok(())
|
|
}
|
|
|
|
async fn upsert_metrics_bucket(
|
|
&self,
|
|
step: ProxyNodeMetricsStep,
|
|
node_id: &str,
|
|
expected_tunnel_generation: Option<&str>,
|
|
bucket_start: u64,
|
|
sample: &TunnelMetricsSample,
|
|
) -> Result<(), DataLayerError> {
|
|
let sql = match step {
|
|
ProxyNodeMetricsStep::OneMinute => UPSERT_PROXY_NODE_METRICS_1M_SQL,
|
|
ProxyNodeMetricsStep::OneHour => UPSERT_PROXY_NODE_METRICS_1H_SQL,
|
|
};
|
|
sqlx::query(sql)
|
|
.bind(node_id)
|
|
.bind(i64::try_from(bucket_start).unwrap_or(i64::MAX))
|
|
.bind(sample.samples)
|
|
.bind(sample.uptime_samples)
|
|
.bind(sample.active_connections_sum)
|
|
.bind(sample.active_connections_max)
|
|
.bind(sample.heartbeat_rtt_ms_sum)
|
|
.bind(sample.heartbeat_rtt_ms_max)
|
|
.bind(sample.connect_errors_delta)
|
|
.bind(sample.disconnects_delta)
|
|
.bind(sample.error_events_delta)
|
|
.bind(sample.ws_in_bytes_delta)
|
|
.bind(sample.ws_out_bytes_delta)
|
|
.bind(sample.ws_in_frames_delta)
|
|
.bind(sample.ws_out_frames_delta)
|
|
.bind(expected_tunnel_generation)
|
|
.execute(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
Ok(())
|
|
}
|
|
|
|
fn registration_lock_key(ip: &str, port: i32) -> i64 {
|
|
let mut hasher = Sha256::new();
|
|
hasher.update(ip.as_bytes());
|
|
hasher.update(b":");
|
|
hasher.update(port.to_string().as_bytes());
|
|
let digest = hasher.finalize();
|
|
let mut bytes = [0u8; 8];
|
|
bytes.copy_from_slice(&digest[..8]);
|
|
i64::from_be_bytes(bytes)
|
|
}
|
|
|
|
fn normalize_remote_config(
|
|
mutation: &ProxyNodeRemoteConfigMutation,
|
|
existing: Option<&serde_json::Value>,
|
|
) -> Option<serde_json::Value> {
|
|
let mut config = match existing {
|
|
Some(serde_json::Value::Object(map)) => map.clone(),
|
|
_ => serde_json::Map::new(),
|
|
};
|
|
|
|
if let Some(node_name) = mutation.node_name.as_ref() {
|
|
config.insert(
|
|
"node_name".to_string(),
|
|
serde_json::Value::String(node_name.clone()),
|
|
);
|
|
}
|
|
if let Some(allowed_ports) = mutation.allowed_ports.as_ref() {
|
|
config.insert(
|
|
"allowed_ports".to_string(),
|
|
serde_json::json!(allowed_ports),
|
|
);
|
|
}
|
|
if let Some(log_level) = mutation.log_level.as_ref() {
|
|
config.insert(
|
|
"log_level".to_string(),
|
|
serde_json::Value::String(log_level.clone()),
|
|
);
|
|
}
|
|
if let Some(heartbeat_interval) = mutation.heartbeat_interval {
|
|
config.insert(
|
|
"heartbeat_interval".to_string(),
|
|
serde_json::json!(heartbeat_interval),
|
|
);
|
|
}
|
|
if let Some(scheduling_state) = mutation.scheduling_state.as_ref() {
|
|
match scheduling_state {
|
|
Some(state) => {
|
|
config.insert(
|
|
"scheduling_state".to_string(),
|
|
serde_json::Value::String(state.clone()),
|
|
);
|
|
}
|
|
None => {
|
|
config.remove("scheduling_state");
|
|
}
|
|
}
|
|
}
|
|
if let Some(upgrade_to) = mutation.upgrade_to.as_ref() {
|
|
match upgrade_to {
|
|
Some(version) => {
|
|
config.insert(
|
|
"upgrade_to".to_string(),
|
|
serde_json::Value::String(version.clone()),
|
|
);
|
|
}
|
|
None => {
|
|
config.remove("upgrade_to");
|
|
}
|
|
}
|
|
}
|
|
|
|
(!config.is_empty()).then_some(serde_json::Value::Object(config))
|
|
}
|
|
|
|
fn duplicate_proxy_node_detail(name: &str, ip: &str, port: i32) -> String {
|
|
format!("已存在相同地址的代理节点: {name} ({ip}:{port})")
|
|
}
|
|
|
|
async fn find_duplicate_proxy_node_locked(
|
|
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
|
ip: &str,
|
|
port: i32,
|
|
exclude_node_id: Option<&str>,
|
|
) -> Result<Option<(String, String, i32)>, DataLayerError> {
|
|
let row = if let Some(exclude_node_id) = exclude_node_id {
|
|
sqlx::query(FIND_DUPLICATE_PROXY_NODE_EXCLUDING_ID_SQL)
|
|
.bind(ip)
|
|
.bind(port)
|
|
.bind(exclude_node_id)
|
|
.fetch_optional(&mut **tx)
|
|
.await
|
|
.map_postgres_err()?
|
|
} else {
|
|
sqlx::query(FIND_DUPLICATE_PROXY_NODE_SQL)
|
|
.bind(ip)
|
|
.bind(port)
|
|
.fetch_optional(&mut **tx)
|
|
.await
|
|
.map_postgres_err()?
|
|
};
|
|
|
|
row.map(|row| {
|
|
Ok((
|
|
row.try_get("name").map_postgres_err()?,
|
|
row.try_get("ip").map_postgres_err()?,
|
|
row.try_get("port").map_postgres_err()?,
|
|
))
|
|
})
|
|
.transpose()
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl ProxyNodeReadRepository for SqlxProxyNodeRepository {
|
|
async fn list_proxy_nodes(&self) -> Result<Vec<StoredProxyNode>, DataLayerError> {
|
|
let mut builder = QueryBuilder::<Postgres>::new(proxy_node_columns());
|
|
builder.push(" ORDER BY name ASC, id ASC");
|
|
let mut rows = builder.build().fetch(&self.pool);
|
|
let mut items = Vec::new();
|
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
|
items.push(Self::row_to_stored(&row)?);
|
|
}
|
|
Ok(items)
|
|
}
|
|
|
|
async fn find_proxy_node(
|
|
&self,
|
|
node_id: &str,
|
|
) -> Result<Option<StoredProxyNode>, DataLayerError> {
|
|
let mut builder = QueryBuilder::<Postgres>::new(proxy_node_columns());
|
|
let mut where_clause = WhereClause::new();
|
|
push_eq(&mut builder, &mut where_clause, "id", node_id.to_string());
|
|
push_limit(&mut builder, 1);
|
|
let row = builder
|
|
.build()
|
|
.fetch_optional(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
row.map(|row| Self::row_to_stored(&row)).transpose()
|
|
}
|
|
|
|
async fn list_proxy_node_events(
|
|
&self,
|
|
node_id: &str,
|
|
limit: usize,
|
|
) -> Result<Vec<StoredProxyNodeEvent>, DataLayerError> {
|
|
let mut builder = QueryBuilder::<Postgres>::new(proxy_node_event_columns());
|
|
let mut where_clause = WhereClause::new();
|
|
push_eq(
|
|
&mut builder,
|
|
&mut where_clause,
|
|
"node_id",
|
|
node_id.to_string(),
|
|
);
|
|
builder.push(" ORDER BY created_at DESC, id DESC");
|
|
push_limit(&mut builder, i64::try_from(limit).unwrap_or(i64::MAX));
|
|
let mut rows = builder.build().fetch(&self.pool);
|
|
let mut items = Vec::new();
|
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
|
items.push(Self::row_to_event(&row)?);
|
|
}
|
|
Ok(items)
|
|
}
|
|
|
|
async fn list_proxy_node_events_filtered(
|
|
&self,
|
|
node_id: &str,
|
|
query: &ProxyNodeEventQuery,
|
|
) -> Result<Vec<StoredProxyNodeEvent>, DataLayerError> {
|
|
let mut builder = QueryBuilder::<Postgres>::new(proxy_node_event_columns());
|
|
let mut where_clause = WhereClause::new();
|
|
push_eq(
|
|
&mut builder,
|
|
&mut where_clause,
|
|
"node_id",
|
|
node_id.to_string(),
|
|
);
|
|
if let Some(from_unix_secs) = query.from_unix_secs {
|
|
where_clause.push_next(&mut builder);
|
|
builder
|
|
.push("created_at >= TO_TIMESTAMP(")
|
|
.push_bind(from_unix_secs as f64)
|
|
.push("::double precision)");
|
|
}
|
|
if let Some(to_unix_secs) = query.to_unix_secs {
|
|
where_clause.push_next(&mut builder);
|
|
builder
|
|
.push("created_at <= TO_TIMESTAMP(")
|
|
.push_bind(to_unix_secs as f64)
|
|
.push("::double precision)");
|
|
}
|
|
if let Some(event_type) = query.event_type.as_deref() {
|
|
where_clause.push_next(&mut builder);
|
|
builder
|
|
.push("LOWER(CAST(event_type AS TEXT)) = LOWER(")
|
|
.push_bind(event_type.to_string())
|
|
.push("::text)");
|
|
}
|
|
builder.push(" ORDER BY created_at DESC, id DESC");
|
|
push_limit(&mut builder, i64::try_from(query.limit).unwrap_or(i64::MAX));
|
|
let mut rows = builder.build().fetch(&self.pool);
|
|
let mut items = Vec::new();
|
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
|
items.push(Self::row_to_event(&row)?);
|
|
}
|
|
Ok(items)
|
|
}
|
|
|
|
async fn list_proxy_node_metrics(
|
|
&self,
|
|
node_id: &str,
|
|
step: ProxyNodeMetricsStep,
|
|
from_unix_secs: u64,
|
|
to_unix_secs: u64,
|
|
limit: usize,
|
|
) -> Result<Vec<StoredProxyNodeMetricsBucket>, DataLayerError> {
|
|
let sql = match step {
|
|
ProxyNodeMetricsStep::OneMinute => LIST_PROXY_NODE_METRICS_1M_SQL,
|
|
ProxyNodeMetricsStep::OneHour => LIST_PROXY_NODE_METRICS_1H_SQL,
|
|
};
|
|
let mut rows = sqlx::query(sql)
|
|
.bind(node_id)
|
|
.bind(i64::try_from(from_unix_secs).unwrap_or(i64::MAX))
|
|
.bind(i64::try_from(to_unix_secs).unwrap_or(i64::MAX))
|
|
.bind(i64::try_from(limit).unwrap_or(i64::MAX))
|
|
.fetch(&self.pool);
|
|
let mut items = Vec::new();
|
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
|
items.push(Self::row_to_node_metric(&row)?);
|
|
}
|
|
Ok(items)
|
|
}
|
|
|
|
async fn list_proxy_fleet_metrics(
|
|
&self,
|
|
step: ProxyNodeMetricsStep,
|
|
from_unix_secs: u64,
|
|
to_unix_secs: u64,
|
|
limit: usize,
|
|
) -> Result<Vec<StoredProxyFleetMetricsBucket>, DataLayerError> {
|
|
let sql = match step {
|
|
ProxyNodeMetricsStep::OneMinute => LIST_PROXY_FLEET_METRICS_1M_SQL,
|
|
ProxyNodeMetricsStep::OneHour => LIST_PROXY_FLEET_METRICS_1H_SQL,
|
|
};
|
|
let mut rows = sqlx::query(sql)
|
|
.bind(i64::try_from(from_unix_secs).unwrap_or(i64::MAX))
|
|
.bind(i64::try_from(to_unix_secs).unwrap_or(i64::MAX))
|
|
.bind(i64::try_from(limit).unwrap_or(i64::MAX))
|
|
.fetch(&self.pool);
|
|
let mut items = Vec::new();
|
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
|
items.push(Self::row_to_fleet_metric(&row)?);
|
|
}
|
|
Ok(items)
|
|
}
|
|
}
|
|
|
|
fn proxy_node_columns() -> &'static str {
|
|
FIND_PROXY_NODE_SQL
|
|
.split_once("WHERE id = $1")
|
|
.map(|(prefix, _)| prefix)
|
|
.unwrap_or(FIND_PROXY_NODE_SQL)
|
|
}
|
|
|
|
fn proxy_node_event_columns() -> &'static str {
|
|
LIST_PROXY_NODE_EVENTS_SQL
|
|
.split_once("WHERE node_id = $1")
|
|
.map(|(prefix, _)| prefix)
|
|
.unwrap_or(LIST_PROXY_NODE_EVENTS_SQL)
|
|
}
|
|
|
|
#[async_trait]
|
|
impl ProxyNodeWriteRepository for SqlxProxyNodeRepository {
|
|
async fn reset_stale_tunnel_statuses(&self) -> Result<usize, DataLayerError> {
|
|
let result = sqlx::query(RESET_STALE_TUNNEL_STATUSES_SQL)
|
|
.execute(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
Ok(result.rows_affected() as usize)
|
|
}
|
|
|
|
async fn compare_and_set_proxy_password(
|
|
&self,
|
|
node_id: &str,
|
|
expected: &str,
|
|
replacement: &str,
|
|
) -> Result<bool, DataLayerError> {
|
|
let result = sqlx::query(
|
|
r#"
|
|
UPDATE proxy_nodes
|
|
SET proxy_password = $2, updated_at = NOW()
|
|
WHERE id = $1 AND proxy_password = $3
|
|
"#,
|
|
)
|
|
.bind(node_id)
|
|
.bind(replacement)
|
|
.bind(expected)
|
|
.execute(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
Ok(result.rows_affected() == 1)
|
|
}
|
|
|
|
async fn compare_and_set_proxy_metadata(
|
|
&self,
|
|
node_id: &str,
|
|
expected: &serde_json::Value,
|
|
replacement: &serde_json::Value,
|
|
) -> Result<bool, DataLayerError> {
|
|
let expected = serde_json::to_string(expected).map_err(|err| {
|
|
DataLayerError::InvalidInput(format!("proxy_nodes.proxy_metadata is invalid: {err}"))
|
|
})?;
|
|
let replacement = serde_json::to_string(replacement).map_err(|err| {
|
|
DataLayerError::InvalidInput(format!("proxy_nodes.proxy_metadata is invalid: {err}"))
|
|
})?;
|
|
let result = sqlx::query(
|
|
r#"
|
|
UPDATE proxy_nodes
|
|
SET proxy_metadata = $2::json, updated_at = NOW()
|
|
WHERE id = $1 AND proxy_metadata::jsonb = $3::jsonb
|
|
"#,
|
|
)
|
|
.bind(node_id)
|
|
.bind(replacement)
|
|
.bind(expected)
|
|
.execute(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
Ok(result.rows_affected() == 1)
|
|
}
|
|
|
|
async fn create_manual_node(
|
|
&self,
|
|
mutation: &ProxyNodeManualCreateMutation,
|
|
) -> Result<StoredProxyNode, DataLayerError> {
|
|
let lock_key = Self::registration_lock_key(&mutation.ip, mutation.port);
|
|
let mut tx = self.pool.begin().await.map_postgres_err()?;
|
|
sqlx::query("SELECT pg_advisory_xact_lock($1)")
|
|
.bind(lock_key)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
|
|
if let Some((name, ip, port)) =
|
|
Self::find_duplicate_proxy_node_locked(&mut tx, &mutation.ip, mutation.port, None)
|
|
.await?
|
|
{
|
|
return Err(DataLayerError::InvalidInput(
|
|
Self::duplicate_proxy_node_detail(&name, &ip, port),
|
|
));
|
|
}
|
|
|
|
let node_id = requested_proxy_node_id(mutation.node_id.as_deref())?
|
|
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
|
|
let tunnel_generation = uuid::Uuid::new_v4().to_string();
|
|
if let Some((ip, port)) = proxy_node_id_owner_locked(&mut tx, &node_id).await? {
|
|
return Err(proxy_node_id_in_use_error(&node_id, &ip, port));
|
|
}
|
|
sqlx::query(INSERT_MANUAL_PROXY_NODE_SQL)
|
|
.bind(&node_id)
|
|
.bind(&mutation.name)
|
|
.bind(&mutation.ip)
|
|
.bind(mutation.port)
|
|
.bind(mutation.region.as_deref())
|
|
.bind(&mutation.proxy_url)
|
|
.bind(mutation.proxy_username.as_deref())
|
|
.bind(mutation.proxy_password.as_deref())
|
|
.bind(mutation.registered_by.as_deref())
|
|
.bind(&tunnel_generation)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
|
|
tx.commit().await.map_err(postgres_error)?;
|
|
self.find_proxy_node(&node_id).await?.ok_or_else(|| {
|
|
DataLayerError::UnexpectedValue("created manual proxy node missing".to_string())
|
|
})
|
|
}
|
|
|
|
async fn update_manual_node(
|
|
&self,
|
|
mutation: &ProxyNodeManualUpdateMutation,
|
|
) -> Result<Option<StoredProxyNode>, DataLayerError> {
|
|
let existing = self.find_proxy_node(&mutation.node_id).await?;
|
|
let Some(existing) = existing else {
|
|
return Ok(None);
|
|
};
|
|
if !existing.is_manual {
|
|
return Err(DataLayerError::InvalidInput(
|
|
"只能编辑手动添加的代理节点".to_string(),
|
|
));
|
|
}
|
|
|
|
let next_ip = mutation.ip.as_deref().unwrap_or(existing.ip.as_str());
|
|
let next_port = mutation.port.unwrap_or(existing.port);
|
|
let lock_key = Self::registration_lock_key(next_ip, next_port);
|
|
let mut tx = self.pool.begin().await.map_postgres_err()?;
|
|
sqlx::query("SELECT pg_advisory_xact_lock($1)")
|
|
.bind(lock_key)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
|
|
if let Some((name, ip, port)) = Self::find_duplicate_proxy_node_locked(
|
|
&mut tx,
|
|
next_ip,
|
|
next_port,
|
|
Some(&mutation.node_id),
|
|
)
|
|
.await?
|
|
{
|
|
return Err(DataLayerError::InvalidInput(
|
|
Self::duplicate_proxy_node_detail(&name, &ip, port),
|
|
));
|
|
}
|
|
|
|
let result = sqlx::query(UPDATE_MANUAL_PROXY_NODE_SQL)
|
|
.bind(&mutation.node_id)
|
|
.bind(mutation.name.as_deref())
|
|
.bind(mutation.ip.as_deref())
|
|
.bind(mutation.port)
|
|
.bind(mutation.region.as_deref())
|
|
.bind(mutation.proxy_url.as_deref())
|
|
.bind(mutation.proxy_username.as_deref())
|
|
.bind(mutation.proxy_password.as_deref())
|
|
.bind(&existing.tunnel_generation)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
|
|
if result.rows_affected() == 0 {
|
|
tx.rollback().await.map_postgres_err()?;
|
|
return Ok(None);
|
|
}
|
|
|
|
tx.commit().await.map_err(postgres_error)?;
|
|
self.find_proxy_node(&mutation.node_id).await
|
|
}
|
|
|
|
async fn register_node(
|
|
&self,
|
|
mutation: &ProxyNodeRegistrationMutation,
|
|
) -> Result<StoredProxyNode, DataLayerError> {
|
|
let normalized_proxy_metadata = merge_proxy_metadata_for_registration(
|
|
None,
|
|
normalize_proxy_metadata(
|
|
mutation.proxy_metadata.as_ref(),
|
|
mutation.proxy_version.as_deref(),
|
|
),
|
|
);
|
|
let lock_key = Self::registration_lock_key(&mutation.ip, mutation.port);
|
|
let mut tx = self.pool.begin().await.map_postgres_err()?;
|
|
|
|
sqlx::query("SELECT pg_advisory_xact_lock($1)")
|
|
.bind(lock_key)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
|
|
let existing = sqlx::query(FIND_EXISTING_TUNNEL_NODE_SQL)
|
|
.bind(&mutation.ip)
|
|
.bind(mutation.port)
|
|
.fetch_optional(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
|
|
let node_id = if let Some(row) = existing.as_ref() {
|
|
let existing = Self::row_to_stored(row)?;
|
|
if let Some(requested_id) = requested_proxy_node_id(mutation.node_id.as_deref())? {
|
|
if requested_id != existing.id {
|
|
return Err(proxy_node_registration_identity_error(
|
|
&requested_id,
|
|
&existing.id,
|
|
));
|
|
}
|
|
}
|
|
let registration_proxy_metadata = merge_proxy_metadata_for_registration(
|
|
existing.proxy_metadata.as_ref(),
|
|
normalized_proxy_metadata,
|
|
);
|
|
sqlx::query(UPDATE_PROXY_NODE_REGISTRATION_SQL)
|
|
.bind(&existing.id)
|
|
.bind(&mutation.name)
|
|
.bind(&mutation.ip)
|
|
.bind(mutation.port)
|
|
.bind(mutation.region.as_deref())
|
|
.bind(mutation.registered_by.as_deref())
|
|
.bind(mutation.heartbeat_interval)
|
|
.bind(mutation.active_connections)
|
|
.bind(mutation.total_requests)
|
|
.bind(mutation.avg_latency_ms)
|
|
.bind(mutation.hardware_info.as_ref())
|
|
.bind(mutation.estimated_max_concurrency)
|
|
.bind(mutation.tunnel_mode)
|
|
.bind(registration_proxy_metadata.as_ref())
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
existing.id
|
|
} else {
|
|
let node_id = requested_proxy_node_id(mutation.node_id.as_deref())?
|
|
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
|
|
let tunnel_generation = uuid::Uuid::new_v4().to_string();
|
|
if let Some((ip, port)) = proxy_node_id_owner_locked(&mut tx, &node_id).await? {
|
|
return Err(proxy_node_id_in_use_error(&node_id, &ip, port));
|
|
}
|
|
sqlx::query(INSERT_PROXY_NODE_SQL)
|
|
.bind(&node_id)
|
|
.bind(&mutation.name)
|
|
.bind(&mutation.ip)
|
|
.bind(mutation.port)
|
|
.bind(mutation.region.as_deref())
|
|
.bind(mutation.registered_by.as_deref())
|
|
.bind(mutation.heartbeat_interval)
|
|
.bind(mutation.active_connections)
|
|
.bind(mutation.total_requests)
|
|
.bind(mutation.avg_latency_ms)
|
|
.bind(mutation.hardware_info.as_ref())
|
|
.bind(mutation.estimated_max_concurrency)
|
|
.bind(mutation.tunnel_mode)
|
|
.bind(normalized_proxy_metadata.as_ref())
|
|
.bind(&tunnel_generation)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
node_id
|
|
};
|
|
|
|
tx.commit().await.map_err(postgres_error)?;
|
|
self.find_proxy_node(&node_id).await?.ok_or_else(|| {
|
|
DataLayerError::UnexpectedValue("registered proxy node missing".to_string())
|
|
})
|
|
}
|
|
|
|
async fn apply_heartbeat(
|
|
&self,
|
|
mutation: &ProxyNodeHeartbeatMutation,
|
|
) -> Result<Option<StoredProxyNode>, DataLayerError> {
|
|
let existing = self.find_proxy_node(&mutation.node_id).await?;
|
|
let Some(existing) = existing else {
|
|
return Ok(None);
|
|
};
|
|
if mutation
|
|
.expected_tunnel_generation
|
|
.as_deref()
|
|
.is_some_and(|expected| expected != existing.tunnel_generation)
|
|
{
|
|
return Ok(None);
|
|
}
|
|
if !existing.tunnel_mode {
|
|
return Err(DataLayerError::InvalidInput(
|
|
"non-tunnel mode is no longer supported, please upgrade aether-tunnel to use tunnel mode"
|
|
.to_string(),
|
|
));
|
|
}
|
|
|
|
let tunnel_generation = existing.tunnel_generation.clone();
|
|
let has_proxy_metadata_update = normalize_heartbeat_proxy_metadata(
|
|
None,
|
|
mutation.proxy_metadata.as_ref(),
|
|
mutation.proxy_version.as_deref(),
|
|
)
|
|
.is_some();
|
|
|
|
let result = sqlx::query(APPLY_HEARTBEAT_SQL)
|
|
.bind(&mutation.node_id)
|
|
.bind(mutation.heartbeat_interval)
|
|
.bind(mutation.active_connections)
|
|
.bind(mutation.avg_latency_ms)
|
|
.bind(mutation.total_requests_delta)
|
|
.bind(mutation.failed_requests_delta)
|
|
.bind(mutation.dns_failures_delta)
|
|
.bind(mutation.stream_errors_delta)
|
|
.bind(&tunnel_generation)
|
|
.execute(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
if result.rows_affected() == 0 {
|
|
return Ok(None);
|
|
}
|
|
|
|
let mut updated = None;
|
|
let mut tunnel_metrics_sample = None;
|
|
if has_proxy_metadata_update {
|
|
for _ in 0..8 {
|
|
let Some(current) = self.find_proxy_node(&mutation.node_id).await? else {
|
|
return Ok(None);
|
|
};
|
|
if current.tunnel_generation != tunnel_generation {
|
|
return Ok(None);
|
|
}
|
|
let Some(replacement) = normalize_heartbeat_proxy_metadata(
|
|
current.proxy_metadata.as_ref(),
|
|
mutation.proxy_metadata.as_ref(),
|
|
mutation.proxy_version.as_deref(),
|
|
) else {
|
|
break;
|
|
};
|
|
if current.proxy_metadata.as_ref() == Some(&replacement) {
|
|
tunnel_metrics_sample = build_tunnel_metrics_sample(
|
|
current.proxy_metadata.as_ref(),
|
|
Some(&replacement),
|
|
current.active_connections,
|
|
current.tunnel_connected,
|
|
);
|
|
updated = Some(current);
|
|
break;
|
|
}
|
|
|
|
let result = sqlx::query(CAS_HEARTBEAT_PROXY_METADATA_SQL)
|
|
.bind(&mutation.node_id)
|
|
.bind(&replacement)
|
|
.bind(&tunnel_generation)
|
|
.bind(current.proxy_metadata.as_ref())
|
|
.execute(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
if result.rows_affected() == 0 {
|
|
continue;
|
|
}
|
|
let Some(after_cas) = self.find_proxy_node(&mutation.node_id).await? else {
|
|
return Ok(None);
|
|
};
|
|
if after_cas.tunnel_generation != tunnel_generation {
|
|
return Ok(None);
|
|
}
|
|
tunnel_metrics_sample = build_tunnel_metrics_sample(
|
|
current.proxy_metadata.as_ref(),
|
|
after_cas.proxy_metadata.as_ref(),
|
|
after_cas.active_connections,
|
|
after_cas.tunnel_connected,
|
|
);
|
|
updated = Some(after_cas);
|
|
break;
|
|
}
|
|
}
|
|
let updated = if let Some(updated) = updated {
|
|
updated
|
|
} else {
|
|
let Some(current) = self.find_proxy_node(&mutation.node_id).await? else {
|
|
return Ok(None);
|
|
};
|
|
if current.tunnel_generation != tunnel_generation {
|
|
return Ok(None);
|
|
}
|
|
current
|
|
};
|
|
let now_unix_secs = updated
|
|
.last_heartbeat_at_unix_secs
|
|
.unwrap_or_else(|| chrono::Utc::now().timestamp().max(0) as u64);
|
|
|
|
if let Some(sample) = tunnel_metrics_sample.as_ref() {
|
|
self.upsert_metrics_bucket(
|
|
ProxyNodeMetricsStep::OneMinute,
|
|
&updated.id,
|
|
Some(tunnel_generation.as_str()),
|
|
bucket_start_unix_secs(now_unix_secs, ProxyNodeMetricsStep::OneMinute),
|
|
sample,
|
|
)
|
|
.await?;
|
|
self.upsert_metrics_bucket(
|
|
ProxyNodeMetricsStep::OneHour,
|
|
&updated.id,
|
|
Some(tunnel_generation.as_str()),
|
|
bucket_start_unix_secs(now_unix_secs, ProxyNodeMetricsStep::OneHour),
|
|
sample,
|
|
)
|
|
.await?;
|
|
|
|
for error in &sample.recent_error_events {
|
|
log_reported_tunnel_error_event(&updated.id, error, now_unix_secs);
|
|
let detail = build_tunnel_error_event_detail(error);
|
|
let event_metadata = serde_json::json!({
|
|
"source": "heartbeat",
|
|
"category": error.category,
|
|
"message": error.message,
|
|
"severity": error.severity.as_deref(),
|
|
"component": error.component.as_deref(),
|
|
"summary": error.summary.as_deref(),
|
|
"operator_action": error.operator_action.as_deref(),
|
|
"timestamp_unix_secs": error.timestamp_unix_secs,
|
|
"timestamp_unix_ms": error.timestamp_unix_ms,
|
|
});
|
|
self.insert_event(
|
|
&updated.id,
|
|
Some(tunnel_generation.as_str()),
|
|
PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR,
|
|
Some(detail.as_str()),
|
|
Some(&event_metadata),
|
|
Some(if error.timestamp_unix_secs == 0 {
|
|
now_unix_secs
|
|
} else {
|
|
error.timestamp_unix_secs
|
|
}),
|
|
)
|
|
.await?;
|
|
}
|
|
}
|
|
|
|
if reconcile_remote_config_after_heartbeat(
|
|
updated.remote_config.as_ref(),
|
|
mutation.proxy_version.as_deref(),
|
|
) != updated.remote_config
|
|
{
|
|
return self
|
|
.update_remote_config(&ProxyNodeRemoteConfigMutation {
|
|
node_id: mutation.node_id.clone(),
|
|
expected_tunnel_generation: Some(tunnel_generation),
|
|
node_name: None,
|
|
allowed_ports: None,
|
|
log_level: None,
|
|
heartbeat_interval: None,
|
|
scheduling_state: None,
|
|
upgrade_to: Some(None),
|
|
})
|
|
.await;
|
|
}
|
|
|
|
Ok(Some(updated))
|
|
}
|
|
|
|
async fn record_traffic(
|
|
&self,
|
|
mutation: &ProxyNodeTrafficMutation,
|
|
) -> Result<bool, DataLayerError> {
|
|
let mut tx = self.pool.begin().await.map_postgres_err()?;
|
|
let row = sqlx::query(&format!("{FIND_PROXY_NODE_SQL} FOR UPDATE"))
|
|
.bind(&mutation.node_id)
|
|
.fetch_optional(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
let Some(row) = row else {
|
|
tx.rollback().await.map_postgres_err()?;
|
|
return Ok(false);
|
|
};
|
|
let generation = Self::row_to_stored(&row)?.tunnel_generation;
|
|
let Some(expected_generation) = mutation.expected_tunnel_generation.as_deref() else {
|
|
tx.rollback().await.map_postgres_err()?;
|
|
return Ok(false);
|
|
};
|
|
if expected_generation != generation {
|
|
tx.rollback().await.map_postgres_err()?;
|
|
return Ok(false);
|
|
}
|
|
let result = sqlx::query(RECORD_PROXY_NODE_TRAFFIC_SQL)
|
|
.bind(&mutation.node_id)
|
|
.bind(mutation.total_requests_delta)
|
|
.bind(mutation.failed_requests_delta)
|
|
.bind(mutation.dns_failures_delta)
|
|
.bind(mutation.stream_errors_delta)
|
|
.bind(expected_generation)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
let applied = result.rows_affected() > 0;
|
|
tx.commit().await.map_err(postgres_error)?;
|
|
Ok(applied)
|
|
}
|
|
|
|
async fn update_tunnel_status(
|
|
&self,
|
|
mutation: &ProxyNodeTunnelStatusMutation,
|
|
) -> Result<Option<StoredProxyNode>, DataLayerError> {
|
|
let existing = self.find_proxy_node(&mutation.node_id).await?;
|
|
let Some(existing) = existing else {
|
|
return Ok(None);
|
|
};
|
|
if mutation
|
|
.expected_tunnel_generation
|
|
.as_deref()
|
|
.is_some_and(|expected| expected != existing.tunnel_generation)
|
|
{
|
|
return Ok(None);
|
|
}
|
|
let tunnel_generation = existing.tunnel_generation.clone();
|
|
|
|
let observed_at_unix_secs = mutation.observed_at_unix_secs;
|
|
let event_type = if mutation.connected {
|
|
"connected"
|
|
} else {
|
|
"disconnected"
|
|
};
|
|
let event_detail = mutation.detail.clone().unwrap_or_else(|| {
|
|
format!(
|
|
"[tunnel_node_status] conn_count={}",
|
|
i32::max(mutation.conn_count, 0)
|
|
)
|
|
});
|
|
|
|
let mut tx = self.pool.begin().await.map_postgres_err()?;
|
|
let result = sqlx::query(UPDATE_TUNNEL_STATUS_SQL)
|
|
.bind(&mutation.node_id)
|
|
.bind(mutation.connected)
|
|
.bind(observed_at_unix_secs.map(|value| value as f64))
|
|
.bind(&tunnel_generation)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
let stale = result.rows_affected() == 0;
|
|
let persisted_detail = stale
|
|
.then(|| format!("[stale_ignored] {event_detail}"))
|
|
.unwrap_or(event_detail);
|
|
|
|
sqlx::query(INSERT_PROXY_NODE_EVENT_SQL)
|
|
.bind(&mutation.node_id)
|
|
.bind(event_type)
|
|
.bind(persisted_detail)
|
|
.bind(None::<serde_json::Value>)
|
|
.bind(
|
|
(!stale)
|
|
.then_some(observed_at_unix_secs)
|
|
.flatten()
|
|
.map(|value| value as f64),
|
|
)
|
|
.bind(&tunnel_generation)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
|
|
tx.commit().await.map_err(postgres_error)?;
|
|
let current = self.find_proxy_node(&mutation.node_id).await?;
|
|
Ok(current.filter(|node| node.tunnel_generation == tunnel_generation))
|
|
}
|
|
|
|
async fn unregister_node(
|
|
&self,
|
|
node_id: &str,
|
|
) -> Result<Option<StoredProxyNode>, DataLayerError> {
|
|
let mut tx = self.pool.begin().await.map_postgres_err()?;
|
|
let row = sqlx::query(&format!("{FIND_PROXY_NODE_SQL} FOR UPDATE"))
|
|
.bind(node_id)
|
|
.fetch_optional(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
let Some(row) = row else {
|
|
tx.rollback().await.map_postgres_err()?;
|
|
return Ok(None);
|
|
};
|
|
let generation = Self::row_to_stored(&row)?.tunnel_generation;
|
|
sqlx::query(UNREGISTER_PROXY_NODE_SQL)
|
|
.bind(node_id)
|
|
.bind(&generation)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
let updated_row = sqlx::query(&format!("{FIND_PROXY_NODE_SQL} FOR UPDATE"))
|
|
.bind(node_id)
|
|
.fetch_one(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
let updated = Self::row_to_stored(&updated_row)?;
|
|
tx.commit().await.map_err(postgres_error)?;
|
|
Ok(Some(updated))
|
|
}
|
|
|
|
async fn delete_node(&self, node_id: &str) -> Result<Option<StoredProxyNode>, DataLayerError> {
|
|
let mut tx = self.pool.begin().await.map_postgres_err()?;
|
|
let row = sqlx::query(&format!("{FIND_PROXY_NODE_SQL} FOR UPDATE"))
|
|
.bind(node_id)
|
|
.fetch_optional(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
let Some(row) = row else {
|
|
tx.rollback().await.map_postgres_err()?;
|
|
return Ok(None);
|
|
};
|
|
// Reuse the canonical projection so the value returned to callers is
|
|
// identical to find_proxy_node while the row remains locked.
|
|
let existing = Self::row_to_stored(&row)?;
|
|
let generation = existing.tunnel_generation.as_str();
|
|
|
|
// Keep child cleanup in this transaction and bind it to the locked
|
|
// generation. The FK cascade is still a final backstop for deployments
|
|
// that have the PostgreSQL constraints enabled.
|
|
sqlx::query(DELETE_PROXY_NODE_EVENTS_SQL)
|
|
.bind(node_id)
|
|
.bind(generation)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
sqlx::query(DELETE_PROXY_NODE_METRICS_1M_SQL)
|
|
.bind(node_id)
|
|
.bind(generation)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
sqlx::query(DELETE_PROXY_NODE_METRICS_1H_SQL)
|
|
.bind(node_id)
|
|
.bind(generation)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
|
|
let deleted = sqlx::query(DELETE_PROXY_NODE_SQL)
|
|
.bind(node_id)
|
|
.bind(generation)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
if deleted.rows_affected() != 1 {
|
|
tx.rollback().await.map_postgres_err()?;
|
|
return Ok(None);
|
|
}
|
|
tx.commit().await.map_err(postgres_error)?;
|
|
if let Err(error) = sqlx::query(RETIRE_PROXY_NODE_PENDING_COUNTERS_SQL)
|
|
.bind(node_id)
|
|
.bind(generation)
|
|
.execute(&self.pool)
|
|
.await
|
|
.map_postgres_err()
|
|
{
|
|
tracing::warn!(
|
|
node_id = %node_id,
|
|
tunnel_generation = %generation,
|
|
error = ?error,
|
|
"failed to retire deleted proxy node counter rows"
|
|
);
|
|
}
|
|
Ok(Some(existing))
|
|
}
|
|
|
|
async fn update_remote_config(
|
|
&self,
|
|
mutation: &ProxyNodeRemoteConfigMutation,
|
|
) -> Result<Option<StoredProxyNode>, DataLayerError> {
|
|
for _ in 0..8 {
|
|
let Some(existing) = self.find_proxy_node(&mutation.node_id).await? else {
|
|
return Ok(None);
|
|
};
|
|
if mutation
|
|
.expected_tunnel_generation
|
|
.as_deref()
|
|
.is_some_and(|expected| expected != existing.tunnel_generation)
|
|
{
|
|
return Ok(None);
|
|
}
|
|
if existing.is_manual {
|
|
return Err(DataLayerError::InvalidInput(
|
|
"手动节点不支持远程配置下发".to_string(),
|
|
));
|
|
}
|
|
|
|
let tunnel_generation = existing.tunnel_generation.clone();
|
|
let remote_config =
|
|
Self::normalize_remote_config(mutation, existing.remote_config.as_ref());
|
|
let result = sqlx::query(UPDATE_PROXY_NODE_REMOTE_CONFIG_SQL)
|
|
.bind(&mutation.node_id)
|
|
.bind(mutation.node_name.as_deref())
|
|
.bind(remote_config.as_ref())
|
|
.bind(&tunnel_generation)
|
|
.bind(existing.config_version)
|
|
.execute(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
if result.rows_affected() == 0 {
|
|
continue;
|
|
}
|
|
|
|
let current = self.find_proxy_node(&mutation.node_id).await?;
|
|
return Ok(current.filter(|node| node.tunnel_generation == tunnel_generation));
|
|
}
|
|
|
|
Err(DataLayerError::UnexpectedValue(
|
|
"proxy node remote config changed during every CAS retry".to_string(),
|
|
))
|
|
}
|
|
|
|
async fn increment_manual_node_requests(
|
|
&self,
|
|
node_id: &str,
|
|
total_delta: i64,
|
|
failed_delta: i64,
|
|
latency_ms: Option<i64>,
|
|
) -> Result<(), DataLayerError> {
|
|
let mut tx = self.pool.begin().await.map_postgres_err()?;
|
|
let row = sqlx::query(&format!("{FIND_PROXY_NODE_SQL} FOR UPDATE"))
|
|
.bind(node_id)
|
|
.fetch_optional(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
let Some(row) = row else {
|
|
tx.rollback().await.map_postgres_err()?;
|
|
return Ok(());
|
|
};
|
|
let generation = Self::row_to_stored(&row)?.tunnel_generation;
|
|
sqlx::query(INCREMENT_MANUAL_PROXY_NODE_REQUESTS_SQL)
|
|
.bind(total_delta)
|
|
.bind(failed_delta)
|
|
.bind(latency_ms)
|
|
.bind(node_id)
|
|
.bind(&generation)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
tx.commit().await.map_err(postgres_error)?;
|
|
Ok(())
|
|
}
|
|
|
|
async fn cleanup_proxy_node_metrics(
|
|
&self,
|
|
retain_1m_from_unix_secs: u64,
|
|
retain_1h_from_unix_secs: u64,
|
|
delete_limit: usize,
|
|
) -> Result<ProxyNodeMetricsCleanupSummary, DataLayerError> {
|
|
let delete_limit_i64 = i64::try_from(delete_limit.max(1)).unwrap_or(i64::MAX);
|
|
let deleted_1m = sqlx::query(
|
|
r#"
|
|
WITH expired AS (
|
|
SELECT node_id, bucket_start_unix_secs
|
|
FROM proxy_node_metrics_1m
|
|
WHERE bucket_start_unix_secs < $1
|
|
ORDER BY bucket_start_unix_secs ASC
|
|
LIMIT $2
|
|
)
|
|
DELETE FROM proxy_node_metrics_1m metrics
|
|
USING expired
|
|
WHERE metrics.node_id = expired.node_id
|
|
AND metrics.bucket_start_unix_secs = expired.bucket_start_unix_secs
|
|
"#,
|
|
)
|
|
.bind(i64::try_from(retain_1m_from_unix_secs).unwrap_or(i64::MAX))
|
|
.bind(delete_limit_i64)
|
|
.execute(&self.pool)
|
|
.await
|
|
.map_postgres_err()?
|
|
.rows_affected() as usize;
|
|
|
|
let deleted_1h = sqlx::query(
|
|
r#"
|
|
WITH expired AS (
|
|
SELECT node_id, bucket_start_unix_secs
|
|
FROM proxy_node_metrics_1h
|
|
WHERE bucket_start_unix_secs < $1
|
|
ORDER BY bucket_start_unix_secs ASC
|
|
LIMIT $2
|
|
)
|
|
DELETE FROM proxy_node_metrics_1h metrics
|
|
USING expired
|
|
WHERE metrics.node_id = expired.node_id
|
|
AND metrics.bucket_start_unix_secs = expired.bucket_start_unix_secs
|
|
"#,
|
|
)
|
|
.bind(i64::try_from(retain_1h_from_unix_secs).unwrap_or(i64::MAX))
|
|
.bind(delete_limit_i64)
|
|
.execute(&self.pool)
|
|
.await
|
|
.map_postgres_err()?
|
|
.rows_affected() as usize;
|
|
|
|
Ok(ProxyNodeMetricsCleanupSummary {
|
|
deleted_1m_rows: deleted_1m,
|
|
deleted_1h_rows: deleted_1h,
|
|
})
|
|
}
|
|
}
|
|
|
|
fn requested_proxy_node_id(value: Option<&str>) -> Result<Option<String>, DataLayerError> {
|
|
let Some(value) = value else {
|
|
return Ok(None);
|
|
};
|
|
if value.is_empty() || value.trim() != value {
|
|
return Err(DataLayerError::InvalidInput(
|
|
"proxy node id must be non-empty and unpadded".to_string(),
|
|
));
|
|
}
|
|
Ok(Some(value.to_string()))
|
|
}
|
|
|
|
async fn proxy_node_id_owner_locked(
|
|
tx: &mut sqlx::Transaction<'_, Postgres>,
|
|
node_id: &str,
|
|
) -> Result<Option<(String, i32)>, DataLayerError> {
|
|
let row = sqlx::query("SELECT ip, port FROM proxy_nodes WHERE id = $1 FOR UPDATE")
|
|
.bind(node_id)
|
|
.fetch_optional(&mut **tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
let Some(row) = row else {
|
|
return Ok(None);
|
|
};
|
|
Ok(Some((
|
|
row.try_get::<String, _>("ip").map_postgres_err()?,
|
|
row.try_get::<i32, _>("port").map_postgres_err()?,
|
|
)))
|
|
}
|
|
|
|
fn proxy_node_registration_identity_error(requested_id: &str, existing_id: &str) -> DataLayerError {
|
|
DataLayerError::InvalidInput(format!(
|
|
"proxy node registration identity changed: requested {requested_id}, existing {existing_id}"
|
|
))
|
|
}
|
|
|
|
fn proxy_node_id_in_use_error(node_id: &str, ip: &str, port: i32) -> DataLayerError {
|
|
DataLayerError::InvalidInput(format!(
|
|
"proxy node id is already in use: {node_id} ({ip}:{port})"
|
|
))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
#[test]
|
|
fn proxy_node_sql_uses_json_casts_for_json_columns() {
|
|
assert!(!super::APPLY_HEARTBEAT_SQL.contains("proxy_metadata"));
|
|
assert!(super::CAS_HEARTBEAT_PROXY_METADATA_SQL.contains("proxy_metadata = $2::json"));
|
|
assert!(super::CAS_HEARTBEAT_PROXY_METADATA_SQL
|
|
.contains("proxy_metadata::jsonb IS NOT DISTINCT FROM $4::jsonb"));
|
|
assert!(super::INSERT_PROXY_NODE_SQL
|
|
.contains("\n $11::json,\n $12,\n $13,\n FALSE,\n $14::json,\n $15\n"));
|
|
assert!(super::UPDATE_PROXY_NODE_REGISTRATION_SQL
|
|
.contains("hardware_info = COALESCE($11::json, hardware_info)"));
|
|
assert!(super::UPDATE_PROXY_NODE_REGISTRATION_SQL
|
|
.contains("proxy_metadata = COALESCE($14::json, proxy_metadata)"));
|
|
assert!(!super::UPDATE_PROXY_NODE_REGISTRATION_SQL.contains("tunnel_generation"));
|
|
assert!(super::UPDATE_PROXY_NODE_REMOTE_CONFIG_SQL.contains("remote_config = $3::json"));
|
|
}
|
|
|
|
#[test]
|
|
fn proxy_node_sql_does_not_use_jsonb_casts() {
|
|
assert!(!super::APPLY_HEARTBEAT_SQL.contains("::jsonb"));
|
|
assert!(!super::INSERT_PROXY_NODE_SQL.contains("::jsonb"));
|
|
assert!(!super::UPDATE_PROXY_NODE_REGISTRATION_SQL.contains("::jsonb"));
|
|
assert!(!super::UPDATE_PROXY_NODE_REMOTE_CONFIG_SQL.contains("::jsonb"));
|
|
}
|
|
|
|
#[test]
|
|
fn tunnel_status_sql_rejects_out_of_order_transitions_atomically() {
|
|
assert!(super::UPDATE_TUNNEL_STATUS_SQL.contains("tunnel_connected_at <= CASE"));
|
|
assert!(super::UPDATE_TUNNEL_STATUS_SQL.contains("tunnel_generation = $4"));
|
|
assert!(super::APPLY_HEARTBEAT_SQL.contains("total_requests = total_requests + GREATEST"));
|
|
assert!(super::APPLY_HEARTBEAT_SQL.contains("tunnel_generation = $9"));
|
|
assert!(!super::APPLY_HEARTBEAT_SQL.contains("remote_config ="));
|
|
assert!(!super::APPLY_HEARTBEAT_SQL.contains("config_version ="));
|
|
assert!(super::UPDATE_PROXY_NODE_REMOTE_CONFIG_SQL.contains("config_version = $5"));
|
|
}
|
|
}
|