Add query and stream features; enhance replication routes and logic

Author: JackGuslerGit

Cargo.lock

@@ -66,6 +66,7 @@ features = [
 	"http2",
 	"json",
 	"matched-path",
+	"query",
 	"tokio",
 	"tracing",
 ]
@@ -322,6 +323,7 @@ features = [
 	"rustls",
 	"rustls-native-certs",
 	"socks",
+	"stream",
 ]
 
 [workspace.dependencies.ring]

Cargo.toml

@@ -4,18 +4,21 @@
 //! - `GET  /_tuwunel/replication/status`        — current sequence number + role
 //! - `GET  /_tuwunel/replication/wal?since=N`   — streaming WAL frame feed
 //! - `GET  /_tuwunel/replication/checkpoint`    — full database checkpoint as tar
+//! - `POST /_tuwunel/replication/promote`       — promote secondary to primary
+//! - `POST /_tuwunel/replication/demote`        — demote primary back to secondary
 
 use std::time::Duration;
 
 use axum::{
+	Json,
 	body::Body,
 	extract::{Query, State},
 	http::{StatusCode, header},
 	response::{IntoResponse, Response},
 };
 use bytes::Bytes;
-use futures::StreamExt;
-use serde::Deserialize;
+use futures::{SinkExt, StreamExt};
+use serde::{Deserialize, Serialize};
 use tuwunel_core::Result;
 use tuwunel_database::{WalFrame, is_wal_gap_error};
 
@@ -38,8 +41,16 @@ pub(crate) async fn replication_status(
 		.await
 		.unwrap_or(0);
 
+	let role = if services.server.config.rocksdb_primary_url.is_some()
+		&& !services.replication.is_promoted()
+	{
+		"secondary"
+	} else {
+		"primary"
+	};
+
 	axum::Json(serde_json::json!({
-		"role": "primary",
+		"role": role,
 		"latest_sequence": seq,
 	}))
 }
@@ -214,6 +225,90 @@ pub(crate) async fn replication_checkpoint(
 	}
 }
 
+/// `POST /_tuwunel/replication/promote`
+///
+/// Promotes this secondary to a standalone primary by stopping the replication
+/// worker. After this call returns the instance accepts writes and no longer
+/// tails the primary's WAL. The caller is responsible for updating the VIP or
+/// load balancer to route client traffic to this node.
+///
+/// Returns:
+/// - `200 OK` with `{"status":"promoted"}` on success.
+/// - `409 Conflict` if this instance is already a primary (no `rocksdb_primary_url`
+///   was configured, or it was already promoted).
+pub(crate) async fn replication_promote(
+	State(services): State<crate::State>,
+) -> impl IntoResponse {
+	if services.replication.is_promoted() {
+		return (
+			StatusCode::CONFLICT,
+			axum::Json(serde_json::json!({"error": "already promoted"})),
+		)
+			.into_response();
+	}
+
+	if services.server.config.rocksdb_primary_url.is_none() {
+		return (
+			StatusCode::CONFLICT,
+			axum::Json(serde_json::json!({"error": "not a secondary; no rocksdb_primary_url configured"})),
+		)
+			.into_response();
+	}
+
+	services.replication.promote();
+
+	axum::Json(serde_json::json!({"status": "promoted"})).into_response()
+}
+
+/// Request body for `POST /_tuwunel/replication/demote`.
+#[derive(Debug, Deserialize, Serialize)]
+pub(crate) struct DemoteBody {
+	/// URL of the new primary to replicate from (e.g. `http://host:8008`).
+	pub primary_url: String,
+}
+
+/// `POST /_tuwunel/replication/demote`
+///
+/// Demotes this promoted primary back to a secondary that replicates from
+/// `primary_url`. Resets the resume cursor and triggers a fresh checkpoint
+/// bootstrap from the new primary — the worker restarts replication without
+/// requiring a process restart.
+///
+/// Typical use case: the original primary comes back online after a failover
+/// and needs to re-join the cluster as a secondary under the newly promoted
+/// node.
+///
+/// Returns:
+/// - `200 OK` with `{"status":"demoted","primary_url":"..."}` on success.
+/// - `400 Bad Request` if `primary_url` is missing or empty.
+/// - `409 Conflict` if this instance is not currently promoted (i.e. it is
+///   already actively replicating or was never a secondary).
+pub(crate) async fn replication_demote(
+	State(services): State<crate::State>,
+	Json(body): Json<DemoteBody>,
+) -> impl IntoResponse {
+	if body.primary_url.is_empty() {
+		return (
+			StatusCode::BAD_REQUEST,
+			axum::Json(serde_json::json!({"error": "primary_url is required"})),
+		)
+			.into_response();
+	}
+
+	match services.replication.demote(body.primary_url.clone()).await {
+		| Ok(()) => axum::Json(serde_json::json!({
+			"status": "demoted",
+			"primary_url": body.primary_url,
+		}))
+		.into_response(),
+		| Err(e) => (
+			StatusCode::CONFLICT,
+			axum::Json(serde_json::json!({"error": e.to_string()})),
+		)
+			.into_response(),
+	}
+}
+
 /// Creates a temporary directory that is automatically removed on drop.
 ///
 /// We use a simple wrapper around `std::fs::create_dir_all` on a

src/api/client/replication.rs

@@ -27,7 +27,7 @@ pub(super) use self::{
 };
 use crate::{client, server};
 
