Replication!

This commit is contained in:
Kalzu Rekku committed 2026-10-25 18:54:13 +02:00
1 parent 427a824e45
commit 3d819c7b10
6 files changed
+1004 -71

No files matched your search

+657 -43
View File
@@ -1,27 +1,72 @@
use axum::{
extract::{rejection::JsonRejection, State},
http::StatusCode,
extract::{rejection::JsonRejection, Query, State},
http::{header::AUTHORIZATION, HeaderMap, StatusCode},
routing::get,
Json, Router,
};
use chrono::Utc;
use rusqlite::{params, Connection};
use rusqlite::{params, Connection, OptionalExtension};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::{
fs,
sync::{Arc, Mutex, RwLock},
time::Duration,
};
use tracing::{error, info, warn};
use uuid::Uuid;
// ---------------------------------------------------------------------------
// State
// ---------------------------------------------------------------------------
#[derive(Clone)]
struct AppState {
db: Arc<Mutex<Connection>>,
schema: Arc<RwLock<SchemaConfig>>,
config_path: Arc<String>,
/// Permanent per-instance identity (stored in the `meta` table).
instance_id: Uuid,
role: Role,
primary_url: Option<Arc<String>>,
/// Shared secret for /api/replication/*. None = replication endpoint disabled.
repl_token: Option<Arc<String>>,
sync: Arc<RwLock<SyncStatus>>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Role {
Primary,
Replica,
}
impl Role {
fn as_str(self) -> &'static str {
match self {
Role::Primary => "primary",
Role::Replica => "replica",
}
}
}
#[derive(Debug, Clone, Serialize)]
struct SyncStatus {
/// "primary" | "idle" | "syncing" | "retrying" | "halted"
state: String,
/// Last global tick (seq) applied locally. Replica's cursor.
last_applied_seq: i64,
source_instance_id: Option<Uuid>,
/// Primary's head seq as of the last successful pull.
primary_head_seq: Option<i64>,
last_success: Option<String>,
last_error: Option<String>,
}
// ---------------------------------------------------------------------------
// Schema validation (unchanged)
// ---------------------------------------------------------------------------
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
enum ValidationMode {
@@ -58,6 +103,9 @@ struct SchemaConfig {
#[derive(Debug, Serialize, Deserialize)]
struct ApiResponse {
/// Global tick. Absent on endpoints that aren't records.
#[serde(default, skip_serializing_if = "Option::is_none")]
seq: Option<i64>,
uuid: Uuid,
date: String,
data: Value,
@@ -82,6 +130,11 @@ fn internal_error(context: &str) -> ApiError {
)
}
fn db_error(context: &str, e: rusqlite::Error) -> ApiError {
error!(error = %e, context, "Database error");
api_error(StatusCode::INTERNAL_SERVER_ERROR, "Database error")
}
fn matches_type(value: &Value, expected: &DataType) -> bool {
match expected {
DataType::String => value.is_string(),
@@ -160,10 +213,107 @@ fn load_config(path: &str) -> Result<SchemaConfig, String> {
Ok(config)
}
// ---------------------------------------------------------------------------
// Database init / meta
// ---------------------------------------------------------------------------
fn meta_get(db: &Connection, key: &str) -> rusqlite::Result<Option<String>> {
db.query_row("SELECT value FROM meta WHERE key = ?1", [key], |r| {
r.get::<_, String>(0)
})
.optional()
}
fn meta_set(db: &Connection, key: &str, value: &str) -> rusqlite::Result<()> {
db.execute(
"INSERT INTO meta (key, value) VALUES (?1, ?2)
ON CONFLICT(key) DO UPDATE SET value = excluded.value",
params![key, value],
)?;
Ok(())
}
/// Creates tables and migrates a pre-replication `records` table
/// (uuid PK, no seq) to the journal layout (seq AUTOINCREMENT + uuid UNIQUE).
///
/// `records` is insert-only, so it *is* the append-only journal; `seq` is the
/// global tick. AUTOINCREMENT guarantees seq never gets reused, even after
/// a rollback.
fn init_db(db: &mut Connection) -> rusqlite::Result<()> {
db.execute_batch(
"CREATE TABLE IF NOT EXISTS meta (
key TEXT PRIMARY KEY,
value TEXT NOT NULL
);",
)?;
let exists: bool = db.query_row(
"SELECT EXISTS(SELECT 1 FROM sqlite_master
WHERE type = 'table' AND name = 'records')",
[],
|r| r.get(0),
)?;
if exists {
let has_seq: bool = db.query_row(
"SELECT EXISTS(SELECT 1 FROM pragma_table_info('records')
WHERE name = 'seq')",
[],
|r| r.get(0),
)?;
if !has_seq {
info!("Migrating records table to journal layout");
let tx = db.transaction()?;
tx.execute_batch(
"ALTER TABLE records RENAME TO records_old;
CREATE TABLE records (
seq INTEGER PRIMARY KEY AUTOINCREMENT,
uuid TEXT NOT NULL UNIQUE,
date TEXT NOT NULL,
data TEXT NOT NULL
);
INSERT INTO records (uuid, date, data)
SELECT uuid, date, data FROM records_old ORDER BY rowid;
DROP TABLE records_old;",
)?;
tx.commit()?;
}
}
db.execute_batch(
"CREATE TABLE IF NOT EXISTS records (
seq INTEGER PRIMARY KEY AUTOINCREMENT,
uuid TEXT NOT NULL UNIQUE,
date TEXT NOT NULL,
data TEXT NOT NULL
);",
)?;
Ok(())
}
fn load_or_create_instance_id(db: &Connection) -> rusqlite::Result<Uuid> {
if let Some(existing) = meta_get(db, "instance_id")? {
if let Ok(id) = Uuid::parse_str(&existing) {
return Ok(id);
}
warn!("Stored instance_id is invalid; generating a new one");
}
let id = Uuid::new_v4();
meta_set(db, "instance_id", &id.to_string())?;
Ok(id)
}
// ---------------------------------------------------------------------------
// Regular handlers
// ---------------------------------------------------------------------------
async fn get_info() -> Json<ApiResponse> {
info!(endpoint = "/api/now", "Request received");
Json(ApiResponse {
seq: None,
uuid: Uuid::new_v4(),
date: Utc::now().to_rfc3339(),
data: serde_json::json!({
@@ -179,6 +329,16 @@ async fn create_record(
) -> Result<(StatusCode, Json<ApiResponse>), ApiError> {
info!(endpoint = "/api/records", "Create request received");
if state.role == Role::Replica {
return Err(api_error(
StatusCode::FORBIDDEN,
format!(
"This instance is a read-only replica; write to {}",
state.primary_url.as_deref().map(String::as_str).unwrap_or("the primary")
),
));
}
let Json(request) = payload.map_err(|e| {
warn!(error = %e, "Invalid JSON request");
api_error(StatusCode::BAD_REQUEST, "Invalid JSON request body")
@@ -196,7 +356,8 @@ async fn create_record(
return Err(api_error(StatusCode::BAD_REQUEST, message));
}
let record = ApiResponse {
let mut record = ApiResponse {
seq: None,
uuid: Uuid::new_v4(),
date: Utc::now().to_rfc3339(),
data: request.data,
@@ -227,9 +388,11 @@ async fn create_record(
"Failed to store record",
)
})?;
record.seq = Some(db.last_insert_rowid());
}
info!(uuid = %record.uuid, "Record stored");
info!(uuid = %record.uuid, seq = ?record.seq, "Record stored");
Ok((StatusCode::CREATED, Json(record)))
}
@@ -246,45 +409,29 @@ async fn list_records(
let mut statement = db
.prepare(
"SELECT uuid, date, data
"SELECT seq, uuid, date, data
FROM records
ORDER BY date DESC
ORDER BY seq DESC
LIMIT 100",
)
.map_err(|e| {
error!(error = %e, "Failed to prepare query");
api_error(
StatusCode::INTERNAL_SERVER_ERROR,
"Database query failed",
)
})?;
.map_err(|e| db_error("prepare list", e))?;
let rows = statement
.query_map([], |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, i64>(0)?,
row.get::<_, String>(1)?,
row.get::<_, String>(2)?,
row.get::<_, String>(3)?,
))
})
.map_err(|e| {
error!(error = %e, "Failed to query records");
api_error(
StatusCode::INTERNAL_SERVER_ERROR,
"Database query failed",
)
})?;
.map_err(|e| db_error("query list", e))?;
let mut records = Vec::new();
for row in rows {
let (uuid, date, data_json) = row.map_err(|e| {
error!(error = %e, "Failed to read database row");
api_error(
StatusCode::INTERNAL_SERVER_ERROR,
"Database read failed",
)
})?;
let (seq, uuid, date, data_json) =
row.map_err(|e| db_error("read row", e))?;
let uuid = Uuid::parse_str(&uuid)
.map_err(|_| internal_error("invalid stored UUID"))?;
@@ -292,7 +439,12 @@ async fn list_records(
let data: Value = serde_json::from_str(&data_json)
.map_err(|_| internal_error("invalid stored JSON"))?;
records.push(ApiResponse { uuid, date, data });
records.push(ApiResponse {
seq: Some(seq),
uuid,
date,
data,
});
}
info!(count = records.len(), "Records returned");
@@ -358,6 +510,408 @@ async fn reload_config(
})))
}
// ---------------------------------------------------------------------------
// Replication: serving side (primary, or a replica acting as a relay)
// ---------------------------------------------------------------------------
#[derive(Debug, Serialize, Deserialize)]
struct Change {
seq: i64,
uuid: Uuid,
date: String,
data: Value,
}
#[derive(Debug, Serialize, Deserialize)]
struct ChangesResponse {
instance_id: Uuid,
head_seq: i64,
has_more: bool,
changes: Vec<Change>,
}
#[derive(Debug, Deserialize)]
struct ChangesQuery {
after: Option<i64>,
limit: Option<i64>,
}
fn ct_eq(a: &[u8], b: &[u8]) -> bool {
a.len() == b.len() && a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
}
fn check_token(
headers: &HeaderMap,
expected: &Option<Arc<String>>,
) -> Result<(), ApiError> {
let Some(expected) = expected else {
return Err(api_error(
StatusCode::FORBIDDEN,
"Replication is disabled (no token configured)",
));
};
let provided = headers
.get(AUTHORIZATION)
.and_then(|v| v.to_str().ok())
.and_then(|v| v.strip_prefix("Bearer "))
.unwrap_or("");
if !ct_eq(provided.as_bytes(), expected.as_bytes()) {
return Err(api_error(StatusCode::UNAUTHORIZED, "Invalid token"));
}
Ok(())
}
/// Cursor-based pull. Stateless on the server: the replica says "give me
/// everything after seq N", so dropped connections, duplicate requests and
/// restarts on either side are all harmless.
async fn replication_changes(
State(state): State<AppState>,
headers: HeaderMap,
Query(query): Query<ChangesQuery>,
) -> Result<Json<ChangesResponse>, ApiError> {
check_token(&headers, &state.repl_token)?;
let after = query.after.unwrap_or(0).max(0);
let limit = query.limit.unwrap_or(200).clamp(1, 1000);
let db = state
.db
.lock()
.map_err(|_| internal_error("database lock poisoned"))?;
let head_seq: i64 = db
.query_row("SELECT COALESCE(MAX(seq), 0) FROM records", [], |r| {
r.get(0)
})
.map_err(|e| db_error("head seq", e))?;
let mut statement = db
.prepare(
"SELECT seq, uuid, date, data
FROM records
WHERE seq > ?1
ORDER BY seq ASC
LIMIT ?2",
)
.map_err(|e| db_error("prepare changes", e))?;
let rows = statement
.query_map(params![after, limit], |row| {
Ok((
row.get::<_, i64>(0)?,
row.get::<_, String>(1)?,
row.get::<_, String>(2)?,
row.get::<_, String>(3)?,
))
})
.map_err(|e| db_error("query changes", e))?;
let mut changes = Vec::new();
for row in rows {
let (seq, uuid, date, data_json) =
row.map_err(|e| db_error("read change", e))?;
changes.push(Change {
seq,
uuid: Uuid::parse_str(&uuid)
.map_err(|_| internal_error("invalid stored UUID"))?,
date,
data: serde_json::from_str(&data_json)
.map_err(|_| internal_error("invalid stored JSON"))?,
});
}
let last = changes.last().map(|c| c.seq).unwrap_or(after);
Ok(Json(ChangesResponse {
instance_id: state.instance_id,
head_seq,
has_more: last < head_seq,
changes,
}))
}
async fn replication_status(
State(state): State<AppState>,
) -> Result<Json<Value>, ApiError> {
let status = state
.sync
.read()
.map_err(|_| internal_error("sync status lock poisoned"))?
.clone();
let local_head: i64 = {
let db = state
.db
.lock()
.map_err(|_| internal_error("database lock poisoned"))?;
db.query_row("SELECT COALESCE(MAX(seq), 0) FROM records", [], |r| {
r.get(0)
})
.map_err(|e| db_error("head seq", e))?
};
let lag = status
.primary_head_seq
.map(|head| (head - status.last_applied_seq).max(0));
Ok(Json(serde_json::json!({
"instance_id": state.instance_id,
"role": state.role.as_str(),
"primary_url": state.primary_url.as_deref(),
"local_head_seq": local_head,
"lag_records": lag,
"sync": status,
})))
}
// ---------------------------------------------------------------------------
// Replication: replica side
// ---------------------------------------------------------------------------
#[derive(Debug, Clone)]
struct SyncConfig {
poll_interval: Duration,
max_backoff: Duration,
batch: i64,
}
enum SyncError {
/// Network trouble, 5xx, truncated body... retry with backoff.
Transient(String),
/// Identity/history mismatch. Retrying won't help; needs an operator.
Fatal(String),
}
fn jitter(d: Duration) -> Duration {
// Uniform in [d/2, d]; avoids a fleet of replicas retrying in lockstep.
let ms = d.as_millis() as u64;
let extra = (Uuid::new_v4().as_u128() as u64) % (ms / 2 + 1);
Duration::from_millis(ms / 2 + extra)
}
fn update_status(state: &AppState, f: impl FnOnce(&mut SyncStatus)) {
if let Ok(mut s) = state.sync.write() {
f(&mut s);
}
}
/// Applies a batch atomically: rows + cursor + source identity commit together,
/// so a crash can never leave the cursor ahead of (or behind) the data.
/// INSERT OR IGNORE makes re-delivery idempotent.
fn apply_batch(
db: &Arc<Mutex<Connection>>,
source: Uuid,
changes: &[Change],
) -> Result<i64, String> {
let mut conn = db.lock().map_err(|_| "database lock poisoned".to_string())?;
let tx = conn.transaction().map_err(|e| e.to_string())?;
let mut max_seq = 0;
{
let mut insert = tx
.prepare_cached(
"INSERT OR IGNORE INTO records (seq, uuid, date, data)
VALUES (?1, ?2, ?3, ?4)",
)
.map_err(|e| e.to_string())?;
for c in changes {
let data = serde_json::to_string(&c.data).map_err(|e| e.to_string())?;
insert
.execute(params![c.seq, c.uuid.to_string(), c.date, data])
.map_err(|e| e.to_string())?;
max_seq = max_seq.max(c.seq);
}
}
meta_set(&tx, "last_applied_seq", &max_seq.to_string())
.map_err(|e| e.to_string())?;
meta_set(&tx, "source_instance_id", &source.to_string())
.map_err(|e| e.to_string())?;
tx.commit().map_err(|e| e.to_string())?;
Ok(max_seq)
}
async fn sync_once(
client: &reqwest::Client,
state: &AppState,
primary_url: &str,
token: &str,
batch: i64,
) -> Result<bool, SyncError> {
let (last, known_source) = {
let s = state
.sync
.read()
.map_err(|_| SyncError::Fatal("sync status lock poisoned".into()))?;
(s.last_applied_seq, s.source_instance_id)
};
let url = format!(
"{}/api/replication/changes?after={last}&limit={batch}",
primary_url.trim_end_matches('/')
);
let resp = client
.get(&url)
.bearer_auth(token)
.send()
.await
.map_err(|e| SyncError::Transient(format!("request failed: {e}")))?;
if !resp.status().is_success() {
return Err(SyncError::Transient(format!(
"primary returned {}",
resp.status()
)));
}
// A truncated or garbled body on a flaky link fails to decode here and is
// retried; a partial batch is never applied.
let body: ChangesResponse = resp
.json()
.await
.map_err(|e| SyncError::Transient(format!("bad response body: {e}")))?;
// --- identity / history checks ---
if body.instance_id == state.instance_id {
return Err(SyncError::Fatal(
"source has the same instance_id as this replica (cloned database?)".into(),
));
}
if let Some(known) = known_source {
if known != body.instance_id {
return Err(SyncError::Fatal(format!(
"source instance changed: expected {known}, got {}",
body.instance_id
)));
}
}
if body.head_seq < last {
return Err(SyncError::Fatal(format!(
"source head_seq {} is behind our cursor {last} (source restored from older backup?)",
body.head_seq
)));
}
if body.changes.is_empty() {
update_status(state, |s| {
s.state = "idle".into();
s.primary_head_seq = Some(body.head_seq);
s.last_success = Some(Utc::now().to_rfc3339());
s.last_error = None;
s.source_instance_id = Some(body.instance_id);
});
return Ok(false);
}
// Sequence must be strictly increasing and past our cursor.
let mut prev = last;
for c in &body.changes {
if c.seq <= prev {
return Err(SyncError::Transient(format!(
"non-monotonic seq {} after {prev}",
c.seq
)));
}
prev = c.seq;
}
let count = body.changes.len();
let db = state.db.clone();
let source = body.instance_id;
let changes = body.changes;
let max_seq = tokio::task::spawn_blocking(move || apply_batch(&db, source, &changes))
.await
.map_err(|e| SyncError::Transient(format!("apply task failed: {e}")))?
.map_err(|e| SyncError::Transient(format!("apply failed: {e}")))?;
update_status(state, |s| {
s.state = if body.has_more { "syncing" } else { "idle" }.into();
s.last_applied_seq = max_seq;
s.source_instance_id = Some(source);
s.primary_head_seq = Some(body.head_seq);
s.last_success = Some(Utc::now().to_rfc3339());
s.last_error = None;
});
info!(
applied = count,
last_applied_seq = max_seq,
head = body.head_seq,
"Replication batch applied"
);
Ok(body.has_more)
}
async fn replica_loop(
state: AppState,
primary_url: String,
token: String,
cfg: SyncConfig,
) {
let client = match reqwest::Client::builder()
.connect_timeout(Duration::from_secs(10))
.timeout(Duration::from_secs(60))
.build()
{
Ok(c) => c,
Err(e) => {
error!(error = %e, "Cannot build HTTP client; replication halted");
return;
}
};
let mut backoff = Duration::from_secs(1);
info!(primary = %primary_url, "Replica sync loop started");
loop {
match sync_once(&client, &state, &primary_url, &token, cfg.batch).await {
Ok(more) => {
backoff = Duration::from_secs(1);
if !more {
tokio::time::sleep(jitter(cfg.poll_interval)).await;
}
}
Err(SyncError::Transient(message)) => {
let wait = jitter(backoff);
warn!(reason = %message, retry_in = ?wait, "Replication pull failed");
update_status(&state, |s| {
s.state = "retrying".into();
s.last_error = Some(message);
});
tokio::time::sleep(wait).await;
backoff = (backoff * 2).min(cfg.max_backoff);
}
Err(SyncError::Fatal(message)) => {
error!(reason = %message, "Replication halted; operator action required");
update_status(&state, |s| {
s.state = "halted".into();
s.last_error = Some(message);
});
return;
}
}
}
}
// ---------------------------------------------------------------------------
// main
// ---------------------------------------------------------------------------
fn env_parse<T: std::str::FromStr>(name: &str, default: T) -> T {
std::env::var(name)
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(default)
}
#[tokio::main]
async fn main() {
tracing_subscriber::fmt()
@@ -377,24 +931,80 @@ async fn main() {
"Configuration loaded"
);
let db = Connection::open("messages.db")
// Replication settings
let role = match std::env::var("APP_ROLE").as_deref() {
Ok("replica") => Role::Replica,
Ok("primary") | Err(_) => Role::Primary,
Ok(other) => panic!("APP_ROLE must be 'primary' or 'replica', got '{other}'"),
};
let primary_url = std::env::var("APP_PRIMARY_URL").ok();
let repl_token = std::env::var("APP_REPL_TOKEN").ok().filter(|t| !t.is_empty());
if role == Role::Replica && (primary_url.is_none() || repl_token.is_none()) {
panic!("Replica requires APP_PRIMARY_URL and APP_REPL_TOKEN");
}
if repl_token.is_none() {
warn!("APP_REPL_TOKEN not set: /api/replication/changes is disabled");
}
let sync_cfg = SyncConfig {
poll_interval: Duration::from_secs(env_parse("APP_SYNC_INTERVAL_SECS", 5)),
max_backoff: Duration::from_secs(env_parse("APP_SYNC_MAX_BACKOFF_SECS", 60)),
batch: env_parse("APP_SYNC_BATCH", 200i64).clamp(1, 1000),
};
let mut db = Connection::open("messages.db")
.expect("Failed to open SQLite database");
db.execute_batch(
"CREATE TABLE IF NOT EXISTS records (
uuid TEXT PRIMARY KEY,
date TEXT NOT NULL,
data TEXT NOT NULL
);",
)
.expect("Failed to initialize database");
init_db(&mut db).expect("Failed to initialize database");
let instance_id =
load_or_create_instance_id(&db).expect("Failed to load instance id");
let last_applied_seq = meta_get(&db, "last_applied_seq")
.ok()
.flatten()
.and_then(|v| v.parse::<i64>().ok())
.unwrap_or(0);
let source_instance_id = meta_get(&db, "source_instance_id")
.ok()
.flatten()
.and_then(|v| Uuid::parse_str(&v).ok());
info!(
instance_id = %instance_id,
role = role.as_str(),
last_applied_seq,
"Instance initialised"
);
let state = AppState {
db: Arc::new(Mutex::new(db)),
schema: Arc::new(RwLock::new(schema)),
config_path: Arc::new(config_path),
instance_id,
role,
primary_url: primary_url.clone().map(Arc::new),
repl_token: repl_token.clone().map(Arc::new),
sync: Arc::new(RwLock::new(SyncStatus {
state: if role == Role::Replica { "idle" } else { "primary" }.into(),
last_applied_seq,
source_instance_id,
primary_head_seq: None,
last_success: None,
last_error: None,
})),
};
if role == Role::Replica {
tokio::spawn(replica_loop(
state.clone(),
primary_url.expect("checked above"),
repl_token.expect("checked above"),
sync_cfg,
));
}
let app = Router::new()
.route("/", get(get_info))
.route("/api/now", get(get_info))
@@ -404,16 +1014,20 @@ async fn main() {
"/api/admin/reload-config",
axum::routing::post(reload_config),
)
.route("/api/replication/changes", get(replication_changes))
.route("/api/replication/status", get(replication_status))
// gzip matters on slow links; reqwest's "gzip" feature negotiates it.
.layer(tower_http::compression::CompressionLayer::new())
.with_state(state);
let addr = "127.0.0.1:3000";
let listener = tokio::net::TcpListener::bind(addr)
let addr = std::env::var("APP_BIND").unwrap_or_else(|_| "127.0.0.1:3000".into());
let listener = tokio::net::TcpListener::bind(&addr)
.await
.expect("Failed to bind address");
info!(address = addr, "HTTP server started");
info!(address = %addr, "HTTP server started");
axum::serve(listener, app)
.await
.expect("HTTP server failed");
}
}