diff --git a/Cargo.lock b/Cargo.lock
index 1f659a2cc..5d563353b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2011,6 +2011,7 @@ dependencies = [
"axum",
"base64",
"fluxer_common",
+ "futures-util",
"hex",
"rand 0.10.1",
"reqwest",
@@ -2039,6 +2040,7 @@ dependencies = [
"reqwest",
"serde_json",
"sha2 0.11.0",
+ "tempfile",
"thiserror",
"time",
"tracing",
diff --git a/fluxer_app_proxy/Cargo.toml b/fluxer_app_proxy/Cargo.toml
index 76a9b6213..fb0a7ca1e 100644
--- a/fluxer_app_proxy/Cargo.toml
+++ b/fluxer_app_proxy/Cargo.toml
@@ -10,6 +10,7 @@ anyhow = "1.0.104"
axum = { version = "0.8.9", features = ["macros"] }
base64 = "0.22"
fluxer_common = { path = "../fluxer_common" }
+futures-util = { version = "0.3.32", default-features = false, features = ["std"] }
hex = "0.4"
rand = "0.10"
reqwest = { version = "0.13.4", default-features = false, features = ["json", "rustls", "stream", "gzip", "brotli", "deflate"] }
diff --git a/fluxer_app_proxy/src/bootstrap.rs b/fluxer_app_proxy/src/bootstrap.rs
index 13cc60107..7f2a5d8ce 100644
--- a/fluxer_app_proxy/src/bootstrap.rs
+++ b/fluxer_app_proxy/src/bootstrap.rs
@@ -132,7 +132,7 @@ pub fn inject_bootstrap(
let media = media_endpoint.trim_end_matches('/');
let nonced = html.replace("{{CSP_NONCE_PLACEHOLDER}}", nonce);
- let nonced = apply_static_preconnect(&nonced, static_cdn);
+ let nonced = apply_static_preconnect(nonced, static_cdn);
let nonced = nonced.replace("{{STATIC_CDN_ENDPOINT}}", static_cdn);
let nonced = apply_media_preconnect(&nonced, media, static_cdn);
@@ -143,20 +143,14 @@ pub fn inject_bootstrap(
return nonced.replace("{{FLUXER_BOOTSTRAP}}", script_tag);
}
- if let Some(pos) = nonced.find("
") {
- let insert_at = pos + "".len();
- let mut result = String::with_capacity(nonced.len() + script_tag.len() + 3);
- result.push_str(&nonced[..insert_at]);
- result.push_str("\n\t\t");
- result.push_str(script_tag);
- result.push_str(&nonced[insert_at..]);
- return result;
- }
-
- if let Some(pos) = nonced.find("')
- {
- let insert_at = pos + close + 1;
+ let insert_at = nonced
+ .find("")
+ .map(|pos| pos + "".len())
+ .or_else(|| {
+ let pos = nonced.find("').map(|close| pos + close + 1)
+ });
+ if let Some(insert_at) = insert_at {
let mut result = String::with_capacity(nonced.len() + script_tag.len() + 3);
result.push_str(&nonced[..insert_at]);
result.push_str("\n\t\t");
@@ -168,15 +162,14 @@ pub fn inject_bootstrap(
nonced
}
-fn apply_static_preconnect(html: &str, static_cdn: &str) -> String {
+fn apply_static_preconnect(mut html: String, static_cdn: &str) -> String {
if !static_cdn.is_empty() {
- return html.to_owned();
+ return html;
}
- let mut stripped = html.to_owned();
for tag in STATIC_PRECONNECT_TAGS {
- stripped = stripped.replace(&format!("{tag}\n"), "").replace(tag, "");
+ html = html.replace(&format!("{tag}\n"), "").replace(tag, "");
}
- stripped
+ html
}
fn apply_media_preconnect(html: &str, media: &str, static_cdn: &str) -> String {
diff --git a/fluxer_app_proxy/src/config.rs b/fluxer_app_proxy/src/config.rs
index cdb215743..fc4de7430 100644
--- a/fluxer_app_proxy/src/config.rs
+++ b/fluxer_app_proxy/src/config.rs
@@ -309,30 +309,20 @@ impl fmt::Display for CspReportUri {
}
}
-fn warn_invalid(error: InvalidAppProxyEnvironmentError) {
+fn warn_invalid(error: &InvalidAppProxyEnvironmentError) {
tracing::warn!(%error, "ignoring invalid app proxy environment value");
}
fn parse_optional_http_url(name: &'static str, value: Option) -> Option {
let value = value?;
- match HttpUrl::parse(name, &value) {
- Ok(url) => Some(url),
- Err(error) => {
- warn_invalid(error);
- None
- }
- }
+ HttpUrl::parse(name, &value).inspect_err(warn_invalid).ok()
}
fn parse_optional_http_endpoint(name: &'static str, value: Option) -> Option {
let value = value?;
- match HttpEndpoint::parse(name, &value) {
- Ok(endpoint) => Some(endpoint),
- Err(error) => {
- warn_invalid(error);
- None
- }
- }
+ HttpEndpoint::parse(name, &value)
+ .inspect_err(warn_invalid)
+ .ok()
}
fn parse_env_or_warn(name: &str, raw: &str, default: T) -> T {
@@ -434,25 +424,19 @@ fn read_csp_sources(name: &'static str) -> Vec {
.split([',', ' ', '\t', '\n'])
.map(str::trim)
.filter(|source| !source.is_empty())
- .filter_map(|source| match CspSource::parse(name, source) {
- Ok(source) => Some(source),
- Err(error) => {
- warn_invalid(error);
- None
- }
+ .filter_map(|source| {
+ CspSource::parse(name, source)
+ .inspect_err(warn_invalid)
+ .ok()
})
.collect()
}
fn read_csp_report_uri(name: &'static str) -> Option {
let value = cfg::non_empty_env(name)?;
- match CspReportUri::parse(name, &value) {
- Ok(report_uri) => Some(report_uri),
- Err(error) => {
- warn_invalid(error);
- None
- }
- }
+ CspReportUri::parse(name, &value)
+ .inspect_err(warn_invalid)
+ .ok()
}
impl AppProxyConfig {
@@ -474,13 +458,10 @@ impl AppProxyConfig {
);
let s3_uploads_bucket = cfg::read_env("FLUXER_S3_BUCKET_UPLOADS", "fluxer-uploads");
let s3_uploads_endpoint = s3_public_endpoint.as_ref().and_then(|endpoint| {
- match endpoint.with_host_prefix("FLUXER_S3_BUCKET_UPLOADS", s3_uploads_bucket.trim()) {
- Ok(endpoint) => Some(endpoint),
- Err(error) => {
- warn_invalid(error);
- None
- }
- }
+ endpoint
+ .with_host_prefix("FLUXER_S3_BUCKET_UPLOADS", s3_uploads_bucket.trim())
+ .inspect_err(warn_invalid)
+ .ok()
});
Self {
@@ -499,7 +480,7 @@ impl AppProxyConfig {
"FLUXER_STATIC_CDN_ENDPOINT",
cfg::non_empty_env("FLUXER_STATIC_CDN_ENDPOINT"),
),
- s3_public_endpoint: s3_public_endpoint.clone(),
+ s3_public_endpoint,
s3_uploads_endpoint,
discovery_upstream_url: resolve_discovery_upstream_url_from_env(),
discovery_refresh_interval_ms: parse_env_or_warn(
diff --git a/fluxer_app_proxy/src/csp.rs b/fluxer_app_proxy/src/csp.rs
index 9d0c6d7ba..b40e316a5 100644
--- a/fluxer_app_proxy/src/csp.rs
+++ b/fluxer_app_proxy/src/csp.rs
@@ -287,15 +287,11 @@ fn extend_runtime_s3_sources(target: &mut Vec, runtime_sources: &Runtime
}
fn extend_from(target: &mut Vec, extra: &[CspSource], defaults: &[&str]) {
- for source in defaults {
- if target.iter().any(|existing| existing == source) {
- continue;
- }
- target.push((*source).to_owned());
- }
-
- for source in extra {
- let source = source.as_str();
+ for source in defaults
+ .iter()
+ .copied()
+ .chain(extra.iter().map(CspSource::as_str))
+ {
if target.iter().any(|existing| existing == source) {
continue;
}
diff --git a/fluxer_app_proxy/src/routes/android_association.rs b/fluxer_app_proxy/src/routes/android_association.rs
index 33733846e..0da75daf1 100644
--- a/fluxer_app_proxy/src/routes/android_association.rs
+++ b/fluxer_app_proxy/src/routes/android_association.rs
@@ -2,47 +2,33 @@
use axum::{
Json,
- http::{HeaderValue, header},
+ http::header,
response::{IntoResponse, Response},
};
use serde_json::json;
pub async fn assetlinks() -> Response {
- let body = json!([
- {
+ let body = ["com.fluxer", "com.fluxer.canary"].map(|package_name| {
+ json!({
"relation": [
"delegate_permission/common.handle_all_urls",
"delegate_permission/common.get_login_creds"
],
"target": {
"namespace": "android_app",
- "package_name": "com.fluxer",
+ "package_name": package_name,
"sha256_cert_fingerprints": [
"91:E4:98:E1:B8:A6:C8:BA:99:41:5E:DB:29:78:29:6B:6C:58:BA:A5:E2:D2:A6:49:CE:C6:2D:A7:A8:29:C7:BC"
]
}
- },
- {
- "relation": [
- "delegate_permission/common.handle_all_urls",
- "delegate_permission/common.get_login_creds"
- ],
- "target": {
- "namespace": "android_app",
- "package_name": "com.fluxer.canary",
- "sha256_cert_fingerprints": [
- "91:E4:98:E1:B8:A6:C8:BA:99:41:5E:DB:29:78:29:6B:6C:58:BA:A5:E2:D2:A6:49:CE:C6:2D:A7:A8:29:C7:BC"
- ]
- }
- }
- ]);
+ })
+ });
- let mut response = Json(body).into_response();
- response.headers_mut().insert(
- header::CACHE_CONTROL,
- HeaderValue::from_static("public, max-age=1800"),
- );
- response
+ (
+ [(header::CACHE_CONTROL, "public, max-age=1800")],
+ Json(body),
+ )
+ .into_response()
}
#[cfg(test)]
diff --git a/fluxer_app_proxy/src/routes/apple_association.rs b/fluxer_app_proxy/src/routes/apple_association.rs
index b38c6c7cf..8645c5623 100644
--- a/fluxer_app_proxy/src/routes/apple_association.rs
+++ b/fluxer_app_proxy/src/routes/apple_association.rs
@@ -2,7 +2,7 @@
use axum::{
Json,
- http::{HeaderValue, header},
+ http::header,
response::{IntoResponse, Response},
};
use serde_json::json;
@@ -19,10 +19,9 @@ pub async fn apple_app_site_association() -> Response {
}
});
- let mut response = Json(body).into_response();
- response.headers_mut().insert(
- header::CACHE_CONTROL,
- HeaderValue::from_static("public, max-age=1800"),
- );
- response
+ (
+ [(header::CACHE_CONTROL, "public, max-age=1800")],
+ Json(body),
+ )
+ .into_response()
}
diff --git a/fluxer_app_proxy/src/routes/assets_proxy.rs b/fluxer_app_proxy/src/routes/assets_proxy.rs
index 38c583d14..d572f2b6d 100644
--- a/fluxer_app_proxy/src/routes/assets_proxy.rs
+++ b/fluxer_app_proxy/src/routes/assets_proxy.rs
@@ -9,9 +9,7 @@ use axum::{
};
use std::path::{Path as FsPath, PathBuf};
use std::time::Duration;
-use tokio::io::{AsyncWriteExt, DuplexStream};
use tokio::sync::{OwnedSemaphorePermit, TryAcquireError};
-use tokio_util::io::ReaderStream;
use super::file_stream::stream_file;
use super::spa_static::{CORS_ALLOW_ANY_VALUE, asset_cache_control, guess_mime, is_font_mime};
@@ -19,7 +17,6 @@ use super::spa_static::{CORS_ALLOW_ANY_VALUE, asset_cache_control, guess_mime, i
const ASSET_REQUEST_TIMEOUT: Duration = Duration::from_secs(15);
const PRECOMPRESSED_VARIANTS: &[(&str, &str)] = &[("br", "br"), ("gzip", "gz")];
const MAX_ASSET_SIZE_BYTES: u64 = 100 * 1024 * 1024;
-const UPSTREAM_ASSET_PUMP_BUFFER_BYTES: usize = 64 * 1024;
const UPSTREAM_FAILURE_CACHE_CONTROL: &str = "no-store";
const UPSTREAM_FAILURE_STRIPPED_HEADERS: &[&str] = &[
"cdn-cache-control",
@@ -148,38 +145,59 @@ pub async fn proxy_assets(
response_headers.insert(header::CONTENT_SECURITY_POLICY, state.csp.asset_header());
response_headers.remove("content-security-policy-report-only");
- let body = Body::from_stream(upstream_asset_body(upstream_response, upstream_slot));
+ let body = upstream_asset_body(upstream_response, upstream_slot);
let mut response = Response::new(body);
*response.status_mut() = status;
*response.headers_mut() = response_headers;
response
}
+struct UpstreamAssetReadState {
+ response: reqwest::Response,
+ _permit: OwnedSemaphorePermit,
+ remaining_bytes: u64,
+}
+
fn upstream_asset_body(
- mut upstream_response: reqwest::Response,
+ upstream_response: reqwest::Response,
upstream_slot: OwnedSemaphorePermit,
-) -> ReaderStream {
- let (writer, reader) = tokio::io::duplex(UPSTREAM_ASSET_PUMP_BUFFER_BYTES);
- tokio::spawn(async move {
- let _upstream_slot = upstream_slot;
- let mut writer = writer;
- loop {
- match upstream_response.chunk().await {
- Ok(Some(chunk)) => {
- if writer.write_all(&chunk).await.is_err() {
- return;
- }
- }
- Ok(None) => break,
- Err(err) => {
+) -> Body {
+ let state = UpstreamAssetReadState {
+ response: upstream_response,
+ _permit: upstream_slot,
+ remaining_bytes: MAX_ASSET_SIZE_BYTES,
+ };
+ Body::from_stream(futures_util::stream::try_unfold(
+ state,
+ |mut state| async move {
+ let Some(chunk) = state
+ .response
+ .chunk()
+ .await
+ .inspect_err(|err| {
tracing::warn!(%err, "upstream asset body ended early");
- return;
- }
+ })
+ .map_err(axum::Error::new)?
+ else {
+ return Ok(None);
+ };
+ let chunk_bytes = chunk.len() as u64;
+ if chunk_bytes > state.remaining_bytes {
+ tracing::warn!(
+ maximum_bytes = MAX_ASSET_SIZE_BYTES,
+ remaining_bytes = state.remaining_bytes,
+ chunk_bytes,
+ "upstream asset body exceeds size cap"
+ );
+ return Err(axum::Error::new(std::io::Error::new(
+ std::io::ErrorKind::InvalidData,
+ format!("upstream asset body exceeds {MAX_ASSET_SIZE_BYTES} bytes"),
+ )));
}
- }
- let _ = writer.shutdown().await;
- });
- ReaderStream::new(reader)
+ state.remaining_bytes -= chunk_bytes;
+ Ok(Some((chunk, state)))
+ },
+ ))
}
pub(super) async fn serve_local_asset(
diff --git a/fluxer_app_proxy/src/routes/mod.rs b/fluxer_app_proxy/src/routes/mod.rs
index 29e971eed..3fd8266f5 100644
--- a/fluxer_app_proxy/src/routes/mod.rs
+++ b/fluxer_app_proxy/src/routes/mod.rs
@@ -13,7 +13,7 @@ use axum::{
Router,
extract::Request,
http::{HeaderName, HeaderValue, header},
- middleware::{Next, from_fn, from_fn_with_state},
+ middleware::{Next, from_fn},
response::{IntoResponse, Response},
routing::get,
};
@@ -56,10 +56,7 @@ pub fn build_router(state: AppState) -> Router {
.fallback(get(spa_index::spa_catch_all))
.layer(from_fn(request_id_middleware))
.layer(from_fn(cache_headers_middleware))
- .layer(from_fn_with_state(
- state.clone(),
- security_headers_middleware,
- ))
+ .layer(from_fn(security_headers_middleware))
.layer(
CompressionLayer::new()
.compress_when(DefaultPredicate::new().and(NotForContentType::const_new("font/"))),
@@ -68,14 +65,13 @@ pub fn build_router(state: AppState) -> Router {
.with_state(state)
}
-async fn security_headers_middleware(
- axum::extract::State(_state): axum::extract::State,
- request: Request,
- next: Next,
-) -> Response {
+async fn security_headers_middleware(request: Request, next: Next) -> Response {
let mut response = next.run(request).await;
- let headers = response.headers_mut();
+ set_security_headers(response.headers_mut());
+ response
+}
+fn set_security_headers(headers: &mut axum::http::HeaderMap) {
set_static_header(
headers,
header::STRICT_TRANSPORT_SECURITY,
@@ -89,8 +85,6 @@ async fn security_headers_middleware(
HeaderName::from_static("permissions-policy"),
PERMISSIONS_POLICY_VALUE,
);
-
- response
}
async fn cache_headers_middleware(request: Request, next: Next) -> Response {
diff --git a/fluxer_app_proxy/src/routes/spa_index.rs b/fluxer_app_proxy/src/routes/spa_index.rs
index 0339e7250..48dce480e 100644
--- a/fluxer_app_proxy/src/routes/spa_index.rs
+++ b/fluxer_app_proxy/src/routes/spa_index.rs
@@ -410,19 +410,7 @@ fn build_spa_response(
} else {
headers.insert(header::CACHE_CONTROL, HeaderValue::from_static("no-cache"));
}
- headers.insert(
- header::STRICT_TRANSPORT_SECURITY,
- HeaderValue::from_static("max-age=31536000; includeSubDomains; preload"),
- );
- headers.insert(
- header::X_CONTENT_TYPE_OPTIONS,
- HeaderValue::from_static("nosniff"),
- );
- headers.insert(header::X_FRAME_OPTIONS, HeaderValue::from_static("DENY"));
- headers.insert(
- header::REFERRER_POLICY,
- HeaderValue::from_static("strict-origin-when-cross-origin"),
- );
+ super::set_security_headers(headers);
headers.insert(
axum::http::HeaderName::from_static("accept-ch"),
HeaderValue::from_static(ACCEPT_CH_VALUE),
@@ -431,11 +419,6 @@ fn build_spa_response(
axum::http::HeaderName::from_static("critical-ch"),
HeaderValue::from_static(CRITICAL_CH_VALUE),
);
- headers.insert(
- axum::http::HeaderName::from_static("permissions-policy"),
- HeaderValue::from_static(super::PERMISSIONS_POLICY_VALUE),
- );
-
#[cfg(feature = "time-freeze")]
{
if let Some(tf) = time_freeze_header
diff --git a/fluxer_common/Cargo.toml b/fluxer_common/Cargo.toml
index b9e16e48d..ead7414f7 100644
--- a/fluxer_common/Cargo.toml
+++ b/fluxer_common/Cargo.toml
@@ -13,6 +13,7 @@ maxminddb = "0.28.1"
moka = { version = "0.12.15", features = ["sync"] }
reqwest = { version = "0.13.4", default-features = false, features = ["blocking", "rustls"] }
serde_json = "1.0.150"
+tempfile = "3.27.0"
time = { version = "0.3.47", features = ["formatting", "macros", "parsing"] }
tracing = "0.1.44"
base64 = "0.22"
diff --git a/fluxer_common/src/config.rs b/fluxer_common/src/config.rs
index aefbc77ba..74fe343df 100644
--- a/fluxer_common/src/config.rs
+++ b/fluxer_common/src/config.rs
@@ -203,12 +203,12 @@ fn is_default_port(scheme: &str, port: u16) -> bool {
)
}
-fn strip_trailing_dot(host: &str) -> &str {
- host.strip_suffix('.').unwrap_or(host)
-}
-
fn canonicalize_domain(value: &str) -> String {
- strip_trailing_dot(value.trim().to_lowercase().as_str()).to_owned()
+ let mut domain = value.trim().to_lowercase();
+ if domain.ends_with('.') {
+ domain.pop();
+ }
+ domain
}
fn default_port(scheme: &str) -> u16 {
@@ -285,10 +285,10 @@ where
}
pub fn normalize_public_endpoint(url: &str, base_domain: &str, public_port: Option) -> String {
- let domain = canonicalize_domain(base_domain);
let Some(port) = public_port.filter(|port| *port != 0) else {
return url.to_owned();
};
+ let domain = canonicalize_domain(base_domain);
if domain.is_empty() {
return url.to_owned();
}
diff --git a/fluxer_common/src/geoip.rs b/fluxer_common/src/geoip.rs
index be0e21c59..5aa30ae30 100644
--- a/fluxer_common/src/geoip.rs
+++ b/fluxer_common/src/geoip.rs
@@ -19,7 +19,6 @@ use std::{
path::Path,
time::{Duration, SystemTime},
};
-use time::OffsetDateTime;
const DEFAULT_COUNTRY_CODE: &str = "US";
const GEOIP_CACHE_TTL: Duration = Duration::from_secs(10 * 60);
@@ -242,7 +241,6 @@ fn download_s3_object(
);
};
fs::create_dir_all(parent)?;
- let temp_path = temporary_download_path(destination);
let request = signed_s3_get_request(config, bucket, key)?;
let mut response = reqwest::blocking::Client::new()
.get(request.url)
@@ -253,29 +251,14 @@ fn download_s3_object(
})?;
if !response.status().is_success() {
let status = response.status();
- let _ = fs::remove_file(&temp_path);
anyhow::bail!("failed to download GeoIP database from s3://{bucket}/{key}: HTTP {status}");
}
- {
- let mut file = File::create(&temp_path)?;
- io::copy(&mut response, &mut file)?;
- }
- fs::rename(&temp_path, destination).inspect_err(|_err| {
- let _ = fs::remove_file(&temp_path);
- })?;
+ let mut file = tempfile::Builder::new().make_in(parent, |path| File::create_new(path))?;
+ io::copy(&mut response, &mut file)?;
+ file.persist(destination).map_err(|err| err.error)?;
Ok(())
}
-fn temporary_download_path(destination: &Path) -> std::path::PathBuf {
- let pid = std::process::id();
- let now = OffsetDateTime::now_utc().unix_timestamp_nanos();
- let file_name = destination
- .file_name()
- .and_then(|value| value.to_str())
- .unwrap_or("geoip.mmdb");
- destination.with_file_name(format!("{file_name}.tmp-{pid}-{now}"))
-}
-
struct SignedS3Request {
url: reqwest::Url,
headers: reqwest::header::HeaderMap,
diff --git a/fluxer_snowflakes/src/generator.rs b/fluxer_snowflakes/src/generator.rs
index ea01cafcd..b13422dc5 100644
--- a/fluxer_snowflakes/src/generator.rs
+++ b/fluxer_snowflakes/src/generator.rs
@@ -66,20 +66,15 @@ where
pub async fn generate(&mut self) -> anyhow::Result {
let mut timestamp = self.current_relative_timestamp()?;
- if let Some(last_timestamp) = self.last_timestamp {
- if timestamp < last_timestamp {
+ match self.last_timestamp {
+ Some(last_timestamp) if timestamp <= last_timestamp => {
timestamp = last_timestamp;
- }
- if timestamp == last_timestamp {
self.sequence = (self.sequence + 1) & MAX_SEQUENCE;
if self.sequence == 0 {
timestamp = self.wait_until_next_timestamp(last_timestamp).await?;
}
- } else {
- self.sequence = 0;
}
- } else {
- self.sequence = 0;
+ _ => self.sequence = 0,
}
self.last_timestamp = Some(timestamp);
Ok(create_snowflake_from_relative_timestamp(
diff --git a/fluxer_snowflakes/src/router_impl.rs b/fluxer_snowflakes/src/router_impl.rs
index dda8c995a..82c21bb6b 100644
--- a/fluxer_snowflakes/src/router_impl.rs
+++ b/fluxer_snowflakes/src/router_impl.rs
@@ -4,10 +4,10 @@ use crate::types::{SERVICE_NAME, SnowflakeRequest, SnowflakeResponse};
use fluxer_svc::config::ServiceConfig;
use fluxer_svc::hash_ring::HashRing;
use fluxer_svc::metrics::ServiceMetrics;
-use fluxer_svc::transport::{NatsTransport, TransportMessage, TransportSubscriber};
+use fluxer_svc::transport::{NatsMessage, NatsTransport, TransportMessage, TransportSubscriber};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
-use std::time::Duration;
+use std::time::{Duration, Instant};
use tokio::task::JoinSet;
use tracing::{debug, info, warn};
@@ -61,8 +61,10 @@ pub async fn run_round_robin_router(
config: &ServiceConfig,
transport: NatsTransport,
) -> anyhow::Result<()> {
- let request_subject = format!("svc.{SERVICE_NAME}");
- let queue_group = format!("{SERVICE_NAME}-router");
+ anyhow::ensure!(
+ config.max_concurrent_requests > 0,
+ "snowflake router request concurrency must be positive"
+ );
let picker = Arc::new(SnowflakeShardPicker::new(config.shard_count));
let mut tasks = JoinSet::new();
let health_addr = config.listen_addr;
@@ -82,84 +84,12 @@ pub async fn run_round_robin_router(
.await
}
});
- let req_transport = transport.clone();
- tasks.spawn(async move {
- loop {
- let mut sub = req_transport
- .subscribe_queue(&request_subject, &queue_group)
- .await?;
- info!(subject = request_subject, "snowflake router listening for requests");
- loop {
- let msg = tokio::select! {
- msg_opt = sub.next() => {
- let Some(msg) = msg_opt else {
- warn!("snowflake router request subscription stream ended, will re-subscribe");
- break;
- };
- msg
- }
- _ = req_transport.wait_for_reconnect() => {
- info!("NATS reconnected, re-subscribing snowflake router request listener");
- break;
- }
- };
- let transport = req_transport.clone();
- let picker = picker.clone();
- if msg.payload().len() > MAX_ROUTER_REQUEST_BYTES {
- warn!(
- payload_bytes = msg.payload().len(),
- max_payload_bytes = MAX_ROUTER_REQUEST_BYTES,
- "rejecting oversized snowflake request"
- );
- reply_json_error(&msg, &transport, "request_too_large").await;
- continue;
- }
- let request: SnowflakeRequest = match serde_json::from_slice(msg.payload()) {
- Ok(request) => request,
- Err(error) => {
- warn!(error = %error, "failed to decode snowflake request");
- reply_json_error(&msg, &transport, "decode_error").await;
- continue;
- }
- };
- tokio::spawn(async move {
- let shard_id = picker.pick_shard(&request);
- let shard_subject = format!("svc.{SERVICE_NAME}.shard.{shard_id}");
- let payload = match rmp_serde::to_vec(&request) {
- Ok(payload) => payload,
- Err(error) => {
- warn!(error = %error, "failed to encode snowflake shard request");
- reply_json_error(&msg, &transport, "encode_error").await;
- return;
- }
- };
- let response_bytes = match transport
- .request(&shard_subject, &payload, SHARD_REQUEST_TIMEOUT)
- .await
- {
- Ok(response_bytes) => response_bytes,
- Err(error) => {
- debug!(error = %error, shard_id, "snowflake shard request failed");
- reply_json_error(&msg, &transport, "shard_unavailable").await;
- return;
- }
- };
- let response: SnowflakeResponse = match rmp_serde::from_slice(&response_bytes) {
- Ok(response) => response,
- Err(error) => {
- debug!(error = %error, shard_id, "failed to decode snowflake shard response");
- reply_json_error(&msg, &transport, "shard_decode_error").await;
- return;
- }
- };
- if msg.has_reply() {
- let response_json = serde_json::to_vec(&response).unwrap_or_default();
- let _ = msg.reply(&transport, &response_json).await;
- }
- });
- }
- }
- });
+ tasks.spawn(serve_requests(
+ transport,
+ picker,
+ http_metrics,
+ config.max_concurrent_requests,
+ ));
tokio::select! {
result = tasks.join_next() => {
match result {
@@ -177,15 +107,136 @@ pub async fn run_round_robin_router(
}
}
-async fn reply_json_error(
- msg: &fluxer_svc::transport::NatsMessage,
- transport: &NatsTransport,
- error: &str,
+async fn serve_requests(
+ transport: NatsTransport,
+ picker: Arc,
+ metrics: Arc,
+ max_concurrent_requests: usize,
+) -> anyhow::Result<()> {
+ let request_subject = format!("svc.{SERVICE_NAME}");
+ let queue_group = format!("{SERVICE_NAME}-router");
+ let mut requests = JoinSet::new();
+ loop {
+ let mut sub = transport
+ .subscribe_queue(&request_subject, &queue_group)
+ .await?;
+ info!(
+ subject = request_subject,
+ max_concurrent_requests, "snowflake router listening for requests"
+ );
+ loop {
+ let msg = tokio::select! {
+ result = requests.join_next(), if !requests.is_empty() => {
+ if let Err(error) = result.expect("nonempty snowflake request set") {
+ warn!(error = %error, "snowflake request task failed");
+ }
+ continue;
+ }
+ msg = sub.next() => {
+ let Some(msg) = msg else {
+ warn!("snowflake router request subscription stream ended, will re-subscribe");
+ break;
+ };
+ msg
+ }
+ };
+ while let Some(result) = requests.try_join_next() {
+ if let Err(error) = result {
+ warn!(error = %error, "snowflake request task failed");
+ }
+ }
+ metrics.record_request();
+ if requests.len() >= max_concurrent_requests {
+ debug!(
+ max_concurrent_requests,
+ "shedding overloaded snowflake router request"
+ );
+ metrics.record_request_error();
+ reply_json_error(&msg, &transport, "overloaded").await;
+ continue;
+ }
+ requests.spawn(handle_request(
+ msg,
+ transport.clone(),
+ picker.clone(),
+ metrics.clone(),
+ ));
+ }
+ }
+}
+
+async fn handle_request(
+ msg: NatsMessage,
+ transport: NatsTransport,
+ picker: Arc,
+ metrics: Arc,
) {
+ let started = Instant::now();
+ match forward_request(msg.payload(), &transport, &picker, &metrics).await {
+ Ok(response) => {
+ if msg.has_reply() {
+ let payload = serde_json::to_vec(&response)
+ .expect("snowflake responses contain only JSON-serializable strings");
+ reply(&msg, &transport, &payload).await;
+ }
+ }
+ Err(error) => {
+ metrics.record_request_error();
+ reply_json_error(&msg, &transport, error).await;
+ }
+ }
+ metrics.record_request_duration(started.elapsed().as_millis() as u64);
+}
+
+async fn forward_request(
+ payload: &[u8],
+ transport: &NatsTransport,
+ picker: &SnowflakeShardPicker,
+ metrics: &ServiceMetrics,
+) -> Result {
+ if payload.len() > MAX_ROUTER_REQUEST_BYTES {
+ warn!(
+ payload_bytes = payload.len(),
+ max_payload_bytes = MAX_ROUTER_REQUEST_BYTES,
+ "rejecting oversized snowflake request"
+ );
+ return Err("request_too_large");
+ }
+ let request: SnowflakeRequest = serde_json::from_slice(payload).map_err(|error| {
+ warn!(error = %error, "failed to decode snowflake request");
+ "decode_error"
+ })?;
+ let shard_id = picker.pick_shard(&request);
+ let shard_subject = format!("svc.{SERVICE_NAME}.shard.{shard_id}");
+ let payload = rmp_serde::to_vec(&request).map_err(|error| {
+ warn!(error = %error, "failed to encode snowflake shard request");
+ "encode_error"
+ })?;
+ metrics.record_shard_forward();
+ let response_bytes = transport
+ .request(&shard_subject, &payload, SHARD_REQUEST_TIMEOUT)
+ .await
+ .map_err(|error| {
+ debug!(error = %error, shard_id, "snowflake shard request failed");
+ "shard_unavailable"
+ })?;
+ rmp_serde::from_slice(&response_bytes).map_err(|error| {
+ debug!(error = %error, shard_id, "failed to decode snowflake shard response");
+ "shard_decode_error"
+ })
+}
+
+async fn reply_json_error(msg: &NatsMessage, transport: &NatsTransport, error: &str) {
if msg.has_reply() {
- let error_response =
- serde_json::to_vec(&serde_json::json!({ "error": error })).unwrap_or_default();
- let _ = msg.reply(transport, &error_response).await;
+ let payload = serde_json::to_vec(&serde_json::json!({ "error": error }))
+ .expect("snowflake error responses contain only a JSON-serializable string");
+ reply(msg, transport, &payload).await;
+ }
+}
+
+async fn reply(msg: &NatsMessage, transport: &NatsTransport, payload: &[u8]) {
+ if let Err(error) = msg.reply(transport, payload).await {
+ warn!(error = %error, "failed to reply to snowflake request");
}
}
diff --git a/fluxer_svc/src/config.rs b/fluxer_svc/src/config.rs
index 038a18072..b2d22d003 100644
--- a/fluxer_svc/src/config.rs
+++ b/fluxer_svc/src/config.rs
@@ -75,8 +75,9 @@ impl ServiceConfig {
};
let mode = match optional_from(&get, "FLUXER_SVC_MODE").as_deref() {
+ None | Some("router") => Mode::Router,
Some("shard") => Mode::Shard,
- _ => Mode::Router,
+ Some(other) => anyhow::bail!("unsupported FLUXER_SVC_MODE: {other}"),
};
let shard_count = optional_from(&get, "FLUXER_SVC_SHARD_COUNT")
diff --git a/fluxer_svc/src/hash_ring.rs b/fluxer_svc/src/hash_ring.rs
index 276b94d08..4412e9c78 100644
--- a/fluxer_svc/src/hash_ring.rs
+++ b/fluxer_svc/src/hash_ring.rs
@@ -9,6 +9,10 @@ impl HashRing {
Self { shard_count }
}
+ pub fn shard_count(&self) -> u32 {
+ self.shard_count
+ }
+
pub fn owner(&self, route_key: &str) -> u32 {
let mut best_shard = 0u32;
let mut best_score = 0u64;
diff --git a/fluxer_svc/src/router.rs b/fluxer_svc/src/router.rs
index 82f7e3d4e..14e29651c 100644
--- a/fluxer_svc/src/router.rs
+++ b/fluxer_svc/src/router.rs
@@ -2,11 +2,15 @@
use crate::config::ServiceConfig;
use crate::hash_ring::HashRing;
-use crate::metrics::{ServiceMetrics, now_ms};
-use crate::transport::{Transport, TransportMessage, TransportSubscriber, reply_message};
+use crate::metrics::ServiceMetrics;
+use crate::transport::{
+ Transport, TransportMessage, TransportSubscriber, reply_bytes, reply_json_error,
+};
+use anyhow::Context;
+use futures::stream::{FuturesUnordered, StreamExt};
use moka::future::Cache;
use std::sync::Arc;
-use std::time::Duration;
+use std::time::{Duration, Instant};
use tokio::sync::{Semaphore, TryAcquireError};
use tokio::task::JoinSet;
use tracing::{debug, info, warn};
@@ -14,6 +18,7 @@ use tracing::{debug, info, warn};
pub(crate) const SHARD_REQUEST_TIMEOUT: Duration = Duration::from_secs(5);
const INFLIGHT_TTL: Duration = Duration::from_millis(200);
const INFLIGHT_MAX_ENTRIES: u64 = 10_000;
+const MAX_BROADCAST_CONCURRENCY: usize = 32;
const MAX_ROUTER_REQUEST_BYTES: usize = 2 * 1024 * 1024;
const LEGACY_SHARD_DECODE_ERROR: &[u8] = br#"{"error":"shard_request_decode_error"}"#;
type InflightKey = (String, String);
@@ -30,6 +35,13 @@ pub trait RouterService: Send + Sync + 'static {
None
}
+ fn is_broadcast_request(_request: &Self::Request) -> bool {
+ false
+ }
+ fn is_broadcast_acknowledgement(_response: &Self::Response) -> bool {
+ false
+ }
+
fn l1_lookup(&self, _req: &Self::Request) -> Option {
None
}
@@ -111,6 +123,88 @@ async fn dispatch_to_shard(
}
}
+async fn forward_to_all_shards(
+ transport: &impl Transport,
+ service: &S,
+ ring: &HashRing,
+ request: &S::Request,
+) -> anyhow::Result> {
+ let payload = rmp_serde::to_vec_named(request)
+ .context("failed to encode broadcast request as msgpack")?;
+ let payload = payload.as_slice();
+ let service_name = service.service_name();
+ let deadline = tokio::time::Instant::now() + SHARD_REQUEST_TIMEOUT;
+ let request_shard = |shard_id| async move {
+ let timeout = deadline.saturating_duration_since(tokio::time::Instant::now());
+ if timeout.is_zero() {
+ anyhow::bail!("broadcast deadline expired before shard {shard_id}");
+ }
+ let subject = format!("svc.{service_name}.shard.{shard_id}");
+ let response_bytes = transport
+ .request(&subject, payload, timeout)
+ .await
+ .with_context(|| format!("broadcast request to shard {shard_id} failed"))?;
+ let response = rmp_serde::from_slice::(&response_bytes)
+ .with_context(|| format!("invalid broadcast response from shard {shard_id}"))?;
+ if !S::is_broadcast_acknowledgement(&response) {
+ anyhow::bail!("unexpected broadcast acknowledgement from shard {shard_id}");
+ }
+ anyhow::Ok(response)
+ };
+
+ let mut shard_ids = 0..ring.shard_count();
+ let mut pending = FuturesUnordered::new();
+ for shard_id in shard_ids.by_ref().take(MAX_BROADCAST_CONCURRENCY) {
+ pending.push(request_shard(shard_id));
+ }
+
+ let mut acknowledgement = None;
+ let mut first_error = None;
+ while let Some(result) = pending.next().await {
+ match result {
+ Ok(response) => acknowledgement = Some(response),
+ Err(error) => {
+ first_error.get_or_insert(error);
+ }
+ }
+ if first_error.is_none()
+ && let Some(shard_id) = shard_ids.next()
+ {
+ pending.push(request_shard(shard_id));
+ }
+ }
+ if let Some(error) = first_error {
+ return Err(error);
+ }
+ let acknowledgement = acknowledgement.context("broadcast request has no configured shards")?;
+ if S::CACHES_RESPONSES {
+ rmp_serde::to_vec_named(&acknowledgement)
+ .context("failed to encode broadcast acknowledgement as msgpack")
+ } else {
+ serde_json::to_vec(&acknowledgement)
+ .context("failed to encode broadcast acknowledgement as json")
+ }
+}
+
+async fn reply_json_response(
+ message: &impl TransportMessage,
+ transport: &impl Transport,
+ response: &impl serde::Serialize,
+ metrics: &ServiceMetrics,
+) {
+ if !message.has_reply() {
+ return;
+ }
+ match serde_json::to_vec(response) {
+ Ok(payload) => reply_bytes(message, transport, &payload).await,
+ Err(error) => {
+ warn!(error = %error, subject = message.subject(), "failed to encode router response");
+ metrics.record_request_error();
+ reply_json_error(message, transport, "encode_error").await;
+ }
+ }
+}
+
async fn handle_router_request(
msg: T::Message,
transport: T,
@@ -122,7 +216,7 @@ async fn handle_router_request(
S: RouterService,
T: Transport,
{
- let request_start = now_ms();
+ let request_start = Instant::now();
metrics.record_request();
if msg.payload().len() > MAX_ROUTER_REQUEST_BYTES {
warn!(
@@ -131,12 +225,7 @@ async fn handle_router_request(
"rejecting oversized router request"
);
metrics.record_request_error();
- if msg.has_reply() {
- let error_response =
- serde_json::to_vec(&serde_json::json!({"error": "request_too_large"}))
- .unwrap_or_default();
- let _ = reply_message(&msg, &transport, &error_response).await;
- }
+ reply_json_error(&msg, &transport, "request_too_large").await;
return;
}
let request: S::Request = match serde_json::from_slice(msg.payload()) {
@@ -144,25 +233,20 @@ async fn handle_router_request(
Err(err) => {
warn!(error = %err, "failed to decode incoming request");
metrics.record_request_error();
- if msg.has_reply() {
- let error_response =
- serde_json::to_vec(&serde_json::json!({"error": "decode_error"}))
- .unwrap_or_default();
- let _ = reply_message(&msg, &transport, &error_response).await;
- }
+ reply_json_error(&msg, &transport, "decode_error").await;
return;
}
};
let request = Arc::new(request);
+ let broadcast = S::is_broadcast_request(&request);
- if let Some(cached) = service.l1_lookup(&request) {
+ if S::CACHES_RESPONSES
+ && !broadcast
+ && let Some(cached) = service.l1_lookup(&request)
+ {
metrics.record_cache_hit();
- let elapsed = (now_ms() - request_start).max(0) as u64;
- metrics.record_request_duration(elapsed);
- if msg.has_reply() {
- let response_bytes = serde_json::to_vec(&cached).unwrap_or_default();
- let _ = reply_message(&msg, &transport, &response_bytes).await;
- }
+ metrics.record_request_duration(request_start.elapsed().as_millis() as u64);
+ reply_json_response(&msg, &transport, &cached, &metrics).await;
return;
}
@@ -170,7 +254,9 @@ async fn handle_router_request(
metrics.record_cache_miss();
metrics.record_shard_forward();
- let coalesce_result = if let Some(coalesce_key) = S::coalesce_key(&request) {
+ let coalesce_result = if broadcast {
+ forward_to_all_shards::(&transport, service.as_ref(), ring.as_ref(), &request).await
+ } else if let Some(coalesce_key) = S::coalesce_key(&request) {
let forward_transport = transport.clone();
let forward_service = service.clone();
let forward_ring = ring.clone();
@@ -208,36 +294,31 @@ async fn handle_router_request(
.await
};
- let elapsed = (now_ms() - request_start).max(0) as u64;
- metrics.record_request_duration(elapsed);
+ metrics.record_request_duration(request_start.elapsed().as_millis() as u64);
match coalesce_result {
Ok(response_bytes) => {
if !S::CACHES_RESPONSES {
if msg.has_reply() {
- let _ = reply_message(&msg, &transport, &response_bytes).await;
+ reply_bytes(&msg, &transport, &response_bytes).await;
}
return;
}
match rmp_serde::from_slice::(&response_bytes) {
Ok(response) => {
- service.l1_insert(&request, &response);
- if msg.has_reply() {
- let json = serde_json::to_vec(&response).unwrap_or_default();
- let _ = reply_message(&msg, &transport, &json).await;
+ if !broadcast {
+ service.l1_insert(&request, &response);
}
+ reply_json_response(&msg, &transport, &response, &metrics).await;
}
Err(err) => {
debug!(error = %err, "failed to decode shard response");
if msg.has_reply() {
if serde_json::from_slice::(&response_bytes).is_ok() {
- let _ = reply_message(&msg, &transport, &response_bytes).await;
+ reply_bytes(&msg, &transport, &response_bytes).await;
return;
}
- let error_response =
- serde_json::to_vec(&serde_json::json!({"error": "shard_decode_error"}))
- .unwrap_or_default();
- let _ = reply_message(&msg, &transport, &error_response).await;
+ reply_json_error(&msg, &transport, "shard_decode_error").await;
}
}
}
@@ -245,12 +326,7 @@ async fn handle_router_request(
Err(err) => {
debug!(error = %err, "shard request failed (coalesced)");
metrics.record_request_error();
- if msg.has_reply() {
- let error_response =
- serde_json::to_vec(&serde_json::json!({"error": "shard_unavailable"}))
- .unwrap_or_default();
- let _ = reply_message(&msg, &transport, &error_response).await;
- }
+ reply_json_error(&msg, &transport, "shard_unavailable").await;
}
}
}
@@ -267,7 +343,6 @@ where
let ring = Arc::new(HashRing::new(config.shard_count));
let name = service.service_name().to_owned();
let request_subject = format!("svc.{name}");
- let invalidate_subject = format!("svc.{name}.invalidate.>");
let queue_group = format!("{name}-router");
let metrics = Arc::new(ServiceMetrics::default());
@@ -295,6 +370,7 @@ where
let req_metrics = metrics.clone();
let req_permits = Arc::new(Semaphore::new(config.max_concurrent_requests));
tasks.spawn(async move {
+ let mut requests = JoinSet::new();
loop {
let mut sub = req_transport
.subscribe_queue(&request_subject, &req_queue)
@@ -307,6 +383,12 @@ where
loop {
let msg = tokio::select! {
+ result = requests.join_next(), if !requests.is_empty() => {
+ if let Err(err) = result.expect("nonempty router request set") {
+ warn!(error = %err, "router request task failed");
+ }
+ continue;
+ }
msg_opt = sub.next() => {
let Some(msg) = msg_opt else {
warn!("router request subscription stream ended, will re-subscribe");
@@ -314,24 +396,20 @@ where
};
msg
}
- _ = req_transport.wait_for_reconnect() => {
- info!("NATS reconnected, re-subscribing router request listener");
- break;
- }
};
+ while let Some(result) = requests.try_join_next() {
+ if let Err(err) = result {
+ warn!(error = %err, "router request task failed");
+ }
+ }
let permit = match req_permits.clone().try_acquire_owned() {
Ok(permit) => permit,
Err(TryAcquireError::NoPermits) => {
debug!("shedding router request, no permits available");
req_metrics.record_request();
req_metrics.record_request_error();
- if msg.has_reply() {
- let error_response =
- serde_json::to_vec(&serde_json::json!({"error": "overloaded"}))
- .unwrap_or_default();
- let _ = reply_message(&msg, &req_transport, &error_response).await;
- }
+ reply_json_error(&msg, &req_transport, "overloaded").await;
continue;
}
Err(TryAcquireError::Closed) => return anyhow::Ok(()),
@@ -341,7 +419,7 @@ where
let ring = ring.clone();
let inflight = inflight.clone();
let metrics = req_metrics.clone();
- tokio::spawn(async move {
+ requests.spawn(async move {
let _permit = permit;
handle_router_request::(msg, transport, service, ring, inflight, metrics)
.await;
@@ -350,39 +428,32 @@ where
}
});
- let inv_transport = transport.clone();
- let inv_service = service.clone();
- tasks.spawn(async move {
- loop {
- let mut sub = inv_transport.subscribe(&invalidate_subject).await?;
- info!(
- subject = invalidate_subject,
- "router listening for cache invalidations"
- );
-
+ if S::CACHES_RESPONSES {
+ let inv_transport = transport.clone();
+ let inv_service = service.clone();
+ let invalidate_prefix = format!("svc.{name}.invalidate.");
+ let invalidate_subject = format!("{invalidate_prefix}>");
+ tasks.spawn(async move {
loop {
- tokio::select! {
- msg_opt = sub.next() => {
- let Some(msg) = msg_opt else {
- warn!("router invalidation subscription stream ended, will re-subscribe");
- break;
- };
- let subject = msg.subject().to_owned();
- let key = subject
- .strip_prefix(&format!("svc.{name}.invalidate."))
- .unwrap_or("");
- if !key.is_empty() {
- inv_service.l1_invalidate(key);
- }
- }
- _ = inv_transport.wait_for_reconnect() => {
- info!("NATS reconnected, re-subscribing router invalidation listener");
- break;
+ let mut sub = inv_transport.subscribe(&invalidate_subject).await?;
+ info!(
+ subject = invalidate_subject,
+ "router listening for cache invalidations"
+ );
+
+ while let Some(msg) = sub.next().await {
+ if let Some(key) = msg
+ .subject()
+ .strip_prefix(&invalidate_prefix)
+ .filter(|key| !key.is_empty())
+ {
+ inv_service.l1_invalidate(key);
}
}
+ warn!("router invalidation subscription stream ended, will re-subscribe");
}
- }
- });
+ });
+ }
tokio::select! {
result = tasks.join_next() => {
diff --git a/fluxer_svc/src/shard.rs b/fluxer_svc/src/shard.rs
index 405f55bf7..99d519ee2 100644
--- a/fluxer_svc/src/shard.rs
+++ b/fluxer_svc/src/shard.rs
@@ -1,11 +1,15 @@
// SPDX-License-Identifier: AGPL-3.0-or-later
use crate::config::ServiceConfig;
-use crate::metrics::{ServiceMetrics, now_ms};
-use crate::transport::{Transport, TransportMessage, TransportSubscriber, reply_message};
+use crate::metrics::ServiceMetrics;
+use crate::transport::{
+ Transport, TransportMessage, TransportSubscriber, reply_bytes, reply_json_error,
+};
+use anyhow::Context;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
-use tokio::sync::{Semaphore, TryAcquireError};
+use std::time::Instant;
+use tokio::sync::{Semaphore, TryAcquireError, oneshot};
use tokio::task::JoinSet;
use tracing::{debug, info, warn};
@@ -93,9 +97,14 @@ where
let shard_is_serving = is_serving.clone();
let shard_permits = request_permits.clone();
let shard_metrics = metrics.clone();
- tasks.spawn(async move {
- loop {
- let mut sub = shard_transport.subscribe(&shard_subject).await?;
+ let (stop_requests, mut stop_requests_rx) = oneshot::channel();
+ let request_task = tasks.spawn(async move {
+ let mut requests = JoinSet::new();
+ 'listening: loop {
+ let mut sub = tokio::select! {
+ _ = &mut stop_requests_rx => break 'listening,
+ result = shard_transport.subscribe(&shard_subject) => result?,
+ };
info!(
subject = shard_subject,
shard_id,
@@ -105,12 +114,23 @@ where
loop {
tokio::select! {
+ _ = &mut stop_requests_rx => break 'listening,
+ result = requests.join_next(), if !requests.is_empty() => {
+ if let Err(err) = result.expect("nonempty shard request set") {
+ warn!(error = %err, "shard request task failed");
+ }
+ }
msg_opt = sub.next() => {
let Some(msg) = msg_opt else {
warn!("shard subscription stream ended, will re-subscribe");
break;
};
+ while let Some(result) = requests.try_join_next() {
+ if let Err(err) = result {
+ warn!(error = %err, "shard request task failed");
+ }
+ }
if !shard_is_serving.load(Ordering::SeqCst) {
continue;
}
@@ -133,27 +153,25 @@ where
debug!("shedding shard request, no permits available");
shard_metrics.record_request();
shard_metrics.record_request_error();
- reply_shard_error(&msg, &transport, "overloaded").await;
+ reply_json_error(&msg, &transport, "overloaded").await;
continue;
}
Err(TryAcquireError::Closed) => return anyhow::Ok(()),
};
- let raw_payload = msg.payload().to_vec();
-
- tokio::spawn(async move {
+ requests.spawn(async move {
let _permit = permit;
if !is_serving.load(Ordering::SeqCst) {
return;
}
- let request_start = now_ms();
+ let request_start = Instant::now();
metrics.record_request();
- let encoding = WireEncoding::detect(&raw_payload);
- let request: S::Request = match encoding.decode(&raw_payload) {
+ let encoding = WireEncoding::detect(msg.payload());
+ let request: S::Request = match encoding.decode(msg.payload()) {
Ok(r) => r,
Err(err) => {
warn!(error = %err, ?encoding, "failed to decode shard request");
metrics.record_request_error();
- reply_shard_error(&msg, &transport, "shard_request_decode_error")
+ reply_json_error(&msg, &transport, "shard_request_decode_error")
.await;
return;
}
@@ -161,25 +179,20 @@ where
match service.handle(request).await {
Ok(response) => {
- let elapsed = (now_ms() - request_start).max(0) as u64;
- metrics.record_request_duration(elapsed);
+ metrics.record_request_duration(request_start.elapsed().as_millis() as u64);
if msg.has_reply() {
match encoding.encode(&response) {
Ok(response_bytes) => {
- if let Err(err) =
- reply_message(&msg, &transport, &response_bytes).await
- {
- debug!(
- error = %err,
- "failed to send shard reply"
- );
- }
+ reply_bytes(&msg, &transport, &response_bytes).await;
}
Err(err) => {
warn!(
error = %err,
+ ?encoding,
"failed to encode shard response"
);
+ metrics.record_request_error();
+ reply_json_error(&msg, &transport, "encode_error").await;
}
}
}
@@ -187,20 +200,19 @@ where
Err(err) => {
warn!(error = %err, "shard handler returned error");
metrics.record_request_error();
- let elapsed = (now_ms() - request_start).max(0) as u64;
- metrics.record_request_duration(elapsed);
- reply_shard_error(&msg, &transport, "shard_handler_error").await;
+ metrics.record_request_duration(request_start.elapsed().as_millis() as u64);
+ reply_json_error(&msg, &transport, "shard_handler_error").await;
}
}
});
}
- _ = shard_transport.wait_for_reconnect() => {
- info!("NATS reconnected, re-subscribing shard listener");
- break;
- }
}
}
}
+ while let Some(result) = requests.join_next().await {
+ result.context("shard request task failed while draining")?;
+ }
+ anyhow::Ok(())
});
tokio::select! {
@@ -217,20 +229,34 @@ where
is_serving.store(false, Ordering::SeqCst);
- let max_permits = config.max_concurrent_requests;
- let drain_permits = request_permits.clone();
- crate::shutdown::drain_with_timeout(
- async move {
- if let Ok(_permit) = drain_permits.acquire_many(max_permits as u32).await {
- info!(
- max_concurrent_requests = max_permits,
- "all in-flight requests drained"
- );
+ let _ = stop_requests.send(());
+ let drain = async {
+ loop {
+ let (task_id, result) = tasks.join_next_with_id().await
+ .expect("running shard listener while draining")
+ .context("shard service task failed while draining")?;
+ result?;
+ if task_id == request_task.id() {
+ return anyhow::Ok(());
}
- },
- crate::shutdown::DEFAULT_DRAIN_TIMEOUT,
- )
- .await;
+ }
+ };
+ match tokio::time::timeout(crate::shutdown::DEFAULT_DRAIN_TIMEOUT, drain).await {
+ Ok(result) => {
+ result?;
+ info!(
+ max_concurrent_requests = config.max_concurrent_requests,
+ "all in-flight requests drained"
+ );
+ info!("graceful drain completed");
+ }
+ Err(_) => {
+ warn!(
+ timeout_secs = crate::shutdown::DEFAULT_DRAIN_TIMEOUT.as_secs(),
+ "drain timeout exceeded, proceeding with shutdown"
+ );
+ }
+ }
info!("shard shutdown complete");
Ok(())
@@ -238,16 +264,6 @@ where
}
}
-async fn reply_shard_error(msg: &impl TransportMessage, transport: &impl Transport, code: &str) {
- if !msg.has_reply() {
- return;
- }
- let response = serde_json::to_vec(&serde_json::json!({ "error": code })).unwrap_or_default();
- if let Err(err) = reply_message(msg, transport, &response).await {
- debug!(error = %err, "failed to send shard error reply");
- }
-}
-
#[cfg(test)]
mod tests {
use super::*;
diff --git a/fluxer_svc/src/transport.rs b/fluxer_svc/src/transport.rs
index 4b9611f0b..eb63dacf8 100644
--- a/fluxer_svc/src/transport.rs
+++ b/fluxer_svc/src/transport.rs
@@ -6,7 +6,7 @@ use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use tokio::sync::Notify;
-use tracing::{info, warn};
+use tracing::{debug, info, warn};
const NATS_SUBSCRIPTION_CAPACITY: usize = 8_192;
const SLOW_CONSUMER_LOG_INTERVAL_MS: u64 = 1_000;
@@ -62,6 +62,29 @@ where
Ok(())
}
+pub(crate) async fn reply_bytes(
+ message: &impl TransportMessage,
+ transport: &impl Transport,
+ payload: &[u8],
+) {
+ if let Err(error) = reply_message(message, transport, payload).await {
+ debug!(error = %error, subject = message.subject(), "failed to send service reply");
+ }
+}
+
+pub(crate) async fn reply_json_error(
+ message: &impl TransportMessage,
+ transport: &impl Transport,
+ code: &str,
+) {
+ if !message.has_reply() {
+ return;
+ }
+ let payload = serde_json::to_vec(&serde_json::json!({ "error": code }))
+ .expect("service error responses contain only a JSON-serializable string");
+ reply_bytes(message, transport, &payload).await;
+}
+
#[derive(Clone)]
pub struct NatsTransport {
client: async_nats::Client,
diff --git a/fluxer_users/Cargo.toml b/fluxer_users/Cargo.toml
index 2e8372187..ab3664536 100644
--- a/fluxer_users/Cargo.toml
+++ b/fluxer_users/Cargo.toml
@@ -11,7 +11,7 @@ chrono = { version = "0.4", default-features = false }
fluxer_common = { path = "../fluxer_common" }
fluxer-svc = { path = "../fluxer_svc" }
futures = "0.3.32"
-moka = { version = "0.12.15", features = ["future", "sync"] }
+moka = { version = "0.12.15", features = ["future"] }
rmp-serde = "1.3"
scylla = { version = "1.6.0", features = ["chrono-04"], optional = true }
serde = { version = "1.0.228", features = ["derive"] }
diff --git a/fluxer_users/src/main.rs b/fluxer_users/src/main.rs
index fe644ae7f..5ca4df749 100644
--- a/fluxer_users/src/main.rs
+++ b/fluxer_users/src/main.rs
@@ -26,10 +26,7 @@ async fn main() -> anyhow::Result<()> {
);
match config.mode {
- Mode::Router => {
- let router = UsersRouter::new(config.cache_max_entries, config.cache_ttl);
- fluxer_svc::router::run_router(&config, router, transport).await
- }
+ Mode::Router => fluxer_svc::router::run_router(&config, UsersRouter, transport).await,
Mode::Shard => {
let shard = match config.database_backend {
DatabaseBackend::Postgres => {
@@ -37,12 +34,7 @@ async fn main() -> anyhow::Result<()> {
fluxer_svc::postgres::PostgresConfig::from_service_config(&config);
let pool = fluxer_svc::postgres::connect(&postgres_config).await?;
let kv = fluxer_svc::postgres::KvClient::new(pool, &postgres_config)?;
- UsersShard::new_postgres(
- kv,
- transport.clone(),
- config.cache_max_entries,
- config.cache_ttl,
- )?
+ UsersShard::new_postgres(kv, config.cache_max_entries, config.cache_ttl)
}
DatabaseBackend::Cassandra => {
#[cfg(feature = "scylla")]
@@ -50,13 +42,8 @@ async fn main() -> anyhow::Result<()> {
let scylla_config =
fluxer_svc::scylla::ScyllaConfig::from_service_config(&config);
let db = fluxer_svc::scylla::connect(&scylla_config).await?;
- UsersShard::new_scylla(
- db,
- transport.clone(),
- config.cache_max_entries,
- config.cache_ttl,
- )
- .await?
+ UsersShard::new_scylla(db, config.cache_max_entries, config.cache_ttl)
+ .await?
}
#[cfg(not(feature = "scylla"))]
{
diff --git a/fluxer_users/src/router_impl.rs b/fluxer_users/src/router_impl.rs
index f8283aabd..aaf3e064b 100644
--- a/fluxer_users/src/router_impl.rs
+++ b/fluxer_users/src/router_impl.rs
@@ -2,36 +2,24 @@
use crate::types::{UserRequest, UserResponse};
use fluxer_svc::router::RouterService;
-use moka::sync::Cache;
-use std::time::Duration;
-pub struct UsersRouter {
- l1: Cache,
-}
-
-impl UsersRouter {
- pub fn new(max_entries: u64, ttl: Duration) -> Self {
- Self {
- l1: Cache::builder()
- .max_capacity(max_entries)
- .time_to_live(ttl)
- .build(),
- }
- }
-}
+pub struct UsersRouter;
impl RouterService for UsersRouter {
type Request = UserRequest;
type Response = UserResponse;
+ const CACHES_RESPONSES: bool = false;
+
fn service_name(&self) -> &str {
"users"
}
fn route_key(req: &UserRequest) -> String {
match req {
- UserRequest::GetById { user_id } => user_id.to_string(),
- UserRequest::GetPartialById { user_id } => user_id.to_string(),
+ UserRequest::GetById { user_id }
+ | UserRequest::GetPartialById { user_id }
+ | UserRequest::Invalidate { user_id } => user_id.to_string(),
UserRequest::GetPartialsByIds { user_ids } => user_ids
.iter()
.min()
@@ -43,7 +31,6 @@ impl RouterService for UsersRouter {
.min()
.cloned()
.unwrap_or_else(|| "0".to_owned()),
- UserRequest::Invalidate { user_id } => user_id.to_string(),
}
}
@@ -74,117 +61,12 @@ impl RouterService for UsersRouter {
}
}
- fn l1_lookup(&self, req: &UserRequest) -> Option {
- match req {
- UserRequest::GetById { user_id } => self.l1.get(&user_id.to_string()),
- UserRequest::GetPartialById { user_id } => {
- let cached = self.l1.get(&user_id.to_string())?;
- match cached {
- UserResponse::Found(ref user) => {
- Some(UserResponse::FoundPartial(user.to_partial()))
- }
- UserResponse::FoundPartial(_) => Some(cached),
- UserResponse::NotFound => Some(UserResponse::NotFound),
- _ => None,
- }
- }
- UserRequest::GetPartialsByIds { user_ids } => {
- let mut partials = Vec::with_capacity(user_ids.len());
- for user_id in user_ids {
- let cached = self.l1.get(&user_id.to_string())?;
- match cached {
- UserResponse::Found(ref user) => partials.push(user.to_partial()),
- UserResponse::FoundPartial(partial) => partials.push(partial),
- UserResponse::NotFound => {}
- _ => return None,
- }
- }
- Some(UserResponse::FoundPartials(partials))
- }
- UserRequest::GetApiPartialById { user_id } => {
- let cached = self.l1.get(user_id)?;
- match cached {
- UserResponse::Found(ref user) => {
- Some(UserResponse::FoundApiPartial(user.to_api_partial()))
- }
- UserResponse::FoundPartial(ref partial) => {
- Some(UserResponse::FoundApiPartial(partial.to_api_partial()))
- }
- UserResponse::FoundApiPartial(partial) => {
- Some(UserResponse::FoundApiPartial(partial))
- }
- UserResponse::NotFound => Some(UserResponse::NotFound),
- _ => None,
- }
- }
- UserRequest::GetApiPartialsByIds { user_ids } => {
- let mut partials = Vec::with_capacity(user_ids.len());
- for user_id in user_ids {
- let cached = self.l1.get(user_id)?;
- match cached {
- UserResponse::Found(ref user) => partials.push(user.to_api_partial()),
- UserResponse::FoundPartial(ref partial) => {
- partials.push(partial.to_api_partial())
- }
- UserResponse::FoundApiPartial(partial) => partials.push(partial),
- UserResponse::NotFound => {}
- _ => return None,
- }
- }
- Some(UserResponse::FoundApiPartials(partials))
- }
- UserRequest::Invalidate { .. } => None,
- }
+ fn is_broadcast_request(req: &UserRequest) -> bool {
+ matches!(req, UserRequest::Invalidate { .. })
}
- fn l1_insert(&self, req: &UserRequest, resp: &UserResponse) {
- match req {
- UserRequest::GetById { user_id } => {
- self.l1.insert(user_id.to_string(), resp.clone());
- }
- UserRequest::GetPartialById { user_id } => {
- if !matches!(
- self.l1.get(&user_id.to_string()),
- Some(UserResponse::Found(_))
- ) {
- self.l1.insert(user_id.to_string(), resp.clone());
- }
- }
- UserRequest::GetPartialsByIds { .. } => {
- if let UserResponse::FoundPartials(partials) = resp {
- for partial in partials {
- let key = partial.user_id.to_string();
- if !matches!(self.l1.get(&key), Some(UserResponse::Found(_))) {
- self.l1
- .insert(key, UserResponse::FoundPartial(partial.clone()));
- }
- }
- }
- }
- UserRequest::GetApiPartialById { .. } => {
- if let UserResponse::FoundApiPartial(partial) = resp {
- self.l1.insert(
- partial.id.clone(),
- UserResponse::FoundApiPartial(partial.clone()),
- );
- }
- }
- UserRequest::GetApiPartialsByIds { .. } => {
- if let UserResponse::FoundApiPartials(partials) = resp {
- for partial in partials {
- self.l1.insert(
- partial.id.clone(),
- UserResponse::FoundApiPartial(partial.clone()),
- );
- }
- }
- }
- UserRequest::Invalidate { .. } => {}
- }
- }
-
- fn l1_invalidate(&self, key: &str) {
- self.l1.invalidate(key);
+ fn is_broadcast_acknowledgement(response: &UserResponse) -> bool {
+ matches!(response, UserResponse::Invalidated)
}
}
@@ -192,72 +74,39 @@ impl RouterService for UsersRouter {
mod tests {
use super::*;
- fn api_partial(id: &str) -> crate::types::ApiUserPartial {
- crate::types::ApiUserPartial {
- id: id.to_owned(),
- username: "Ada".to_owned(),
- discriminator: "0007".to_owned(),
- global_name: Some("Ada Lovelace".to_owned()),
- avatar: Some("avatar_hash".to_owned()),
- avatar_color: Some(0x336699),
- bot: None,
- system: None,
- flags: 1,
- mention_flags: None,
- }
+ #[test]
+ fn coalesce_key_ignores_batch_id_order() {
+ let forward = UserRequest::GetPartialsByIds {
+ user_ids: vec![7, 42, 7],
+ };
+ let reversed = UserRequest::GetPartialsByIds {
+ user_ids: vec![42, 7],
+ };
+ assert_eq!(
+ UsersRouter::coalesce_key(&forward),
+ UsersRouter::coalesce_key(&reversed)
+ );
}
#[test]
- fn l1_caches_api_partial_batch_responses() {
- let router = UsersRouter::new(100, Duration::from_secs(30));
- let partial = api_partial("9223372036854775807");
- router.l1_insert(
- &UserRequest::GetApiPartialsByIds {
- user_ids: vec![partial.id.clone()],
- },
- &UserResponse::FoundApiPartials(vec![partial.clone()]),
+ fn coalesce_key_separates_batches_with_different_ids() {
+ let left = UserRequest::GetApiPartialsByIds {
+ user_ids: vec!["7".to_owned(), "42".to_owned()],
+ };
+ let right = UserRequest::GetApiPartialsByIds {
+ user_ids: vec!["7".to_owned()],
+ };
+ assert_ne!(
+ UsersRouter::coalesce_key(&left),
+ UsersRouter::coalesce_key(&right)
);
-
- let cached = router.l1_lookup(&UserRequest::GetApiPartialsByIds {
- user_ids: vec![partial.id.clone()],
- });
-
- match cached {
- Some(UserResponse::FoundApiPartials(partials)) => {
- assert_eq!(partials.len(), 1);
- assert_eq!(partials[0].id, partial.id);
- assert_eq!(partials[0].username, partial.username);
- }
- other => panic!("unexpected cached response: {other:?}"),
- }
}
#[test]
- fn l1_invalidates_cached_api_partials_by_user_id() {
- let router = UsersRouter::new(100, Duration::from_secs(30));
- let partial = api_partial("42");
- router.l1_insert(
- &UserRequest::GetApiPartialById {
- user_id: partial.id.clone(),
- },
- &UserResponse::FoundApiPartial(partial.clone()),
- );
-
- assert!(
- router
- .l1_lookup(&UserRequest::GetApiPartialById {
- user_id: partial.id.clone(),
- })
- .is_some()
- );
- router.l1_invalidate(&partial.id);
-
- assert!(
- router
- .l1_lookup(&UserRequest::GetApiPartialById {
- user_id: partial.id,
- })
- .is_none()
+ fn invalidations_are_never_coalesced() {
+ assert_eq!(
+ UsersRouter::coalesce_key(&UserRequest::Invalidate { user_id: 42 }),
+ None
);
}
}
diff --git a/fluxer_users/src/shard_impl.rs b/fluxer_users/src/shard_impl.rs
index e9a03ac1a..b913db51e 100644
--- a/fluxer_users/src/shard_impl.rs
+++ b/fluxer_users/src/shard_impl.rs
@@ -4,7 +4,6 @@ use crate::types::{ApiUserPartial, User, UserPartial, UserRequest, UserResponse}
#[cfg(feature = "scylla")]
use chrono::{DateTime, NaiveDate, Utc};
use fluxer_svc::shard::ShardService;
-use fluxer_svc::transport::NatsTransport;
use fluxer_svc::{postgres, postgres::KeyPart};
use futures::stream::{self, StreamExt};
use moka::future::Cache;
@@ -17,8 +16,12 @@ use scylla::statement::prepared::PreparedStatement;
#[cfg(feature = "scylla")]
use scylla::value::MaybeEmpty;
use serde::Deserialize;
+use std::collections::{HashMap, hash_map::DefaultHasher};
+use std::fmt::Write;
use std::future::Future;
+use std::hash::{Hash, Hasher};
use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
#[cfg(feature = "scylla")]
@@ -60,6 +63,8 @@ const PARTIAL_USER_COLUMNS: &str = "\
mention_flags";
const USER_BATCH_SIZE: usize = 128;
const USER_BATCH_CONCURRENCY: usize = 8;
+const USER_CACHE_MIN_GENERATION_STRIPES: usize = 4096;
+const USER_CACHE_MAX_GENERATION_STRIPES: usize = 1 << 20;
const FLUXER_SYSTEM_USER_ID: i64 = 0;
const FLUXER_SYSTEM_USERNAME: &str = "Fluxer";
const FLUXER_SYSTEM_DISCRIMINATOR: i32 = 0;
@@ -68,12 +73,19 @@ const USER_FLAG_STAFF: i64 = 1;
pub struct UsersShard {
storage: UsersStorage,
caches: UserCaches,
- transport: NatsTransport,
}
struct UserCaches {
- full: Cache>,
- partial: Cache>,
+ full: Cache>,
+ partial: Cache>,
+ generations: Box<[AtomicU64]>,
+ generation_bumps: AtomicU64,
+}
+
+#[derive(Clone, Copy, Eq, Hash, PartialEq)]
+struct UserCacheKey {
+ user_id: i64,
+ generation: u64,
}
#[derive(Clone)]
@@ -159,26 +171,9 @@ struct FullUserDbRow {
timezone_privacy_flags: Option,
}
-#[cfg(feature = "scylla")]
-#[derive(Debug, DeserializeRow)]
-struct PartialUserDbRow {
- user_id: i64,
- username: String,
- discriminator: i32,
- global_name: Option,
- avatar_hash: Option,
- bot: Option,
- system: Option,
- flags: Option,
- banner_hash: Option,
- banner_color: Option,
- accent_color: Option,
- avatar_color: Option,
- mention_flags: Option,
-}
-
#[derive(Debug, Deserialize)]
-struct PartialUserKvRow {
+#[cfg_attr(feature = "scylla", derive(DeserializeRow))]
+struct PartialUserDbRow {
user_id: i64,
username: String,
discriminator: i32,
@@ -256,6 +251,16 @@ struct FullUserKvRow {
timezone_privacy_flags: Option,
}
+fn generation_stripes(max_entries: u64) -> usize {
+ usize::try_from(max_entries)
+ .unwrap_or(USER_CACHE_MAX_GENERATION_STRIPES)
+ .clamp(
+ USER_CACHE_MIN_GENERATION_STRIPES,
+ USER_CACHE_MAX_GENERATION_STRIPES,
+ )
+ .next_power_of_two()
+}
+
impl UserCaches {
fn new(max_entries: u64, ttl: Duration) -> Self {
Self {
@@ -267,6 +272,23 @@ impl UserCaches {
.max_capacity(max_entries)
.time_to_live(ttl)
.build(),
+ generations: (0..generation_stripes(max_entries))
+ .map(|_| AtomicU64::new(0))
+ .collect(),
+ generation_bumps: AtomicU64::new(0),
+ }
+ }
+
+ fn generation(&self, user_id: i64) -> &AtomicU64 {
+ let mut hasher = DefaultHasher::new();
+ user_id.hash(&mut hasher);
+ &self.generations[hasher.finish() as usize % self.generations.len()]
+ }
+
+ fn key(&self, user_id: i64) -> UserCacheKey {
+ UserCacheKey {
+ user_id,
+ generation: self.generation(user_id).load(Ordering::SeqCst),
}
}
@@ -274,13 +296,14 @@ impl UserCaches {
where
F: Future