-pub fn build(router: Router<State>, server: &Server) -> Router<State> {
+pub fn build(router: Router<State>, server: &Server, state: State) -> Router<State> {
 	let config = &server.config;
 	let mut router = router
         .ruma_route(&client::get_timezone_key_route)
@@ -263,7 +263,15 @@ pub fn build(router: Router<State>, server: &Server) -> Router<State> {
 				"/_tuwunel/replication/checkpoint",
 				get(client::replication_checkpoint),
 			)
-			.layer(middleware::from_fn(check_replication_token)),
+			.route(
+				"/_tuwunel/replication/promote",
+				post(client::replication_promote),
+			)
+			.route(
+				"/_tuwunel/replication/demote",
+				post(client::replication_demote),
+			)
+			.layer(middleware::from_fn_with_state(state, check_replication_token)),
 	);
 
 	if config.allow_legacy_media {

src/api/router.rs

@@ -1,6 +1,6 @@
 use axum::{
 	body::Body,
-	extract::Request,
+	extract::{Request, State},
 	http::StatusCode,
 	middleware::Next,
 	response::{IntoResponse, Response},
@@ -16,10 +16,11 @@ pub(super) const TOKEN_HEADER: &str = "x-tuwunel-replication-token";
 /// - `401 Unauthorized` if the token is missing or incorrect.
 /// - Passes through to the handler if the token matches.
 pub(crate) async fn check_replication_token(
-	axum::extract::State(services): axum::extract::State<super::State>,
+	State(services): State<super::State>,
 	request: Request<Body>,
 	next: Next,
 ) -> Response {
+
 	let Some(ref expected) = services.server.config.rocksdb_replication_token else {
 		return (
 			StatusCode::NOT_IMPLEMENTED,

src/api/router/replication_auth.rs

@@ -44,7 +44,7 @@ pub(crate) async fn open(ctx: Arc<Context>, desc: &[Descriptor]) -> Result<Arc<S
 	} else if config.rocksdb_secondary {
 		let secondary_path = config
 			.rocksdb_secondary_path
-			.as_deref()
+			.as_ref()
 			.unwrap_or(path);
 		Db::open_cf_descriptors_as_secondary(&db_opts, path, secondary_path, cfds)
 	} else {

src/database/engine/open.rs

@@ -8,7 +8,7 @@ use std::{
 	time::{SystemTime, UNIX_EPOCH},
 };
 
-use rocksdb::Checkpoint;
+use rocksdb::checkpoint::Checkpoint;
 use tuwunel_core::{Err, Result, implement};
 
 use super::Engine;
@@ -219,6 +219,35 @@ pub fn wal_updates_since(&self, since: u64) -> Result<rocksdb::DBWALIterator> {
 	self.db.get_updates_since(since).map_err(map_err)
 }
 
+/// Newtype wrapper making `DBWALIterator` safe to send across threads.
+///
+/// `DBWALIterator` holds a `*mut rocksdb_wal_iterator_t` raw pointer which
+/// is not auto-`Send`. RocksDB WAL iterators are not concurrently shared;
+/// this iterator is consumed by exactly one thread at a time, so sending
+/// ownership across a thread boundary is safe.
+struct SendWalIter(rocksdb::DBWALIterator);
+
+// SAFETY: DBWALIterator is not auto-Send due to its raw pointer, but the
+// underlying RocksDB iterator is safe to use from whichever single thread
+// owns it at any given time. We never share it across threads simultaneously.
+unsafe impl Send for SendWalIter {}
+
+impl Iterator for SendWalIter {
+	type Item = Result<WalFrame>;
+
+	fn next(&mut self) -> Option<Self::Item> {
+		self.0.next().map(|result| {
+			result
+				.map(|(seq, batch)| {
+					let data = batch.data().to_vec();
+					let count = batch_count_from_bytes(&data);
+					WalFrame::data(seq, count, data)
+				})
+				.map_err(map_err)
+		})
+	}
+}
+
 /// Return a higher-level iterator of [`WalFrame`]s starting at `since`.
 ///
 /// Wraps `wal_updates_since` and maps each rocksdb batch into a `WalFrame`,
@@ -233,15 +262,7 @@ pub fn wal_frame_iter(
 	since: u64,
 ) -> Result<Box<dyn Iterator<Item = Result<WalFrame>> + Send>> {
 	let iter = self.db.get_updates_since(since).map_err(map_err)?;
-	Ok(Box::new(iter.map(|result| {
-		result
-			.map(|(seq, batch)| {
-				let data = batch.data().to_vec();
-				let count = batch_count_from_bytes(&data);
-				WalFrame::data(seq, count, data)
-			})
-			.map_err(map_err)
-	})))
+	Ok(Box::new(SendWalIter(iter)))
 }
 
 /// Returns `true` if `err` indicates the requested WAL sequence is older

src/database/engine/replication.rs

@@ -10,7 +10,7 @@ use tuwunel_service::Services;
 pub(crate) fn build(services: &Arc<Services>) -> (Router, Guard) {
 	let router = Router::<state::State>::new();
 	let (state, guard) = state::create(services.clone());
-	let router = tuwunel_api::router::build(router, &services.server)
+	let router = tuwunel_api::router::build(router, &services.server, state)
 		.route("/", get(it_works))
 		.fallback(not_found)
 		.with_state(state);