diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 65bf668..c5d6e27 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -231,10 +231,10 @@ replication is a primary-to-replica stream built on top of the same persistence 1. a replica connects and performs a handshake 2. the primary sends one in-memory snapshot per shard plus the current offset -3. the replica loads those snapshots -4. the primary streams incremental `AofRecord` updates +3. the replica loads those snapshots and acknowledges their offsets +4. the primary streams incremental `AofRecord` updates, and the replica acknowledges them as it applies them -successful mutations inside a shard emit `ReplicationEvent`s with monotonically increasing per-shard offsets. the replication server uses those offsets for streaming and for the `WAIT` command's acknowledgement tracking. +each record a shard writes gets the next per-shard offset. a shard collects its records during a burst of requests and sends them as one `ReplicationEvent` on a broadcast channel, so the shards lock the channel once per burst. for each replica, the primary moves those events into a queue of its own, up to 256 MiB. records written during a full sync wait there until the replica has loaded it. a replica that falls further behind is told to resync and starts over with a full sync. the replica applies records in batches, one message per shard. the acknowledged offsets drive the `WAIT` command, which counts a replica only once it has loaded its full sync. ### slot migration diff --git a/CHANGELOG.md b/CHANGELOG.md index 52363b9..8891d5c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,9 +11,15 @@ - raft and the `openraft` dependency. raft only ever ran on the bootstrap node, and gossip with config epochs already decided slot ownership and failover, as in Redis Cluster. `raft-port-offset` and `--cluster-raft-port-offset` are still accepted and ignored, so existing configs keep working (#443) ### performance +- replication keeps up with the primary under heavy writes. on a 32-core machine a replica loads a full sync of 1.26M keys in about 1 second instead of 13.5, and stays in step with 1.1M writes a second. a replica that fell behind, as every one did under load, was resynced from scratch in a loop and served a partial dataset each time (#455): + - the primary queues each replica's records itself, up to 256 MiB, so writes during a full sync no longer overrun the 65,536-event broadcast channel. a replica further behind gets a resync at once instead of first receiving the stale backlog + - shards send their records once per burst instead of once per record. every send locks the broadcast channel, which held the primary to about 600k writes a second with a replica attached + - the replica applies records and snapshot keys in batches, one message per shard, instead of waiting for each one - cluster mode is about as fast as standalone without pipelining: about 860k requests a second against 320k on a 32-core machine. with a pipeline of 16 it serves about 2.3M against 320k, 15% below standalone. every command locked the migration state, an async mutex shared by every connection, and read the cluster state up to three times. it now reads each once, and the migration state takes a shared lock (#453) ### fixes +- WAIT counts a replica only once it has loaded its full sync. it counted the replica as soon as the primary had sent the snapshot, while the replica was still loading it, or had just flushed its data to start over (#455) +- `INFO replication` on a primary counts the replicas streaming from it in `connected_replicas`. it counted the replicas gossip had reported, which could read 0 while one streamed. a new `sync_full` field counts the full syncs sent (#455) - failover safety (#452): - an election needs a majority of all primaries, the failed one included, as in Redis Cluster. it counted only the other primaries, so with two primaries a replica cut off with one of them won with a single vote while its primary kept serving - a primary votes for one replica per failed primary for two failure timeouts, so two replicas of the same primary can't both win. a replica stops its election once another node owns the failed primary's slots diff --git a/crates/ember-core/src/shard/aof.rs b/crates/ember-core/src/shard/aof.rs index 0703c0b..492d735 100644 --- a/crates/ember-core/src/shard/aof.rs +++ b/crates/ember-core/src/shard/aof.rs @@ -886,24 +886,60 @@ impl AofLog { } } -/// Broadcasts a single replication event, used by blocking pop operations -/// that bypass the normal process_message path. -pub(super) fn broadcast_replication( - record: AofRecord, - replication_tx: &Option>, - replication_offset: &AtomicU64, +/// Collects a shard's replication records during a burst and sends them +/// to the replicas as one event at its end. +/// +/// Every shard sends into one broadcast channel, and each send locks it, +/// so one send per record caps the primary's write rate. +pub(super) struct ReplicationLog { + tx: broadcast::Sender, + /// Offset of the last record taken. Shared with the `ShardHandle`, so + /// WAIT can read it as soon as the write is applied. + offset: Arc, + pending: Vec, shard_id: u16, -) { - if let Some(ref tx) = *replication_tx { - let offset = replication_offset.fetch_add(1, Ordering::Release) + 1; - if tx - .send(ReplicationEvent { - shard_id, - offset, - record, - }) - .is_err() - { +} + +impl ReplicationLog { + pub(super) fn new( + tx: broadcast::Sender, + offset: Arc, + shard_id: u16, + ) -> Self { + Self { + tx, + offset, + pending: Vec::new(), + shard_id, + } + } + + /// The offset of the last record taken. + pub(super) fn offset(&self) -> u64 { + self.offset.load(Ordering::Acquire) + } + + /// Takes records applied by a write, numbering each with the next + /// offset. + pub(super) fn push(&mut self, records: impl IntoIterator) { + for record in records { + self.offset.fetch_add(1, Ordering::Release); + self.pending.push(record); + } + } + + /// Sends the records taken since the last call as one event. + pub(super) fn flush(&mut self) { + if self.pending.is_empty() { + return; + } + let records: Arc<[AofRecord]> = std::mem::take(&mut self.pending).into(); + let event = ReplicationEvent { + shard_id: self.shard_id, + offset: self.offset() + 1 - records.len() as u64, + records, + }; + if self.tx.send(event).is_err() { // no replicas are currently connected — normal during startup // or after a replica disconnects, but tracked so operators // can alert on unexpected drops in a replicated deployment. diff --git a/crates/ember-core/src/shard/blocking.rs b/crates/ember-core/src/shard/blocking.rs index 06bb788..2f78924 100644 --- a/crates/ember-core/src/shard/blocking.rs +++ b/crates/ember-core/src/shard/blocking.rs @@ -115,12 +115,9 @@ fn pop_for( if let Some(ref mut aof) = *ctx.aof { aof.append(std::slice::from_ref(&record)); } - aof::broadcast_replication( - record, - ctx.replication_tx, - ctx.replication_offset, - ctx.shard_id, - ); + if let Some(ref mut replication) = *ctx.replication { + replication.push([record]); + } true } diff --git a/crates/ember-core/src/shard/mod.rs b/crates/ember-core/src/shard/mod.rs index b8909a7..c62f620 100644 --- a/crates/ember-core/src/shard/mod.rs +++ b/crates/ember-core/src/shard/mod.rs @@ -78,19 +78,20 @@ use crate::types::sorted_set::{ScoreBound, ZAddFlags}; use crate::types::Value; use ember_protocol::command::{BitOpKind, BitRange}; -/// A mutation event broadcast to replication subscribers. +/// Mutations broadcast to replication subscribers, one event per burst of +/// requests a shard handles. /// -/// Published after every successful mutation on the hot path. The -/// `offset` is per-shard and monotonically increasing — replicas use it -/// to detect gaps and trigger re-sync when they fall behind. +/// Each record has a per-shard offset that increases by one per record. +/// The first record's offset is `offset`, the next one `offset + 1`, and +/// so on. #[derive(Debug, Clone)] pub struct ReplicationEvent { /// The shard that produced this event. pub shard_id: u16, - /// Monotonically increasing per-shard offset. + /// Offset of the first record. pub offset: u64, - /// The mutation record, ready to replay on a replica. - pub record: AofRecord, + /// The mutation records, in order, ready to replay on a replica. + pub records: Arc<[AofRecord]>, } /// Optional persistence configuration for a shard. @@ -312,6 +313,8 @@ async fn run_shard(prepared: PreparedShard) { schema_registry, } = prepared; let shard_id = config.shard_id; + let mut replication = + replication_tx.map(|tx| aof::ReplicationLog::new(tx, replication_offset, shard_id)); // tokio's interval panics on zero, so hold both to at least 1ms let expiry_interval = config.expiry_interval.max(Duration::from_millis(1)); let fsync_interval = config.fsync_interval.max(Duration::from_millis(1)); @@ -432,8 +435,7 @@ async fn run_shard(prepared: PreparedShard) { persistence: &persistence, drop_handle: &drop_handle, shard_id, - replication_tx: &replication_tx, - replication_offset: &replication_offset, + replication: &mut replication, lpop_waiters: &mut lpop_waiters, rpop_waiters: &mut rpop_waiters, #[cfg(feature = "protobuf")] @@ -444,6 +446,9 @@ async fn run_shard(prepared: PreparedShard) { if let Some(ref mut aof) = aof { aof.flush(); } + if let Some(ref mut replication) = replication { + replication.flush(); + } for (reply, response) in held_replies.drain(..) { reply.send(response); } @@ -487,10 +492,8 @@ struct ProcessCtx<'a> { persistence: &'a Option, drop_handle: &'a Option, shard_id: u16, - replication_tx: &'a Option>, - /// Offset of the last replication event sent. Shared with the - /// `ShardHandle` so WAIT can read it. - replication_offset: &'a AtomicU64, + /// Where replication records wait for the end of the burst. + replication: &'a mut Option, /// Waiters for BLPOP — keyed by list name, FIFO order. lpop_waiters: &'a mut HashMap>>, /// Waiters for BRPOP — keyed by list name, FIFO order. @@ -666,16 +669,8 @@ fn process_single(mut request: ShardRequest, reply: ReplySender, ctx: &mut Proce } // broadcast mutation events to replication subscribers - if let Some(ref tx) = *ctx.replication_tx { - for record in records { - let offset = ctx.replication_offset.fetch_add(1, Ordering::Release) + 1; - // ignore send errors — no subscribers or lagged consumers - let _ = tx.send(ReplicationEvent { - shard_id, - offset, - record, - }); - } + if let Some(ref mut replication) = *ctx.replication { + replication.push(records); } if let Some(key) = pushed_key { @@ -685,9 +680,10 @@ fn process_single(mut request: ShardRequest, reply: ReplySender, ctx: &mut Proce // handle special requests that need access to persistence state match request_kind { RequestKind::SerializeSnapshot => { - // the snapshot covers every event up to the current offset, - // since this task sends all of them and is busy doing this now - let offset = ctx.replication_offset.load(Ordering::Acquire); + // the snapshot covers every record up to the current offset, + // including those not yet sent: they go out later with offsets + // the replica knows to skip + let offset = ctx.replication.as_ref().map_or(0, |r| r.offset()); let resp = persistence::handle_serialize_snapshot(ctx.keyspace, shard_id, offset); ctx.respond(reply, resp); return; diff --git a/crates/ember-server/src/cluster/replication.rs b/crates/ember-server/src/cluster/replication.rs index fe37e4e..fb4d038 100644 --- a/crates/ember-server/src/cluster/replication.rs +++ b/crates/ember-server/src/cluster/replication.rs @@ -11,8 +11,6 @@ pub struct ReplicationInfo { pub role: NodeRole, /// Address of the primary this node replicates from (replica only). pub primary_addr: Option, - /// Number of connected replicas (primary only). - pub replica_count: usize, } impl ClusterCoordinator { @@ -136,8 +134,6 @@ impl ClusterCoordinator { } /// Returns replication status for the `INFO replication` section. - /// - /// Returns `(role, Option)`. pub async fn replication_info(&self) -> ReplicationInfo { let state = self.state.read().await; let local = state.nodes.get(&self.local_id); @@ -150,15 +146,6 @@ impl ClusterCoordinator { } else { None }; - let replica_count = if role == NodeRole::Primary { - local.map(|n| n.replicas.len()).unwrap_or(0) - } else { - 0 - }; - ReplicationInfo { - role, - primary_addr, - replica_count, - } + ReplicationInfo { role, primary_addr } } } diff --git a/crates/ember-server/src/connection/exec/server.rs b/crates/ember-server/src/connection/exec/server.rs index 9320ba6..02b5eee 100644 --- a/crates/ember-server/src/connection/exec/server.rs +++ b/crates/ember-server/src/connection/exec/server.rs @@ -379,7 +379,14 @@ async fn render_info( match info.role { NodeRole::Primary => { out.push_str("role:primary\r\n"); - out.push_str(&format!("connected_replicas:{}\r\n", info.replica_count)); + out.push_str(&format!( + "connected_replicas:{}\r\n", + ctx.replica_tracker.connected_count() + )); + out.push_str(&format!( + "sync_full:{}\r\n", + ctx.replica_tracker.full_syncs() + )); } NodeRole::Replica => { out.push_str("role:replica\r\n"); diff --git a/crates/ember-server/src/replication.rs b/crates/ember-server/src/replication.rs index cff02ab..0e9b9de 100644 --- a/crates/ember-server/src/replication.rs +++ b/crates/ember-server/src/replication.rs @@ -30,31 +30,32 @@ //! // Incremental records (unbounded stream): //! [MSG_RECORD: 1B][shard_id: 2B][offset: 8B][record_len: 4B][record_bytes] //! -//! // Replica → primary, after applying each record: +//! // Replica → primary, after loading the snapshots and after applying +//! // records. offset is the newest one applied on that shard: //! [MSG_ACK: 1B][shard_id: 2B][offset: 8B] //! -//! // When replica falls behind (broadcast lag): +//! // When the replica falls too far behind (see REPLICA_QUEUE_LIMIT): //! [MSG_RESYNC: 1B] primary closes the connection; replica reconnects //! ``` //! -//! Offsets count replication events per shard. A snapshot's offset is the -//! last event it contains, so the replica skips records at or below it. +//! Offsets count replication records per shard. A snapshot's offset is the +//! last record it contains, so the replica skips records at or below it. use std::collections::HashMap; use std::io; use std::net::SocketAddr; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::Duration; -use bytes::Bytes; +use bytes::{BufMut, Bytes, BytesMut}; use ember_cluster::ClusterSecret; use ember_core::{Engine, ShardRequest, ShardResponse}; use ember_persistence::aof::AofRecord; use ember_persistence::snapshot; -use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufWriter}; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader, BufWriter}; use tokio::net::{TcpListener, TcpStream}; -use tokio::sync::broadcast; +use tokio::sync::{broadcast, mpsc}; use tracing::{debug, error, info, warn}; // -- protocol constants -- @@ -72,10 +73,19 @@ const MSG_RECORD: u8 = 4; const MSG_RESYNC: u8 = 5; const MSG_ACK: u8 = 6; +/// The most record bytes the primary queues for one replica. Writes made +/// while the replica loads its full sync wait in this queue. A replica +/// that falls further behind is sent MSG_RESYNC. +const REPLICA_QUEUE_LIMIT: usize = 256 * 1024 * 1024; + +/// The most requests a replica collects before sending them to its shards. +const MAX_BATCH: usize = 4096; + /// Tracks each replica's acknowledged offset per shard for WAIT. /// -/// A replica starts at the offsets of the snapshot it received and moves -/// forward as it acknowledges records. WAIT compares these against the +/// A replica starts at offset 0 on every shard. It acknowledges the +/// snapshots' offsets once it has loaded them, then each record it +/// applies. WAIT compares these against the /// shards' current offsets from [`Engine::replication_offsets`]. #[derive(Debug)] pub struct ReplicaTracker { @@ -83,6 +93,8 @@ pub struct ReplicaTracker { /// assigned when the replica connects. acked: Mutex>>, next_id: AtomicU64, + /// Full syncs started since the server started. + full_syncs: AtomicU64, } impl ReplicaTracker { @@ -90,15 +102,16 @@ impl ReplicaTracker { Self { acked: Mutex::new(HashMap::new()), next_id: AtomicU64::new(0), + full_syncs: AtomicU64::new(0), } } - /// Registers a replica that loaded snapshots at `snapshot_offsets`. - /// Returns its id. - pub fn register(&self, snapshot_offsets: Vec) -> u64 { + /// Registers a replica of a primary with `shards` shards. Returns its + /// id. + pub fn register(&self, shards: usize) -> u64 { let id = self.next_id.fetch_add(1, Ordering::Relaxed); if let Ok(mut map) = self.acked.lock() { - map.insert(id, snapshot_offsets); + map.insert(id, vec![0; shards]); } id } @@ -132,6 +145,16 @@ impl ReplicaTracker { .unwrap_or(0) } + /// Counts a full sync sent to a replica. + pub fn full_sync_started(&self) { + self.full_syncs.fetch_add(1, Ordering::Relaxed); + } + + /// Full syncs started since the server started. + pub fn full_syncs(&self) -> u64 { + self.full_syncs.load(Ordering::Relaxed) + } + /// Returns the total number of currently connected replicas. pub fn connected_count(&self) -> usize { self.acked.lock().map(|map| map.len()).unwrap_or(0) @@ -343,17 +366,20 @@ impl ReplicationServer { // subscribe before snapshotting so we don't miss events that // happen between snapshot and incremental stream start - let mut rx = match self.engine.subscribe_replication() { - Some(rx) => rx, - None => { - return Err(std::io::Error::other( - "replication channel not configured on this engine", - )); - } + let Some(rx) = self.engine.subscribe_replication() else { + return Err(std::io::Error::other( + "replication channel not configured on this engine", + )); }; + // move events into this replica's own queue as they happen. the + // broadcast ring holds a fixed number of events, too few to cover a + // full sync under write load + let (queue_tx, mut queue) = mpsc::unbounded_channel(); + let queued = Arc::new(AtomicUsize::new(0)); + let forwarder = tokio::spawn(forward_events(rx, queue_tx, Arc::clone(&queued))); // --- full sync --- - let mut snapshot_offsets = Vec::with_capacity(self.engine.shard_count()); + self.tracker.full_sync_started(); for shard_idx in 0..self.engine.shard_count() { let resp = self .engine @@ -388,7 +414,6 @@ impl ReplicationServer { write_u8(&mut stream, MSG_SHARD_OFFSET).await?; write_u16_le(&mut stream, shard_id).await?; write_u64_le(&mut stream, offset).await?; - snapshot_offsets.push(offset); } stream.flush().await?; @@ -401,7 +426,8 @@ impl ReplicationServer { let mut writer = BufWriter::with_capacity(65536, write_inner); // Register this replica and spawn an ACK reader task. - let replica_id = self.tracker.register(snapshot_offsets); + // WAIT counts the replica once it acknowledges having loaded these + let replica_id = self.tracker.register(self.engine.shard_count()); let tracker = Arc::clone(&self.tracker); let mut ack_reader = read_half; @@ -426,52 +452,86 @@ impl ReplicationServer { }); // --- incremental stream --- - let result = self.stream_records(&mut writer, &mut rx).await; + let result = stream_records(&mut writer, &mut queue, &queued).await; // clean up regardless of result + forwarder.abort(); ack_task.abort(); self.tracker.remove(replica_id); result } +} - /// Streams replication records to the replica until the connection closes - /// or the broadcast channel is exhausted. - async fn stream_records( - &self, - writer: &mut BufWriter>, - rx: &mut broadcast::Receiver, - ) -> std::io::Result<()> { - loop { - match rx.recv().await { - Ok(event) => { - let record_bytes = event.record.to_bytes().map_err(|e| { - std::io::Error::other(format!("record serialization failed: {e}")) - })?; - let record_len = u32::try_from(record_bytes.len()).map_err(|_| { - std::io::Error::new(std::io::ErrorKind::InvalidData, "record too large") - })?; - - write_u8(writer, MSG_RECORD).await?; - write_u16_le(writer, event.shard_id).await?; - write_u64_le(writer, event.offset).await?; - write_u32_le(writer, record_len).await?; - writer.write_all(&record_bytes).await?; - writer.flush().await?; - } - Err(broadcast::error::RecvError::Lagged(count)) => { - warn!("replication stream lagged by {count} events; triggering resync"); - let _ = write_u8(writer, MSG_RESYNC).await; - let _ = writer.flush().await; - return Ok(()); - } - Err(broadcast::error::RecvError::Closed) => { - info!("replication broadcast channel closed; disconnecting replica"); - return Ok(()); - } +/// Encodes the records of each replication event as MSG_RECORD frames and +/// queues them for one replica. Stops, closing the queue, when the replica +/// falls more than [`REPLICA_QUEUE_LIMIT`] bytes behind or the broadcast +/// ring overruns. +async fn forward_events( + mut rx: broadcast::Receiver, + queue: mpsc::UnboundedSender, + queued: Arc, +) { + loop { + let event = match rx.recv().await { + Ok(event) => event, + Err(broadcast::error::RecvError::Lagged(count)) => { + warn!("replication stream lagged by {count} events; triggering resync"); + return; } + Err(broadcast::error::RecvError::Closed) => return, + }; + let mut frames = BytesMut::new(); + for (offset, record) in (event.offset..).zip(event.records.iter()) { + let bytes = match record.to_bytes() { + Ok(bytes) => bytes, + Err(e) => { + error!("replication record serialization failed: {e}"); + return; + } + }; + let Ok(len) = u32::try_from(bytes.len()) else { + error!("replication record too large: {} bytes", bytes.len()); + return; + }; + frames.put_u8(MSG_RECORD); + frames.put_u16_le(event.shard_id); + frames.put_u64_le(offset); + frames.put_u32_le(len); + frames.put_slice(&bytes); + } + + let size = frames.len(); + if queued.fetch_add(size, Ordering::Relaxed) + size > REPLICA_QUEUE_LIMIT { + warn!("replica is over {REPLICA_QUEUE_LIMIT} bytes behind; triggering resync"); + return; + } + if queue.send(frames.freeze()).is_err() { + return; + } + } +} + +/// Writes queued records to the replica, flushing whenever the queue runs +/// dry. Once the forwarder closes the queue, the rest of it is stale: it +/// sends MSG_RESYNC instead, so the replica starts over with a full sync. +async fn stream_records( + writer: &mut BufWriter>, + queue: &mut mpsc::UnboundedReceiver, + queued: &AtomicUsize, +) -> std::io::Result<()> { + while let Some(frame) = queue.recv().await { + if queue.is_closed() { + break; + } + queued.fetch_sub(frame.len(), Ordering::Relaxed); + writer.write_all(&frame).await?; + if queue.is_empty() { + writer.flush().await?; } } + write_u8(writer, MSG_RESYNC).await?; + writer.flush().await } // -- replica-side client -- @@ -540,7 +600,9 @@ impl ReplicationClient { /// Performs full sync + incremental stream for one connection session. /// Resets `backoff` once the handshake succeeds, so a replica that ran /// fine for hours does not wait the maximum delay after one drop. - async fn sync(&self, mut stream: TcpStream, backoff: &mut Duration) -> std::io::Result<()> { + async fn sync(&self, stream: TcpStream, backoff: &mut Duration) -> std::io::Result<()> { + // reads come through a buffer; writes go straight to the socket + let mut stream = BufReader::with_capacity(65536, stream); let our_shards = self.engine.shard_count() as u16; // send handshake @@ -599,6 +661,7 @@ impl ReplicationClient { .map_err(|e| std::io::Error::other(format!("flush before sync failed: {e:?}")))?; // receive per-shard snapshots and the offset each one covers + let mut batch = Batch::new(usize::from(primary_shards)); let mut snapshot_offsets = vec![0u64; usize::from(primary_shards)]; let mut snap_buf = Vec::new(); for _ in 0..primary_shards { @@ -609,7 +672,7 @@ impl ReplicationClient { snap_buf.clear(); snap_buf.resize(snap_len, 0); stream.read_exact(&mut snap_buf).await?; - self.apply_snapshot(shard_id, &snap_buf).await?; + self.apply_snapshot(&mut batch, shard_id, &snap_buf).await?; expect_tag(&mut stream, MSG_SHARD_OFFSET, "MSG_SHARD_OFFSET").await?; let offset_shard = read_u16_le(&mut stream).await?; @@ -618,10 +681,19 @@ impl ReplicationClient { *slot = offset; } } + batch.apply(&self.engine).await; + let mut acks = Vec::new(); + for (shard, offset) in snapshot_offsets.iter().enumerate() { + acks.push(MSG_ACK); + acks.extend_from_slice(&(shard as u16).to_le_bytes()); + acks.extend_from_slice(&offset.to_le_bytes()); + } + stream.write_all(&acks).await?; info!("full sync applied, starting incremental replay"); - // incremental stream + // incremental stream. records are collected into a batch, which + // is applied and acknowledged once the buffered input is used up loop { let msg = read_u8(&mut stream).await?; match msg { @@ -649,20 +721,10 @@ impl ReplicationClient { // the handshake checked that both sides have the same shard // count, so the primary's shard id picks the same shard a // key lookup would, and it also covers keyless records - if !in_snapshot { - if let Some(request) = ember_core::from_aof_record(&record) { - if let Err(e) = - self.engine.send_to_shard(shard_id.into(), request).await - { - warn!("replication apply failed: {e:?}"); - } - } - } - - // acknowledge the record so WAIT can count this replica - write_u8(&mut stream, MSG_ACK).await?; - write_u16_le(&mut stream, shard_id).await?; - write_u64_le(&mut stream, offset).await?; + let request = (!in_snapshot) + .then(|| ember_core::from_aof_record(&record)) + .flatten(); + batch.push(shard_id, request, Some(offset)); } MSG_RESYNC => { info!("primary requested resync; reconnecting"); @@ -675,6 +737,18 @@ impl ReplicationClient { )); } } + + // acknowledge once the records are applied, so WAIT can count + // this replica + if stream.buffer().is_empty() || batch.is_full() { + let mut acks = Vec::new(); + for (shard, offset) in batch.apply(&self.engine).await { + acks.push(MSG_ACK); + acks.extend_from_slice(&shard.to_le_bytes()); + acks.extend_from_slice(&offset.to_le_bytes()); + } + stream.write_all(&acks).await?; + } } } @@ -682,7 +756,12 @@ impl ReplicationClient { /// /// Each entry goes in through RESTORE, which replaces the key with the /// snapshot's value as is, for every value type. - async fn apply_snapshot(&self, shard_id: u16, data: &[u8]) -> std::io::Result<()> { + async fn apply_snapshot( + &self, + batch: &mut Batch, + shard_id: u16, + data: &[u8], + ) -> std::io::Result<()> { let (_, entries) = snapshot::read_snapshot_from_bytes(data).map_err(|e| { std::io::Error::new( std::io::ErrorKind::InvalidData, @@ -708,14 +787,83 @@ impl ReplicationClient { data: Bytes::from(data), replace: true, }; - if let Err(e) = self.engine.send_to_shard(shard_id.into(), request).await { - warn!("snapshot restore failed: {e:?}"); + batch.push(shard_id, Some(request), None); + if batch.is_full() { + batch.apply(&self.engine).await; } } Ok(()) } } +/// Requests a replica has collected for its shards, and the offsets to +/// acknowledge once they are applied. +/// +/// Each shard gets its requests as one message, and the shards apply them +/// in parallel. +struct Batch { + requests: Vec>, + /// The newest offset collected for each shard. + acks: Vec>, + len: usize, +} + +impl Batch { + fn new(shards: usize) -> Self { + Self { + requests: (0..shards).map(|_| Vec::new()).collect(), + acks: vec![None; shards], + len: 0, + } + } + + /// Adds `request` for `shard`, if there is one, and `offset` to + /// acknowledge once it is applied. + fn push(&mut self, shard: u16, request: Option, offset: Option) { + let shard = usize::from(shard); + if let (Some(request), Some(requests)) = (request, self.requests.get_mut(shard)) { + requests.push(request); + self.len += 1; + } + if let Some(ack) = self.acks.get_mut(shard) { + *ack = offset.or(*ack); + } + } + + fn is_full(&self) -> bool { + self.len >= MAX_BATCH + } + + /// Sends the requests to their shards and waits until all are applied. + /// Returns the `(shard, offset)` pairs this covers. + async fn apply(&mut self, engine: &Engine) -> Vec<(u16, u64)> { + let mut replies = Vec::with_capacity(self.len); + for (shard, requests) in self.requests.iter_mut().enumerate() { + if requests.is_empty() { + continue; + } + match engine + .dispatch_batch_to_shard(shard, std::mem::take(requests)) + .await + { + Ok(batch) => replies.extend(batch), + Err(e) => warn!("replication apply failed: {e:?}"), + } + } + self.len = 0; + for reply in replies { + if reply.await.is_err() { + warn!("replication apply failed: shard dropped the request"); + } + } + self.acks + .iter_mut() + .enumerate() + .filter_map(|(shard, ack)| Some((shard as u16, ack.take()?))) + .collect() + } +} + #[cfg(test)] mod tests { use super::*; @@ -723,8 +871,11 @@ mod tests { #[test] fn tracker_counts_replicas_that_reached_every_shard() { let tracker = ReplicaTracker::new(); - let a = tracker.register(vec![3, 0]); - let b = tracker.register(vec![3, 0]); + let a = tracker.register(2); + let b = tracker.register(2); + // both loaded snapshots at offsets [3, 0] + tracker.update(a, 0, 3); + tracker.update(b, 0, 3); tracker.update(a, 1, 2); // a replica that is ahead on one shard but behind on another does not count tracker.update(b, 0, 9); diff --git a/tests/integration/src/cluster.rs b/tests/integration/src/cluster.rs index c18de63..5abfde3 100644 --- a/tests/integration/src/cluster.rs +++ b/tests/integration/src/cluster.rs @@ -1026,3 +1026,83 @@ fn cluster_on_a_public_address_needs_a_cluster_secret() { let stderr = String::from_utf8_lossy(&output.stderr); assert!(stderr.contains("cluster-auth-pass"), "{stderr}"); } + +/// Sends the commands `make(0)`, `make(1)`, … `make(count - 1)` in +/// pipelined chunks and checks that none fails. +async fn pipeline( + c: &mut crate::helpers::TestClient, + count: usize, + make: impl Fn(usize) -> Vec, +) { + for start in (0..count).step_by(1000) { + let end = (start + 1000).min(count); + let mut out = Vec::new(); + for i in start..end { + let args = make(i); + out.extend_from_slice(format!("*{}\r\n", args.len()).as_bytes()); + for arg in args { + out.extend_from_slice(format!("${}\r\n{arg}\r\n", arg.len()).as_bytes()); + } + } + c.write_raw(&out).await; + for _ in start..end { + let reply = c.read_frame().await; + assert!(!matches!(reply, Frame::Error(_)), "{reply:?}"); + } + } +} + +#[tokio::test] +async fn a_replica_keeps_up_with_writes_during_its_full_sync() { + use std::time::{Duration, Instant}; + + let opts = || ServerOptions { + cluster_enabled: true, + ..Default::default() + }; + // a broadcast channel smaller than the writes made during the full + // sync, which a replica that waited on it would miss + let primary = TestServer::start_with(ServerOptions { + cluster_bootstrap: true, + replication_broadcast_capacity: Some(1024), + ..opts() + }); + let replica = TestServer::start_with(opts()); + let mut c0 = primary.connect().await; + let mut c1 = replica.connect().await; + + // enough data that the full sync takes a while + pipeline(&mut c0, 50_000, |i| { + vec!["SET".into(), format!("key:{i}"), "x".repeat(64)] + }) + .await; + + let primary_id = c0.get_bulk(&["CLUSTER", "MYID"]).await.unwrap(); + c0.ok(&["CLUSTER", "MEET", "127.0.0.1", &replica.port.to_string()]) + .await; + let deadline = Instant::now() + Duration::from_secs(30); + while !matches!( + c1.cmd(&["CLUSTER", "REPLICATE", &primary_id]).await, + Frame::Simple(_) + ) { + assert!( + Instant::now() <= deadline, + "replica never learned the primary" + ); + tokio::time::sleep(Duration::from_millis(250)).await; + } + + // writes while the replica loads the full sync. they wait for it in + // the replica's queue + pipeline(&mut c0, 100_000, |_| vec!["INCR".into(), "counter".into()]).await; + + assert_eq!(c0.get_int(&["WAIT", "1", "30000"]).await, 1); + assert_eq!( + c1.get_bulk(&["GET", "counter"]).await, + Some("100000".into()) + ); + assert_eq!(c1.get_int(&["DBSIZE"]).await, 50_001); + let info = c0.get_bulk(&["INFO", "replication"]).await.unwrap(); + assert!(info.contains("connected_replicas:1"), "{info}"); + assert!(info.contains("sync_full:1"), "the replica resynced: {info}"); +} diff --git a/tests/integration/src/helpers.rs b/tests/integration/src/helpers.rs index eb3e6c9..bd99f51 100644 --- a/tests/integration/src/helpers.rs +++ b/tests/integration/src/helpers.rs @@ -59,6 +59,9 @@ pub struct ServerOptions { /// Start the gRPC listener on a free port. When false, gRPC is disabled /// so test servers don't compete for the default port. pub grpc: bool, + /// Events the replication broadcast channel holds + /// (`replication-broadcast-capacity`, set through a config file). + pub replication_broadcast_capacity: Option, } impl TestServer { @@ -128,6 +131,15 @@ impl TestServer { if let Some(secs) = opts.idle_timeout_secs { cmd.arg("--idle-timeout-secs").arg(secs.to_string()); } + if let Some(capacity) = opts.replication_broadcast_capacity { + let path = std::env::temp_dir().join(format!("ember-test-{port}.toml")); + std::fs::write( + &path, + format!("[engine]\nreplication-broadcast-capacity = {capacity}\n"), + ) + .expect("failed to write the test config"); + cmd.arg("--config").arg(path); + } let data_dir = if opts.appendonly || opts.cluster_enabled { if opts.appendonly {