- Add BUILD.bazel across tidal, tidal-net, tidal-server, tidalctl for bzlmod build - Add tidal/ crate docs (README, CHANGELOG, CONTRIBUTING, AGENTS, CLAUDE, API, ARCHITECTURE) and ai-lookup reference - Add docker standalone/cluster/deploy images, compose, and prometheus config - Harden WAL (batch format, writer, dedup, diagnostics), text syncer/collectors, and vector registry - Expand tidalctl CLI and tests; restructure WAL/visibility integration test suites - Refine tidal-net transport/client/server and tidal-server cluster/scatter-gather
1060 lines
38 KiB
Rust
1060 lines
38 KiB
Rust
//! Cluster mode: multi-region tidalDB behind a single HTTP surface.
|
|
//!
|
|
//! Wraps [`SimulatedCluster`] with region name ↔ [`RegionId`] mapping and
|
|
//! exposes cluster management HTTP routes matching the runbook.
|
|
//!
|
|
//! # Transport (m8p8)
|
|
//!
|
|
//! Each follower region is wired to a real [`GrpcTransport`] from `tidal-net`
|
|
//! (not in-process crossbeam channels): leader writes are encoded as WAL
|
|
//! segments and shipped to each follower over gRPC, where the follower's
|
|
//! segment-receiver thread applies them. The transport is a self-loop — the
|
|
//! follower's own gRPC server is both the ship target and the receive source —
|
|
//! so replication traverses a real gRPC/TCP hop (serialization, circuit
|
|
//! breaker, HTTP/2) on the loopback interface.
|
|
//!
|
|
//! All regions still run inside this single process; true multi-process /
|
|
//! multi-host deployment with process isolation is tracked separately (m8p10).
|
|
|
|
use std::{
|
|
collections::HashMap,
|
|
net::{SocketAddr, TcpStream},
|
|
path::Path,
|
|
sync::{
|
|
Arc,
|
|
atomic::{AtomicBool, Ordering},
|
|
},
|
|
time::{Duration, Instant},
|
|
};
|
|
|
|
use axum::{
|
|
Json, Router,
|
|
extract::{Query, Request, State},
|
|
http::StatusCode,
|
|
middleware::{self, Next},
|
|
response::{IntoResponse, Response},
|
|
routing::{get, post},
|
|
};
|
|
use serde::{Deserialize, Serialize};
|
|
use tidal_net::{GrpcTransport, config::GrpcTransportConfig};
|
|
use tidaldb::{
|
|
query::{retrieve::Retrieve, search::Search},
|
|
replication::{
|
|
shard::{RegionId, ShardId},
|
|
transport::Transport,
|
|
},
|
|
schema::EntityId,
|
|
testing::{ClusterConfig, SimulatedCluster},
|
|
};
|
|
|
|
use crate::{
|
|
dto::{
|
|
EmbeddingRequest, FeedItem, FeedQuery, FeedResponse, ItemRequest, SearchItem,
|
|
SearchQueryParams, SearchResponse, SignalRequest, default_limit, default_profile,
|
|
feed_items, search_items,
|
|
},
|
|
error::{Result, ServerError},
|
|
};
|
|
|
|
/// How long to wait for each follower's gRPC server to bind before serving.
|
|
const GRPC_READY_TIMEOUT: Duration = Duration::from_secs(5);
|
|
/// Poll interval while waiting for a gRPC server to become connectable.
|
|
const GRPC_READY_POLL: Duration = Duration::from_millis(20);
|
|
/// Max attempts to bring a follower's gRPC server up on an auto-allocated port,
|
|
/// retrying with a fresh port to self-heal a transient bind race.
|
|
const GRPC_BUILD_ATTEMPTS: u32 = 3;
|
|
|
|
// ── Topology YAML ──────────────────────────────────────────────────────────
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
pub struct TopologySpec {
|
|
pub regions: Vec<RegionSpec>,
|
|
pub leader: String,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
pub struct RegionSpec {
|
|
pub name: String,
|
|
/// Address this region's gRPC replication server binds to, e.g.
|
|
/// `"127.0.0.1:9601"`. When omitted, an OS-assigned loopback port is used —
|
|
/// the right default for a single-process cluster where peers never need a
|
|
/// stable address. Explicit addresses are required only for a future
|
|
/// multi-process deployment (m8p10).
|
|
#[serde(default)]
|
|
pub grpc_addr: Option<String>,
|
|
}
|
|
|
|
const DEFAULT_TOPOLOGY_YAML: &str = include_str!("../config/default-cluster.yaml");
|
|
|
|
/// Load the cluster topology spec from YAML (or the compiled-in default when
|
|
/// `path` is `None`).
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns [`ServerError`] if the file cannot be read or the YAML fails to parse.
|
|
pub fn load_topology(path: Option<&Path>) -> Result<TopologySpec> {
|
|
let raw = match path {
|
|
Some(p) => std::fs::read_to_string(p).map_err(|e| ServerError::io(p, e))?,
|
|
None => DEFAULT_TOPOLOGY_YAML.to_string(),
|
|
};
|
|
serde_yml::from_str(&raw)
|
|
.map_err(|e| ServerError::SchemaConfig(format!("parse topology yaml: {e}")))
|
|
}
|
|
|
|
// ── ClusterState ───────────────────────────────────────────────────────────
|
|
|
|
/// HTTP-facing cluster state wrapping the simulated cluster fabric.
|
|
///
|
|
/// Maps human-readable region names (e.g. "us-east") to [`RegionId`] values
|
|
/// used by the distributed fabric. All cluster operations go through this
|
|
/// layer.
|
|
///
|
|
/// # Experimental
|
|
///
|
|
/// Replication between regions runs over the **real `tidal-net` gRPC
|
|
/// transport** (see the module-level docs), but every region still lives inside
|
|
/// this single process: there is no process or host isolation, so a crash takes
|
|
/// the whole "cluster" down. It provides faithful multi-region replication
|
|
/// semantics, not production high-availability. Cluster mode is therefore gated
|
|
/// behind an explicit operator opt-in — see [`ensure_experimental_enabled`].
|
|
pub struct ClusterState {
|
|
/// `Some` for the lifetime of the server; consumed by [`shutdown`] /
|
|
/// [`Drop`] so every node's `TidalDb` is dropped (checkpoint + WAL fsync +
|
|
/// thread join) deterministically on shutdown.
|
|
///
|
|
/// [`shutdown`]: ClusterState::shutdown
|
|
///
|
|
/// Held behind an [`Arc`] so scatter-gather can hand each shard's blocking
|
|
/// query to a detached `'static` worker thread (see
|
|
/// [`crate::scatter_gather`]) without borrowing `&self` for the worker's
|
|
/// lifetime. Cloning the `Arc` is cheap and the underlying
|
|
/// [`SimulatedCluster`] is `Sync`, so concurrent reads are sound.
|
|
cluster: Option<Arc<SimulatedCluster>>,
|
|
name_to_id: HashMap<String, RegionId>,
|
|
id_to_name: HashMap<RegionId, String>,
|
|
shutting_down: AtomicBool,
|
|
}
|
|
|
|
/// Environment variable that opts in to the experimental cluster mode.
|
|
pub const EXPERIMENTAL_CLUSTER_ENV: &str = "TIDAL_ALLOW_EXPERIMENTAL_CLUSTER";
|
|
|
|
/// Gate cluster mode behind an explicit operator opt-in.
|
|
///
|
|
/// Replication runs over the real `tidal-net` gRPC transport, but every region
|
|
/// lives in this single process — there is no process/host isolation, so it is
|
|
/// not production HA and must not be started by accident and mistaken for it.
|
|
/// Starting is permitted only when either:
|
|
///
|
|
/// * the `--experimental-cluster` CLI flag is passed (`flag_set == true`), or
|
|
/// * the [`EXPERIMENTAL_CLUSTER_ENV`] env var is set to a truthy value
|
|
/// (`1`, `true`, `yes`).
|
|
///
|
|
/// On success a loud `WARN` is emitted describing exactly what cluster mode is
|
|
/// and is not.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns [`ServerError::ExperimentalDisabled`] when neither opt-in is present.
|
|
pub fn ensure_experimental_enabled(flag_set: bool) -> Result<()> {
|
|
let env_set = std::env::var(EXPERIMENTAL_CLUSTER_ENV)
|
|
.map(|v| matches!(v.trim().to_ascii_lowercase().as_str(), "1" | "true" | "yes"))
|
|
.unwrap_or(false);
|
|
|
|
if !(flag_set || env_set) {
|
|
return Err(ServerError::ExperimentalDisabled(format!(
|
|
"cluster mode replicates over real gRPC but runs all regions in ONE process \
|
|
(no host/process isolation) — it is not production HA. To run it anyway, \
|
|
pass --experimental-cluster or set {EXPERIMENTAL_CLUSTER_ENV}=1."
|
|
)));
|
|
}
|
|
|
|
tracing::warn!(
|
|
"EXPERIMENTAL cluster mode enabled. Regions replicate over the real tidal-net gRPC \
|
|
transport (loopback), but all regions run in a SINGLE process — there is no host or \
|
|
process isolation, so this provides NO production high-availability. True \
|
|
multi-process deployment is tracked as m8p10. Do NOT use this for production traffic."
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
// ── gRPC transport wiring (m8p8) ─────────────────────────────────────────────
|
|
|
|
/// Build one real [`GrpcTransport`] per follower region.
|
|
///
|
|
/// Each follower's transport is a **self-loop**: its gRPC server binds the
|
|
/// follower's address and its single peer entry points at that same address, so
|
|
/// a leader ship (`send_segment(follower_shard, …)` inside
|
|
/// [`SimulatedCluster::write_signal`]) crosses a real gRPC/TCP hop into the
|
|
/// follower's own inbound queue, which the follower's segment-receiver thread
|
|
/// drains and applies. This is exactly the bidirectional pipe the crossbeam
|
|
/// `ChannelTransport` provided, but over the wire. The leader region receives
|
|
/// nothing, so it gets no transport.
|
|
///
|
|
/// Returns a map keyed by follower [`RegionId`] suitable for
|
|
/// [`ClusterConfig::transports`].
|
|
///
|
|
/// # Blocking
|
|
///
|
|
/// Each [`GrpcTransport::new`] blocks on its own runtime and is then verified
|
|
/// ready — must run on a non-async thread.
|
|
fn build_grpc_transports(
|
|
topology: &TopologySpec,
|
|
name_to_id: &HashMap<String, RegionId>,
|
|
leader_id: RegionId,
|
|
) -> Result<HashMap<RegionId, Arc<dyn Transport>>> {
|
|
let mut transports: HashMap<RegionId, Arc<dyn Transport>> = HashMap::new();
|
|
|
|
for region in &topology.regions {
|
|
let region_id = *name_to_id
|
|
.get(®ion.name)
|
|
.expect("region name was just inserted into name_to_id");
|
|
if region_id == leader_id {
|
|
continue; // the leader applies writes locally and receives nothing
|
|
}
|
|
|
|
let shard = ShardId(region_id.0);
|
|
let transport =
|
|
build_ready_follower_transport(shard, region.grpc_addr.as_deref(), ®ion.name)?;
|
|
transports.insert(region_id, Arc::new(transport));
|
|
}
|
|
|
|
Ok(transports)
|
|
}
|
|
|
|
/// Build a follower's self-loop gRPC transport and block until its server is
|
|
/// connectable.
|
|
///
|
|
/// A transport whose server does not come up within [`GRPC_READY_TIMEOUT`] is
|
|
/// dropped (freeing its runtime and port) and retried on a freshly-allocated
|
|
/// loopback port — this self-heals the rare bind race where the port probed by
|
|
/// [`free_loopback_addr`] is taken before tonic rebinds it, which otherwise
|
|
/// surfaces only as a silent serve failure and a startup timeout. An explicit
|
|
/// operator-provided `grpc_addr` is tried once: reallocating would silently
|
|
/// ignore the operator's chosen address.
|
|
fn build_ready_follower_transport(
|
|
shard: ShardId,
|
|
addr_spec: Option<&str>,
|
|
region_name: &str,
|
|
) -> Result<GrpcTransport> {
|
|
let attempts = if addr_spec.is_none() {
|
|
GRPC_BUILD_ATTEMPTS
|
|
} else {
|
|
1
|
|
};
|
|
let mut last = String::new();
|
|
|
|
for attempt in 1..=attempts {
|
|
let addr = resolve_grpc_addr(addr_spec, region_name)?;
|
|
let mut peers = HashMap::new();
|
|
peers.insert(shard, addr); // self-loop: ship target == own server
|
|
|
|
let transport = GrpcTransport::new(GrpcTransportConfig {
|
|
local_shard: shard,
|
|
listen_addr: addr,
|
|
peers,
|
|
insecure: true,
|
|
..GrpcTransportConfig::default()
|
|
})
|
|
.map_err(|e| {
|
|
ServerError::Cluster(format!(
|
|
"build gRPC transport for region '{region_name}' on {addr}: {e}"
|
|
))
|
|
})?;
|
|
|
|
if grpc_server_ready(addr) {
|
|
tracing::info!(region = %region_name, %addr, attempt, "follower gRPC transport ready");
|
|
return Ok(transport);
|
|
}
|
|
|
|
last = format!(
|
|
"gRPC server {addr} for region '{region_name}' did not become ready within \
|
|
{GRPC_READY_TIMEOUT:?}"
|
|
);
|
|
tracing::warn!(region = %region_name, %addr, attempt, max = attempts, "{last}; retrying");
|
|
drop(transport); // free the runtime + port before the next attempt
|
|
}
|
|
|
|
Err(ServerError::Cluster(last))
|
|
}
|
|
|
|
/// Block until `addr` accepts a TCP connection, or [`GRPC_READY_TIMEOUT`]
|
|
/// elapses. A successful connect proves the gRPC listener is bound, which is
|
|
/// enough for the lazily-connected client's first `ship_segment`.
|
|
fn grpc_server_ready(addr: SocketAddr) -> bool {
|
|
let deadline = Instant::now() + GRPC_READY_TIMEOUT;
|
|
loop {
|
|
if TcpStream::connect_timeout(&addr, GRPC_READY_POLL).is_ok() {
|
|
return true;
|
|
}
|
|
if Instant::now() > deadline {
|
|
return false;
|
|
}
|
|
std::thread::sleep(GRPC_READY_POLL);
|
|
}
|
|
}
|
|
|
|
/// Resolve a follower's gRPC bind address: parse the topology `grpc_addr`, or
|
|
/// allocate an OS-assigned loopback port when unset.
|
|
fn resolve_grpc_addr(spec: Option<&str>, region_name: &str) -> Result<SocketAddr> {
|
|
spec.map_or_else(
|
|
|| {
|
|
free_loopback_addr().map_err(|e| {
|
|
ServerError::Cluster(format!(
|
|
"allocate loopback gRPC port for region '{region_name}': {e}"
|
|
))
|
|
})
|
|
},
|
|
|s| {
|
|
s.parse().map_err(|e| {
|
|
ServerError::Cluster(format!(
|
|
"invalid grpc_addr '{s}' for region '{region_name}': {e}"
|
|
))
|
|
})
|
|
},
|
|
)
|
|
}
|
|
|
|
/// Reserve a free loopback port by binding `127.0.0.1:0` and reading back the
|
|
/// assigned address. The probe listener is dropped immediately so tonic can
|
|
/// bind the same port; the rare race where another bind steals the port in
|
|
/// between is retried by [`build_ready_follower_transport`].
|
|
fn free_loopback_addr() -> std::io::Result<SocketAddr> {
|
|
let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
|
|
listener.local_addr()
|
|
}
|
|
|
|
impl ClusterState {
|
|
/// Build from topology, schema, and ranking profiles.
|
|
///
|
|
/// Wires each follower region to a real [`GrpcTransport`] (m8p8) and waits
|
|
/// for every follower's gRPC server to bind before returning, so the first
|
|
/// HTTP write always finds a connectable peer.
|
|
///
|
|
/// # Blocking / threading
|
|
///
|
|
/// MUST be called from a non-async thread. [`GrpcTransport::new`] starts a
|
|
/// gRPC server by blocking on its own tokio runtime and asserts it is not
|
|
/// inside another runtime; calling this from the axum/main reactor would
|
|
/// panic. The caller (`main::run_cluster`) hops to a dedicated `std::thread`.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns [`ServerError::SchemaConfig`] when the topology has no regions,
|
|
/// contains duplicate region names, or names a leader that is not one of the
|
|
/// declared regions; or [`ServerError::Cluster`] when a follower's gRPC
|
|
/// transport cannot be built or its server does not become ready in time.
|
|
pub fn new(
|
|
topology: &TopologySpec,
|
|
schema: tidaldb::schema::Schema,
|
|
profiles: Vec<tidaldb::ranking::profile::RankingProfile>,
|
|
) -> Result<Self> {
|
|
if topology.regions.is_empty() {
|
|
return Err(ServerError::SchemaConfig(
|
|
"topology must declare at least one region".into(),
|
|
));
|
|
}
|
|
|
|
let mut name_to_id = HashMap::new();
|
|
let mut id_to_name = HashMap::new();
|
|
let mut region_ids = Vec::new();
|
|
|
|
for (i, region) in topology.regions.iter().enumerate() {
|
|
let id = RegionId(u16::try_from(i).map_err(|_| {
|
|
ServerError::SchemaConfig("topology declares more than 65535 regions".into())
|
|
})?);
|
|
if name_to_id.insert(region.name.clone(), id).is_some() {
|
|
return Err(ServerError::SchemaConfig(format!(
|
|
"duplicate region name '{}' in topology",
|
|
region.name
|
|
)));
|
|
}
|
|
id_to_name.insert(id, region.name.clone());
|
|
region_ids.push(id);
|
|
}
|
|
|
|
let leader_id = *name_to_id.get(&topology.leader).ok_or_else(|| {
|
|
ServerError::SchemaConfig(format!("leader '{}' not found in regions", topology.leader))
|
|
})?;
|
|
|
|
// m8p8: wire a real gRPC transport per follower region.
|
|
let transports = build_grpc_transports(topology, &name_to_id, leader_id)?;
|
|
tracing::info!(
|
|
followers = transports.len(),
|
|
"cluster replication wired over gRPC transport"
|
|
);
|
|
|
|
let config = ClusterConfig {
|
|
regions: region_ids,
|
|
leader_region: leader_id,
|
|
schema,
|
|
profiles,
|
|
transports: Some(transports),
|
|
};
|
|
|
|
let cluster = SimulatedCluster::build(config);
|
|
|
|
Ok(Self {
|
|
cluster: Some(Arc::new(cluster)),
|
|
name_to_id,
|
|
id_to_name,
|
|
shutting_down: AtomicBool::new(false),
|
|
})
|
|
}
|
|
|
|
pub fn set_shutting_down(&self) {
|
|
self.shutting_down.store(true, Ordering::Release);
|
|
}
|
|
|
|
pub fn is_shutting_down(&self) -> bool {
|
|
self.shutting_down.load(Ordering::Acquire)
|
|
}
|
|
|
|
/// Cleanly shut down every node in the cluster.
|
|
///
|
|
/// Drops the wrapped [`SimulatedCluster`], which drops each node's
|
|
/// `TidalDb`. `TidalDb::Drop` runs the full shutdown path (checkpoint
|
|
/// in-memory signal state → flush storage → write WAL checkpoint marker +
|
|
/// fsync → join the WAL writer, sweeper, checkpoint, and text-syncer
|
|
/// threads). The drop is idempotent, so a later [`Drop`] of `ClusterState`
|
|
/// is a no-op.
|
|
///
|
|
/// Each follower's gRPC segment-receiver thread is detached and parks in
|
|
/// `recv_segment`; it (and the [`GrpcTransport`]'s tokio runtime it holds)
|
|
/// is reclaimed when the process exits, which is immediately after a server
|
|
/// shutdown. [`GrpcTransport`]'s `Drop` uses `shutdown_background`, so if a
|
|
/// transport ever does drop here (on the reactor thread) it does not block
|
|
/// or panic.
|
|
pub fn shutdown(&mut self) {
|
|
self.set_shutting_down();
|
|
if self.cluster.take().is_some() {
|
|
tracing::info!("cluster shutdown: closing all nodes (checkpoint + WAL fsync)");
|
|
}
|
|
}
|
|
|
|
/// Access the underlying cluster, panicking if it has been shut down.
|
|
///
|
|
/// The cluster is present for the entire request-serving lifetime; it is
|
|
/// only taken during [`shutdown`](Self::shutdown), after axum has stopped
|
|
/// accepting requests, so no live handler can observe `None`.
|
|
fn cluster_ref(&self) -> &SimulatedCluster {
|
|
self.cluster
|
|
.as_ref()
|
|
.expect("cluster accessed after shutdown")
|
|
}
|
|
|
|
/// Clone the cluster `Arc` for scatter-gather fan-out.
|
|
///
|
|
/// Scatter-gather spawns one detached `'static` worker per shard, so each
|
|
/// worker needs an owned handle that can outlive the request future if the
|
|
/// shard's blocking query exceeds the deadline. See
|
|
/// [`crate::scatter_gather`].
|
|
fn cluster_arc(&self) -> Arc<SimulatedCluster> {
|
|
Arc::clone(
|
|
self.cluster
|
|
.as_ref()
|
|
.expect("cluster accessed after shutdown"),
|
|
)
|
|
}
|
|
|
|
fn resolve_region(&self, name: &str) -> Result<RegionId> {
|
|
self.name_to_id
|
|
.get(name)
|
|
.copied()
|
|
.ok_or_else(|| ServerError::BadRequest(format!("unknown region '{name}'")))
|
|
}
|
|
|
|
fn region_name(&self, id: RegionId) -> &str {
|
|
self.id_to_name.get(&id).map_or_else(
|
|
|| {
|
|
tracing::warn!(region_id = id.0, "unknown region ID in name lookup");
|
|
"unknown"
|
|
},
|
|
String::as_str,
|
|
)
|
|
}
|
|
|
|
/// Default region for reads when no `?region=` is specified: the leader.
|
|
fn default_read_region(&self) -> RegionId {
|
|
self.cluster_ref().leader_region()
|
|
}
|
|
|
|
fn read_region(&self, region_name: Option<&str>) -> Result<RegionId> {
|
|
region_name.map_or_else(
|
|
|| Ok(self.default_read_region()),
|
|
|name| self.resolve_region(name),
|
|
)
|
|
}
|
|
|
|
/// All region IDs (shard list for scatter-gather).
|
|
pub fn shard_ids(&self) -> Vec<RegionId> {
|
|
self.cluster_ref().regions()
|
|
}
|
|
|
|
/// Access the underlying cluster.
|
|
pub fn cluster(&self) -> &SimulatedCluster {
|
|
self.cluster_ref()
|
|
}
|
|
|
|
/// Region name mapping for scatter-gather metadata.
|
|
pub const fn id_to_name_map(&self) -> &HashMap<RegionId, String> {
|
|
&self.id_to_name
|
|
}
|
|
}
|
|
|
|
impl Drop for ClusterState {
|
|
/// Backstop for the explicit [`shutdown`](ClusterState::shutdown): if the
|
|
/// server exited without calling it, dropping the cluster here still runs
|
|
/// each node's `TidalDb::Drop` (checkpoint + WAL fsync + thread join).
|
|
fn drop(&mut self) {
|
|
self.shutdown();
|
|
}
|
|
}
|
|
|
|
// ── Cluster HTTP routes ────────────────────────────────────────────────────
|
|
|
|
/// Build the cluster-mode router.
|
|
///
|
|
/// Exposes the same data routes as standalone (items, embeddings, signals,
|
|
/// feed, search) plus cluster management endpoints.
|
|
pub fn build_cluster_router(state: Arc<ClusterState>, api_key: Option<Arc<str>>) -> Router {
|
|
let public = Router::new()
|
|
.route("/health", get(cluster_health))
|
|
.route("/health/startup", get(health_startup))
|
|
.route("/health/live", get(health_live))
|
|
.route("/cluster/status", get(cluster_status))
|
|
.with_state(Arc::clone(&state));
|
|
|
|
let protected = Router::new()
|
|
.route("/items", post(create_item))
|
|
.route("/embeddings", post(write_embedding))
|
|
.route("/signals", post(write_signal))
|
|
.route("/feed", get(feed))
|
|
.route("/search", get(search))
|
|
.route("/cluster/promote", post(cluster_promote))
|
|
.route("/cluster/partition", post(cluster_partition))
|
|
.route("/cluster/heal", post(cluster_heal))
|
|
// Sharded (scatter-gather) routes.
|
|
.route("/sharded/items", post(sharded_create_item))
|
|
.route("/sharded/embeddings", post(sharded_write_embedding))
|
|
.route("/sharded/signals", post(sharded_write_signal))
|
|
.route("/sharded/feed", get(sharded_feed))
|
|
.route("/sharded/search", get(sharded_search))
|
|
.layer(axum::extract::DefaultBodyLimit::max(2 * 1024 * 1024))
|
|
.with_state(state);
|
|
|
|
let protected = match api_key {
|
|
Some(key) => protected.layer(middleware::from_fn(move |req: Request, next: Next| {
|
|
let key = key.clone();
|
|
async move { crate::router::bearer_auth(req, next, &key).await }
|
|
})),
|
|
None => protected,
|
|
};
|
|
|
|
public.merge(protected)
|
|
}
|
|
|
|
// ── Health ──────────────────────────────────────────────────────────────────
|
|
|
|
async fn health_startup() -> Json<serde_json::Value> {
|
|
Json(serde_json::json!({ "ok": true, "service": "tidaldb" }))
|
|
}
|
|
|
|
async fn health_live() -> Json<serde_json::Value> {
|
|
Json(serde_json::json!({ "ok": true, "service": "tidaldb" }))
|
|
}
|
|
|
|
async fn cluster_health(
|
|
State(state): State<Arc<ClusterState>>,
|
|
) -> std::result::Result<(StatusCode, Json<serde_json::Value>), ClusterAppError> {
|
|
if state.is_shutting_down() {
|
|
return Ok((
|
|
StatusCode::SERVICE_UNAVAILABLE,
|
|
Json(serde_json::json!({
|
|
"ok": false,
|
|
"service": "tidaldb",
|
|
"cause": "shutting down"
|
|
})),
|
|
));
|
|
}
|
|
|
|
let leader = state.cluster().leader_region();
|
|
let items = state.cluster().item_count(leader);
|
|
|
|
Ok((
|
|
StatusCode::OK,
|
|
Json(serde_json::json!({
|
|
"ok": true,
|
|
"service": "tidaldb",
|
|
"mode": "cluster",
|
|
"leader": state.region_name(leader),
|
|
"items": items,
|
|
})),
|
|
))
|
|
}
|
|
|
|
// ── Cluster management ─────────────────────────────────────────────────────
|
|
|
|
#[derive(Serialize)]
|
|
struct ClusterStatusResponse {
|
|
leader: String,
|
|
relay_log_len: u64,
|
|
regions: Vec<RegionStatus>,
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
struct RegionStatus {
|
|
name: String,
|
|
applied_events: u64,
|
|
lag_events: u64,
|
|
partitioned: bool,
|
|
}
|
|
|
|
async fn cluster_status(State(state): State<Arc<ClusterState>>) -> Json<ClusterStatusResponse> {
|
|
let cluster = state.cluster();
|
|
let leader_region = cluster.leader_region();
|
|
let relay_log_len = cluster.relay_log_len();
|
|
|
|
let regions: Vec<RegionStatus> = cluster
|
|
.regions()
|
|
.into_iter()
|
|
.map(|region| {
|
|
let applied = cluster.applied_count(region);
|
|
let lag = relay_log_len.saturating_sub(applied);
|
|
RegionStatus {
|
|
name: state.region_name(region).to_string(),
|
|
applied_events: applied,
|
|
lag_events: lag,
|
|
partitioned: cluster.is_partitioned(region),
|
|
}
|
|
})
|
|
.collect();
|
|
|
|
Json(ClusterStatusResponse {
|
|
leader: state.region_name(leader_region).to_string(),
|
|
relay_log_len,
|
|
regions,
|
|
})
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct RegionRequest {
|
|
region: String,
|
|
}
|
|
|
|
async fn cluster_promote(
|
|
State(state): State<Arc<ClusterState>>,
|
|
Json(req): Json<RegionRequest>,
|
|
) -> std::result::Result<Json<serde_json::Value>, ClusterAppError> {
|
|
let region_id = state.resolve_region(&req.region).map_err(ClusterAppError)?;
|
|
state.cluster().promote_leader(region_id);
|
|
Ok(Json(serde_json::json!({
|
|
"ok": true,
|
|
"leader": req.region,
|
|
})))
|
|
}
|
|
|
|
async fn cluster_partition(
|
|
State(state): State<Arc<ClusterState>>,
|
|
Json(req): Json<RegionRequest>,
|
|
) -> std::result::Result<Json<serde_json::Value>, ClusterAppError> {
|
|
let region_id = state.resolve_region(&req.region).map_err(ClusterAppError)?;
|
|
state.cluster().partition_region(region_id);
|
|
Ok(Json(serde_json::json!({
|
|
"ok": true,
|
|
"partitioned": req.region,
|
|
})))
|
|
}
|
|
|
|
async fn cluster_heal(
|
|
State(state): State<Arc<ClusterState>>,
|
|
Json(req): Json<RegionRequest>,
|
|
) -> std::result::Result<Json<serde_json::Value>, ClusterAppError> {
|
|
let region_id = state.resolve_region(&req.region).map_err(ClusterAppError)?;
|
|
// heal_region re-ships missed segments over gRPC (blocking), so offload it.
|
|
let cluster = state.cluster_arc();
|
|
offload_cluster_write(move || {
|
|
cluster.heal_region(region_id);
|
|
Ok(())
|
|
})
|
|
.await?;
|
|
Ok(Json(serde_json::json!({
|
|
"ok": true,
|
|
"healed": req.region,
|
|
})))
|
|
}
|
|
|
|
// ── Data routes (cluster mode) ─────────────────────────────────────────────
|
|
|
|
async fn create_item(
|
|
State(state): State<Arc<ClusterState>>,
|
|
Json(req): Json<ItemRequest>,
|
|
) -> std::result::Result<StatusCode, ClusterAppError> {
|
|
state
|
|
.cluster()
|
|
.write_item_with_metadata(EntityId::new(req.entity_id), &req.metadata)
|
|
.map_err(|e| ClusterAppError(ServerError::from(e)))?;
|
|
Ok(StatusCode::CREATED)
|
|
}
|
|
|
|
async fn write_embedding(
|
|
State(state): State<Arc<ClusterState>>,
|
|
Json(req): Json<EmbeddingRequest>,
|
|
) -> std::result::Result<StatusCode, ClusterAppError> {
|
|
state
|
|
.cluster()
|
|
.write_item_embedding(EntityId::new(req.entity_id), &req.values)
|
|
.map_err(|e| ClusterAppError(ServerError::from(e)))?;
|
|
Ok(StatusCode::NO_CONTENT)
|
|
}
|
|
|
|
async fn write_signal(
|
|
State(state): State<Arc<ClusterState>>,
|
|
Json(req): Json<SignalRequest>,
|
|
) -> std::result::Result<StatusCode, ClusterAppError> {
|
|
// write_signal ships to followers over gRPC (a blocking `runtime.block_on`),
|
|
// so it must run off the async reactor — see [`offload_cluster_write`].
|
|
let cluster = state.cluster_arc();
|
|
let signal = req.signal;
|
|
let entity = EntityId::new(req.entity_id);
|
|
let weight = req.weight;
|
|
offload_cluster_write(move || {
|
|
cluster
|
|
.write_signal(&signal, entity, weight)
|
|
.map_err(ServerError::from)
|
|
})
|
|
.await?;
|
|
Ok(StatusCode::NO_CONTENT)
|
|
}
|
|
|
|
async fn feed(
|
|
State(state): State<Arc<ClusterState>>,
|
|
Query(query): Query<FeedQuery>,
|
|
) -> std::result::Result<Json<FeedResponse>, ClusterAppError> {
|
|
let region = state
|
|
.read_region(query.region.as_deref())
|
|
.map_err(ClusterAppError)?;
|
|
|
|
let mut builder = Retrieve::builder()
|
|
.profile(&query.profile)
|
|
.limit(query.limit as usize);
|
|
if let Some(user_id) = query.user_id {
|
|
builder = builder.for_user(user_id);
|
|
}
|
|
let retrieve = builder
|
|
.build()
|
|
.map_err(|e| ClusterAppError(ServerError::Tidal(e.into())))?;
|
|
|
|
// `retrieve` is a blocking `TidalDb` query; running it directly on the axum
|
|
// reactor would stall every other in-flight request behind one slow shard.
|
|
// Offload it to the blocking pool — see [`offload_cluster_read`].
|
|
let cluster = state.cluster_arc();
|
|
let result = offload_cluster_read(move || {
|
|
cluster
|
|
.retrieve(region, &retrieve)
|
|
.map_err(ServerError::from)
|
|
})
|
|
.await?;
|
|
|
|
Ok(Json(FeedResponse {
|
|
items: feed_items(&result.items),
|
|
total_candidates: result.total_candidates,
|
|
region: query.region,
|
|
}))
|
|
}
|
|
|
|
async fn search(
|
|
State(state): State<Arc<ClusterState>>,
|
|
Query(query): Query<SearchQueryParams>,
|
|
) -> std::result::Result<Json<SearchResponse>, ClusterAppError> {
|
|
let region = state
|
|
.read_region(query.region.as_deref())
|
|
.map_err(ClusterAppError)?;
|
|
|
|
let mut builder = Search::builder().query(&query.query).limit(query.limit);
|
|
if let Some(user_id) = query.user_id {
|
|
builder = builder.for_user(user_id);
|
|
}
|
|
let search_query = builder
|
|
.build()
|
|
.map_err(|e| ClusterAppError(ServerError::Tidal(e.into())))?;
|
|
|
|
// The text-index reload and the search are both blocking `TidalDb` calls;
|
|
// running them on the axum reactor would stall the whole node. Offload the
|
|
// pair to the blocking pool so the reactor stays free — see
|
|
// [`offload_cluster_read`].
|
|
let cluster = state.cluster_arc();
|
|
let result = offload_cluster_read(move || {
|
|
// Reload text index on the target region before searching.
|
|
cluster
|
|
.node(region)
|
|
.db
|
|
.reload_text_index()
|
|
.map_err(ServerError::from)?;
|
|
cluster
|
|
.search(region, &search_query)
|
|
.map_err(ServerError::from)
|
|
})
|
|
.await?;
|
|
|
|
Ok(Json(SearchResponse {
|
|
items: search_items(&result.items),
|
|
total_candidates: result.total_candidates,
|
|
region: query.region,
|
|
}))
|
|
}
|
|
|
|
// ── Sharded (scatter-gather) routes ─────────────────────────────────────────
|
|
|
|
async fn sharded_create_item(
|
|
State(state): State<Arc<ClusterState>>,
|
|
Json(req): Json<ItemRequest>,
|
|
) -> std::result::Result<StatusCode, ClusterAppError> {
|
|
let shards = state.shard_ids();
|
|
crate::scatter_gather::sharded_write_item(
|
|
state.cluster(),
|
|
EntityId::new(req.entity_id),
|
|
&req.metadata,
|
|
&shards,
|
|
)
|
|
.map_err(ClusterAppError)?;
|
|
Ok(StatusCode::CREATED)
|
|
}
|
|
|
|
async fn sharded_write_embedding(
|
|
State(state): State<Arc<ClusterState>>,
|
|
Json(req): Json<EmbeddingRequest>,
|
|
) -> std::result::Result<StatusCode, ClusterAppError> {
|
|
let shards = state.shard_ids();
|
|
crate::scatter_gather::sharded_write_embedding(
|
|
state.cluster(),
|
|
EntityId::new(req.entity_id),
|
|
&req.values,
|
|
&shards,
|
|
)
|
|
.map_err(ClusterAppError)?;
|
|
Ok(StatusCode::NO_CONTENT)
|
|
}
|
|
|
|
async fn sharded_write_signal(
|
|
State(state): State<Arc<ClusterState>>,
|
|
Json(req): Json<SignalRequest>,
|
|
) -> std::result::Result<StatusCode, ClusterAppError> {
|
|
let shards = state.shard_ids();
|
|
crate::scatter_gather::sharded_write_signal(
|
|
state.cluster(),
|
|
&req.signal,
|
|
EntityId::new(req.entity_id),
|
|
req.weight,
|
|
&shards,
|
|
)
|
|
.map_err(ClusterAppError)?;
|
|
Ok(StatusCode::NO_CONTENT)
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct ShardedFeedQuery {
|
|
#[serde(default)]
|
|
user_id: Option<u64>,
|
|
#[serde(default = "default_profile")]
|
|
profile: String,
|
|
#[serde(default = "default_limit")]
|
|
limit: u32,
|
|
#[serde(default)]
|
|
deadline_ms: Option<u64>,
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
struct ShardedFeedResponse {
|
|
items: Vec<FeedItem>,
|
|
total_candidates: usize,
|
|
scatter_gather: ScatterGatherInfo,
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
struct ScatterGatherInfo {
|
|
degraded: bool,
|
|
#[serde(skip_serializing_if = "Vec::is_empty")]
|
|
unavailable_shards: Vec<String>,
|
|
shards_queried: usize,
|
|
elapsed_ms: u64,
|
|
shard_deadline_ms: u64,
|
|
}
|
|
|
|
async fn sharded_feed(
|
|
State(state): State<Arc<ClusterState>>,
|
|
Query(query): Query<ShardedFeedQuery>,
|
|
) -> std::result::Result<Json<ShardedFeedResponse>, ClusterAppError> {
|
|
let mut builder = Retrieve::builder()
|
|
.profile(&query.profile)
|
|
.limit(query.limit as usize);
|
|
if let Some(user_id) = query.user_id {
|
|
builder = builder.for_user(user_id);
|
|
}
|
|
let retrieve = builder
|
|
.build()
|
|
.map_err(|e| ClusterAppError(ServerError::Tidal(e.into())))?;
|
|
|
|
// The scatter-gather coordinator merges shard results and reads creator
|
|
// metadata for coordinator-level diversity — all blocking. Offload the
|
|
// whole coordination off the reactor so a slow shard never stalls it.
|
|
let cluster = state.cluster_arc();
|
|
let shards = state.shard_ids();
|
|
let region_names = state.id_to_name_map().clone();
|
|
let deadline_ms = query.deadline_ms;
|
|
let (result, meta) = offload_cluster_read(move || {
|
|
crate::scatter_gather::scatter_gather_retrieve(
|
|
&cluster,
|
|
&retrieve,
|
|
&shards,
|
|
®ion_names,
|
|
deadline_ms,
|
|
)
|
|
})
|
|
.await?;
|
|
|
|
Ok(Json(ShardedFeedResponse {
|
|
items: feed_items(&result.items),
|
|
total_candidates: result.total_candidates,
|
|
scatter_gather: ScatterGatherInfo {
|
|
degraded: meta.degraded,
|
|
unavailable_shards: meta.unavailable_shards,
|
|
shards_queried: meta.shards_queried,
|
|
elapsed_ms: meta.elapsed_ms,
|
|
shard_deadline_ms: meta.shard_deadline_ms,
|
|
},
|
|
}))
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct ShardedSearchQuery {
|
|
query: String,
|
|
#[serde(default)]
|
|
user_id: Option<u64>,
|
|
#[serde(default = "default_limit")]
|
|
limit: u32,
|
|
#[serde(default)]
|
|
deadline_ms: Option<u64>,
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
struct ShardedSearchResponse {
|
|
items: Vec<SearchItem>,
|
|
total_candidates: usize,
|
|
scatter_gather: ScatterGatherInfo,
|
|
}
|
|
|
|
async fn sharded_search(
|
|
State(state): State<Arc<ClusterState>>,
|
|
Query(query): Query<ShardedSearchQuery>,
|
|
) -> std::result::Result<Json<ShardedSearchResponse>, ClusterAppError> {
|
|
let mut builder = Search::builder().query(&query.query).limit(query.limit);
|
|
if let Some(user_id) = query.user_id {
|
|
builder = builder.for_user(user_id);
|
|
}
|
|
let search_query = builder
|
|
.build()
|
|
.map_err(|e| ClusterAppError(ServerError::Tidal(e.into())))?;
|
|
|
|
// Offload the blocking scatter-gather coordination off the reactor (see
|
|
// [`sharded_feed`] / [`offload_cluster_read`]).
|
|
let cluster = state.cluster_arc();
|
|
let shards = state.shard_ids();
|
|
let region_names = state.id_to_name_map().clone();
|
|
let deadline_ms = query.deadline_ms;
|
|
let (result, meta) = offload_cluster_read(move || {
|
|
crate::scatter_gather::scatter_gather_search(
|
|
&cluster,
|
|
&search_query,
|
|
&shards,
|
|
®ion_names,
|
|
deadline_ms,
|
|
)
|
|
})
|
|
.await?;
|
|
|
|
Ok(Json(ShardedSearchResponse {
|
|
items: search_items(&result.items),
|
|
total_candidates: result.total_candidates,
|
|
scatter_gather: ScatterGatherInfo {
|
|
degraded: meta.degraded,
|
|
unavailable_shards: meta.unavailable_shards,
|
|
shards_queried: meta.shards_queried,
|
|
elapsed_ms: meta.elapsed_ms,
|
|
shard_deadline_ms: meta.shard_deadline_ms,
|
|
},
|
|
}))
|
|
}
|
|
|
|
// ── Blocking-work offload ────────────────────────────────────────────────────
|
|
|
|
/// Run a blocking READ- only cluster query (RETRIEVE / SEARCH) on the tokio
|
|
/// blocking pool, off the async reactor, and await its result.
|
|
///
|
|
/// `TidalDb::retrieve`/`search` (and `reload_text_index`) are synchronous,
|
|
/// CPU/IO-bound calls. Running them directly inside an axum handler would block
|
|
/// the reactor thread for the whole query, so one slow shard would stall every
|
|
/// other in-flight request on that worker. Reads never call
|
|
/// [`GrpcTransport::send_segment`], so unlike [`offload_cluster_write`] they do
|
|
/// NOT need a runtime-free OS thread — `spawn_blocking` (the purpose-built
|
|
/// blocking pool) is the right tool and keeps the reactor free.
|
|
async fn offload_cluster_read<F, T>(f: F) -> std::result::Result<T, ClusterAppError>
|
|
where
|
|
F: FnOnce() -> Result<T> + Send + 'static,
|
|
T: Send + 'static,
|
|
{
|
|
tokio::task::spawn_blocking(f)
|
|
.await
|
|
.map_err(|e| {
|
|
ClusterAppError(ServerError::Cluster(format!(
|
|
"cluster read worker panicked: {e}"
|
|
)))
|
|
})?
|
|
.map_err(ClusterAppError)
|
|
}
|
|
|
|
/// Run a cluster operation that may ship WAL segments over gRPC on a dedicated
|
|
/// OS thread, off the async reactor, and await its result.
|
|
///
|
|
/// [`GrpcTransport::send_segment`] blocks on its own tokio runtime and asserts
|
|
/// it is not invoked from within a runtime (`Handle::try_current().is_err()`),
|
|
/// so it must not run on an axum worker — nor on a `spawn_blocking` thread,
|
|
/// which still carries a runtime handle and would trip that assert. A fresh
|
|
/// `std::thread` has no ambient runtime; the result is bridged back through a
|
|
/// oneshot the handler awaits, so the reactor is never blocked.
|
|
async fn offload_cluster_write<F, T>(f: F) -> std::result::Result<T, ClusterAppError>
|
|
where
|
|
F: FnOnce() -> Result<T> + Send + 'static,
|
|
T: Send + 'static,
|
|
{
|
|
let (tx, rx) = tokio::sync::oneshot::channel();
|
|
std::thread::Builder::new()
|
|
.name("cluster-write".into())
|
|
.spawn(move || {
|
|
let _ = tx.send(f());
|
|
})
|
|
.map_err(|e| {
|
|
ClusterAppError(ServerError::Cluster(format!(
|
|
"spawn cluster write worker: {e}"
|
|
)))
|
|
})?;
|
|
rx.await
|
|
.map_err(|_| {
|
|
ClusterAppError(ServerError::Cluster(
|
|
"cluster write worker dropped without responding".into(),
|
|
))
|
|
})?
|
|
.map_err(ClusterAppError)
|
|
}
|
|
|
|
// ── Error handling ─────────────────────────────────────────────────────────
|
|
|
|
struct ClusterAppError(ServerError);
|
|
|
|
impl IntoResponse for ClusterAppError {
|
|
fn into_response(self) -> Response {
|
|
let status = crate::router::status_from_error(&self.0);
|
|
let body = serde_json::json!({ "error": self.0.to_string() });
|
|
(status, Json(body)).into_response()
|
|
}
|
|
}
|