//! PostgreSQL logical replication stream management
//!
//! This module provides high-level management of logical replication streams,
//! including connection management, slot creation, and message processing.
//!
//! ## Async Stream API
//!
//! The `EventStream` and `EventStreamRef` types provide a native `.next().await` API
//! for simple usage without requiring trait imports. This design prioritizes ergonomics
//! and ease of use.
//!
//! If you need to use this with the `futures::Stream` trait for stream combinators,
//! you can easily wrap it using `futures::stream::unfold`. See the EventStream
//! documentation for an example.
use crate::error::{ReplicationError, Result};
use crate::lsn::SharedLsnFeedback;
use crate::types::{ChangeEvent, EventType, Lsn, ReplicationSlotOptions, SlotType, TimestampTz};
use crate::{
format_lsn, parse_keepalive_message, LogicalReplicationParser, PgReplicationConnection,
ReplicationConnectionRetry, ReplicationState, RetryConfig, XLogRecPtr, INVALID_XLOG_REC_PTR,
};
// `RelationInfo` is referenced only by the unit tests in this file, via `super::*`.
#[cfg(test)]
use crate::RelationInfo;
use bytes::{Buf, Bytes};
use std::sync::Arc;
use std::future::Future;
use std::time::{Duration, Instant};
use tokio_util::sync::CancellationToken;
use tracing::{debug, error, info, warn};
/// PostgreSQL logical replication stream
pub struct LogicalReplicationStream {
connection: PgReplicationConnection,
parser: LogicalReplicationParser,
state: ReplicationState,
config: ReplicationStreamConfig,
/// What this stream knows about its replication slot. See [`SlotState`].
slot: SlotState,
retry_handler: ReplicationConnectionRetry,
last_health_check: Instant,
/// Shared LSN feedback for communication with consumer. This allows the consumer to update flushed/applied LSN after commits
pub shared_lsn_feedback: Arc<SharedLsnFeedback>,
/// `IDENTIFY_SYSTEM`'s identity from the first successful connection, re-checked on every reconnect. See [`ServerIdentity`].
identity: Option<ServerIdentity>,
/// Counter that throttles per-event time checks inside `next_event`. We only consult `Instant::now()` once every `FEEDBACK_CHECK_EVENT_INTERVAL` events so the hot path avoids a syscall on every message.
feedback_check_counter: u32,
/// Bounded replay terminal. See [`ReplayStop`].
replay_stop: ReplayStop,
}
/// What this stream knows about the replication slot it is using.
///
/// Replaces a `slot_created: bool` plus two loose `Option` fields. The bool
/// could not carry the distinction that matters: it was set to `true` both when
/// *we* created the slot and when `CREATE_REPLICATION_SLOT` failed with "already
/// exists", so nothing downstream could tell whose slot it was. That is the
/// difference between dropping our own failed attempt and destroying a slot the
/// operator created — `discard_replication_slot` now asks the type instead of
/// trusting the caller to have checked an unrelated field first.
///
/// Snapshot state is deliberately *nested inside* ownership rather than sitting
/// beside it: an exported snapshot exists only because we created the slot, and
/// modelling it as a sibling field made "both or neither" an unchecked
/// convention.
#[derive(Debug, Default)]
enum SlotState {
/// No slot yet, or ours was discarded after a failed snapshot.
#[default]
Absent,
/// A slot with this name already existed. We did not create it, so we must never drop it, and it exported no snapshot — that is the ordinary "resume an existing subscription" path.
PreExisting,
/// We created it in this session.
Ours {
/// The LSN the exported snapshot is consistent at. Also the position
/// `START_REPLICATION` resumes from, so it is the LSN stamped on
/// snapshot rows.
///
/// `None` only when the server returned an unparsable value, which is deliberately non-fatal: it makes the snapshot helper unavailable, notreplication.
consistent_point: Option<Lsn>,
/// The name exported by `EXPORT_SNAPSHOT`.
///
/// Cleared on a successful `START_REPLICATION`, which destroys the snapshot server-side — *independently* of `consistent_point`, which outlives it and stays valid as the resume position. Those differing lifetimes are why this is an inner `Option` rather than a second variant.
exported_snapshot: Option<String>,
},
}
impl SlotState {
/// True once a usable slot exists, however it got there. This is the old
/// `slot_created` predicate, and the only question `ensure_replication_slot`
/// needs to answer.
#[inline]
fn is_present(&self) -> bool {
!matches!(self, SlotState::Absent)
}
/// The consistent point, if this stream created the slot.
#[inline]
fn consistent_point(&self) -> Option<Lsn> {
match self {
SlotState::Ours {
consistent_point, ..
} => *consistent_point,
_ => None,
}
}
/// The exported snapshot name, while it is still alive server-side.
#[inline]
fn exported_snapshot(&self) -> Option<&str> {
match self {
SlotState::Ours {
exported_snapshot, ..
} => exported_snapshot.as_deref(),
_ => None,
}
}
/// Forget the exported snapshot name, keeping everything else.
///
/// Called after a successful `START_REPLICATION`: the server has run
/// `SnapBuildClearExportedSnapshot`, so the name is a handle to nothing.
#[inline]
fn clear_exported_snapshot(&mut self) {
if let SlotState::Ours {
exported_snapshot, ..
} = self
{
*exported_snapshot = None;
}
}
}
/// How far bounded replay (`stop_at_lsn`) has progressed.
///
/// Replaces `stop_at_reached: Option<Lsn>` + `copy_done_sent: bool`, which
/// between them could spell `(None, true)` — "CopyDone already sent for a
/// terminal that was never armed". Nothing rejected it: `copy_done_sent` was
/// read at exactly one site, nested inside the `Some` check, so the pairing was
/// an unstated rule spread across three arming sites.
///
/// As an enum the one-shot guard *is* the transition: `Armed -> Stopped` can
/// happen only once, because `Stopped` no longer matches `Armed`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ReplayStop {
/// Bounded replay is off, or the target has not been reached yet.
Running,
/// The terminal has been reached; the next `next_event` owes the server a
/// CopyDone. Holds the crossing commit's `end_lsn` on the typed path, a
/// keepalive's `wal_end` when the target is past the last commit, and the
/// transport-level `wal_end` on the raw path.
Armed(Lsn),
/// CopyDone has been sent. Every later call reports `StreamStopped` again.
Stopped(Lsn),
}
impl ReplayStop {
/// The LSN the terminal was armed at, once it has been reached.
#[inline]
fn reached(self) -> Option<Lsn> {
match self {
ReplayStop::Running => None,
ReplayStop::Armed(lsn) | ReplayStop::Stopped(lsn) => Some(lsn),
}
}
/// Arm the terminal, unless it is already armed or stopped.
#[inline]
fn arm(&mut self, lsn: Lsn) {
if matches!(self, ReplayStop::Running) {
*self = ReplayStop::Armed(lsn);
}
}
}
/// A raw XLogData (`'w'`) message: the undecoded pgoutput payload plus its WAL
/// positions. Bring your own decoder.
///
/// Returned by [`LogicalReplicationStream::next_raw_event`]. `data` is a zero-copy slice of the connection read buffer and is valid until the next call to `next_raw_event` or `next_event`.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RawXLogData {
/// Server WAL start position for this message (`start_lsn`).
pub wal_start: Lsn,
/// WAL end position — the next byte after this message. Ack with this.
pub wal_end: Lsn,
/// Server send time (Postgres-epoch microseconds).
pub server_time: TimestampTz,
/// Undecoded pgoutput message bytes (everything after the 25-byte header).
pub data: Bytes,
}
/// How often `next_event` consults `Instant::now()` to decide whether to send
/// feedback. Must be a power of two so the modulo folds to a bitmask.
const FEEDBACK_CHECK_EVENT_INTERVAL: u32 = 128;
/// How often `next_event_with_retry` performs the connection health check.
/// Must be a power of two. The health check itself is time-gated, but its
/// `Instant::now()` call is a vDSO syscall, so we additionally amortize it
/// across this many events on the hot path.
const HEALTH_CHECK_EVENT_INTERVAL: u32 = 1024;
/// The `systemid` and `timeline` reported by `IDENTIFY_SYSTEM`.
///
/// Captured on the first connection and compared after every reconnect. Both
/// values come from a row the connection already fetches, so the check costs
/// nothing extra.
#[derive(Debug, Clone, PartialEq, Eq)]
struct ServerIdentity {
system_id: String,
timeline: u32,
}
impl ServerIdentity {
/// Read the first row of an `IDENTIFY_SYSTEM` result.
///
/// `None` when a column is missing or the timeline is not a number — an
/// unexpected result shape disables the check rather than failing the
/// connection, since the check is a safety net, not a requirement.
fn from_identify_system(result: &crate::connection::PgResult) -> Option<Self> {
Some(Self {
system_id: result.get_value(0, 0)?,
timeline: result.get_value(0, 1)?.parse().ok()?,
})
}
}
/// Configuration for the replication stream
///
/// `#[non_exhaustive]`: build one with [`builder`](Self::builder) or
/// [`new`](Self::new) plus the `with_*` setters, not with a struct literal.
/// Closing literal construction is what makes adding a field a non-breaking
/// change — otherwise every new field breaks every downstream literal again.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct ReplicationStreamConfig {
pub slot_name: String,
/// Publication to replicate from.
///
/// To stream from several publications at once, pass them comma-separated with each name quoted — `r#"pub_a","pub_b"#` — which produces the `publication_names '"pub_a","pub_b"'` list pgoutput expects. One stream and one slot cover all of them.
pub publication_name: String,
pub protocol_version: u32,
pub streaming_mode: StreamingMode,
pub messages: bool,
pub binary: bool,
pub two_phase: bool,
pub origin: Option<OriginFilter>,
pub feedback_interval: Duration,
pub connection_timeout: Duration,
pub health_check_interval: Duration,
pub retry_config: RetryConfig,
pub slot_options: ReplicationSlotOptions,
pub slot_type: SlotType,
/// Bounded replay: stop after the transaction whose commit `end_lsn` reaches this LSN. `None` (the default) streams unbounded.
pub stop_at_lsn: Option<Lsn>,
}
/// Streaming mode for logical replication (pgoutput)
///
/// From PostgreSQL docs:
/// - `off` is the default
/// - `on` enables streaming of in-progress transactions (protocol v2+)
/// - `parallel` enables extra per-message data for parallelization (protocol v4+)
///
/// See: <https://www.postgresql.org/docs/current/protocol-logical-replication.html>
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StreamingMode {
Off,
On,
Parallel,
}
impl StreamingMode {
#[inline]
pub fn as_str(&self) -> &'static str {
match self {
StreamingMode::Off => "off",
StreamingMode::On => "on",
StreamingMode::Parallel => "parallel",
}
}
}
/// Origin filter option for logical replication
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OriginFilter {
None,
Any,
}
impl OriginFilter {
#[inline]
pub fn as_str(&self) -> &'static str {
match self {
OriginFilter::None => "none",
OriginFilter::Any => "any",
}
}
}
impl ReplicationStreamConfig {
/// Create a new replication stream configuration
///
/// # Arguments
///
/// * `slot_name` - Name of the replication slot to use (will be created if it doesn't exist)
/// * `publication_name` - Name of the PostgreSQL publication to replicate from
/// * `protocol_version` - Protocol version (1-4). Version 2+ supports streaming transactions
/// * `streaming_mode` - Streaming mode for logical decoding (off/on/parallel).
/// `on` requires protocol version >= 2; `parallel` requires protocol version >= 4.
/// * `feedback_interval` - How often to send status feedback to PostgreSQL
/// * `connection_timeout` - Maximum time to wait for connection establishment
/// * `health_check_interval` - How often to check connection health
/// * `retry_config` - Configuration for connection retry behavior
///
/// # Example
///
/// ```
/// use pg_walstream::{ReplicationStreamConfig, RetryConfig, StreamingMode};
/// use std::time::Duration;
///
/// let config = ReplicationStreamConfig::new(
/// "my_slot".to_string(),
/// "my_publication".to_string(),
/// 2,
/// StreamingMode::On,
/// Duration::from_secs(10),
/// Duration::from_secs(30),
/// Duration::from_secs(60),
/// RetryConfig::default(),
/// );
/// ```
#[allow(clippy::too_many_arguments)]
pub fn new(
slot_name: String,
publication_name: String,
protocol_version: u32,
streaming_mode: StreamingMode,
feedback_interval: Duration,
connection_timeout: Duration,
health_check_interval: Duration,
retry_config: RetryConfig,
) -> Self {
Self {
slot_name,
publication_name,
protocol_version,
streaming_mode,
messages: false,
binary: false,
two_phase: false,
origin: None,
feedback_interval,
connection_timeout,
health_check_interval,
retry_config,
slot_options: ReplicationSlotOptions {
snapshot: Some("nothing".to_string()),
..Default::default()
},
slot_type: SlotType::Logical,
stop_at_lsn: None,
}
}
/// Enable or disable logical replication messages (pg_logical_emit_message)
#[inline]
pub fn with_messages(mut self, enabled: bool) -> Self {
self.messages = enabled;
self
}
/// Enable or disable binary mode for pgoutput
#[inline]
pub fn with_binary(mut self, enabled: bool) -> Self {
self.binary = enabled;
self
}
/// Enable or disable two-phase commit decoding (protocol v3+)
///
/// Two-phase decoding is **two** wire options, not one: `TWO_PHASE` on
/// `CREATE_REPLICATION_SLOT` and `two_phase 'on'` on `START_REPLICATION`.
/// The server accepts the latter on a slot created without the former
/// *without an error* — it simply never decodes `PREPARE`, so the failure is
/// silent. This setter therefore sets both.
///
/// Call it **after** [`with_slot_options`](Self::with_slot_options), which
/// replaces the whole options struct and would otherwise clobber the flag.
#[inline]
pub fn with_two_phase(mut self, enabled: bool) -> Self {
self.two_phase = enabled;
self.slot_options.two_phase = enabled;
self
}
/// Set origin filtering behavior (none|any)
#[inline]
pub fn with_origin(mut self, origin: Option<OriginFilter>) -> Self {
self.origin = origin;
self
}
/// Set streaming mode (off|on|parallel)
#[inline]
pub fn with_streaming_mode(mut self, mode: StreamingMode) -> Self {
self.streaming_mode = mode;
self
}
/// Set replication slot creation options
///
/// This allows customizing how the replication slot is created, including:
/// - `temporary`: Create a temporary slot (dropped when connection ends)
/// - `snapshot`: Control snapshot export behavior (`"export"`, `"use"`, `"nothing"`)
/// - `two_phase`: Enable two-phase commit support for the slot
/// - `failover`: Enable failover synchronization
///
/// # Replaces, does not merge
///
/// This assigns the whole struct, so it **clobbers
/// [`with_two_phase`](Self::with_two_phase)** if called afterwards — that
/// setter writes `slot_options.two_phase`, and the failure is silent (the
/// slot is simply created without `TWO_PHASE`, and `PREPARE` is never
/// decoded). Call `with_two_phase` last, or set `two_phase` in the struct
/// here.
///
/// The `snapshot` field is the exception: `snapshot()` sets `"export"` itself
/// immediately before creating the slot, so whatever is written here cannot
/// stop an initial snapshot from working.
///
/// # Example: Temporary slot with exported snapshot
///
/// ```
/// use pg_walstream::{ReplicationStreamConfig, ReplicationSlotOptions, RetryConfig, StreamingMode};
/// use std::time::Duration;
///
/// let config = ReplicationStreamConfig::new(
/// "my_temp_slot".to_string(),
/// "my_publication".to_string(),
/// 2,
/// StreamingMode::On,
/// Duration::from_secs(10),
/// Duration::from_secs(30),
/// Duration::from_secs(60),
/// RetryConfig::default(),
/// ).with_slot_options(ReplicationSlotOptions {
/// temporary: true,
/// snapshot: Some("export".to_string()),
/// ..Default::default()
/// });
/// ```
#[inline]
pub fn with_slot_options(mut self, options: ReplicationSlotOptions) -> Self {
self.slot_options = options;
self
}
/// Set the replication slot type (default: `SlotType::Logical`)
///
/// Use `SlotType::Physical` for physical replication slots that stream raw WAL data without logical decoding.
#[inline]
pub fn with_slot_type(mut self, slot_type: SlotType) -> Self {
self.slot_type = slot_type;
self
}
}
impl Default for ReplicationStreamConfig {
fn default() -> Self {
Self {
slot_name: String::new(),
publication_name: String::new(),
protocol_version: 2,
streaming_mode: StreamingMode::Off,
messages: false,
binary: false,
two_phase: false,
origin: None,
feedback_interval: Duration::from_secs(10),
connection_timeout: Duration::from_secs(30),
health_check_interval: Duration::from_secs(60),
retry_config: RetryConfig::default(),
slot_options: ReplicationSlotOptions {
snapshot: Some("nothing".to_string()),
..Default::default()
},
slot_type: SlotType::Logical,
stop_at_lsn: None,
}
}
}
impl ReplicationStreamConfig {
/// Start a config from the two required fields; everything else defaults.
/// Non-breaking alternative to the 8-arg [`new`](Self::new).
#[inline]
pub fn builder(slot: impl Into<String>, publication: impl Into<String>) -> Self {
Self {
slot_name: slot.into(),
publication_name: publication.into(),
..Default::default()
}
}
/// Set the logical replication protocol version (1-4).
#[inline]
pub fn with_protocol_version(mut self, v: u32) -> Self {
self.protocol_version = v;
self
}
/// Set how often status feedback is sent to PostgreSQL.
#[inline]
pub fn with_feedback_interval(mut self, d: Duration) -> Self {
self.feedback_interval = d;
self
}
/// Set the connection establishment timeout.
#[inline]
pub fn with_connection_timeout(mut self, d: Duration) -> Self {
self.connection_timeout = d;
self
}
/// Set how often connection health is checked.
#[inline]
pub fn with_health_check_interval(mut self, d: Duration) -> Self {
self.health_check_interval = d;
self
}
/// Set the connection retry configuration.
#[inline]
pub fn with_retry_config(mut self, r: RetryConfig) -> Self {
self.retry_config = r;
self
}
/// Stop replication after the first committed transaction whose commit `end_lsn` reaches `lsn`. That transaction is delivered in full, then the stream ends cleanly (client CopyDone) and `next_event` returns [`crate::ReplicationError::StreamStopped`]. `None` (default) is unbounded.
///
/// If `lsn` lies *beyond* the last commit the server has, no commit boundary can ever reach it, so the stream instead ends once a keepalive reports the server's send position past `lsn` — rather than blocking forever waiting for a transaction that may never be written. This mirrors pg_recvlogical's `--endpos` `STREAM_STOP_KEEPALIVE` exit and is safe for the same reason: a keepalive's `wal_end` is `sentPtr`, assigned only *after* the output plugin has run and queued its data on the same socket, so nothing committing at or below it can still be undelivered. The `StreamStopped` LSN is then that keepalive position, not a commit boundary.
#[inline]
pub fn with_stop_at_lsn(mut self, lsn: impl Into<Lsn>) -> Self {
self.stop_at_lsn = Some(lsn.into());
self
}
}
/// Bounded replay: given the configured `stop_at`, return `Some(end_lsn)` when
/// `event` is a committed-transaction boundary (`Commit` / `StreamCommit` /
/// `CommitPrepared`) whose `end_lsn` reaches the target, else `None`.
///
/// Pure and allocation-free. The `stop_at?` short-circuits with zero cost when
/// bounded replay is off, and the match only ever fires on commit boundaries,
/// so the DML hot path is untouched.
#[inline]
fn reached_stop_lsn(stop_at: Option<Lsn>, event: &ChangeEvent) -> Option<Lsn> {
let target = stop_at?;
let end_lsn = match &event.event_type {
EventType::Commit { end_lsn, .. }
| EventType::StreamCommit { end_lsn, .. }
| EventType::CommitPrepared { end_lsn, .. } => *end_lsn,
_ => return None,
};
(end_lsn >= target).then_some(end_lsn)
}
impl LogicalReplicationStream {
/// Create a new logical replication stream
///
/// This establishes a connection to PostgreSQL and prepares the stream for replication.
/// It does not create the replication slot or start replication - call `start()` for that.
///
/// The stream automatically creates a shared LSN feedback tracker accessible via the
/// `shared_lsn_feedback` field. The consumer should use this to update flushed/applied
/// LSN values after committing data to the destination. This allows the stream to send
/// accurate feedback to PostgreSQL, which is crucial for WAL retention management.
///
/// # Arguments
///
/// * `connection_string` - PostgreSQL connection string. Must include `replication=database`
/// parameter. Example: `"postgresql://user:pass@host:5432/dbname?replication=database"`
/// * `config` - Replication stream configuration
///
/// # Returns
///
/// A new `LogicalReplicationStream` instance with an initialized LSN feedback tracker.
///
/// # Errors
///
/// Returns an error if:
/// - Connection to PostgreSQL fails
/// - Connection string is invalid
/// - PostgreSQL version is too old (< 14.0)
/// - Authentication fails
///
/// # Example
///
/// ```no_run
/// use pg_walstream::{LogicalReplicationStream, ReplicationStreamConfig, RetryConfig, StreamingMode};
/// use std::time::Duration;
///
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let config = ReplicationStreamConfig::new(
/// "my_slot".to_string(),
/// "my_publication".to_string(),
/// 2,
/// StreamingMode::On,
/// Duration::from_secs(10),
/// Duration::from_secs(30),
/// Duration::from_secs(60),
/// RetryConfig::default(),
/// );
///
/// let mut stream = LogicalReplicationStream::new(
/// "postgresql://postgres:password@localhost:5432/mydb?replication=database",
/// config,
/// ).await?;
///
/// // Access LSN feedback directly
/// stream.shared_lsn_feedback.update_applied_lsn(12345);
/// # Ok(())
/// # }
/// ```
pub async fn new(connection_string: &str, config: ReplicationStreamConfig) -> Result<Self> {
info!("Creating logical replication stream with retry support");
let retry_handler = ReplicationConnectionRetry::new(
config.retry_config,
with_connect_timeout(connection_string, config.connection_timeout),
);
// Establish initial connection with retry and enforce timeout
let connection = timeout_or_error(
config.connection_timeout,
retry_handler.connect_with_retry(),
)
.await?;
let parser = LogicalReplicationParser::with_protocol_version(config.protocol_version);
let state = ReplicationState::new();
let last_health_check = Instant::now();
// Create shared LSN feedback for consumer
let shared_lsn_feedback = SharedLsnFeedback::new_shared();
Ok(Self {
connection,
parser,
state,
config,
slot: SlotState::Absent,
retry_handler,
last_health_check,
shared_lsn_feedback,
identity: None,
feedback_check_counter: 0,
replay_stop: ReplayStop::Running,
})
}
/// Initialize the replication stream
async fn initialize(&mut self) -> Result<()> {
info!("Initializing replication stream");
// Identify the system, and remember who it is: a later reconnect that
// lands on a different cluster or a forked timeline is rejected.
let identity = self.connection.identify_system()?;
self.identity = ServerIdentity::from_identify_system(&identity);
info!("System identification successful");
// Create replication slot if it doesn't exist
self.ensure_replication_slot().await?;
info!("Replication stream initialized");
Ok(())
}
/// Ensure the replication slot exists, creating it if necessary.
///
/// When the slot is created with `snapshot: Some("export")` (i.e. `EXPORT_SNAPSHOT`),
/// PostgreSQL returns a snapshot identifier that can be retrieved via
/// [`exported_snapshot_name()`](Self::exported_snapshot_name). This snapshot can be
/// imported by a **separate regular connection** using
/// `SET TRANSACTION SNAPSHOT '<snapshot_id>'` to perform a consistent read of the
/// database at the slot's starting LSN.
///
/// **Important:** The exported snapshot is only valid between this call and
/// [`start()`](Self::start). Once `start()` issues `START_REPLICATION`, PostgreSQL
/// destroys the exported snapshot. You **must** read the snapshot on a separate
/// connection before calling `start()`.
///
/// If the slot already exists, this is a no-op.
///
/// # Initial Snapshot Workflow
///
/// ```text
/// Connection 1 (replication): Connection 2 (snapshot read):
/// ensure_replication_slot()
/// → CREATE_REPLICATION_SLOT ...
/// EXPORT_SNAPSHOT
/// → snapshot_id available
/// BEGIN TRANSACTION
/// ISOLATION LEVEL REPEATABLE READ;
/// SET TRANSACTION SNAPSHOT '<snapshot_id>';
/// -- read initial table state
/// COMMIT;
/// start(None)
/// → START_REPLICATION ...
pub async fn ensure_replication_slot(&mut self) -> Result<()> {
if self.slot.is_present() {
return Ok(());
}
info!("Creating replication slot: {}", self.config.slot_name);
let output_plugin = match self.config.slot_type {
SlotType::Logical => Some("pgoutput"),
SlotType::Physical => None,
};
match self.connection.create_replication_slot_with_options(
&self.config.slot_name,
self.config.slot_type,
output_plugin,
&self.config.slot_options,
) {
Ok(result) => {
// CREATE_REPLICATION_SLOT returns: slot_name(0), consistent_point(1), snapshot_name(2), output_plugin(3)
//
// The consistent point is the LSN the exported snapshot is taken
// at, and equals the slot's `confirmed_flush` — so it is both the
// LSN to stamp on snapshot rows and the position `start(None)`
// resumes from. A malformed value is not fatal: it only means the
// initial-snapshot helper is unavailable, not that replication is.
let exported_snapshot = result
.get_value(0, 2)
.filter(|name| !name.is_empty())
.inspect(|name| info!("Exported snapshot name: {}", name));
info!("Replication slot created successfully");
self.slot = SlotState::Ours {
consistent_point: parse_consistent_point(&result),
exported_snapshot,
};
}
Err(e) => {
let error_msg = e.to_string();
if error_msg.contains("already exists") {
warn!("Replication slot already exists, continuing");
// NOT `Ours`: we did not create it, so it exported no snapshot and we must never drop it.
self.slot = SlotState::PreExisting;
} else {
return Err(e);
}
}
}
Ok(())
}
/// Start the replication stream
///
/// This initializes the replication slot (creating it if necessary) and begins
/// streaming changes from PostgreSQL. ///
/// # Arguments
///
/// * `start_lsn` - Optional LSN to start replication from. If `None`, starts from
/// the current WAL position. Use this to resume replication from a known position.
///
/// # Returns
///
/// Returns `Ok(())` if replication started successfully.
///
/// # Errors
///
/// Returns an error if:
/// - System identification fails
/// - Replication slot creation fails (if it doesn't exist)
/// - Starting replication command fails
///
/// # Example
///
/// ```no_run
/// use pg_walstream::{LogicalReplicationStream, ReplicationStreamConfig, RetryConfig, StreamingMode};
/// use std::time::Duration;
///
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let config = ReplicationStreamConfig::new(
/// "my_slot".to_string(),
/// "my_publication".to_string(),
/// 2, StreamingMode::On,
/// Duration::from_secs(10),
/// Duration::from_secs(30),
/// Duration::from_secs(60),
/// RetryConfig::default(),
/// );
///
/// let mut stream = LogicalReplicationStream::new("connection_string", config).await?;
///
/// // Start from the beginning
/// stream.start(None).await?;
///
/// // Or resume from a specific LSN
/// // stream.start(Some(0x16B374D848)).await?;
/// # Ok(())
/// # }
/// ```
pub async fn start(&mut self, start_lsn: Option<XLogRecPtr>) -> Result<()> {
info!("Starting logical replication stream");
self.initialize().await?;
let start_lsn = start_lsn.unwrap_or(INVALID_XLOG_REC_PTR);
// Seed `applied` with the caller's startpoint: `start(Some(X))` asserts
// everything below X is processed, which is what `applied` means. Without
// it, `recover_connection` (which resumes from `applied`) would rewind an
// explicit startpoint to `confirmed_flush`. Monotonic and 0-ignoring, so
// `start(None)` is a no-op and a later ack past X wins.
self.shared_lsn_feedback.update_applied_lsn(start_lsn);
let options = self.build_replication_options()?;
// Start replication
let options_ref = options
.iter()
.map(|(k, v)| (k.as_str(), v.as_str()))
.collect::<Vec<_>>();
self.connection
.start_replication(&self.config.slot_name, start_lsn, &options_ref)?;
self.slot.clear_exported_snapshot();
info!(
"Logical replication started with LSN: {}",
format_lsn(start_lsn)
);
Ok(())
}
/// Pump the COPY stream until the next XLogData (`'w'`) frame arrives.
///
/// Handles cancellation, the bounded-replay terminal (CopyDone-once + `StreamStopped`), throttled feedback, and keepalives internally. Returns the full `'w'` CopyData buffer (1-byte tag + 24-byte header + payload).
/// Shared by [`next_event`](Self::next_event) and [`next_raw_event`](Self::next_raw_event).
async fn next_wal_frame(&mut self, cancellation_token: &CancellationToken) -> Result<Bytes> {
loop {
// Check cancellation before processing.
if cancellation_token.is_cancelled() {
return Err(ReplicationError::Cancelled(
"Operation cancelled".to_string(),
));
}
// Bounded replay terminal: once a stop boundary has been delivered, end the COPY stream cleanly (once) and report StreamStopped on this and every subsequent call. Runs before any socket read so no further WAL is consumed past the target.
match self.replay_stop {
ReplayStop::Running => {}
ReplayStop::Armed(reached) => {
// The `Armed -> Stopped` transition IS the one-shot guard:
debug!("Bounded replay: reached stop_at_lsn at {reached}, sending CopyDone");
self.replay_stop = ReplayStop::Stopped(reached);
if let Err(e) = self.connection.end_copy().await {
warn!("Failed to send CopyDone at stop_at_lsn: {e}");
}
return Err(ReplicationError::stream_stopped(reached));
}
ReplayStop::Stopped(reached) => {
return Err(ReplicationError::stream_stopped(reached));
}
}
// Throttle feedback checks: only consult `Instant::now()` and the
// shared atomic once every FEEDBACK_CHECK_EVENT_INTERVAL iterations.
self.feedback_check_counter = self.feedback_check_counter.wrapping_add(1);
if self.feedback_check_counter & (FEEDBACK_CHECK_EVENT_INTERVAL - 1) == 0 {
self.maybe_send_feedback().await;
}
// Get data from replication stream (returns Bytes for zero-copy).
let data = self
.connection
.get_copy_data_async(cancellation_token)
.await?;
if data.is_empty() {
continue;
}
match data[0] {
b'w' => return Ok(data),
b'k' => {
// Keepalive message — handled internally.
self.process_keepalive_message(&data).await?;
}
other => {
warn!("Received unknown message type: {}", other as char);
}
}
}
}
/// Process the next single replication event with cancellation support
///
/// This is the core method for retrieving individual change events from PostgreSQL.
/// It handles WAL messages, keepalive messages, and automatically sends feedback
/// to PostgreSQL when appropriate.
///
/// For production use, consider using `next_event_with_retry()` which adds automatic
/// retry and recovery logic, or use `into_stream()` for an iterator-like interface.
///
/// # Arguments
/// * `cancellation_token` - Optional cancellation token to abort the operation
///
/// # Returns
/// * `Ok(event)` - Successfully received a change event
/// * `Err(ReplicationError::Cancelled(_))` - Operation was cancelled
/// * `Err(_)` - Connection or protocol errors occurred
///
/// # Example
///
/// ```no_run
/// use pg_walstream::{LogicalReplicationStream, ReplicationStreamConfig, RetryConfig, StreamingMode};
/// use tokio_util::sync::CancellationToken;
/// use std::time::Duration;
///
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let config = ReplicationStreamConfig::new(
/// "my_slot".to_string(),
/// "my_publication".to_string(),
/// 2, StreamingMode::On,
/// Duration::from_secs(10),
/// Duration::from_secs(30),
/// Duration::from_secs(60),
/// RetryConfig::default(),
/// );
///
/// let mut stream = LogicalReplicationStream::new("connection_string", config).await?;
/// stream.start(None).await?;
///
/// let cancel_token = CancellationToken::new();
///
/// // Simple polling loop
/// loop {
/// match stream.next_event(&cancel_token).await {
/// Ok(event) => {
/// println!("Event: {:?}", event);
/// }
/// Err(_) => break,
/// }
/// }
/// # Ok(())
/// # }
/// ```
pub async fn next_event(
&mut self,
cancellation_token: &CancellationToken,
) -> Result<ChangeEvent> {
// Loop until we get a valid event or encounter an error/cancellation
loop {
let data = self.next_wal_frame(cancellation_token).await?;
// 'w' WAL data message (zero-copy: Bytes slicing avoids copies).
if let Some(event) = self.process_wal_message(data)? {
// Bounded replay: if this event is a commit boundary that reaches
// stop_at_lsn, remember it — the full crossing transaction is
// still delivered now; the NEXT call terminates.
if let Some(reached) = reached_stop_lsn(self.config.stop_at_lsn, &event) {
self.replay_stop.arm(reached);
}
return Ok(event);
}
// Control message with no ChangeEvent (e.g. Relation) → next frame.
}
}
/// Like [`next_event`](Self::next_event), but returns the **undecoded** XLogData payload instead of a parsed `ChangeEvent`. Keepalives, feedback throttling, cancellation, and the bounded-replay terminal are handled identically; only pgoutput decoding is skipped — bring your own decoder.
///
/// Ack progress after durable processing by advancing the applied LSN: `stream.shared_lsn_feedback.update_applied_lsn(raw.wal_end.value())`. There is no auto-ack and no retry/recovery on this path (that is the point — you own restart semantics), so wrap reconnection yourself if you need it.
///
/// In raw mode `config.stop_at_lsn` is a **transport-level** cutoff: the stream stops after the first message whose `wal_end >= stop_at_lsn`, which may fall mid-transaction. Terminates with [`ReplicationError::StreamStopped`].
///
/// **Do not interleave** `next_raw_event` with `next_event` on the same stream: raw mode bypasses the parser's relation cache.
///
/// # Example
///
/// ```no_run
/// use pg_walstream::{LogicalReplicationStream, ReplicationStreamConfig, RetryConfig, StreamingMode};
/// use tokio_util::sync::CancellationToken;
/// use std::time::Duration;
///
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let config = ReplicationStreamConfig::new(
/// "my_slot".to_string(),
/// "my_publication".to_string(),
/// 2, StreamingMode::On,
/// Duration::from_secs(10),
/// Duration::from_secs(30),
/// Duration::from_secs(60),
/// RetryConfig::default(),
/// );
/// let mut stream = LogicalReplicationStream::new("connection_string", config).await?;
/// stream.start(None).await?;
/// let cancel = CancellationToken::new();
///
/// loop {
/// match stream.next_raw_event(&cancel).await {
/// Ok(raw) => {
/// // Decode `raw.data` (pgoutput bytes) yourself, then ack:
/// stream.shared_lsn_feedback.update_applied_lsn(raw.wal_end.value());
/// }
/// Err(e) if matches!(
/// e,
/// pg_walstream::ReplicationError::Cancelled(_)
/// | pg_walstream::ReplicationError::StreamStopped(_)
/// ) => break,
/// Err(e) => return Err(e.into()),
/// }
/// }
/// # Ok(())
/// # }
/// ```
pub async fn next_raw_event(
&mut self,
cancellation_token: &CancellationToken,
) -> Result<RawXLogData> {
let data = self.next_wal_frame(cancellation_token).await?;
let raw = self.decode_raw_frame(data)?;
// Transport-level bounded replay: arm the terminal when this message's wal_end reaches the configured stop LSN. The next `next_wal_frame` call sends CopyDone once and returns StreamStopped. Armed here (not in the decoder) so the decoder stays policy-free and the typed path can reuse it; mirrors how `next_event` arms the commit-boundary stop.
if let Some(stop) = self.config.stop_at_lsn {
if raw.wal_end >= stop {
self.replay_stop.arm(raw.wal_end);
}
}
Ok(raw)
}
/// Check connection health and attempt recovery if needed
pub async fn check_connection_health(&mut self) -> Result<()> {
let now = Instant::now();
if now.duration_since(self.last_health_check) < self.config.health_check_interval {
return Ok(()); // Skip health check if not enough time has passed
}
self.last_health_check = now;
debug!("Performing connection health check");
if !self.connection.is_alive() {
warn!("Connection health check failed, attempting recovery");
match self.recover_connection().await {
Ok(_) => {
info!("Connection recovered successfully");
}
Err(e) => {
error!("Failed to recover connection: {}", e);
return Err(e);
}
}
} else {
debug!("Connection health check passed");
}
Ok(())
}
/// Reject a reconnect that landed on a different cluster or a forked timeline.
///
/// An LSN only identifies a WAL record within one cluster on one timeline. After a promotion (failover, switchover, or point-in-time recovery) the server runs on a new timeline whose WAL diverges from the old one at the switchpoint; records the consumer already saw past that point no longer exist. Resuming silently would stream a different history under the same LSNs, and nothing downstream could tell.
///
/// A promoted standby keeps the primary's `systemid`, so the timeline is the only signal for a failover; the `systemid` check catches the separate mistake of pointing the consumer at an unrelated cluster that happens to have a slot by the same name.
///
/// Returns [`ReplicationError::ReplicationSlot`], which [`is_permanent`](ReplicationError::is_permanent) reports as `true`, so the retry loop stops instead of reconnecting into the fork. Deciding whether the stored LSN survived requires the switchpoint from `TIMELINE_HISTORY <tli>`, which is a policy call left to the caller.
fn check_server_identity(&self, current: &ServerIdentity) -> Result<()> {
let Some(previous) = self.identity.as_ref() else {
return Ok(());
};
if previous.system_id != current.system_id {
return Err(ReplicationError::replication_slot(format!(
"reconnected to a different cluster: system identifier was {}, now {}; \
the stored LSN belongs to the old cluster",
previous.system_id, current.system_id
)));
}
if previous.timeline != current.timeline {
return Err(ReplicationError::replication_slot(format!(
"server switched timeline {} -> {} (failover or point-in-time recovery); \
WAL diverges at the switchpoint, so the stored LSN may name a record that \
no longer exists. Check the confirmed LSN against the switchpoint reported \
by `TIMELINE_HISTORY {}` before resuming, or re-sync from a new snapshot",
previous.timeline, current.timeline, current.timeline
)));
}
Ok(())
}
/// Recover connection after a failure
async fn recover_connection(&mut self) -> Result<()> {
info!("Attempting to recover replication connection");
// Attempt reconnection with retry logic and enforce timeout
self.connection = timeout_or_error(
self.config.connection_timeout,
self.retry_handler.connect_with_retry(),
)
.await?;
// Re-initialize the connection. A reconnect can land somewhere other than where we left off, and the stored LSN is only meaningful on the original cluster and timeline.
let identity = self.connection.identify_system()?;
match ServerIdentity::from_identify_system(&identity) {
Some(current) => {
self.check_server_identity(¤t)?;
self.identity = Some(current);
}
// Fail closed. An unreadable IDENTIFY_SYSTEM used to skip the check
// entirely, which is precisely backwards: we are about to resume from a
// stored LSN that is only meaningful on the same cluster and timeline,
// and we just lost the ability to confirm either. A conforming server
// always returns systemid and timeline, so getting here means the
// response was malformed — treat that as the mismatch it might be.
//
// `permanent_connection`, not `replication_slot`: the server never
// looked at the slot, and `ReplicationSlot` is documented as
// "invalidated / does not exist — drop it and re-sync", which would
// send an operator to destroy a healthy slot over a bad response.
None if self.identity.is_some() => {
return Err(ReplicationError::permanent_connection(
"IDENTIFY_SYSTEM returned no usable system identifier/timeline on reconnect; \
refusing to resume from the stored LSN without confirming the cluster and \
timeline are unchanged"
.to_string(),
));
}
None => {
warn!("IDENTIFY_SYSTEM returned no usable identity; cluster/timeline changes cannot be detected");
}
}
// Temporary slots are dropped when the connection dies, so we must
// recreate them on recovery.
if self.config.slot_options.temporary {
self.slot = SlotState::Absent;
}
// Ensure replication slot still exists (recreate if temporary)
self.ensure_replication_slot().await?;
// Restart position for the reconnect.
//
// NEVER `last_received_lsn`. That advances at frame-decode time
// (`parse_xlogdata_header`, and keepalives in `process_keepalive_message`),
// i.e. the moment bytes arrive — long before the consumer has processed or
// acked them. PostgreSQL does not clamp a startpoint that is *ahead* of the
// slot: `CreateDecodingContext` (logical.c) forwards a startpoint that is
// *behind* `confirmed_flush` up to it, but has no symmetric branch for one
// that is ahead, and `SnapBuildXactNeedsSkip` (snapbuild.c,
// `ptr < start_decoding_at`) then makes `DecodeTXNNeedSkip` drop every
// transaction committing below it. Restarting from `last_received_lsn`
// therefore *permanently discards* everything received-but-not-yet-applied.
//
// Use the consumer's applied position instead. `start(Some(X))` seeds it,
// and every ack raises it, so it is the single authority for "where may we
// safely resume". Zero means "nothing acked and no explicit startpoint",
// which is `INVALID_XLOG_REC_PTR`, so the server resumes from
// `confirmed_flush` — exactly what a cold `start(None)` does. Recovery is
// now at-least-once like cold start, instead of at-most-once. Replaying a
// transaction is safe; skipping one is not.
let last_lsn = self.shared_lsn_feedback.get_feedback_lsn().1;
if last_lsn == INVALID_XLOG_REC_PTR && self.state.last_received_lsn != INVALID_XLOG_REC_PTR
{
warn!(
"Reconnecting with no applied LSN ever acked; replaying from the slot's \
confirmed_flush. Call `shared_lsn_feedback.update_applied_lsn(..)` after \
durable processing to bound replay."
);
}
let options = self.build_replication_options()?;
let options_ref = options
.iter()
.map(|(k, v)| (k.as_str(), v.as_str()))
.collect::<Vec<_>>();
// On failure, mark the connection we just installed as dead.
//
// `start_replication` reports a server error without touching liveness, so
// the freshly-connected socket stays `is_alive() == true` while never
// having entered COPY mode. `next_event_with_retry` skips recovery when
// the connection looks alive, so every later read failed
// `ensure_replication_mode` with `Protocol(..)`, which the live-connection
// decode guard then treats as terminal *before* the backoff sleep — a
// zero-delay busy-spin for `futures::Stream` consumers, and a hard exit for
// `for_each_event`. The canonical trigger is 55006 "replication slot is
// already active for PID" after a network blip, which clears on its own
// once `wal_sender_timeout` (default 60 s) reaps the old walsender: a
// routine, self-healing condition that must not kill the stream.
if let Err(e) =
self.connection
.start_replication(&self.config.slot_name, last_lsn, &options_ref)
{
self.connection.mark_dead();
return Err(e);
}
// The server resumes at a transaction boundary, so a StreamStart seen
// before the disconnect has no matching StreamStop coming.
self.parser.reset();
info!("Replication connection recovered and restarted");
Ok(())
}
/// Enhanced next_event with automatic retry, recovery and cancellation support
///
/// This is the recommended method for production use. It wraps `next_event()` with:
/// - Automatic retry logic for transient errors (up to 3 attempts)
/// - Connection health checks and automatic recovery
/// - Exponential backoff between retries
/// - Proper cancellation handling
///
/// For an even more ergonomic API, consider using `into_stream()` which provides
/// an async iterator interface that automatically calls this method.
///
/// # Arguments
/// * `cancellation_token` - Token to signal cancellation/shutdown
///
/// # Returns
/// * `Ok(event)` - Successfully received a change event
/// * `Err(ReplicationError::Cancelled(_))` - Operation was cancelled
/// * `Err(_)` - Permanent error after exhausting retries
///
/// # Example
///
/// ```no_run
/// use pg_walstream::{LogicalReplicationStream, ReplicationStreamConfig, RetryConfig, StreamingMode};
/// use tokio_util::sync::CancellationToken;
/// use std::time::Duration;
///
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let config = ReplicationStreamConfig::new(
/// "my_slot".to_string(),
/// "my_publication".to_string(),
/// 2, StreamingMode::On,
/// Duration::from_secs(10),
/// Duration::from_secs(30),
/// Duration::from_secs(60),
/// RetryConfig::default(),
/// );
///
/// let mut stream = LogicalReplicationStream::new("connection_string", config).await?;
/// stream.start(None).await?;
///
/// let cancel_token = CancellationToken::new();
///
/// // Robust polling loop with automatic retry
/// loop {
/// match stream.next_event_with_retry(&cancel_token).await {
/// Ok(event) => {
/// println!("Event: {:?}", event);
/// }
/// Err(e) => {
/// eprintln!("Fatal error: {}", e);
/// break;
/// }
/// }
/// }
/// # Ok(())
/// # }
/// ```
pub async fn next_event_with_retry(
&mut self,
cancellation_token: &CancellationToken,
) -> Result<ChangeEvent> {
// Bounded-replay terminal is authoritative: once a stop boundary has been
// reached, delegate straight to `next_event` (which sends CopyDone once and
// returns StreamStopped) and skip the health check below. Otherwise a dead
// connection could be reconnected by `recover_connection` and the COPY
// stream resurrected past `stop_at_lsn`.
if self.replay_stop.reached().is_some() {
return self.next_event(cancellation_token).await;
}
// Perform periodic health check.
//
// The check itself is cheap when nothing is wrong (it just compares `Instant::now()` against `last_health_check`), but `Instant::now()` is a vDSO syscall on Linux. At 100k+ rps that adds up, so we gate the health check on the per-event feedback counter and let it amortize across `HEALTH_CHECK_EVENT_INTERVAL` events. The actual time-based interval inside `check_connection_health()` still applies, this just skips the syscall on the inner-loop hot path.
if self.feedback_check_counter & (HEALTH_CHECK_EVENT_INTERVAL - 1) == 0 {
if let Err(e) = self.check_connection_health().await {
warn!("Health check failed: {}", e);
// Don't fail immediately, try to continue
}
}
// Try to get the next event with retry logic for transient connection errors
const MAX_ATTEMPTS: u32 = 3;
let mut attempt = 0;
loop {
attempt += 1;
match self.next_event(cancellation_token).await {
Ok(event) => return Ok(event),
Err(e) => {
// Cancellation and bounded-replay stop are terminal — never retry.
if matches!(
e,
ReplicationError::Cancelled(_) | ReplicationError::StreamStopped(_)
) {
debug!("Terminal signal, ending stream: {}", e);
return Err(e);
}
if e.is_permanent() {
error!("Permanent error in event processing: {}", e);
return Err(e);
}
// A decode failure on a *live* connection is not recoverable by
// retrying: `next_event` has already consumed the offending frame,
// so looping just reads the NEXT frame and returns that — silently
// dropping a change the consumer never saw, while the received LSN
// has already advanced past it. Surface it instead.
//
// Server-reported conditions do NOT reach this guard as
// `Protocol`: `from_sqlstate` maps the transient ones (57P01
// admin shutdown, 55006 slot busy, …) to `TransientConnection`,
// which is not in the match below. And any error out of the COPY
// read path clears the backend's liveness first — the native
// worker stores `alive = false` for every non-cancelled error,
// and the libpq backend reports `is_alive() == false` once
// `copy_end` is latched — so `is_alive()` is already false by the
// time a server error arrives here. What this guard actually
// catches is parse-stage failures on a connection that is still
// healthy. Unclassified SQLSTATEs still fall back to `Protocol`;
// those genuinely are terminal here for the reason above.
if self.connection.is_alive()
&& matches!(
e,
ReplicationError::Protocol(_)
| ReplicationError::Buffer(_)
| ReplicationError::Deserialize(_)
)
{
error!("Unrecoverable decode error on a live connection: {}", e);
return Err(e);
}
// Check if we've exhausted retry attempts
if attempt >= MAX_ATTEMPTS {
error!(
"Exhausted retry attempts ({}) for event processing: {}",
MAX_ATTEMPTS, e
);
return Err(e);
}
warn!(
"Transient error in event processing (attempt {}/{}): {}",
attempt, MAX_ATTEMPTS, e
);
// Attempt connection recovery if connection is dead
if !self.connection.is_alive() {
if let Err(recovery_err) = self.recover_connection().await {
// Only a *permanent* recovery failure ends the stream.
// Returning on a transient one skipped the backoff
// below and handed the consumer an error it would
// immediately retry, at zero delay. The connection is
// marked dead on this path, so the next attempt
// re-enters recovery rather than reading from a socket
// that never entered COPY mode.
if recovery_err.is_permanent() || attempt >= MAX_ATTEMPTS {
error!("Failed to recover connection: {}", recovery_err);
return Err(recovery_err);
}
warn!(
"Connection recovery failed (attempt {}/{}), backing off: {}",
attempt, MAX_ATTEMPTS, recovery_err
);
} else {
info!("Connection recovered successfully");
}
}
// Exponential backoff before retry
let delay = Duration::from_millis(1000 * (1 << (attempt - 1)));
debug!("Waiting {:?} before retry attempt {}", delay, attempt + 1);
// Use tokio::select! to allow cancellation during backoff
tokio::select! {
biased;
_ = tokio::time::sleep(delay) => {},
_ = cancellation_token.cancelled() => {
return Err(ReplicationError::cancelled(
"Operation cancelled during retry backoff"
));
}
}
}
}
}
}
/// Parse and validate the fixed 25-byte XLogData (`'w'`) header, advancing
/// the received LSN. Returns `(wal_start, wal_end, send_time)`.
///
/// Shared by [`process_wal_message`](Self::process_wal_message) (typed path)
/// and [`decode_raw_frame`](Self::decode_raw_frame) (raw path); each caller
/// then handles the message body (or its absence) its own way. Layout:
/// `'w'` (1) + start_lsn (8) + end_lsn (8) + send_time (8) = 25 bytes.
/// `#[inline]`, so callers keep identical codegen to a hand-inlined header
/// decode.
#[inline]
fn parse_xlogdata_header(&mut self, data: &Bytes) -> Result<(u64, u64, i64)> {
// 'w' (1) + start_lsn (8) + end_lsn (8) + send_time (8) = 25 bytes.
if data.len() < 25 {
return Err(ReplicationError::protocol(
"WAL message too short".to_string(),
));
}
let mut header = &data[1..25];
let wal_start = header.get_u64();
let wal_end = header.get_u64();
let send_time = header.get_i64();
// Update LSN tracking using the server's WAL end position for this message.
if wal_end > 0 {
self.state.update_received_lsn(wal_end);
}
Ok((wal_start, wal_end, send_time))
}
/// Process a WAL data message (zero-copy: uses Bytes slicing)
///
/// The XLogData header is fixed-layout (25 bytes), so we decode it directly
/// from the slice instead of constructing a `BufferReader` — mirroring the
/// keepalive fast path. This avoids per-event bounds-check overhead and an
/// extra `Bytes` slice for the header region.
#[inline]
fn process_wal_message(&mut self, data: impl Into<Bytes>) -> Result<Option<ChangeEvent>> {
// Decode the XLogData envelope with the shared, policy-free decoder, then decode its pgoutput body. A header-only frame has an empty body → no event; the raw path surfaces that same frame instead (each caller owns the skip-vs-surface choice).
let raw = self.decode_raw_frame(data.into())?;
if raw.data.is_empty() {
return Ok(None);
}
let replication_message = self.parser.parse_wal_message_bytes(raw.data)?;
crate::protocol::message_to_change_event(
&mut self.state,
replication_message,
raw.wal_start.value(),
)
}
/// Decode a `'w'` CopyData frame into a [`RawXLogData`] without touching the
/// pgoutput parser.
///
/// Policy-free: it parses the fixed 25-byte header (advancing the received
/// LSN, in `parse_xlogdata_header`) and returns the zero-copy body slice — it
/// arms NO stop. Callers own that: [`next_raw_event`](Self::next_raw_event)
/// arms the transport-level stop, `next_event` arms the commit-boundary stop.
/// That separation is what lets the typed `process_wal_message` reuse this
/// decoder without inheriting the raw path's stop semantics.
///
/// A header-only frame (len == 25) yields a `RawXLogData` with an empty
/// `data` payload — a conformant server never sends one (it sends a keepalive
/// when idle), so this is only reached by a malformed/edge frame.
#[inline]
fn decode_raw_frame(&mut self, data: Bytes) -> Result<RawXLogData> {
let (wal_start, wal_end, server_time) = self.parse_xlogdata_header(&data)?;
// Zero-copy slice of the message body — refcount only, no memcpy, `slice(25..)` is an empty `Bytes` for a header-only frame.
Ok(RawXLogData {
wal_start: Lsn::new(wal_start),
wal_end: Lsn::new(wal_end),
server_time,
data: data.slice(25..),
})
}
/// Process a keepalive message
async fn process_keepalive_message(&mut self, data: &[u8]) -> Result<()> {
let keepalive = parse_keepalive_message(data)?;
debug!(
"Received keepalive: wal_end={}, reply_requested={}",
format_lsn(keepalive.wal_end),
keepalive.reply_requested
);
self.state.update_received_lsn(keepalive.wal_end);
// Bounded replay, idle-server terminal.
//
// On the typed path `stop_at_reached` is otherwise armed only from a
// delivered commit boundary (`next_event`), so a target set past the last
// commit never fires: the server has nothing more to send, `next_wal_frame`
// loops on keepalives forever (there is no read timeout), and the "catch up
// to now, then exit" consumer simply never exits. (`next_raw_event` also
// arms it, from the transport-level `wal_end`; that path is unaffected here.)
//
// Safe despite `wal_end` not being a commit boundary. A keepalive's
// `wal_end` is the walsender's `sentPtr`, and `XLogSendLogical` assigns
// `sentPtr = reader->EndRecPtr` only *after* `LogicalDecodingProcessRecord`
// has run the output plugin, whose `pq_putmessage_noblock` lands in the same
// FIFO socket buffer as the keepalive's. So everything committing at or
// below `wal_end` is already on the wire ahead of this keepalive; a
// transaction still open across `sentPtr` necessarily commits *above* it and
// is out of range by construction. This is pg_recvlogical's
// `STREAM_STOP_KEEPALIVE` exit, whose in-tree comment makes the same
// argument. (The "--endpos may truncate partway through a transaction"
// caveat in the docs is about its *per-record* endpos check, not this one.)
if let Some(stop) = self.config.stop_at_lsn {
if keepalive.wal_end >= stop.value() {
self.replay_stop.arm(Lsn::new(keepalive.wal_end));
}
}
// pg_recvlogical also flushes feedback when endpos is reached, so the
// server records final progress before the stream is torn down.
if keepalive.reply_requested || self.replay_stop.reached().is_some() {
self.send_feedback().await?;
}
Ok(())
}
/// Fetch the consumer's flushed/applied LSNs from the shared tracker, each capped
/// to `last_received_lsn` — a consumer must never report progress ahead of what we
/// have actually received. A zero from the tracker means "unset" and stays zero.
#[inline]
fn capped_feedback_lsns(&self) -> (u64, u64) {
let (f, a) = self.shared_lsn_feedback.get_feedback_lsn();
let cap = |v: u64| {
if v > 0 {
v.min(self.state.last_received_lsn)
} else {
0
}
};
(cap(f), cap(a))
}
/// Check if feedback should be sent and send it
///
/// Implements intelligent throttling: only sends feedback if the configured
/// interval has elapsed AND the LSN values have actually changed. This avoids
/// unnecessary status updates and reduces network overhead.
#[inline]
pub async fn maybe_send_feedback(&mut self) {
// Check if enough time has elapsed
if !self
.state
.should_send_feedback(self.config.feedback_interval)
{
return;
}
// Get current LSN values that would be sent
let (flushed_lsn, applied_lsn) = self.capped_feedback_lsns();
// Only send feedback if LSN values have changed
if self.state.lsn_has_changed(flushed_lsn, applied_lsn) {
if let Err(e) = self.send_feedback().await {
warn!("Failed to send feedback: {}", e);
}
}
}
/// Send feedback to the server
///
/// This method sends a standby status update to PostgreSQL with three LSN values:
/// - write_lsn (received): Updated when data is received from the replication stream
/// - flush_lsn: Updated when data is written to the destination
/// - replay_lsn (applied): Updated when data is committed to the destination
///
/// If a shared LSN feedback tracker is configured, it will use the flushed/applied
/// values from there (updated by the consumer). Otherwise, it falls back to the
/// local state values.
///
/// The method tracks the last sent LSN values to enable intelligent throttling
/// and avoid sending redundant status updates.
pub async fn send_feedback(&mut self) -> Result<()> {
if self.state.last_received_lsn == 0 {
return Ok(());
}
// Fetch the consumer's committed progress (capped to what we've received).
// This allows the consumer to update these values after committing to destination.
let (flushed_lsn, applied_lsn) = self.capped_feedback_lsns();
self.connection
.send_standby_status_update(
self.state.last_received_lsn,
flushed_lsn,
applied_lsn,
false, // Don't request reply
)
.await?;
// Record the LSN values we sent for intelligent throttling
self.state
.mark_feedback_sent_with_lsn(flushed_lsn, applied_lsn);
debug!(
"Sent feedback: received={}, flushed={}, applied={}",
format_lsn(self.state.last_received_lsn),
format_lsn(flushed_lsn),
format_lsn(applied_lsn)
);
Ok(())
}
/// Stop the replication stream
pub async fn stop(&mut self) -> Result<()> {
if let Err(e) = self.send_feedback().await {
warn!("Failed to send final feedback: {}", e);
}
info!(
"Stopping logical replication stream (last received LSN: {})",
format_lsn(self.current_lsn())
);
Ok(())
}
/// Get the current LSN position
pub fn current_lsn(&self) -> XLogRecPtr {
self.state.last_received_lsn
}
/// Returns the exported snapshot name from slot creation, if available.
///
/// When the replication slot is created with `snapshot: Some("export")` (i.e. `EXPORT_SNAPSHOT`),
/// PostgreSQL returns a snapshot identifier that can be used with
/// `SET TRANSACTION SNAPSHOT` on a separate connection to read a consistent snapshot
/// of the database at the slot's starting LSN.
///
/// **Important:** This snapshot is only valid between
/// [`ensure_replication_slot()`](Self::ensure_replication_slot) and
/// [`start()`](Self::start). Once `START_REPLICATION` is issued, PostgreSQL
/// destroys the snapshot. You must read the snapshot on a separate connection
/// **before** calling `start()`; after a *successful* `start()` this returns
/// `None` rather than a name the server has already discarded.
///
/// Any `start()` that **fails** deliberately leaves the name in place. The
/// caller still owns the stream and may go on to call
/// [`snapshot()`](Self::snapshot), which reads this to choose `Available` vs
/// `Unavailable`; reporting `None` there would give `Unavailable`, which the
/// documented usage reads as "nothing to copy" and streams with no baseline.
/// Keeping it means a failure that never reached the server (the snapshot is
/// intact) still works, and one that did destroy the snapshot fails loudly at
/// `SET TRANSACTION SNAPSHOT` instead of silently.
///
/// Returns `None` if no snapshot was exported (e.g. `NOEXPORT_SNAPSHOT`) or if
/// the slot has not yet been created.
pub fn exported_snapshot_name(&self) -> Option<&str> {
self.slot.exported_snapshot()
}
/// Best-effort drop of the replication slot this stream created.
///
/// Used by the initial-snapshot helper when an attempt fails part-way. A slot
/// created seconds ago retains no WAL worth keeping, and leaving it behind is
/// actively harmful: the next attempt would find it already present,
/// [`ensure_replication_slot`](Self::ensure_replication_slot) would swallow
/// the "already exists" error, no snapshot would be exported, and the caller
/// would silently stream with no baseline.
///
/// Failure to drop is logged, not returned: this runs on an error path and
/// must not mask the error that got us here.
pub(crate) fn discard_replication_slot(&mut self) {
// Only ever drop a slot this stream created. A pre-existing slot belongs
// to whoever made it, and dropping one would destroy their subscription's
// position — the caller's own WAL retention along with it.
//
// This check is the reason `SlotState` exists. The `slot_created: bool`
// it replaced was set to `true` on *both* the create and the
// "already exists" branch, so it could not distinguish the two: guarding
// on it would have passed here and dropped the operator's slot anyway.
// Until now the guarantee lived at the single call site, which happened
// to check an unrelated field first.
if !matches!(self.slot, SlotState::Ours { .. }) {
debug!(
"Not dropping replication slot {}: this stream did not create it",
self.config.slot_name
);
return;
}
let slot_name = self.config.slot_name.clone();
match self.connection.drop_replication_slot(&slot_name, false) {
Ok(()) => {
info!("Dropped replication slot {slot_name} after a failed snapshot attempt");
self.slot = SlotState::Absent;
}
Err(e) => {
warn!(
"Could not drop replication slot {slot_name} after a failed snapshot \
attempt: {e}. Drop it manually, or the next snapshot attempt will \
find it present and silently proceed without a baseline."
);
}
}
}
/// The LSN the exported snapshot is consistent at.
///
/// `None` when the slot was never created by this stream, or pre-existed (see
/// [`ensure_replication_slot`](Self::ensure_replication_slot), which treats
/// "already exists" as success).
///
/// Crate-internal on purpose. Callers who need this value are holding a
/// [`Snapshot`](crate::snapshot::Snapshot), whose `consistent_point()` is not
/// an `Option` — and a caller holding the stream *directly* only gets here
/// down the `SnapshotOutcome::Unavailable` arm, where it is always `None`.
/// A public `Option` accessor would therefore be a trap: unwrappable exactly
/// when it is useless.
#[inline]
pub(crate) fn consistent_point(&self) -> Option<Lsn> {
self.slot.consistent_point()
}
/// The connection string this stream was built from.
#[inline]
pub(crate) fn conninfo(&self) -> &str {
self.retry_handler.connection_string()
}
/// The stream's configuration.
#[inline]
pub(crate) fn config(&self) -> &ReplicationStreamConfig {
&self.config
}
/// Require the slot to be created with `SNAPSHOT 'export'`.
///
/// Called by [`snapshot`](Self::snapshot) so the entry point sets its own
/// precondition, instead of a builder call the caller has to remember — and
/// that [`ReplicationStreamConfig::with_slot_options`] could silently undo.
/// A no-op once the slot exists, which is the `Unavailable` path.
#[inline]
pub(crate) fn require_exported_snapshot(&mut self) {
self.config.slot_options.snapshot = Some("export".to_string());
}
/// Server version number, e.g. 160001 for PG 16.1.
#[inline]
pub(crate) fn server_version(&self) -> i32 {
self.connection.server_version()
}
/// Returns `true` if this stream's replication slot is configured as temporary.
///
/// Temporary slots are automatically dropped when the replication connection
/// is closed. This means that on reconnection, the slot will be recreated.
pub fn is_temporary_slot(&self) -> bool {
self.config.slot_options.temporary
}
/// Convert into an async stream of change events
///
/// This method provides an iterator-like interface for consuming replication events.
/// The returned `EventStream` has a built-in `next()` method that doesn't require
/// importing any traits - simply call `.next().await` to get the next event.
///
/// If you prefer to use the `futures::Stream` trait, enable the `stream` feature
/// in your `Cargo.toml`:
///
/// ```toml
/// [dependencies]
/// pg_walstream = { version = "0.1", features = ["stream"] }
/// ```
///
/// The stream automatically handles retries, reconnections, and cancellation.
/// When cancelled, the stream will cleanly terminate.
///
/// # Arguments
///
/// * `cancellation_token` - Token to signal cancellation/shutdown
///
/// # Returns
///
/// An `EventStream` that yields `Result<ChangeEvent>` items.
///
/// # Example
///
/// ```no_run
/// use pg_walstream::{LogicalReplicationStream, ReplicationStreamConfig, RetryConfig, StreamingMode};
/// use tokio_util::sync::CancellationToken;
/// use std::time::Duration;
///
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let config = ReplicationStreamConfig::new(
/// "my_slot".to_string(),
/// "my_publication".to_string(),
/// 2, StreamingMode::On,
/// Duration::from_secs(10),
/// Duration::from_secs(30),
/// Duration::from_secs(60),
/// RetryConfig::default(),
/// );
///
/// let mut stream = LogicalReplicationStream::new("connection_string", config).await?;
/// stream.start(None).await?;
///
/// let cancel_token = CancellationToken::new();
/// let mut event_stream = stream.into_stream(cancel_token);
///
/// loop {
/// match event_stream.next_event().await {
/// Ok(event) => println!("Event: {:?}", event),
/// Err(_) => break,
/// }
/// }
/// # Ok(())
/// # }
/// ```
pub fn into_stream(self, cancellation_token: CancellationToken) -> EventStream {
let shared_feedback = Arc::clone(&self.shared_lsn_feedback);
EventStream {
state: StreamState::Ready {
stream: Box::new(self),
terminated: false,
},
cancellation_token,
shared_feedback,
}
}
/// Create an async stream of change events (borrows self)
///
/// Similar to `into_stream` but borrows the stream instead of consuming it.
/// This is useful when you want to maintain ownership of the LogicalReplicationStream
/// while still using the Stream interface.
///
/// # Arguments
///
/// * `cancellation_token` - Token to signal cancellation/shutdown
///
/// # Returns
///
/// An `EventStreamRef` that yields `Result<ChangeEvent>` items.
pub fn stream(&mut self, cancellation_token: CancellationToken) -> EventStreamRef<'_> {
EventStreamRef {
inner: self,
cancellation_token,
}
}
fn build_replication_options(&self) -> Result<Vec<(String, String)>> {
self.validate_replication_options()?;
let proto_version = self.config.protocol_version.to_string();
let publication_names = format!("\"{}\"", self.config.publication_name);
let mut options = vec![
("proto_version".to_string(), proto_version),
("publication_names".to_string(), publication_names),
];
if !matches!(self.config.streaming_mode, StreamingMode::Off) {
options.push((
"streaming".to_string(),
self.config.streaming_mode.as_str().to_string(),
));
}
if self.config.messages {
options.push(("messages".to_string(), "on".to_string()));
}
if self.config.binary {
options.push(("binary".to_string(), "on".to_string()));
}
if self.config.two_phase {
options.push(("two_phase".to_string(), "on".to_string()));
}
if let Some(origin) = self.config.origin {
options.push(("origin".to_string(), origin.as_str().to_string()));
}
Ok(options)
}
fn validate_replication_options(&self) -> Result<()> {
match self.config.protocol_version {
1..=4 => {}
version => {
return Err(ReplicationError::config(format!(
"Unsupported protocol version: {version} (expected 1-4)"
)))
}
}
if matches!(self.config.streaming_mode, StreamingMode::On)
&& self.config.protocol_version < 2
{
return Err(ReplicationError::config(
"streaming=on requires protocol version >= 2".to_string(),
));
}
if matches!(self.config.streaming_mode, StreamingMode::Parallel)
&& self.config.protocol_version < 4
{
return Err(ReplicationError::config(
"streaming=parallel requires protocol version >= 4".to_string(),
));
}
if self.config.two_phase && self.config.protocol_version < 3 {
return Err(ReplicationError::config(
"two_phase requires protocol version >= 3".to_string(),
));
}
Ok(())
}
}
/// Async stream of PostgreSQL replication events (owned version)
///
/// This struct implements both a native `.next().await` API and the
/// [`futures_core::Stream`] trait, so you can use it with any stream
/// combinator from `futures` or `tokio-stream`.
///
/// Create an `EventStream` by calling `into_stream()` on `LogicalReplicationStream`.
///
/// # Using as `futures::Stream`
///
/// The stream yields `Result<ChangeEvent>` items. When cancelled or after a
/// permanent error, the stream terminates (yields `None`). It also implements
/// [`futures_core::FusedStream`] so `select!` macros can detect termination.
///
/// ```ignore
/// use futures::StreamExt;
///
/// let mut event_stream = stream.into_stream(cancel_token);
///
/// // Works directly with StreamExt combinators!
/// while let Some(result) = event_stream.next().await {
/// match result {
/// Ok(event) => println!("Event: {:?}", event),
/// Err(e) => {
/// eprintln!("Error: {}", e);
/// break;
/// }
/// }
/// }
/// println!("Stream terminated");
/// ```
///
/// # Using the native API
///
/// ```no_run
/// use pg_walstream::{LogicalReplicationStream, ReplicationStreamConfig, RetryConfig, StreamingMode};
/// use tokio_util::sync::CancellationToken;
/// use std::time::Duration;
///
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let config = ReplicationStreamConfig::new(
/// "my_slot".to_string(),
/// "my_publication".to_string(),
/// 2, StreamingMode::On,
/// Duration::from_secs(10),
/// Duration::from_secs(30),
/// Duration::from_secs(60),
/// RetryConfig::default(),
/// );
///
/// let mut stream = LogicalReplicationStream::new("connection_string", config).await?;
/// stream.start(None).await?;
///
/// let cancel_token = CancellationToken::new();
/// let mut event_stream = stream.into_stream(cancel_token);
///
/// // Native API — no trait imports needed
/// loop {
/// match event_stream.next_event().await {
/// Ok(event) => {
/// println!("Received: {:?}", event);
/// }
/// Err(e) if matches!(e, pg_walstream::ReplicationError::Cancelled(_)) => {
/// println!("Stream cancelled");
/// break;
/// }
/// Err(e) => {
/// eprintln!("Error: {}", e);
/// break;
/// }
/// }
/// }
/// # Ok(())
/// # }
/// ```
pub struct EventStream {
/// Where the replication stream currently lives. See [`StreamState`].
state: StreamState,
cancellation_token: CancellationToken,
/// Cached reference to the shared LSN feedback so callers can always
/// update/read LSN values even while the stream is parked in a poll.
shared_feedback: Arc<SharedLsnFeedback>,
}
/// The future a `poll_next` parks the replication stream inside.
///
/// It **owns** the stream for the duration of the poll and hands it back
/// alongside the result — which is why "polling" and "holding the stream" are
/// mutually exclusive, and why dropping this future drops the connection.
type PollFuture =
std::pin::Pin<Box<dyn Future<Output = (LogicalReplicationStream, Result<ChangeEvent>)> + Send>>;
/// Where an [`EventStream`]'s replication stream is, and whether it is done.
///
/// Replaces `inner: Option<_>` + `inflight: Option<_>` + `terminated: bool`,
/// which spelled three legal states across eight representable ones. The pair
/// `(inner: None, inflight: None)` was reachable: `shutdown()` landed there
/// whenever a poll was in flight, dropping the future that owned the connection
/// and skipping the final LSN feedback while still returning `Ok(())`. Nothing
/// checked the invariant — the `expect`s only asserted it.
///
/// `terminated` stays a field rather than becoming a third variant because the
/// stream is still *there* after termination: `inner()` must keep working once
/// `shutdown()` has run. Inside `Ready` it is an independent binary fact, and
/// `Polling` cannot contradict it, so exactly the three legal states are
/// representable.
enum StreamState {
/// The stream is in hand. `terminated` distinguishes "pollable" from
/// "finished, but still readable through `inner()`".
Ready {
stream: Box<LogicalReplicationStream>,
terminated: bool,
},
/// The stream has been moved into an in-flight poll future.
Polling(PollFuture),
}
impl StreamState {
/// The stream, unless it is currently parked inside a poll future.
#[inline]
fn stream(&self) -> Option<&LogicalReplicationStream> {
match self {
StreamState::Ready { stream, .. } => Some(stream),
StreamState::Polling(_) => None,
}
}
#[inline]
fn stream_mut(&mut self) -> Option<&mut LogicalReplicationStream> {
match self {
StreamState::Ready { stream, .. } => Some(stream),
StreamState::Polling(_) => None,
}
}
#[inline]
fn is_terminated(&self) -> bool {
matches!(
self,
StreamState::Ready {
terminated: true,
..
}
)
}
/// Whether a poll future currently holds the stream.
#[cfg(test)]
fn is_polling(&self) -> bool {
matches!(self, StreamState::Polling(_))
}
/// Take the stream out, leaving a poll future that resolves to `result`.
///
/// Test-only: reproduces the state a `Poll::Pending` leaves behind, without
/// needing a real server to stall on.
#[cfg(test)]
fn park_with_result(&mut self, result: Result<ChangeEvent>) {
let placeholder: PollFuture = Box::pin(core::future::pending());
let StreamState::Ready { stream, .. } =
core::mem::replace(self, StreamState::Polling(placeholder))
else {
panic!("park_with_result expects a stream in hand");
};
let stream = *stream;
*self = StreamState::Polling(Box::pin(async move { (stream, result) }));
}
/// Flip the terminal flag without going through a poll.
#[cfg(test)]
fn mark_terminated(&mut self) {
match self {
StreamState::Ready { terminated, .. } => *terminated = true,
StreamState::Polling(_) => panic!("cannot terminate a parked stream"),
}
}
/// Park the stream in a fresh poll future. A no-op unless idle.
fn start_poll(&mut self, cancellation_token: &CancellationToken) {
if !matches!(
self,
StreamState::Ready {
terminated: false,
..
}
) {
return;
}
// Moving the stream out of `Ready` needs something to leave behind.
// `pending()` is zero-sized, so boxing it does not allocate, and it
// occupies `self` only until the real future is built two lines down —
// the guard above is what stops it ever displacing a live poll.
let placeholder: PollFuture = Box::pin(core::future::pending());
if let StreamState::Ready { stream, .. } =
core::mem::replace(self, StreamState::Polling(placeholder))
{
let mut stream = *stream;
let cancellation_token = cancellation_token.clone();
*self = StreamState::Polling(Box::pin(async move {
let result = stream.next_event_with_retry(&cancellation_token).await;
(stream, result)
}));
}
}
}
/// Extract the `consistent_point` from a `CREATE_REPLICATION_SLOT` result.
///
/// The command returns `slot_name(0), consistent_point(1), snapshot_name(2),
/// output_plugin(3)`. Column 1 is the LSN the exported snapshot is taken at, and
/// equals the slot's `confirmed_flush` — so it is both the LSN to stamp on
/// snapshot rows and the position `start(None)` resumes from.
///
/// A missing or malformed value is **not** fatal: it only means the
/// initial-snapshot helper is unavailable, not that replication is. Returning
/// `None` lets `snapshot()` report `Unavailable` rather than failing a stream
/// that would otherwise work.
///
/// Split out as a free function so it is testable without a live connection.
fn parse_consistent_point(result: &crate::connection::PgResult) -> Option<Lsn> {
let raw = result.get_value(0, 1).filter(|v| !v.is_empty())?;
match crate::types::parse_lsn(&raw) {
Ok(lsn) => Some(Lsn::new(lsn)),
Err(e) => {
warn!("Ignoring unparsable consistent_point {raw:?}: {e}");
None
}
}
}
/// Crate-internal seam so the consume loop is unit-testable without a live
/// PostgreSQL connection. Implemented by [`EventStream`]; mocked in tests.
pub(crate) trait EventSource {
async fn recv(&mut self) -> Result<ChangeEvent>;
fn ack(&self, lsn: XLogRecPtr);
}
impl EventSource for EventStream {
#[inline]
async fn recv(&mut self) -> Result<ChangeEvent> {
self.next_event().await
}
#[inline]
fn ack(&self, lsn: XLogRecPtr) {
self.update_applied_lsn(lsn);
}
}
/// Drive `source`, running `handler` per event and advancing the applied LSN
/// only after the handler resolves `Ok`. `Cancelled` recv → graceful `Ok(())`.
pub(crate) async fn for_each_event_impl<S, F, Fut>(source: &mut S, mut handler: F) -> Result<()>
where
S: EventSource,
F: FnMut(ChangeEvent) -> Fut,
Fut: Future<Output = Result<()>>,
{
loop {
match source.recv().await {
Ok(ev) => {
let lsn = ev.lsn.value();
handler(ev).await?;
source.ack(lsn);
}
Err(ReplicationError::Cancelled(_)) | Err(ReplicationError::StreamStopped(_)) => {
return Ok(())
}
Err(e) => return Err(e),
}
}
}
impl EventStream {
/// Get a reference to the underlying LogicalReplicationStream
///
/// # Panics
///
/// Panics if called while a `Stream::poll_next` future is in flight — which
/// includes the ordinary case of a poll that was *abandoned*, such as
/// `tokio::time::timeout(d, stream.next())` expiring or a `select!` arm
/// losing. Call [`next_event`](Self::next_event) first: it reclaims the
/// parked stream. These accessors are synchronous and cannot.
pub fn inner(&self) -> &LogicalReplicationStream {
self.state
.stream()
.expect("inner stream is temporarily taken during poll")
}
/// Get a mutable reference to the underlying LogicalReplicationStream
///
/// # Panics
///
/// Same as [`inner`](Self::inner): panics while a poll is in flight,
/// abandoned polls included. Call [`next_event`](Self::next_event) first.
pub fn inner_mut(&mut self) -> &mut LogicalReplicationStream {
self.state
.stream_mut()
.expect("inner stream is temporarily taken during poll")
}
/// Get the current LSN position
///
/// This is the **received** LSN (the last WAL byte read off the wire), which
/// is a different value from the flushed/applied pair
/// [`get_feedback_lsn`](Self::get_feedback_lsn) reports.
///
/// # Panics
///
/// Panics if called while a `Stream::poll_next` future is in flight — it
/// reads through [`inner`](Self::inner). Unlike the `*_lsn` feedback methods,
/// this value lives only on the inner stream, so it cannot be served while
/// that stream is parked. Call [`next_event`](Self::next_event) first: it
/// reclaims the parked stream.
pub fn current_lsn(&self) -> XLogRecPtr {
self.inner().current_lsn()
}
/// Update the flushed LSN feedback
///
/// Call this after data has been written/flushed to the destination database,
/// but not yet committed (e.g., during batch writes).
/// This method is always safe to call, even during stream polling.
///
/// Accepts anything convertible to [`Lsn`], so `event.lsn` and a raw `u64`
/// both work without an explicit `.value()`.
#[inline]
pub fn update_flushed_lsn(&self, lsn: impl Into<Lsn>) {
self.shared_feedback.update_flushed_lsn(lsn);
}
/// Update the applied LSN feedback
///
/// Call this after data has been committed to the destination database.
/// This is the most common feedback update in typical replication scenarios.
/// This method is always safe to call, even during stream polling.
///
/// Accepts anything convertible to [`Lsn`], so `event.lsn` and a raw `u64`
/// both work without an explicit `.value()`.
#[inline]
pub fn update_applied_lsn(&self, lsn: impl Into<Lsn>) {
self.shared_feedback.update_applied_lsn(lsn);
}
/// Get the current feedback LSN values
///
/// Returns a tuple of (flushed_lsn, applied_lsn).
/// This method is always safe to call, even during stream polling.
#[inline]
pub fn get_feedback_lsn(&self) -> (XLogRecPtr, XLogRecPtr) {
self.shared_feedback.get_feedback_lsn()
}
/// Returns whether this stream has terminated.
///
/// A terminated stream will always return `None` from `poll_next()`.
/// The stream terminates when:
/// - The cancellation token is triggered
/// - A permanent error occurs
/// - `shutdown()` is called
#[inline]
pub fn is_terminated(&self) -> bool {
self.state.is_terminated()
}
/// Get the next event from the stream (native async API)
///
/// This method provides a native async API without requiring any trait imports.
/// It returns `Ok(event)` when an event is available, and `Err(e)` when an error occurs
/// (including cancellation via `Err(ReplicationError::Cancelled(_))`).
///
/// **Note**: If you also use this type as a `futures::Stream`, prefer calling
/// `StreamExt::next()` instead to avoid confusion. This method is named `next_event()`
/// to avoid conflict with `StreamExt::next()`.
pub async fn next_event(&mut self) -> Result<ChangeEvent> {
// A `poll_next` that returned `Pending` parked the stream inside
// `inflight`, and dropping the `Next`/`select!` future that drove it does
// not clear that field. Resume the parked future rather than asserting it
// cannot exist: abandoning a poll is ordinary use (`tokio::time::timeout`
// around `StreamExt::next`), and its event is still owed to the caller.
if let StreamState::Polling(fut) = &mut self.state {
let (stream, result) = fut.as_mut().await;
self.state = StreamState::Ready {
stream: Box::new(stream),
terminated: false,
};
return result;
}
let inner = self
.state
.stream_mut()
.expect("Polling was handled above, so the stream is in hand");
inner.next_event_with_retry(&self.cancellation_token).await
}
/// Consume the stream to completion, running `handler` on each event and
/// automatically advancing the applied LSN after each successful handler.
///
/// - `handler` returning `Ok(())` acks that event (advances applied LSN).
/// - `handler` returning `Err(e)` stops the loop and returns `e` (no ack).
/// - A cancelled stream returns `Ok(())` for graceful shutdown.
///
/// The handler future must be self-contained (own its captures or clone an
/// `Arc`); it may not borrow the handler's environment across `.await`.
pub async fn for_each_event<F, Fut>(&mut self, handler: F) -> Result<()>
where
F: FnMut(ChangeEvent) -> Fut,
Fut: Future<Output = Result<()>>,
{
for_each_event_impl(self, handler).await
}
/// Gracefully shut down the replication stream.
///
/// This method:
/// 1. Cancels the cancellation token (signals all async operations to stop)
/// 2. Sends a final LSN feedback to PostgreSQL
/// 3. Marks the stream as terminated
///
/// After calling this, `next_event()` and `Stream::poll_next` will return `None`.
pub async fn shutdown(&mut self) -> Result<()> {
self.cancellation_token.cancel();
// Reclaim the stream from an abandoned poll before stopping it. The
// parked future *owns* the stream, so dropping it here would skip step 2
// entirely and tear the COPY down without a CopyDone — while still
// returning `Ok(())`. Cancelling first is what makes this await prompt:
// the parked `next_event_with_retry` honours the token.
if let StreamState::Polling(fut) = &mut self.state {
let (stream, _) = fut.as_mut().await;
self.state = StreamState::Ready {
stream: Box::new(stream),
terminated: false,
};
}
let StreamState::Ready { stream, terminated } = &mut self.state else {
// Unreachable: the block above converted Polling to Ready.
return Ok(());
};
let result = stream.stop().await;
*terminated = true;
result
}
}
impl futures_core::Stream for EventStream {
type Item = Result<ChangeEvent>;
fn poll_next(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
if self.state.is_terminated() {
return std::task::Poll::Ready(None);
}
// Park the stream in a poll future if it is not already parked.
//
// Destructured rather than cloning the token: `poll_next` runs on every
// wakeup, and `start_poll` is a no-op on all but the first, so an
// unconditional `Arc` clone would put an atomic increment/decrement pair
// on every wakeup to serve a branch that rarely runs. `start_poll` clones
// only on the branch that actually parks.
let Self {
state,
cancellation_token,
..
} = &mut *self;
state.start_poll(cancellation_token);
let StreamState::Polling(fut) = &mut self.state else {
// Unreachable: `start_poll` parks anything that is Ready and not
// terminated, and the terminated case returned above.
return std::task::Poll::Ready(None);
};
match fut.as_mut().poll(cx) {
std::task::Poll::Pending => std::task::Poll::Pending,
std::task::Poll::Ready((stream, result)) => {
// Take the stream back out of the resolved future. Whether this
// was the last event decides `terminated`, so the state
// transition is part of each arm rather than a separate step.
let stream = Box::new(stream);
match result {
Ok(event) => {
self.state = StreamState::Ready {
stream,
terminated: false,
};
std::task::Poll::Ready(Some(Ok(event)))
}
Err(ref e) if e.is_cancelled() || e.is_stream_stopped() => {
self.state = StreamState::Ready {
stream,
terminated: true,
};
std::task::Poll::Ready(None)
}
Err(ref e) if e.is_permanent() => {
self.state = StreamState::Ready {
stream,
terminated: true,
};
std::task::Poll::Ready(Some(Err(result.unwrap_err())))
}
Err(e) => {
// Transient error — yield it but keep the stream alive
self.state = StreamState::Ready {
stream,
terminated: false,
};
std::task::Poll::Ready(Some(Err(e)))
}
}
}
}
}
}
impl futures_core::FusedStream for EventStream {
fn is_terminated(&self) -> bool {
self.state.is_terminated()
}
}
/// Async stream of PostgreSQL replication events (borrowed version)
///
/// This struct provides an iterator-like interface for consuming replication events,
/// similar to `EventStream` but borrows the underlying `LogicalReplicationStream`
/// instead of owning it.
///
/// It has a built-in `next()` method that returns `Future<Result<ChangeEvent>>`,
/// allowing you to call `.next().await` without importing any traits.
///
/// Create an `EventStreamRef` by calling `stream()` on `LogicalReplicationStream`.
pub struct EventStreamRef<'a> {
inner: &'a mut LogicalReplicationStream,
cancellation_token: CancellationToken,
}
impl<'a> EventStreamRef<'a> {
/// Get a reference to the underlying LogicalReplicationStream
pub fn inner(&self) -> &LogicalReplicationStream {
self.inner
}
/// Get a mutable reference to the underlying LogicalReplicationStream
pub fn inner_mut(&mut self) -> &mut LogicalReplicationStream {
self.inner
}
/// Get the current LSN position
pub fn current_lsn(&self) -> XLogRecPtr {
self.inner.current_lsn()
}
/// Get the next event from the stream
///
/// This method provides a native async API without requiring any trait imports.
/// It returns `Ok(event)` when an event is available, and `Err(e)` when an error occurs
/// (including cancellation via `Err(ReplicationError::Cancelled(_))`).
///
/// # Example
///
/// ```no_run
/// use pg_walstream::{LogicalReplicationStream, ReplicationStreamConfig, RetryConfig, StreamingMode};
/// use tokio_util::sync::CancellationToken;
/// use std::time::Duration;
///
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let config = ReplicationStreamConfig::new(
/// "my_slot".to_string(),
/// "my_publication".to_string(),
/// 2, StreamingMode::On,
/// Duration::from_secs(10),
/// Duration::from_secs(30),
/// Duration::from_secs(60),
/// RetryConfig::default(),
/// );
///
/// let mut stream = LogicalReplicationStream::new("connection_string", config).await?;
/// stream.start(None).await?;
///
/// let cancel_token = CancellationToken::new();
/// let mut event_stream = stream.stream(cancel_token);
///
/// // No need to import futures::StreamExt!
/// loop {
/// match event_stream.next().await {
/// Ok(event) => println!("Event: {:?}", event),
/// Err(e) if matches!(e, pg_walstream::ReplicationError::Cancelled(_)) => {
/// println!("Stream cancelled, shutting down gracefully");
/// break;
/// }
/// Err(e) => {
/// eprintln!("Error: {}", e);
/// break;
/// }
/// }
/// }
/// # Ok(())
/// # }
/// ```
pub async fn next(&mut self) -> Result<ChangeEvent> {
self.inner
.next_event_with_retry(&self.cancellation_token)
.await
}
}
/// Fold `connection_timeout` into the connection string as `connect_timeout`
/// unless the caller already set one.
///
/// This is what gives [`timeout_or_error`] teeth. `connect_with_retry` calls the
/// *blocking* `PgReplicationConnection::connect` inline, so the `tokio::time::timeout`
/// wrapped around it can only fire at the backoff `await` *between* attempts — a
/// single connect parked against a peer that completes the TCP handshake and then
/// goes silent is unbounded, and no caller-side timeout can reach it.
///
/// `connect_timeout` is the bound that does work, on both backends: the native
/// backend wraps the whole handshake in its own `tokio::time::timeout` driven by
/// the worker's runtime (where the timer *is* polled), and libpq honours the
/// keyword natively. A caller who sets it explicitly wins — this only supplies a
/// default. `Duration::ZERO` means "no timeout", matching libpq, so it adds nothing.
fn with_connect_timeout(conn_str: &str, timeout: Duration) -> String {
if timeout.is_zero() || conn_str.contains("connect_timeout") {
return conn_str.to_string();
}
// libpq's `connect_timeout` is whole seconds, and a sub-second value must not
// round down to 0 — that is the "disabled" sentinel.
let secs = timeout.as_secs().max(1);
if conn_str.starts_with("postgres://") || conn_str.starts_with("postgresql://") {
let sep = if conn_str.contains('?') { '&' } else { '?' };
format!("{conn_str}{sep}connect_timeout={secs}")
} else {
format!("{conn_str} connect_timeout={secs}")
}
}
async fn timeout_or_error<T>(
duration: Duration,
future: impl Future<Output = Result<T>>,
) -> Result<T> {
match tokio::time::timeout(duration, future).await {
Ok(result) => result,
Err(_) => Err(ReplicationError::timeout(format!(
"Connection attempt exceeded timeout of {duration:?}"
))),
}
}
#[cfg(test)]
mod reached_stop_lsn_tests {
use super::*;
use crate::column_value::{ColumnValue, RowData};
fn ts() -> chrono::DateTime<chrono::Utc> {
chrono::DateTime::<chrono::Utc>::from_timestamp(0, 0).unwrap()
}
fn commit(end: u64) -> ChangeEvent {
ChangeEvent {
event_type: EventType::Commit {
commit_timestamp: ts(),
commit_lsn: Lsn::new(end),
end_lsn: Lsn::new(end),
},
lsn: Lsn::new(end),
metadata: None,
}
}
fn stream_commit(end: u64) -> ChangeEvent {
ChangeEvent {
event_type: EventType::StreamCommit {
transaction_id: 7,
commit_lsn: Lsn::new(end),
end_lsn: Lsn::new(end),
commit_timestamp: ts(),
},
lsn: Lsn::new(end),
metadata: None,
}
}
fn commit_prepared(end: u64) -> ChangeEvent {
ChangeEvent {
event_type: EventType::CommitPrepared {
flags: 0,
transaction_id: 7,
commit_lsn: Lsn::new(end),
end_lsn: Lsn::new(end),
commit_timestamp: ts(),
gid: std::sync::Arc::from("g"),
},
lsn: Lsn::new(end),
metadata: None,
}
}
fn insert(lsn: u64) -> ChangeEvent {
let data = RowData::from_pairs(vec![("id", ColumnValue::text("1"))]);
ChangeEvent::insert("public", "t", 1, data, Lsn::new(lsn))
}
#[test]
fn none_when_feature_off() {
assert_eq!(reached_stop_lsn(None, &commit(9999)), None);
}
#[test]
fn stops_at_or_after_target_on_commit() {
assert_eq!(reached_stop_lsn(Some(Lsn::new(100)), &commit(99)), None);
assert_eq!(
reached_stop_lsn(Some(Lsn::new(100)), &commit(100)),
Some(Lsn::new(100))
);
assert_eq!(
reached_stop_lsn(Some(Lsn::new(100)), &commit(101)),
Some(Lsn::new(101))
);
}
#[test]
fn handles_stream_commit_and_commit_prepared() {
assert_eq!(
reached_stop_lsn(Some(Lsn::new(50)), &stream_commit(50)),
Some(Lsn::new(50))
);
assert_eq!(
reached_stop_lsn(Some(Lsn::new(50)), &commit_prepared(60)),
Some(Lsn::new(60))
);
assert_eq!(
reached_stop_lsn(Some(Lsn::new(50)), &stream_commit(49)),
None
);
}
#[test]
fn never_stops_on_non_commit_events() {
assert_eq!(reached_stop_lsn(Some(Lsn::new(1)), &insert(9999)), None);
}
}
#[cfg(test)]
pub(crate) mod tests {
use super::*;
use crate::column_value::{ColumnValue, RowData};
use crate::protocol::TupleData;
use crate::types::{parse_lsn, ReplicaIdentity};
#[allow(dead_code)]
const _: () = {
fn assert_send<T: Send>() {}
fn assertions() {
assert_send::<PgReplicationConnection>();
assert_send::<crate::connection::PgResult>();
assert_send::<LogicalReplicationStream>();
assert_send::<EventStream>();
}
};
/// Helper function to create a test configuration
fn create_test_config() -> ReplicationStreamConfig {
ReplicationStreamConfig::new(
"test_slot".to_string(),
"test_publication".to_string(),
2,
StreamingMode::On,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
)
}
#[test]
fn test_replication_stream_config_creation() {
let config = create_test_config();
assert_eq!(config.slot_name, "test_slot");
assert_eq!(config.publication_name, "test_publication");
assert_eq!(config.protocol_version, 2);
assert_eq!(config.streaming_mode, StreamingMode::On);
assert_eq!(config.feedback_interval, Duration::from_secs(10));
assert_eq!(config.connection_timeout, Duration::from_secs(30));
assert_eq!(config.health_check_interval, Duration::from_secs(60));
assert!(matches!(config.slot_type, SlotType::Logical));
}
#[tokio::test]
async fn test_connection_timeout_exceeded() {
let config = ReplicationStreamConfig::new(
"test_slot".to_string(),
"test_publication".to_string(),
2,
StreamingMode::On,
Duration::from_secs(10),
Duration::from_millis(1),
Duration::from_secs(60),
RetryConfig::default(),
);
let result = timeout_or_error(config.connection_timeout, async {
tokio::time::sleep(Duration::from_millis(50)).await;
Ok::<_, ReplicationError>(())
})
.await;
assert!(matches!(result, Err(ReplicationError::Timeout(_))));
}
#[test]
fn test_replication_stream_config_clone() {
let config1 = create_test_config();
let config2 = config1.clone();
assert_eq!(config1.slot_name, config2.slot_name);
assert_eq!(config1.publication_name, config2.publication_name);
assert_eq!(config1.protocol_version, config2.protocol_version);
}
#[test]
fn test_replication_state_new() {
let state = ReplicationState::new();
assert_eq!(state.last_received_lsn, 0);
}
#[test]
fn test_replication_state_update_lsn() {
let mut state = ReplicationState::new();
// Update to a higher LSN
state.update_received_lsn(1000);
assert_eq!(state.last_received_lsn, 1000);
// Update to an even higher LSN
state.update_received_lsn(2000);
assert_eq!(state.last_received_lsn, 2000);
// Trying to update to a lower LSN should not change it
state.update_received_lsn(500);
assert_eq!(state.last_received_lsn, 2000);
}
#[test]
#[allow(clippy::assertions_on_constants)]
fn test_feedback_check_interval_is_power_of_two() {
// The hot-path gate folds `counter % INTERVAL` into a bitmask via
// `counter & (INTERVAL - 1)`, which is only correct when INTERVAL
// is a power of two and non-zero.
assert!(FEEDBACK_CHECK_EVENT_INTERVAL > 0);
assert_eq!(
FEEDBACK_CHECK_EVENT_INTERVAL & (FEEDBACK_CHECK_EVENT_INTERVAL - 1),
0,
"FEEDBACK_CHECK_EVENT_INTERVAL must be a power of two"
);
}
#[test]
fn test_feedback_check_counter_initialized_zero() {
let config = create_test_config();
let stream = create_test_stream(config);
assert_eq!(stream.feedback_check_counter, 0);
}
#[test]
fn test_feedback_check_gate_semantics() {
// The gate fires when `(counter + 1) & (INTERVAL - 1) == 0`, i.e. at
// event numbers INTERVAL, 2*INTERVAL, 3*INTERVAL, ...
let interval = FEEDBACK_CHECK_EVENT_INTERVAL;
let mut counter: u32 = 0;
let mut fire_events = Vec::new();
for event_idx in 1..=(interval * 3) {
counter = counter.wrapping_add(1);
if counter & (interval - 1) == 0 {
fire_events.push(event_idx);
}
}
assert_eq!(
fire_events,
vec![interval, interval * 2, interval * 3],
"gate should fire exactly on multiples of the interval"
);
}
#[test]
#[allow(clippy::assertions_on_constants)]
fn test_health_check_event_interval_is_power_of_two() {
// `next_event_with_retry` uses the same `counter & (INTERVAL - 1) == 0`
// bitmask trick to gate `Instant::now()` calls, so HEALTH_CHECK_EVENT_INTERVAL
// must be a non-zero power of two.
assert!(HEALTH_CHECK_EVENT_INTERVAL > 0);
assert_eq!(
HEALTH_CHECK_EVENT_INTERVAL & (HEALTH_CHECK_EVENT_INTERVAL - 1),
0,
"HEALTH_CHECK_EVENT_INTERVAL must be a power of two"
);
}
#[test]
#[allow(clippy::assertions_on_constants)]
fn test_health_check_event_interval_at_least_feedback_interval() {
// The health check interval is intentionally coarser than the feedback
// interval (and is a multiple of it) so feedback never gets starved by
// the health-check throttle. Pin the relationship in case either
// constant is ever bumped without consideration.
assert!(
HEALTH_CHECK_EVENT_INTERVAL >= FEEDBACK_CHECK_EVENT_INTERVAL,
"health-check interval should not be tighter than the feedback interval"
);
assert_eq!(
HEALTH_CHECK_EVENT_INTERVAL % FEEDBACK_CHECK_EVENT_INTERVAL,
0,
"HEALTH_CHECK_EVENT_INTERVAL should be a multiple of FEEDBACK_CHECK_EVENT_INTERVAL"
);
}
#[test]
fn test_health_check_gate_fires_on_first_event() {
// The gate is `counter & (INTERVAL - 1) == 0`. The counter starts at
// 0, so the very first call to `next_event_with_retry` must run the
// health check — important for catching dead connections at startup
// even on low-throughput streams.
let counter: u32 = 0;
assert_eq!(counter & (HEALTH_CHECK_EVENT_INTERVAL - 1), 0);
}
#[test]
fn test_health_check_gate_fires_on_multiples() {
// Mirrors `test_feedback_check_gate_semantics` for the health-check gate.
let interval = HEALTH_CHECK_EVENT_INTERVAL;
let mut counter: u32 = 0;
let mut fire_events = Vec::new();
// Walk through 3 intervals' worth of events and record where the gate fires.
for event_idx in 0..(interval * 3) {
if counter & (interval - 1) == 0 {
fire_events.push(event_idx);
}
counter = counter.wrapping_add(1);
}
assert_eq!(
fire_events,
vec![0, interval, interval * 2],
"health-check gate should fire at counter 0, INTERVAL, 2*INTERVAL"
);
}
#[test]
fn test_health_check_gate_handles_counter_wraparound() {
// The counter is `u32` and uses `wrapping_add`. The gate must continue
// to fire correctly when the counter wraps from u32::MAX back to 0.
let interval = HEALTH_CHECK_EVENT_INTERVAL;
// u32::MAX = 0xFFFF_FFFF; with INTERVAL=1024, (INTERVAL - 1) = 0x3FF,
// so u32::MAX & 0x3FF == 0x3FF (gate does NOT fire).
let counter: u32 = u32::MAX;
assert_ne!(counter & (interval - 1), 0);
// After wrapping_add(1) the counter wraps to 0, and the gate fires.
let next = counter.wrapping_add(1);
assert_eq!(next, 0);
assert_eq!(next & (interval - 1), 0);
}
#[test]
fn test_replication_state_should_send_feedback() {
let mut state = ReplicationState::new();
let feedback_interval = Duration::from_millis(50);
// Update LSN so there's something to send feedback about
state.update_received_lsn(1000);
// Should send feedback after receiving data (interval has elapsed since creation)
std::thread::sleep(Duration::from_millis(60));
assert!(state.should_send_feedback(feedback_interval));
// Mark as sent
state.mark_feedback_sent();
// Immediately after, should not send
assert!(!state.should_send_feedback(feedback_interval));
// After waiting long enough, should send again
std::thread::sleep(Duration::from_millis(60));
assert!(state.should_send_feedback(feedback_interval));
}
#[test]
fn test_relation_info_creation() {
let columns = vec![
crate::protocol::ColumnInfo {
flags: 1,
name: Arc::from("id"),
type_id: 23,
type_modifier: -1,
},
crate::protocol::ColumnInfo {
flags: 0,
name: Arc::from("name"),
type_id: 25,
type_modifier: -1,
},
];
let relation = RelationInfo::new(16384, "public", "users", b'd', columns);
assert_eq!(relation.relation_id, 16384);
assert_eq!(&*relation.namespace, "public");
assert_eq!(&*relation.relation_name, "users");
assert_eq!(relation.columns.len(), 2);
}
#[test]
fn test_relation_info_full_name() {
let relation = RelationInfo::new(16384, "public", "users", b'd', vec![]);
assert_eq!(relation.full_name(), "public.users");
}
#[test]
fn test_relation_info_get_column_by_index() {
let columns = vec![
crate::protocol::ColumnInfo {
flags: 1,
name: Arc::from("id"),
type_id: 23,
type_modifier: -1,
},
crate::protocol::ColumnInfo {
flags: 0,
name: Arc::from("name"),
type_id: 25,
type_modifier: -1,
},
];
let relation = RelationInfo::new(16384, "public", "users", b'd', columns);
// Valid index
assert!(relation.get_column_by_index(0).is_some());
assert_eq!(&*relation.get_column_by_index(0).unwrap().name, "id");
// Invalid index
assert!(relation.get_column_by_index(10).is_none());
}
#[test]
fn test_relation_info_get_key_columns() {
let columns = vec![
crate::protocol::ColumnInfo {
flags: 1, // Key column
name: Arc::from("id"),
type_id: 23,
type_modifier: -1,
},
crate::protocol::ColumnInfo {
flags: 0, // Non-key column
name: Arc::from("name"),
type_id: 25,
type_modifier: -1,
},
];
let relation = RelationInfo::new(16384, "public", "users", b'd', columns);
let key_columns = relation.get_key_columns();
assert_eq!(key_columns.len(), 1);
assert_eq!(&*key_columns[0].name, "id");
}
#[test]
fn test_change_event_insert_creation() {
let data = RowData::from_pairs(vec![
("id", ColumnValue::text("1")),
("name", ColumnValue::text("Alice")),
]);
let event = ChangeEvent::insert("public", "users", 16384, data.clone(), Lsn::new(1000));
match event.event_type {
EventType::Insert {
schema,
table,
relation_oid,
data: event_data,
} => {
assert_eq!(&*schema, "public");
assert_eq!(&*table, "users");
assert_eq!(relation_oid, 16384);
assert_eq!(event_data.len(), 2);
}
_ => panic!("Expected Insert event"),
}
assert_eq!(event.lsn.value(), 1000);
}
#[test]
fn test_change_event_update_creation() {
let old_data = RowData::from_pairs(vec![
("id", ColumnValue::text("1")),
("name", ColumnValue::text("Alice")),
]);
let new_data = RowData::from_pairs(vec![
("id", ColumnValue::text("1")),
("name", ColumnValue::text("Bob")),
]);
let event = ChangeEvent::update(
"public",
"users",
16384,
Some(old_data),
new_data,
ReplicaIdentity::Default,
vec![Arc::from("id")],
Lsn::new(2000),
);
match event.event_type {
EventType::Update {
schema,
table,
relation_oid,
old_data,
new_data,
replica_identity,
key_columns,
} => {
assert_eq!(&*schema, "public");
assert_eq!(&*table, "users");
assert_eq!(relation_oid, 16384);
assert!(old_data.is_some());
assert_eq!(new_data.len(), 2);
assert_eq!(replica_identity, ReplicaIdentity::Default);
assert_eq!(key_columns.len(), 1);
}
_ => panic!("Expected Update event"),
}
}
#[test]
fn test_change_event_delete_creation() {
let old_data = RowData::from_pairs(vec![("id", ColumnValue::text("1"))]);
let event = ChangeEvent::delete(
"public",
"users",
16384,
old_data,
ReplicaIdentity::Default,
vec![Arc::from("id")],
Lsn::new(3000),
);
match event.event_type {
EventType::Delete {
schema,
table,
relation_oid,
old_data,
replica_identity,
key_columns,
} => {
assert_eq!(&*schema, "public");
assert_eq!(&*table, "users");
assert_eq!(relation_oid, 16384);
assert_eq!(old_data.len(), 1);
assert_eq!(replica_identity, ReplicaIdentity::Default);
assert_eq!(key_columns.len(), 1);
}
_ => panic!("Expected Delete event"),
}
}
#[test]
fn test_tuple_to_data_skips_unchanged_columns() {
let columns = vec![
crate::protocol::ColumnInfo {
flags: 0,
name: Arc::from("col1"),
type_id: 25,
type_modifier: -1,
},
crate::protocol::ColumnInfo {
flags: 0,
name: Arc::from("col2"),
type_id: 25,
type_modifier: -1,
},
];
let relation = RelationInfo::new(1, "public", "test", b'd', columns);
let tuple = TupleData::new(vec![
crate::protocol::ColumnData::unchanged(),
crate::protocol::ColumnData::text(b"value".to_vec()),
]);
let data = tuple.into_row_data(&relation);
assert_eq!(data.len(), 1);
assert_eq!(data.get("col2").unwrap(), "value");
}
#[test]
fn test_tuple_to_data_empty_text_column() {
let columns = vec![crate::protocol::ColumnInfo {
flags: 0,
name: Arc::from("col1"),
type_id: 25,
type_modifier: -1,
}];
let relation = RelationInfo::new(1, "public", "test", b'd', columns);
let tuple = TupleData::new(vec![crate::protocol::ColumnData::text(Vec::new())]);
let data = tuple.into_row_data(&relation);
assert_eq!(data.get("col1").unwrap(), "");
}
#[test]
fn test_cancellation_token_basic() {
use tokio_util::sync::CancellationToken;
let cancel_token = CancellationToken::new();
// Initially not cancelled
assert!(!cancel_token.is_cancelled());
// Cancel it
cancel_token.cancel();
// Should be cancelled now
assert!(cancel_token.is_cancelled());
}
#[test]
fn test_cancellation_token_clone() {
use tokio_util::sync::CancellationToken;
let cancel_token = CancellationToken::new();
let clone = cancel_token.clone();
// Neither should be cancelled initially
assert!(!cancel_token.is_cancelled());
assert!(!clone.is_cancelled());
// Cancel the original
cancel_token.cancel();
// Both should be cancelled
assert!(cancel_token.is_cancelled());
assert!(clone.is_cancelled());
}
#[test]
fn test_shared_lsn_feedback_integration() {
let feedback = SharedLsnFeedback::new_shared();
// Initial values should be 0
assert_eq!(feedback.get_flushed_lsn(), 0);
assert_eq!(feedback.get_applied_lsn(), 0);
// Update flushed LSN
feedback.update_flushed_lsn(1000);
assert_eq!(feedback.get_flushed_lsn(), 1000);
assert_eq!(feedback.get_applied_lsn(), 0);
// Update applied LSN (should also update flushed)
feedback.update_applied_lsn(2000);
assert_eq!(feedback.get_flushed_lsn(), 2000);
assert_eq!(feedback.get_applied_lsn(), 2000);
// Get both values at once
let (flushed, applied) = feedback.get_feedback_lsn();
assert_eq!(flushed, 2000);
assert_eq!(applied, 2000);
}
#[test]
fn test_lsn_value_operations() {
let lsn1 = Lsn::new(1000);
let lsn2 = Lsn::new(2000);
let lsn3 = Lsn::new(1000);
// Test value retrieval
assert_eq!(lsn1.value(), 1000);
assert_eq!(lsn2.value(), 2000);
// Test equality
assert_eq!(lsn1, lsn3);
assert_ne!(lsn1, lsn2);
// Test ordering
assert!(lsn1 < lsn2);
assert!(lsn2 > lsn1);
}
#[test]
fn test_replica_identity_conversions() {
// Test Default
assert_eq!(
ReplicaIdentity::from_byte(b'd'),
Some(ReplicaIdentity::Default)
);
assert_eq!(ReplicaIdentity::Default.to_byte(), b'd');
// Test Nothing
assert_eq!(
ReplicaIdentity::from_byte(b'n'),
Some(ReplicaIdentity::Nothing)
);
assert_eq!(ReplicaIdentity::Nothing.to_byte(), b'n');
// Test Full
assert_eq!(
ReplicaIdentity::from_byte(b'f'),
Some(ReplicaIdentity::Full)
);
assert_eq!(ReplicaIdentity::Full.to_byte(), b'f');
// Test Index
assert_eq!(
ReplicaIdentity::from_byte(b'i'),
Some(ReplicaIdentity::Index)
);
assert_eq!(ReplicaIdentity::Index.to_byte(), b'i');
// Test invalid byte
assert_eq!(ReplicaIdentity::from_byte(b'x'), None);
}
#[test]
fn test_retry_config_default() {
let config = RetryConfig::default();
assert_eq!(config.max_attempts, 5);
assert_eq!(config.initial_delay, Duration::from_secs(1));
assert_eq!(config.max_delay, Duration::from_secs(60));
assert_eq!(config.multiplier, 2.0);
assert_eq!(config.max_duration, Duration::from_secs(300));
assert!(config.jitter);
}
#[test]
fn test_retry_config_custom() {
let config = RetryConfig {
max_attempts: 10,
initial_delay: Duration::from_millis(500),
max_delay: Duration::from_secs(30),
multiplier: 3.0,
max_duration: Duration::from_secs(600),
jitter: false,
};
assert_eq!(config.max_attempts, 10);
assert_eq!(config.initial_delay, Duration::from_millis(500));
assert!(!config.jitter);
}
#[test]
fn test_exponential_backoff_progression() {
let config = RetryConfig {
max_attempts: 5,
initial_delay: Duration::from_millis(100),
max_delay: Duration::from_secs(10),
multiplier: 2.0,
max_duration: Duration::from_secs(60),
jitter: false,
};
let mut backoff = crate::retry::ExponentialBackoff::new(&config);
// First delay should be initial_delay
let delay1 = backoff.next_delay();
assert!(delay1 >= Duration::from_millis(100));
assert!(delay1 <= Duration::from_millis(100));
// Second delay should be doubled
let delay2 = backoff.next_delay();
assert!(delay2 >= Duration::from_millis(200));
// Third delay should be doubled again
let delay3 = backoff.next_delay();
assert!(delay3 >= Duration::from_millis(400));
}
#[test]
fn test_format_lsn() {
assert_eq!(format_lsn(0), "0/0");
assert_eq!(format_lsn(0x16B374D848), "16/B374D848");
assert_eq!(format_lsn(0x100000000), "1/0");
assert_eq!(format_lsn(0xFFFFFFFFFFFFFFFF), "FFFFFFFF/FFFFFFFF");
}
#[test]
fn test_parse_lsn() {
assert_eq!(parse_lsn("0/0").unwrap(), 0);
assert_eq!(parse_lsn("16/B374D848").unwrap(), 0x16B374D848);
assert_eq!(parse_lsn("1/0").unwrap(), 0x100000000);
// Test invalid formats
assert!(parse_lsn("invalid").is_err());
assert!(parse_lsn("1/2/3").is_err());
assert!(parse_lsn("xyz/abc").is_err());
}
#[test]
fn test_lsn_roundtrip() {
let original_lsn = 0x16B374D848u64;
let formatted = format_lsn(original_lsn);
let parsed = parse_lsn(&formatted).unwrap();
assert_eq!(original_lsn, parsed);
}
#[test]
fn test_event_stream_lsn_feedback_update() {
let feedback = SharedLsnFeedback::new_shared();
// Simulate updating flushed LSN
feedback.update_flushed_lsn(1000);
assert_eq!(feedback.get_flushed_lsn(), 1000);
// Simulate updating applied LSN (should also update flushed)
feedback.update_applied_lsn(2000);
assert_eq!(feedback.get_flushed_lsn(), 2000);
assert_eq!(feedback.get_applied_lsn(), 2000);
}
#[test]
fn test_event_stream_get_feedback_lsn() {
let feedback = SharedLsnFeedback::new_shared();
feedback.update_flushed_lsn(5000);
feedback.update_applied_lsn(10000);
let (flushed, applied) = feedback.get_feedback_lsn();
assert_eq!(flushed, 10000); // Applied also updates flushed
assert_eq!(applied, 10000);
}
#[test]
fn test_event_stream_lsn_monotonic_increase() {
let feedback = SharedLsnFeedback::new_shared();
// LSNs should only increase
feedback.update_applied_lsn(1000);
assert_eq!(feedback.get_applied_lsn(), 1000);
feedback.update_applied_lsn(2000);
assert_eq!(feedback.get_applied_lsn(), 2000);
// Trying to set a lower value should keep the higher one
feedback.update_applied_lsn(500);
assert_eq!(feedback.get_applied_lsn(), 2000);
}
#[test]
fn test_config_with_different_protocol_versions() {
let config_v1 = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
1,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
assert_eq!(config_v1.protocol_version, 1);
assert_eq!(config_v1.streaming_mode, StreamingMode::Off);
let config_v4 = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
4,
StreamingMode::Parallel,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
assert_eq!(config_v4.protocol_version, 4);
assert_eq!(config_v4.streaming_mode, StreamingMode::Parallel);
}
#[test]
fn test_config_with_custom_intervals() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
2,
StreamingMode::On,
Duration::from_millis(500),
Duration::from_secs(5),
Duration::from_secs(15),
RetryConfig::default(),
);
assert_eq!(config.feedback_interval, Duration::from_millis(500));
assert_eq!(config.connection_timeout, Duration::from_secs(5));
assert_eq!(config.health_check_interval, Duration::from_secs(15));
}
#[test]
fn test_replication_state_lsn_tracking() {
let mut state = ReplicationState::new();
// Track multiple LSN updates
state.update_received_lsn(100);
assert_eq!(state.last_received_lsn, 100);
state.update_received_lsn(200);
assert_eq!(state.last_received_lsn, 200);
state.update_received_lsn(300);
assert_eq!(state.last_received_lsn, 300);
// Lower LSN should not update
state.update_received_lsn(150);
assert_eq!(state.last_received_lsn, 300);
}
#[test]
fn test_replication_state_feedback_timing() {
let mut state = ReplicationState::new();
let interval = Duration::from_millis(100);
// Initial state should allow feedback
state.update_received_lsn(1000);
std::thread::sleep(Duration::from_millis(110));
assert!(state.should_send_feedback(interval));
// Mark as sent
state.mark_feedback_sent();
// Should not send immediately after
assert!(!state.should_send_feedback(interval));
// After interval, should send again
std::thread::sleep(Duration::from_millis(110));
assert!(state.should_send_feedback(interval));
}
#[test]
fn test_replication_state_zero_lsn() {
let mut state = ReplicationState::new();
// Test with zero LSN
assert_eq!(state.last_received_lsn, 0);
state.update_received_lsn(0);
assert_eq!(state.last_received_lsn, 0);
// Update to non-zero
state.update_received_lsn(1);
assert_eq!(state.last_received_lsn, 1);
}
// ========================================
// RelationInfo Advanced Tests
// ========================================
#[test]
fn test_relation_info_empty_columns() {
let relation = RelationInfo::new(12345, "schema", "table", b'd', vec![]);
assert_eq!(relation.columns.len(), 0);
assert!(relation.get_column_by_index(0).is_none());
assert_eq!(relation.get_key_columns().len(), 0);
}
#[test]
fn test_relation_info_multiple_key_columns() {
let columns = vec![
crate::protocol::ColumnInfo {
flags: 1,
name: Arc::from("id1"),
type_id: 23,
type_modifier: -1,
},
crate::protocol::ColumnInfo {
flags: 1,
name: Arc::from("id2"),
type_id: 23,
type_modifier: -1,
},
crate::protocol::ColumnInfo {
flags: 0,
name: Arc::from("data"),
type_id: 25,
type_modifier: -1,
},
];
let relation = RelationInfo::new(12345, "public", "composite_key_table", b'd', columns);
let key_columns = relation.get_key_columns();
assert_eq!(key_columns.len(), 2);
assert_eq!(&*key_columns[0].name, "id1");
assert_eq!(&*key_columns[1].name, "id2");
}
#[test]
fn test_relation_info_different_replica_identities() {
// Test Default
let rel_default = RelationInfo::new(1, "public", "t1", b'd', vec![]);
assert_eq!(rel_default.replica_identity, b'd');
// Test Full
let rel_full = RelationInfo::new(2, "public", "t2", b'f', vec![]);
assert_eq!(rel_full.replica_identity, b'f');
// Test Nothing
let rel_nothing = RelationInfo::new(3, "public", "t3", b'n', vec![]);
assert_eq!(rel_nothing.replica_identity, b'n');
// Test Index
let rel_index = RelationInfo::new(4, "public", "t4", b'i', vec![]);
assert_eq!(rel_index.replica_identity, b'i');
}
#[test]
fn test_change_event_with_null_values() {
let data = RowData::from_pairs(vec![
("id", ColumnValue::text("1")),
("nullable_field", ColumnValue::Null),
]);
let event = ChangeEvent::insert(
"public".to_string(),
"users".to_string(),
12345,
data,
Lsn::new(1000),
);
match event.event_type {
EventType::Insert { data, .. } => {
assert_eq!(data.get("id").unwrap(), &ColumnValue::text("1"));
assert_eq!(data.get("nullable_field").unwrap(), &ColumnValue::Null);
}
_ => panic!("Expected Insert event"),
}
}
#[test]
fn test_change_event_update_without_old_data() {
let new_data = RowData::new();
let event = ChangeEvent::update(
"public",
"users",
12345,
None, // No old data
new_data,
ReplicaIdentity::Nothing,
vec![],
Lsn::new(2000),
);
match event.event_type {
EventType::Update {
old_data,
replica_identity,
..
} => {
assert!(old_data.is_none());
assert_eq!(replica_identity, ReplicaIdentity::Nothing);
}
_ => panic!("Expected Update event"),
}
}
#[test]
fn test_change_event_with_large_lsn() {
let data = RowData::new();
let large_lsn = 0xFFFFFFFF_FFFFFFFFu64;
let event = ChangeEvent::insert("public", "test", 1, data, Lsn::new(large_lsn));
assert_eq!(event.lsn.value(), large_lsn);
}
#[test]
fn test_change_event_metadata_field() {
let data = RowData::new();
let event = ChangeEvent::insert("public", "test", 1, data, Lsn::new(1000));
// Verify metadata field exists and is initially None
assert!(event.metadata.is_none());
}
#[test]
fn test_lsn_arithmetic_operations() {
let lsn1 = Lsn::new(100);
let lsn2 = Lsn::new(200);
let lsn3 = Lsn::new(100);
// Test equality
assert_eq!(lsn1, lsn3);
assert_ne!(lsn1, lsn2);
// Test ordering
assert!(lsn1 < lsn2);
assert!(lsn2 > lsn1);
assert!(lsn1 <= lsn3);
assert!(lsn1 >= lsn3);
}
#[test]
fn test_lsn_edge_cases() {
let lsn_zero = Lsn::new(0);
let lsn_max = Lsn::new(u64::MAX);
assert_eq!(lsn_zero.value(), 0);
assert_eq!(lsn_max.value(), u64::MAX);
assert!(lsn_zero < lsn_max);
}
#[test]
fn test_format_lsn_edge_cases() {
// Test zero
assert_eq!(format_lsn(0), "0/0");
// Test max value
assert_eq!(format_lsn(u64::MAX), "FFFFFFFF/FFFFFFFF");
// Test boundary values
assert_eq!(format_lsn(0xFFFFFFFF), "0/FFFFFFFF");
assert_eq!(format_lsn(0x100000000), "1/0");
}
#[test]
fn test_parse_lsn_various_formats() {
// Standard format
assert_eq!(parse_lsn("1/2A").unwrap(), 0x10000002A);
// Lowercase hex
assert_eq!(parse_lsn("a/b").unwrap(), 0xa0000000b);
// Mixed case
assert_eq!(parse_lsn("Ab/Cd").unwrap(), 0xab000000cd);
// Leading zeros
assert_eq!(parse_lsn("00/00").unwrap(), 0);
assert_eq!(parse_lsn("01/00000001").unwrap(), 0x100000001);
}
#[test]
fn test_parse_lsn_invalid_formats() {
assert!(parse_lsn("").is_err());
assert!(parse_lsn("123").is_err());
assert!(parse_lsn("abc").is_err());
assert!(parse_lsn("1/2/3").is_err());
assert!(parse_lsn("xyz/123").is_err());
assert!(parse_lsn("123/xyz").is_err());
assert!(parse_lsn("/").is_err());
assert!(parse_lsn("1/").is_err());
assert!(parse_lsn("/1").is_err());
}
#[test]
fn test_retry_config_with_jitter() {
let config_with_jitter = RetryConfig {
max_attempts: 3,
initial_delay: Duration::from_millis(100),
max_delay: Duration::from_secs(5),
multiplier: 2.0,
max_duration: Duration::from_secs(30),
jitter: true,
};
assert!(config_with_jitter.jitter);
let config_without_jitter = RetryConfig {
jitter: false,
..config_with_jitter
};
assert!(!config_without_jitter.jitter);
}
#[test]
fn test_retry_config_max_duration() {
let config = RetryConfig {
max_attempts: 100,
initial_delay: Duration::from_millis(10),
max_delay: Duration::from_secs(30),
multiplier: 2.0,
max_duration: Duration::from_secs(5),
jitter: false,
};
assert_eq!(config.max_duration, Duration::from_secs(5));
}
#[test]
fn test_exponential_backoff_respects_max_delay() {
let config = RetryConfig {
max_attempts: 10,
initial_delay: Duration::from_millis(100),
max_delay: Duration::from_millis(500),
multiplier: 2.0,
max_duration: Duration::from_secs(60),
jitter: false,
};
let mut backoff = crate::retry::ExponentialBackoff::new(&config);
// First few delays
let _d1 = backoff.next_delay();
let _d2 = backoff.next_delay();
let _d3 = backoff.next_delay();
// After several iterations, should not exceed max_delay
let d4 = backoff.next_delay();
assert!(d4 <= Duration::from_millis(500));
}
#[tokio::test]
async fn test_cancellation_token_async_cancel() {
let token = CancellationToken::new();
let token_clone = token.clone();
// Spawn a task that waits for cancellation
let handle = tokio::spawn(async move {
token_clone.cancelled().await;
"cancelled"
});
// Cancel after a short delay
tokio::time::sleep(Duration::from_millis(10)).await;
token.cancel();
// Task should complete
let result = handle.await.unwrap();
assert_eq!(result, "cancelled");
}
#[tokio::test]
async fn test_cancellation_propagates_to_children() {
let parent = CancellationToken::new();
let child = parent.child_token();
assert!(!parent.is_cancelled());
assert!(!child.is_cancelled());
// Cancel parent
parent.cancel();
// Both should be cancelled
assert!(parent.is_cancelled());
assert!(child.is_cancelled());
}
#[tokio::test]
async fn test_shared_lsn_feedback_concurrent_updates() {
let feedback = SharedLsnFeedback::new_shared();
let feedback_clone = feedback.clone();
// Spawn multiple tasks updating LSNs
let handle1 = tokio::spawn(async move {
for i in 0..100 {
feedback_clone.update_applied_lsn(i * 10);
tokio::task::yield_now().await;
}
});
let feedback_clone2 = feedback.clone();
let handle2 = tokio::spawn(async move {
for i in 0..100 {
feedback_clone2.update_flushed_lsn(i * 10 + 5);
tokio::task::yield_now().await;
}
});
handle1.await.unwrap();
handle2.await.unwrap();
// Final LSN should be consistent
let (flushed, applied) = feedback.get_feedback_lsn();
// Verify both are non-zero after concurrent updates
assert!(
flushed > 0 || applied > 0,
"At least one LSN should be updated"
);
}
#[test]
fn test_shared_lsn_feedback_clone_shares_state() {
let feedback1 = SharedLsnFeedback::new_shared();
let feedback2 = feedback1.clone();
// Update through first reference
feedback1.update_applied_lsn(1000);
// Should be visible through second reference
assert_eq!(feedback2.get_applied_lsn(), 1000);
// Update through second reference
feedback2.update_flushed_lsn(2000);
// Should be visible through first reference
assert_eq!(feedback1.get_flushed_lsn(), 2000);
}
#[test]
fn test_full_change_event_lifecycle() {
// Test complete lifecycle: insert -> update -> delete
let insert_data = RowData::from_pairs(vec![
("id", ColumnValue::text("1")),
("name", ColumnValue::text("Alice")),
("email", ColumnValue::text("alice@example.com")),
]);
let insert_event = ChangeEvent::insert(
"public",
"users",
16384,
insert_data.clone(),
Lsn::new(1000),
);
assert_eq!(insert_event.lsn.value(), 1000);
// Update event
let update_data = RowData::from_pairs(vec![
("id", ColumnValue::text("1")),
("name", ColumnValue::text("Alice")),
("email", ColumnValue::text("alice.new@example.com")),
]);
let update_event = ChangeEvent::update(
"public",
"users",
16384,
Some(insert_data.clone()),
update_data,
ReplicaIdentity::Default,
vec![Arc::from("id")],
Lsn::new(2000),
);
assert_eq!(update_event.lsn.value(), 2000);
// Delete event
let delete_key = RowData::from_pairs(vec![("id", ColumnValue::text("1"))]);
let delete_event = ChangeEvent::delete(
"public",
"users",
16384,
delete_key,
ReplicaIdentity::Default,
vec![Arc::from("id")],
Lsn::new(3000),
);
assert_eq!(delete_event.lsn.value(), 3000);
// Verify LSN progression
assert!(insert_event.lsn < update_event.lsn);
assert!(update_event.lsn < delete_event.lsn);
}
#[test]
fn test_replication_state_feedback_workflow() {
let mut state = ReplicationState::new();
let feedback_interval = Duration::from_millis(50);
// Simulate receiving data
state.update_received_lsn(1000);
// Should send feedback since we have new data
std::thread::sleep(Duration::from_millis(60));
assert!(state.should_send_feedback(feedback_interval));
// Mark as sent
state.mark_feedback_sent();
// Update with more data
state.update_received_lsn(2000);
// Should not send immediately
assert!(!state.should_send_feedback(feedback_interval));
// After interval, should send
std::thread::sleep(Duration::from_millis(60));
assert!(state.should_send_feedback(feedback_interval));
}
#[test]
fn test_relation_info_with_complex_schema() {
let columns = vec![
crate::protocol::ColumnInfo {
flags: 1,
name: Arc::from("user_id"),
type_id: 20, // bigint
type_modifier: -1,
},
crate::protocol::ColumnInfo {
flags: 1,
name: Arc::from("tenant_id"),
type_id: 23, // int4
type_modifier: -1,
},
crate::protocol::ColumnInfo {
flags: 0,
name: Arc::from("created_at"),
type_id: 1184, // timestamptz
type_modifier: -1,
},
crate::protocol::ColumnInfo {
flags: 0,
name: Arc::from("data"),
type_id: 3802, // jsonb
type_modifier: -1,
},
crate::protocol::ColumnInfo {
flags: 0,
name: Arc::from("metadata"),
type_id: 25, // text
type_modifier: -1,
},
];
let relation = RelationInfo::new(
16384,
"tenant_1",
"user_events",
b'i', // index replica identity
columns,
);
assert_eq!(relation.full_name(), "tenant_1.user_events");
assert_eq!(relation.columns.len(), 5);
let key_columns = relation.get_key_columns();
assert_eq!(key_columns.len(), 2);
assert_eq!(&*key_columns[0].name, "user_id");
assert_eq!(&*key_columns[1].name, "tenant_id");
// Test column lookup
assert_eq!(
&*relation.get_column_by_index(2).unwrap().name,
"created_at"
);
assert_eq!(&*relation.get_column_by_index(3).unwrap().name, "data");
assert!(relation.get_column_by_index(10).is_none());
}
#[test]
fn test_lsn_feedback_integration_workflow() {
let feedback = SharedLsnFeedback::new_shared();
// Simulate processing events with LSN progression
let lsns = vec![1000, 1500, 2000, 2500, 3000];
for lsn in lsns {
// Simulate flushing to disk
feedback.update_flushed_lsn(lsn);
assert_eq!(feedback.get_flushed_lsn(), lsn);
// Simulate committing transaction
feedback.update_applied_lsn(lsn);
assert_eq!(feedback.get_applied_lsn(), lsn);
let (flushed, applied) = feedback.get_feedback_lsn();
assert_eq!(flushed, lsn);
assert_eq!(applied, lsn);
}
}
#[test]
fn test_multiple_event_types_with_same_relation() {
let relation_oid = 16384;
let schema = "public".to_string();
let table = "orders".to_string();
// Create multiple event types for same relation
let insert_data = RowData::from_pairs(vec![
("order_id", ColumnValue::text("100")),
("amount", ColumnValue::text("99.99")),
]);
let insert = ChangeEvent::insert(
schema.clone(),
table.clone(),
relation_oid,
insert_data.clone(),
Lsn::new(1000),
);
let update_data = RowData::from_pairs(vec![
("order_id", ColumnValue::text("100")),
("amount", ColumnValue::text("89.99")),
]);
let update = ChangeEvent::update(
schema.clone(),
table.clone(),
relation_oid,
Some(insert_data.clone()),
update_data,
ReplicaIdentity::Full,
vec![Arc::from("order_id")],
Lsn::new(2000),
);
let delete = ChangeEvent::delete(
schema.clone(),
table.clone(),
relation_oid,
insert_data,
ReplicaIdentity::Full,
vec![Arc::from("order_id")],
Lsn::new(3000),
);
// Verify all events reference same relation
match insert.event_type {
EventType::Insert {
relation_oid: oid, ..
} => assert_eq!(oid, relation_oid),
_ => panic!("Expected insert"),
}
match update.event_type {
EventType::Update {
relation_oid: oid, ..
} => assert_eq!(oid, relation_oid),
_ => panic!("Expected update"),
}
match delete.event_type {
EventType::Delete {
relation_oid: oid, ..
} => assert_eq!(oid, relation_oid),
_ => panic!("Expected delete"),
}
}
// ========================================
// Edge Case Tests
// ========================================
#[test]
fn test_empty_table_and_schema_names() {
let data = RowData::new();
let event = ChangeEvent::insert("", "", 0, data, Lsn::new(0));
match event.event_type {
EventType::Insert { schema, table, .. } => {
assert_eq!(&*schema, "");
assert_eq!(&*table, "");
}
_ => panic!("Expected insert"),
}
}
#[test]
fn test_very_long_table_names() {
let long_schema = "a".repeat(100);
let long_table = "b".repeat(100);
let data = RowData::new();
let event = ChangeEvent::insert(
long_schema.clone(),
long_table.clone(),
12345,
data,
Lsn::new(1000),
);
match event.event_type {
EventType::Insert { schema, table, .. } => {
assert_eq!(&*schema, long_schema.as_str());
assert_eq!(&*table, long_table.as_str());
}
_ => panic!("Expected insert"),
}
}
#[test]
fn test_special_characters_in_names() {
let data = RowData::new();
let event = ChangeEvent::insert(
"test-schema_123",
"test_table$with#special@chars",
12345,
data,
Lsn::new(1000),
);
match event.event_type {
EventType::Insert { schema, table, .. } => {
assert!(schema.contains("-"));
assert!(table.contains("$"));
assert!(table.contains("#"));
assert!(table.contains("@"));
}
_ => panic!("Expected insert"),
}
}
#[test]
fn test_large_number_of_columns() {
let mut columns = Vec::new();
for i in 0..100 {
columns.push(crate::protocol::ColumnInfo {
flags: if i < 5 { 1 } else { 0 }, // First 5 are keys
name: Arc::from(format!("col_{i}")),
type_id: 23,
type_modifier: -1,
});
}
let relation = RelationInfo::new(12345, "public", "wide_table", b'd', columns);
assert_eq!(relation.columns.len(), 100);
let key_columns = relation.get_key_columns();
assert_eq!(key_columns.len(), 5);
// Test accessing first and last columns
assert_eq!(&*relation.get_column_by_index(0).unwrap().name, "col_0");
assert_eq!(&*relation.get_column_by_index(99).unwrap().name, "col_99");
assert!(relation.get_column_by_index(100).is_none());
}
#[test]
fn test_retry_config_zero_attempts() {
let config = RetryConfig {
max_attempts: 0,
initial_delay: Duration::from_millis(100),
max_delay: Duration::from_secs(10),
multiplier: 2.0,
max_duration: Duration::from_secs(60),
jitter: false,
};
assert_eq!(config.max_attempts, 0);
}
#[test]
fn test_retry_config_extreme_values() {
let config = RetryConfig {
max_attempts: u32::MAX,
initial_delay: Duration::from_nanos(1),
max_delay: Duration::from_secs(86400), // 1 day
multiplier: 10.0,
max_duration: Duration::from_secs(604800), // 1 week
jitter: true,
};
assert_eq!(config.max_attempts, u32::MAX);
assert_eq!(config.multiplier, 10.0);
}
#[test]
fn test_lsn_with_all_segments() {
// Test LSN with different segment combinations
let test_cases = vec![
(0x00000000_00000000, "0/0"),
(0x00000001_00000000, "1/0"),
(0x00000000_00000001, "0/1"),
(0x12345678_9ABCDEF0, "12345678/9ABCDEF0"),
(0xFFFFFFFF_00000000, "FFFFFFFF/0"),
(0x00000000_FFFFFFFF, "0/FFFFFFFF"),
];
for (lsn_val, expected) in test_cases {
assert_eq!(format_lsn(lsn_val), expected);
assert_eq!(parse_lsn(expected).unwrap(), lsn_val);
}
}
#[test]
fn test_event_with_unicode_data() {
let data = RowData::from_pairs(vec![
("name", ColumnValue::text("Alice 中文 émoji 😀")),
("description", ColumnValue::text("Test with ñ, ü, ö")),
]);
let event = ChangeEvent::insert("public", "users", 12345, data.clone(), Lsn::new(1000));
match event.event_type {
EventType::Insert {
data: event_data, ..
} => {
assert_eq!(
event_data.get("name").unwrap(),
&ColumnValue::text("Alice 中文 émoji 😀")
);
assert_eq!(
event_data.get("description").unwrap(),
&ColumnValue::text("Test with ñ, ü, ö")
);
}
_ => panic!("Expected insert"),
}
}
#[tokio::test]
async fn test_cancellation_token_immediate_cancel() {
let token = CancellationToken::new();
// Cancel immediately
token.cancel();
// Should already be cancelled
assert!(token.is_cancelled());
// cancelled() should complete immediately
tokio::select! {
_ = token.cancelled() => {
// Should reach here immediately
}
_ = tokio::time::sleep(Duration::from_secs(1)) => {
panic!("Cancellation should have been immediate");
}
}
}
#[tokio::test]
async fn test_multiple_child_tokens() {
let parent = CancellationToken::new();
let child1 = parent.child_token();
let child2 = parent.child_token();
let grandchild = child1.child_token();
assert!(!parent.is_cancelled());
assert!(!child1.is_cancelled());
assert!(!child2.is_cancelled());
assert!(!grandchild.is_cancelled());
// Cancel parent
parent.cancel();
// All should be cancelled
assert!(parent.is_cancelled());
assert!(child1.is_cancelled());
assert!(child2.is_cancelled());
assert!(grandchild.is_cancelled());
}
#[test]
fn test_streaming_mode_as_str_off() {
assert_eq!(StreamingMode::Off.as_str(), "off");
}
#[test]
fn test_streaming_mode_as_str_on() {
assert_eq!(StreamingMode::On.as_str(), "on");
}
#[test]
fn test_streaming_mode_as_str_parallel() {
assert_eq!(StreamingMode::Parallel.as_str(), "parallel");
}
#[test]
fn test_streaming_mode_eq_and_copy() {
let a = StreamingMode::On;
let b = a; // Copy
assert_eq!(a, b);
assert_ne!(StreamingMode::Off, StreamingMode::On);
}
#[test]
fn test_origin_filter_as_str_none() {
assert_eq!(OriginFilter::None.as_str(), "none");
}
#[test]
fn test_origin_filter_as_str_any() {
assert_eq!(OriginFilter::Any.as_str(), "any");
}
#[test]
fn test_origin_filter_eq_and_copy() {
let a = OriginFilter::None;
let b = a; // Copy
assert_eq!(a, b);
assert_ne!(OriginFilter::None, OriginFilter::Any);
}
#[test]
fn test_config_with_messages() {
let config = create_test_config().with_messages(true);
assert!(config.messages);
let config2 = config.with_messages(false);
assert!(!config2.messages);
}
#[test]
fn test_config_with_binary() {
let config = create_test_config().with_binary(true);
assert!(config.binary);
}
#[test]
fn test_config_with_two_phase() {
let config = create_test_config().with_two_phase(true);
assert!(config.two_phase);
}
#[test]
fn test_config_with_origin_none() {
let config = create_test_config().with_origin(Some(OriginFilter::None));
assert_eq!(config.origin, Some(OriginFilter::None));
}
#[test]
fn test_config_with_origin_any() {
let config = create_test_config().with_origin(Some(OriginFilter::Any));
assert_eq!(config.origin, Some(OriginFilter::Any));
}
#[test]
fn test_config_with_origin_unset() {
let config = create_test_config().with_origin(None);
assert!(config.origin.is_none());
}
#[test]
fn test_config_with_streaming_mode() {
let config = create_test_config().with_streaming_mode(StreamingMode::Off);
assert_eq!(config.streaming_mode, StreamingMode::Off);
let config2 = config.with_streaming_mode(StreamingMode::Parallel);
assert_eq!(config2.streaming_mode, StreamingMode::Parallel);
}
#[test]
fn test_config_with_slot_type_logical() {
let config = create_test_config().with_slot_type(SlotType::Logical);
assert!(matches!(config.slot_type, SlotType::Logical));
}
#[test]
fn test_config_with_slot_type_physical() {
let config = create_test_config().with_slot_type(SlotType::Physical);
assert!(matches!(config.slot_type, SlotType::Physical));
}
#[test]
fn test_config_builder_chain() {
let config = create_test_config()
.with_messages(true)
.with_binary(true)
.with_two_phase(true)
.with_origin(Some(OriginFilter::None))
.with_streaming_mode(StreamingMode::Parallel)
.with_slot_type(SlotType::Physical);
assert!(config.messages);
assert!(config.binary);
assert!(config.two_phase);
assert_eq!(config.origin, Some(OriginFilter::None));
assert_eq!(config.streaming_mode, StreamingMode::Parallel);
assert!(matches!(config.slot_type, SlotType::Physical));
}
#[test]
fn test_tuple_to_data_with_binary_column() {
use crate::protocol::{ColumnData, ColumnInfo, RelationInfo, TupleData};
let columns = vec![
ColumnInfo::new(0, "binary_col".to_string(), 17, -1), // bytea
];
let relation = RelationInfo::new(1, "public", "t", b'd', columns);
let tuple = TupleData::new(vec![ColumnData::binary(vec![0xDE, 0xAD, 0xBE, 0xEF])]);
let data = tuple.into_row_data(&relation);
let val = data.get("binary_col").unwrap();
assert!(matches!(val, ColumnValue::Binary(_)));
assert_eq!(val.as_bytes(), &[0xDE, 0xAD, 0xBE, 0xEF]);
}
#[test]
fn test_tuple_to_data_with_null_column() {
use crate::protocol::{ColumnData, ColumnInfo, RelationInfo, TupleData};
let columns = vec![ColumnInfo::new(0, "nullable".to_string(), 25, -1)];
let relation = RelationInfo::new(1, "public", "t", b'd', columns);
let tuple = TupleData::new(vec![ColumnData::null()]);
let data = tuple.into_row_data(&relation);
assert_eq!(data.get("nullable").unwrap(), &ColumnValue::Null);
}
#[test]
fn test_tuple_to_data_skips_unchanged_columns_detailed() {
use crate::protocol::{ColumnData, ColumnInfo, RelationInfo, TupleData};
let columns = vec![
ColumnInfo::new(0, "id".to_string(), 23, -1),
ColumnInfo::new(0, "unchanged_col".to_string(), 25, -1),
ColumnInfo::new(0, "updated".to_string(), 25, -1),
];
let relation = RelationInfo::new(1, "public", "t", b'd', columns);
let tuple = TupleData::new(vec![
ColumnData::text(b"1".to_vec()),
ColumnData::unchanged(),
ColumnData::text(b"new_value".to_vec()),
]);
let data = tuple.into_row_data(&relation);
assert_eq!(data.len(), 2);
}
#[test]
fn test_tuple_to_data_empty_tuple() {
use crate::protocol::{ColumnInfo, RelationInfo, TupleData};
let columns: Vec<ColumnInfo> = vec![];
let relation = RelationInfo::new(1, "public", "t", b'd', columns);
let tuple = TupleData::new(vec![]);
let data = tuple.into_row_data(&relation);
assert!(data.is_empty());
}
/// `timeout_or_error` cannot bound the blocking connect inside
/// `connect_with_retry`, so the deadline has to travel in the connection
/// string. Both DSN syntaxes, and an explicit caller value always wins.
#[test]
fn with_connect_timeout_supplies_a_default_in_both_dsn_syntaxes() {
let t = Duration::from_secs(30);
// URI with an existing query string → `&`.
assert_eq!(
with_connect_timeout("postgresql://u@h/db?replication=database", t),
"postgresql://u@h/db?replication=database&connect_timeout=30"
);
// URI with no query string → `?`.
assert_eq!(
with_connect_timeout("postgres://u@h/db", t),
"postgres://u@h/db?connect_timeout=30"
);
// key=value syntax → space-separated.
assert_eq!(
with_connect_timeout("host=h port=5432", t),
"host=h port=5432 connect_timeout=30"
);
}
#[test]
fn with_connect_timeout_never_overrides_the_caller() {
let t = Duration::from_secs(30);
for dsn in [
"postgresql://u@h/db?connect_timeout=5",
"host=h connect_timeout=5",
] {
assert_eq!(with_connect_timeout(dsn, t), dsn, "{dsn}");
}
}
/// `Duration::ZERO` is libpq's "no timeout", so it must add nothing — and a
/// sub-second budget must not round down to that same sentinel.
#[test]
fn with_connect_timeout_handles_the_zero_sentinel() {
assert_eq!(
with_connect_timeout("host=h", Duration::ZERO),
"host=h",
"zero means disabled, not connect_timeout=0"
);
assert_eq!(
with_connect_timeout("host=h", Duration::from_millis(1)),
"host=h connect_timeout=1",
"sub-second must clamp up to 1, not down to the disabled sentinel"
);
}
#[tokio::test]
async fn test_timeout_or_error_success() {
let result = timeout_or_error(Duration::from_secs(5), async {
Ok::<_, ReplicationError>(42)
})
.await;
assert_eq!(result.unwrap(), 42);
}
#[tokio::test]
async fn test_timeout_or_error_inner_error() {
let result = timeout_or_error(Duration::from_secs(5), async {
Err::<i32, _>(ReplicationError::protocol("inner error".to_string()))
})
.await;
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("inner error"));
}
#[tokio::test]
async fn test_timeout_or_error_timeout() {
let result = timeout_or_error(Duration::from_millis(1), async {
tokio::time::sleep(Duration::from_secs(10)).await;
Ok::<_, ReplicationError>(42)
})
.await;
assert!(matches!(result, Err(ReplicationError::Timeout(_))));
}
#[test]
fn test_streaming_mode_debug() {
let debug = format!("{:?}", StreamingMode::On);
assert_eq!(debug, "On");
let debug = format!("{:?}", StreamingMode::Off);
assert_eq!(debug, "Off");
let debug = format!("{:?}", StreamingMode::Parallel);
assert_eq!(debug, "Parallel");
}
#[test]
fn test_origin_filter_debug() {
let debug = format!("{:?}", OriginFilter::None);
assert_eq!(debug, "None");
let debug = format!("{:?}", OriginFilter::Any);
assert_eq!(debug, "Any");
}
#[test]
fn test_config_debug() {
let config = create_test_config();
let debug = format!("{:?}", config);
assert!(debug.contains("test_slot"));
assert!(debug.contains("test_publication"));
}
#[test]
fn test_config_default_optional_fields() {
let config = create_test_config();
assert!(!config.messages);
assert!(!config.binary);
assert!(!config.two_phase);
assert!(config.origin.is_none());
}
/// Helper to create a LogicalReplicationStream for testing without a DB connection.
/// Only safe for testing methods that don't touch self.connection.
pub(crate) fn create_test_stream(config: ReplicationStreamConfig) -> LogicalReplicationStream {
use crate::lsn::SharedLsnFeedback;
LogicalReplicationStream {
connection: PgReplicationConnection::null_for_testing(),
parser: LogicalReplicationParser::with_protocol_version(config.protocol_version),
state: ReplicationState::new(),
config: config.clone(),
slot: SlotState::Absent,
retry_handler: ReplicationConnectionRetry::new(
config.retry_config,
"postgresql://test@localhost/test?replication=database".to_string(),
),
last_health_check: Instant::now(),
shared_lsn_feedback: SharedLsnFeedback::new_shared(),
identity: None,
feedback_check_counter: 0,
replay_stop: ReplayStop::Running,
}
}
/// Build an `IDENTIFY_SYSTEM`-shaped result. Native-only: the libpq
/// `PgResult` wraps a `PGresult` pointer and cannot be fabricated this way.
#[cfg(feature = "rustls-tls")]
fn identify_system_result(cols: &[&str], row: &[Option<&str>]) -> crate::connection::PgResult {
let mut r = crate::connection::PgResult::new();
r.columns = cols.iter().map(|c| c.to_string()).collect();
r.rows = vec![row
.iter()
.map(|v| v.map(|v| v.as_bytes().to_vec()))
.collect()];
r
}
/// `CREATE_REPLICATION_SLOT` returns slot_name(0), consistent_point(1),
/// snapshot_name(2), output_plugin(3). Reading the wrong column would stamp
/// every snapshot row with a bogus LSN, which `update_applied_lsn`'s
/// `fetch_max` would then make permanent.
#[cfg(feature = "rustls-tls")]
#[test]
fn parse_consistent_point_reads_column_one() {
let r = identify_system_result(
&[
"slot_name",
"consistent_point",
"snapshot_name",
"output_plugin",
],
&[
Some("my_slot"),
Some("0/30001A0"),
Some("00000003-0000001B-1"),
Some("pgoutput"),
],
);
assert_eq!(parse_consistent_point(&r), Some(Lsn::new(0x30001A0)));
}
/// A malformed or absent consistent point disables the snapshot helper; it
/// must not fail a replication stream that would otherwise work.
#[cfg(feature = "rustls-tls")]
#[test]
fn parse_consistent_point_tolerates_bad_input() {
let cols = &[
"slot_name",
"consistent_point",
"snapshot_name",
"output_plugin",
];
for bad in [Some(""), Some("not-an-lsn"), Some("0/"), Some("/5"), None] {
let r = identify_system_result(cols, &[Some("s"), bad, Some("snap"), Some("pgoutput")]);
assert_eq!(
parse_consistent_point(&r),
None,
"a bad consistent_point ({bad:?}) must degrade to None, never a wrong LSN"
);
}
}
#[cfg(feature = "rustls-tls")]
#[test]
fn parse_consistent_point_on_an_empty_result() {
let mut r = crate::connection::PgResult::new();
r.columns = vec!["slot_name".to_string(), "consistent_point".to_string()];
assert_eq!(parse_consistent_point(&r), None);
}
/// The slot-discard path runs only on error, so it must never propagate a
/// failure of its own — doing so would mask the error that got us there.
#[cfg(feature = "rustls-tls")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn discard_replication_slot_is_best_effort() {
let mut stream = create_test_stream(create_test_config());
stream.slot = SlotState::Ours {
consistent_point: Some(Lsn::new(42)),
exported_snapshot: Some("snap".to_string()),
};
// The null connection cannot execute DROP_REPLICATION_SLOT, so this
// exercises the failure arm.
stream.discard_replication_slot();
// State is deliberately left alone when the drop did not happen: claiming
// the slot is gone when it is not would be worse than leaving it.
assert!(
matches!(stream.slot, SlotState::Ours { .. }),
"a failed drop must not pretend the slot was removed"
);
assert_eq!(stream.exported_snapshot_name(), Some("snap"));
assert_eq!(stream.consistent_point(), Some(Lsn::new(42)));
}
/// The asymmetry the nested `Option` exists for: `START_REPLICATION` destroys
/// the exported snapshot server-side, but the consistent point survives it as
/// the resume position. Collapsing the two into one `Option` — or into a
/// fourth enum variant — would lose exactly this.
#[test]
fn clearing_the_exported_snapshot_keeps_the_consistent_point() {
let mut slot = SlotState::Ours {
consistent_point: Some(Lsn::new(0x2A)),
exported_snapshot: Some("00000004-00000391-1".to_string()),
};
slot.clear_exported_snapshot();
assert_eq!(
slot.exported_snapshot(),
None,
"the name is dead after START_REPLICATION"
);
assert_eq!(
slot.consistent_point(),
Some(Lsn::new(0x2A)),
"the resume position outlives the snapshot name"
);
assert!(slot.is_present(), "the slot itself is still there");
}
/// A slot we did not create offers no baseline — and, critically, is a
/// *different state* from one we did create. The `slot_created: bool` this
/// replaced was `true` for both, which is why a guard on it could not have
/// protected an operator's slot from `discard_replication_slot`.
#[test]
fn a_pre_existing_slot_is_present_but_exposes_no_baseline() {
let slot = SlotState::PreExisting;
assert!(
slot.is_present(),
"ensure_replication_slot must not try again"
);
assert_eq!(slot.consistent_point(), None);
assert_eq!(slot.exported_snapshot(), None);
assert!(
!matches!(slot, SlotState::Ours { .. }),
"this is the distinction discard_replication_slot guards on"
);
// Absent must stay distinguishable from both.
assert!(!SlotState::Absent.is_present());
}
/// `clear_exported_snapshot` is a no-op on the states that never had one, so
/// a successful `start()` on a pre-existing slot cannot corrupt anything.
#[test]
fn clearing_the_exported_snapshot_is_a_no_op_when_there_is_none() {
for mut slot in [SlotState::Absent, SlotState::PreExisting] {
let was_present = slot.is_present();
slot.clear_exported_snapshot();
assert_eq!(slot.is_present(), was_present);
assert_eq!(slot.exported_snapshot(), None);
}
}
/// Pins the column order of `IDENTIFY_SYSTEM`: systemid is column 0 and
/// timeline is column 1. Reading these transposed would silently compare the
/// wrong values on every reconnect.
#[cfg(feature = "rustls-tls")]
#[test]
fn identity_parses_an_identify_system_row() {
let r = identify_system_result(
&["systemid", "timeline", "xlogpos", "dbname"],
&[
Some("7682669593620992043"),
Some("2"),
Some("0/30001A0"),
Some("test_walstream"),
],
);
assert_eq!(
ServerIdentity::from_identify_system(&r),
Some(identity("7682669593620992043", 2))
);
}
/// An unexpected result shape disables the check rather than failing the
/// connection, so each malformed case must yield `None`, not a panic.
#[cfg(feature = "rustls-tls")]
#[test]
fn identity_returns_none_for_an_unusable_row() {
// No rows at all.
let mut empty = crate::connection::PgResult::new();
empty.columns = vec!["systemid".into(), "timeline".into()];
assert_eq!(ServerIdentity::from_identify_system(&empty), None);
// Timeline column missing.
let short = identify_system_result(&["systemid"], &[Some("7682")]);
assert_eq!(ServerIdentity::from_identify_system(&short), None);
// Timeline present but not a number.
let bad = identify_system_result(&["systemid", "timeline"], &[Some("7682"), Some("x")]);
assert_eq!(ServerIdentity::from_identify_system(&bad), None);
// Timeline is SQL NULL.
let null = identify_system_result(&["systemid", "timeline"], &[Some("7682"), None]);
assert_eq!(ServerIdentity::from_identify_system(&null), None);
}
fn identity(system_id: &str, timeline: u32) -> ServerIdentity {
ServerIdentity {
system_id: system_id.to_string(),
timeline,
}
}
/// The first connection has nothing to compare against, so anything passes.
#[test]
fn identity_check_passes_before_the_first_connection() {
let stream = create_test_stream(create_test_config());
assert!(stream.identity.is_none());
assert!(stream.check_server_identity(&identity("7682", 1)).is_ok());
}
/// Reconnecting to the same cluster on the same timeline is the normal case.
#[test]
fn identity_check_passes_on_an_unchanged_server() {
let mut stream = create_test_stream(create_test_config());
stream.identity = Some(identity("7682", 1));
assert!(stream.check_server_identity(&identity("7682", 1)).is_ok());
}
/// A promoted standby keeps the primary's systemid, so only the timeline
/// moves. This is the failover case, and it must be permanent: retrying
/// would stream the post-fork history under LSNs the consumer already used.
#[test]
fn identity_check_rejects_a_timeline_switch() {
let mut stream = create_test_stream(create_test_config());
stream.identity = Some(identity("7682", 1));
let err = stream
.check_server_identity(&identity("7682", 2))
.unwrap_err();
assert!(
matches!(err, ReplicationError::ReplicationSlot(_)),
"expected ReplicationSlot, got {err:?}"
);
assert!(err.is_permanent(), "a fork must stop the retry loop: {err}");
let msg = err.to_string();
assert!(msg.contains("timeline 1 -> 2"), "{msg}");
assert!(msg.contains("TIMELINE_HISTORY 2"), "{msg}");
}
/// Pointing the consumer at an unrelated cluster (a clone, or staging vs
/// production) that happens to have a slot by the same name.
#[test]
fn identity_check_rejects_a_different_cluster() {
let mut stream = create_test_stream(create_test_config());
stream.identity = Some(identity("7682", 1));
let err = stream
.check_server_identity(&identity("9999", 1))
.unwrap_err();
assert!(err.is_permanent(), "{err}");
let msg = err.to_string();
assert!(msg.contains("different cluster"), "{msg}");
assert!(msg.contains("7682") && msg.contains("9999"), "{msg}");
}
/// systemid is checked first: a wholesale cluster swap is the more
/// actionable diagnosis even when the timeline also differs.
#[test]
fn identity_check_reports_the_cluster_when_both_differ() {
let mut stream = create_test_stream(create_test_config());
stream.identity = Some(identity("7682", 1));
let msg = stream
.check_server_identity(&identity("9999", 3))
.unwrap_err()
.to_string();
assert!(msg.contains("different cluster"), "{msg}");
}
/// `START_REPLICATION` runs `SnapBuildClearExportedSnapshot` server-side, so
/// a name surviving a *successful* `start()` is a handle to something the
/// server already destroyed — `SET TRANSACTION SNAPSHOT` would reject it.
///
/// The converse matters just as much, and is what this pins: a `start()` that
/// failed before reaching the server has **not** destroyed anything, so the
/// name must survive. Clearing it there would send a later `snapshot()` down
/// the `Unavailable` path, which the documented usage reads as "nothing to
/// copy" — a stream with no baseline, failing silently.
///
/// The null connection fails inside `initialize()` (at `identify_system`),
/// which is exactly that case. The success side needs a live server: see
/// `a_direct_start_clears_the_exported_snapshot_name` in
/// `integration-tests/initial_snapshot.rs`.
#[tokio::test]
async fn a_failed_start_keeps_the_exported_snapshot_name() {
let cfg = ReplicationStreamConfig::builder("slot", "pub");
let mut stream = create_test_stream(cfg);
stream.slot = SlotState::Ours {
consistent_point: None,
exported_snapshot: Some("00000004-00000391-1".to_string()),
};
let outcome = stream.start(None).await;
assert!(
outcome.is_err(),
"premise: the null connection must fail the attempt, or this tests nothing"
);
assert_eq!(
stream.exported_snapshot_name(),
Some("00000004-00000391-1"),
"the attempt never reached the server, so the snapshot is still alive \
and must stay reachable"
);
}
#[tokio::test]
async fn next_event_terminates_after_stop_boundary_reached() {
// Simulate a crossing commit already delivered: the terminal guard must
// send CopyDone once (a no-op on the null test connection) and return
// StreamStopped on this and every subsequent call, without reading WAL.
let cfg = ReplicationStreamConfig::builder("slot", "pub").with_stop_at_lsn(0x100u64);
let mut stream = create_test_stream(cfg);
stream.replay_stop = ReplayStop::Armed(Lsn::new(0x100));
let token = CancellationToken::new();
let err = stream.next_event(&token).await.unwrap_err();
assert!(
matches!(err, ReplicationError::StreamStopped(l) if l == Lsn::new(0x100)),
"expected StreamStopped(0x100), got {err:?}"
);
assert_eq!(
stream.replay_stop,
ReplayStop::Stopped(Lsn::new(0x100)),
"the Armed -> Stopped transition is the one-shot CopyDone guard"
);
// Idempotent: a second call still reports StreamStopped without resending.
let err2 = stream.next_event(&token).await.unwrap_err();
assert!(matches!(err2, ReplicationError::StreamStopped(_)));
}
#[tokio::test]
async fn next_event_with_retry_passes_through_stream_stopped() {
// StreamStopped is terminal — returned immediately, never retried.
let cfg = ReplicationStreamConfig::builder("slot", "pub").with_stop_at_lsn(0x200u64);
let mut stream = create_test_stream(cfg);
stream.replay_stop = ReplayStop::Armed(Lsn::new(0x200));
let token = CancellationToken::new();
let err = stream.next_event_with_retry(&token).await.unwrap_err();
assert!(
matches!(err, ReplicationError::StreamStopped(l) if l == Lsn::new(0x200)),
"retry wrapper must pass StreamStopped through unretried, got {err:?}"
);
}
#[tokio::test]
async fn event_stream_terminates_none_on_stream_stopped() {
use futures::StreamExt;
// poll_next maps StreamStopped to a graceful end-of-stream (None + terminated).
let cfg = ReplicationStreamConfig::builder("slot", "pub").with_stop_at_lsn(0x300u64);
let mut stream = create_test_stream(cfg);
stream.replay_stop = ReplayStop::Armed(Lsn::new(0x300));
let token = CancellationToken::new();
let mut es = stream.into_stream(token);
assert!(es.next().await.is_none(), "StreamStopped should yield None");
assert!(es.is_terminated(), "stream should mark itself terminated");
}
#[test]
fn test_validate_replication_options_valid_v1() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
1,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream = create_test_stream(config);
assert!(stream.validate_replication_options().is_ok());
}
#[test]
fn test_validate_replication_options_valid_v2_streaming() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
2,
StreamingMode::On,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream = create_test_stream(config);
assert!(stream.validate_replication_options().is_ok());
}
#[test]
fn test_validate_replication_options_invalid_version() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
5,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream = create_test_stream(config);
let err = stream.validate_replication_options().unwrap_err();
assert!(err.to_string().contains("Unsupported protocol version: 5"));
}
#[test]
fn test_validate_streaming_on_requires_v2() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
1,
StreamingMode::On,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream = create_test_stream(config);
let err = stream.validate_replication_options().unwrap_err();
assert!(err
.to_string()
.contains("streaming=on requires protocol version >= 2"));
}
#[test]
fn test_validate_streaming_parallel_requires_v4() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
3,
StreamingMode::Parallel,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream = create_test_stream(config);
let err = stream.validate_replication_options().unwrap_err();
assert!(err
.to_string()
.contains("streaming=parallel requires protocol version >= 4"));
}
#[test]
fn test_validate_two_phase_requires_v3() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
2,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
)
.with_two_phase(true);
let stream = create_test_stream(config);
let err = stream.validate_replication_options().unwrap_err();
assert!(err
.to_string()
.contains("two_phase requires protocol version >= 3"));
}
#[test]
fn test_validate_two_phase_v3_ok() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
3,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
)
.with_two_phase(true);
let stream = create_test_stream(config);
assert!(stream.validate_replication_options().is_ok());
}
#[test]
fn test_validate_streaming_parallel_v4_ok() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
4,
StreamingMode::Parallel,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream = create_test_stream(config);
assert!(stream.validate_replication_options().is_ok());
}
#[test]
fn test_build_replication_options_basic() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
1,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream = create_test_stream(config);
let options = stream.build_replication_options().unwrap();
assert_eq!(options.len(), 2);
assert_eq!(options[0].0, "proto_version");
assert_eq!(options[0].1, "1");
assert_eq!(options[1].0, "publication_names");
assert_eq!(options[1].1, "\"pub\"");
}
#[test]
fn test_build_replication_options_with_streaming() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
2,
StreamingMode::On,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream = create_test_stream(config);
let options = stream.build_replication_options().unwrap();
assert!(options.iter().any(|(k, v)| k == "streaming" && v == "on"));
}
#[test]
fn test_build_replication_options_with_all_flags() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
4,
StreamingMode::Parallel,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
)
.with_messages(true)
.with_binary(true)
.with_two_phase(true)
.with_origin(Some(OriginFilter::None));
let stream = create_test_stream(config);
let options = stream.build_replication_options().unwrap();
assert!(options
.iter()
.any(|(k, v)| k == "streaming" && v == "parallel"));
assert!(options.iter().any(|(k, v)| k == "messages" && v == "on"));
assert!(options.iter().any(|(k, v)| k == "binary" && v == "on"));
assert!(options.iter().any(|(k, v)| k == "two_phase" && v == "on"));
assert!(options.iter().any(|(k, v)| k == "origin" && v == "none"));
}
#[test]
fn test_build_replication_options_invalid_version() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
99,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream = create_test_stream(config);
assert!(stream.build_replication_options().is_err());
}
#[test]
fn test_get_key_columns_for_relation_k_type() {
use crate::protocol::{ColumnInfo, RelationInfo};
let columns = vec![
ColumnInfo::new(1, "id".to_string(), 23, -1), // key
ColumnInfo::new(0, "name".to_string(), 25, -1), // not key
ColumnInfo::new(1, "email".to_string(), 25, -1), // key
];
let relation = RelationInfo::new(1, "public", "users", b'd', columns);
let keys = crate::protocol::key_columns_for_relation(&relation, Some('K'));
assert_eq!(keys, vec![Arc::from("id"), Arc::from("email")]);
}
#[test]
fn test_get_key_columns_for_relation_o_type() {
use crate::protocol::{ColumnInfo, RelationInfo};
let columns = vec![
ColumnInfo::new(1, "id".to_string(), 23, -1),
ColumnInfo::new(0, "name".to_string(), 25, -1),
];
let relation = RelationInfo::new(1, "public", "users", b'f', columns);
let keys = crate::protocol::key_columns_for_relation(&relation, Some('O'));
// 'O' means full old row, return ALL column names
assert_eq!(keys, vec![Arc::from("id"), Arc::from("name")]);
}
#[test]
fn test_get_key_columns_for_relation_none_with_keys() {
use crate::protocol::{ColumnInfo, RelationInfo};
let columns = vec![
ColumnInfo::new(1, "id".to_string(), 23, -1),
ColumnInfo::new(0, "name".to_string(), 25, -1),
];
let relation = RelationInfo::new(1, "public", "users", b'd', columns);
let keys = crate::protocol::key_columns_for_relation(&relation, None);
// Falls back to key columns from relation
assert_eq!(keys, vec![Arc::from("id")]);
}
#[test]
fn test_get_key_columns_for_relation_none_no_keys() {
use crate::protocol::{ColumnInfo, RelationInfo};
let columns = vec![
ColumnInfo::new(0, "data1".to_string(), 25, -1),
ColumnInfo::new(0, "data2".to_string(), 25, -1),
];
let relation = RelationInfo::new(1, "public", "t", b'n', columns);
let keys = crate::protocol::key_columns_for_relation(&relation, None);
// No key columns, and no flagged key columns, should return empty
assert!(keys.is_empty());
}
#[test]
fn test_get_key_columns_for_relation_unknown_type() {
use crate::protocol::{ColumnInfo, RelationInfo};
let columns = vec![
ColumnInfo::new(1, "id".to_string(), 23, -1),
ColumnInfo::new(0, "name".to_string(), 25, -1),
];
let relation = RelationInfo::new(1, "public", "users", b'd', columns);
let keys = crate::protocol::key_columns_for_relation(&relation, Some('X'));
assert_eq!(keys, vec![Arc::from("id")]);
}
#[test]
fn test_relation_metadata() {
use crate::protocol::{ColumnInfo, RelationInfo};
let config = create_test_config();
let mut stream = create_test_stream(config);
let relation = RelationInfo::new(
100,
"myschema",
"mytable",
b'd',
vec![
ColumnInfo::new(1, "pk_col".to_string(), 23, -1),
ColumnInfo::new(0, "data_col".to_string(), 25, -1),
],
);
stream.state.add_relation(relation);
let meta = crate::protocol::relation_metadata(&stream.state, 100, Some('K')).unwrap();
assert_eq!(&*meta.schema, "myschema");
assert_eq!(&*meta.table, "mytable");
assert_eq!(meta.replica_identity, ReplicaIdentity::Default);
assert_eq!(meta.key_columns, vec![Arc::from("pk_col")]);
}
#[test]
fn test_relation_metadata_missing() {
let stream = create_test_stream(create_test_config());
assert!(crate::protocol::relation_metadata(&stream.state, 999, None).is_none());
}
#[test]
fn test_convert_to_change_event_begin() {
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Begin {
final_lsn: 0x1000,
timestamp: 0,
xid: 42,
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
assert!(result.is_some());
let event = result.unwrap();
match event.event_type {
EventType::Begin { transaction_id, .. } => {
assert_eq!(transaction_id, 42);
}
_ => panic!("Expected Begin event"),
}
}
#[test]
fn test_convert_to_change_event_commit() {
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Commit {
flags: 0,
commit_lsn: 0x2000,
end_lsn: 0x2100,
timestamp: 0,
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
assert!(result.is_some());
let event = result.unwrap();
match event.event_type {
EventType::Commit {
commit_lsn,
end_lsn,
..
} => {
assert_eq!(commit_lsn.0, 0x2000);
assert_eq!(end_lsn.0, 0x2100);
}
_ => panic!("Expected Commit event"),
}
}
#[test]
fn test_convert_to_change_event_relation() {
use crate::protocol::ColumnInfo;
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Relation {
relation_id: 100,
namespace: Arc::from("public"),
relation_name: Arc::from("test_table"),
replica_identity: b'd',
columns: vec![ColumnInfo::new(1, "id".to_string(), 23, -1)],
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
// A first sighting emits a Relation event and caches the schema.
match result.expect("first Relation emits an event").event_type {
EventType::Relation {
relation_id,
relation_name,
..
} => {
assert_eq!(relation_id, 100);
assert_eq!(&*relation_name, "test_table");
}
other => panic!("Expected Relation event, got {other:?}"),
}
assert!(stream.state.get_relation(100).is_some());
}
#[test]
fn test_convert_to_change_event_insert() {
use crate::protocol::{ColumnData, ColumnInfo, RelationInfo, TupleData};
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
// First add a relation
let relation = RelationInfo::new(
100,
"public",
"users",
b'd',
vec![
ColumnInfo::new(1, "id".to_string(), 23, -1),
ColumnInfo::new(0, "name".to_string(), 25, -1),
],
);
stream.state.add_relation(relation);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Insert {
relation_id: 100,
tuple: TupleData::new(vec![
ColumnData::text(b"1".to_vec()),
ColumnData::text(b"Alice".to_vec()),
]),
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
assert!(result.is_some());
let event = result.unwrap();
match event.event_type {
EventType::Insert {
schema,
table,
data,
..
} => {
assert_eq!(&*schema, "public");
assert_eq!(&*table, "users");
assert_eq!(data.get("id").unwrap(), "1");
assert_eq!(data.get("name").unwrap(), "Alice");
}
_ => panic!("Expected Insert event"),
}
}
#[test]
fn test_convert_to_change_event_insert_unknown_relation() {
use crate::protocol::{ColumnData, TupleData};
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Insert {
relation_id: 999, // Unknown
tuple: TupleData::new(vec![ColumnData::text(b"1".to_vec())]),
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
assert!(result.is_none());
}
#[test]
fn test_convert_to_change_event_update_with_old() {
use crate::protocol::{ColumnData, ColumnInfo, RelationInfo, TupleData};
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let relation = RelationInfo::new(
100,
"public",
"users",
b'd',
vec![
ColumnInfo::new(1, "id".to_string(), 23, -1),
ColumnInfo::new(0, "name".to_string(), 25, -1),
],
);
stream.state.add_relation(relation);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Update {
relation_id: 100,
old_tuple: Some(TupleData::new(vec![
ColumnData::text(b"1".to_vec()),
ColumnData::text(b"Old".to_vec()),
])),
new_tuple: TupleData::new(vec![
ColumnData::text(b"1".to_vec()),
ColumnData::text(b"New".to_vec()),
]),
key_type: Some('O'),
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
assert!(result.is_some());
let event = result.unwrap();
match event.event_type {
EventType::Update {
old_data,
new_data,
key_columns,
..
} => {
assert!(old_data.is_some());
assert_eq!(new_data.get("name").unwrap(), "New");
assert_eq!(
key_columns.as_slice(),
&[Arc::from("id"), Arc::from("name")]
);
// 'O' = all columns
}
_ => panic!("Expected Update event"),
}
}
#[test]
fn test_convert_to_change_event_update_unknown_relation() {
use crate::protocol::{ColumnData, TupleData};
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Update {
relation_id: 999,
old_tuple: None,
new_tuple: TupleData::new(vec![ColumnData::text(b"1".to_vec())]),
key_type: None,
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
assert!(result.is_none());
}
#[test]
fn test_convert_to_change_event_delete() {
use crate::protocol::{ColumnData, ColumnInfo, RelationInfo, TupleData};
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let relation = RelationInfo::new(
100,
"public",
"users",
b'd',
vec![
ColumnInfo::new(1, "id".to_string(), 23, -1),
ColumnInfo::new(0, "name".to_string(), 25, -1),
],
);
stream.state.add_relation(relation);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Delete {
relation_id: 100,
old_tuple: TupleData::new(vec![
ColumnData::text(b"1".to_vec()),
ColumnData::text(b"Alice".to_vec()),
]),
key_type: 'K',
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
assert!(result.is_some());
let event = result.unwrap();
match event.event_type {
EventType::Delete {
old_data,
key_columns,
..
} => {
assert_eq!(old_data.get("id").unwrap(), "1");
assert_eq!(key_columns, vec![Arc::from("id")]);
}
_ => panic!("Expected Delete event"),
}
}
#[test]
fn test_convert_to_change_event_delete_unknown_relation() {
use crate::protocol::{ColumnData, TupleData};
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Delete {
relation_id: 999,
old_tuple: TupleData::new(vec![ColumnData::text(b"1".to_vec())]),
key_type: 'K',
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
assert!(result.is_none());
}
#[test]
fn test_convert_to_change_event_truncate() {
use crate::protocol::{ColumnInfo, RelationInfo};
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let relation = RelationInfo::new(
100,
"public",
"users",
b'd',
vec![ColumnInfo::new(1, "id".to_string(), 23, -1)],
);
stream.state.add_relation(relation);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Truncate {
relation_ids: vec![100, 200], // 200 is unknown
flags: 0,
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
assert!(result.is_some());
let event = result.unwrap();
match event.event_type {
EventType::Truncate { tables, .. } => {
assert_eq!(tables.len(), 1); // Only known relation
assert_eq!(&*tables[0], "public.users");
}
_ => panic!("Expected Truncate event"),
}
}
#[test]
fn test_convert_to_change_event_stream_start() {
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::StreamStart {
xid: 42,
first_segment: true,
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
assert!(result.is_some());
match result.unwrap().event_type {
EventType::StreamStart {
transaction_id,
first_segment,
} => {
assert_eq!(transaction_id, 42);
assert!(first_segment);
}
_ => panic!("Expected StreamStart"),
}
}
#[test]
fn test_convert_to_change_event_stream_stop() {
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::StreamStop);
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
assert!(result.is_some());
match result.unwrap().event_type {
EventType::StreamStop => {}
_ => panic!("Expected StreamStop"),
}
}
#[test]
fn test_convert_to_change_event_stream_commit() {
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::StreamCommit {
xid: 42,
flags: 0,
timestamp: 0,
commit_lsn: 0x3000,
end_lsn: 0x3100,
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
assert!(result.is_some());
match result.unwrap().event_type {
EventType::StreamCommit { transaction_id, .. } => {
assert_eq!(transaction_id, 42);
}
_ => panic!("Expected StreamCommit"),
}
}
#[test]
fn test_convert_to_change_event_stream_abort() {
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::StreamAbort {
xid: 42,
subtransaction_xid: 43,
abort_lsn: Some(0x4000),
abort_timestamp: Some(123456),
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
assert!(result.is_some());
match result.unwrap().event_type {
EventType::StreamAbort {
transaction_id,
subtransaction_xid,
abort_lsn,
..
} => {
assert_eq!(transaction_id, 42);
assert_eq!(subtransaction_xid, 43);
assert_eq!(abort_lsn.unwrap().0, 0x4000);
}
_ => panic!("Expected StreamAbort"),
}
}
#[test]
fn test_convert_to_change_event_unknown_message() {
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Origin {
origin_lsn: 0x5000,
origin_name: "test_origin".to_string(),
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
// Origin messages are now surfaced as events
assert!(result.is_some());
let event = result.unwrap();
assert_eq!(event.event_type_str(), "origin");
}
#[test]
fn test_current_lsn() {
let config = create_test_config();
let mut stream = create_test_stream(config);
assert_eq!(stream.current_lsn(), 0);
stream.state.update_received_lsn(12345);
assert_eq!(stream.current_lsn(), 12345);
}
#[tokio::test]
async fn test_stream_stop_method() {
let config = create_test_config();
let mut stream = create_test_stream(config);
// stop() performs async I/O (sends final feedback)
stream.stop().await.unwrap();
}
#[tokio::test]
async fn test_maybe_send_feedback_no_received_lsn() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
1,
StreamingMode::Off,
Duration::from_millis(1), // Very short interval
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let mut stream = create_test_stream(config);
// No received LSN, feedback should be a no-op (won't crash on null conn)
tokio::time::sleep(Duration::from_millis(5)).await;
stream.maybe_send_feedback().await;
// If this didn't panic, the guard condition in send_feedback works
}
#[test]
fn test_relation_metadata_single_part_name() {
use crate::protocol::{ColumnInfo, RelationInfo};
let config = create_test_config();
let mut stream = create_test_stream(config);
// Relation with no namespace (empty string) - creates single-part full_name
let relation = RelationInfo::new(
100,
"",
"just_table",
b'f',
vec![ColumnInfo::new(1, "id".to_string(), 23, -1)],
);
stream.state.add_relation(relation);
let meta = crate::protocol::relation_metadata(&stream.state, 100, None).unwrap();
// With empty namespace, full_name is ".just_table", split gives ["", "just_table"]
assert_eq!(&*meta.table, "just_table");
// Schema should still work
assert_eq!(&*meta.schema, "");
}
/// Build a synthetic WAL message: 'w' + start_lsn(8) + end_lsn(8) + send_time(8) + payload
fn build_wal_message(start_lsn: u64, end_lsn: u64, payload: &[u8]) -> Vec<u8> {
let mut data = Vec::new();
data.push(b'w');
data.extend_from_slice(&start_lsn.to_be_bytes());
data.extend_from_slice(&end_lsn.to_be_bytes());
data.extend_from_slice(&0i64.to_be_bytes()); // send_time
data.extend_from_slice(payload);
data
}
/// Build a Begin message payload: 'B' + final_lsn(8) + timestamp(8) + xid(4)
fn build_begin_payload(final_lsn: u64, xid: u32) -> Vec<u8> {
let mut data = Vec::new();
data.push(b'B');
data.extend_from_slice(&final_lsn.to_be_bytes());
data.extend_from_slice(&0i64.to_be_bytes()); // timestamp
data.extend_from_slice(&xid.to_be_bytes());
data
}
/// Build a Commit message payload: 'C' + flags(1) + commit_lsn(8) + end_lsn(8) + timestamp(8)
fn build_commit_payload(commit_lsn: u64, end_lsn: u64) -> Vec<u8> {
let mut data = Vec::new();
data.push(b'C');
data.push(0u8); // flags
data.extend_from_slice(&commit_lsn.to_be_bytes());
data.extend_from_slice(&end_lsn.to_be_bytes());
data.extend_from_slice(&0i64.to_be_bytes()); // timestamp
data
}
#[test]
fn test_process_wal_message_too_short() {
let mut stream = create_test_stream(create_test_config());
let data = vec![b'w'; 10]; // Too short (< 25 bytes)
let result = stream.process_wal_message(data);
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("too short"));
}
#[test]
fn test_process_wal_message_begin() {
let mut stream = create_test_stream(create_test_config());
let payload = build_begin_payload(0x2000, 42);
let data = build_wal_message(0x1000, 0x1500, &payload);
let result = stream.process_wal_message(data).unwrap();
assert!(result.is_some());
let event = result.unwrap();
match event.event_type {
EventType::Begin { transaction_id, .. } => {
assert_eq!(transaction_id, 42);
}
_ => panic!("Expected Begin event"),
}
// end_lsn should be tracked
assert_eq!(stream.state.last_received_lsn, 0x1500);
}
#[test]
fn test_process_wal_message_commit() {
let mut stream = create_test_stream(create_test_config());
let payload = build_commit_payload(0x2000, 0x2100);
let data = build_wal_message(0x1000, 0x1500, &payload);
let result = stream.process_wal_message(data).unwrap();
assert!(result.is_some());
match result.unwrap().event_type {
EventType::Commit { .. } => {}
_ => panic!("Expected Commit event"),
}
}
#[test]
fn test_process_wal_message_no_payload() {
let mut stream = create_test_stream(create_test_config());
// WAL message with header only (no payload after the 25 header bytes)
let data = build_wal_message(0x1000, 0x1500, &[]);
let result = stream.process_wal_message(data).unwrap();
assert!(result.is_none()); // No message data = None
}
#[test]
fn test_process_wal_message_updates_lsn() {
let mut stream = create_test_stream(create_test_config());
assert_eq!(stream.state.last_received_lsn, 0);
let payload = build_begin_payload(0x2000, 1);
let data = build_wal_message(0x1000, 0x5000, &payload);
let _ = stream.process_wal_message(data);
assert_eq!(stream.state.last_received_lsn, 0x5000);
}
#[test]
fn test_process_wal_message_zero_end_lsn_no_update() {
let mut stream = create_test_stream(create_test_config());
stream.state.update_received_lsn(0x3000);
let payload = build_begin_payload(0x2000, 1);
let data = build_wal_message(0x1000, 0, &payload); // end_lsn = 0
let _ = stream.process_wal_message(data);
assert_eq!(stream.state.last_received_lsn, 0x3000); // unchanged
}
/// Build a keepalive message: 'k' + wal_end(8) + server_time(8) + reply(1)
fn build_keepalive_message(wal_end: u64, reply_requested: bool) -> Vec<u8> {
let mut data = Vec::new();
data.push(b'k');
data.extend_from_slice(&wal_end.to_be_bytes());
data.extend_from_slice(&0i64.to_be_bytes()); // server_time
data.push(if reply_requested { 1 } else { 0 });
data
}
#[tokio::test]
async fn test_process_keepalive_no_reply() {
let mut stream = create_test_stream(create_test_config());
let data = build_keepalive_message(0x4000, false);
let result = stream.process_keepalive_message(&data).await;
assert!(result.is_ok());
assert_eq!(stream.state.last_received_lsn, 0x4000);
}
#[tokio::test]
async fn test_process_keepalive_reply_requested_no_received() {
// When received lsn is 0, send_feedback is a no-op (won't crash on null conn)
let mut stream = create_test_stream(create_test_config());
let data = build_keepalive_message(0, true);
let result = stream.process_keepalive_message(&data).await;
// send_feedback returns early if last_received_lsn == 0
assert!(result.is_ok());
}
#[test]
fn test_event_stream_accessors() {
use tokio_util::sync::CancellationToken;
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let event_stream = stream.into_stream(cancel_token);
assert_eq!(event_stream.current_lsn(), 0);
assert_eq!(event_stream.get_feedback_lsn(), (0, 0));
// Test inner() accessor
let _inner_ref = event_stream.inner();
}
#[test]
fn test_event_stream_update_lsn() {
use tokio_util::sync::CancellationToken;
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let event_stream = stream.into_stream(cancel_token);
event_stream.update_flushed_lsn(1000);
event_stream.update_applied_lsn(2000);
let (flushed, applied) = event_stream.get_feedback_lsn();
// update_applied_lsn(2000) also updates flushed to 2000
assert_eq!(flushed, 2000);
assert_eq!(applied, 2000);
}
#[test]
fn test_event_stream_inner_mut() {
use tokio_util::sync::CancellationToken;
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let mut event_stream = stream.into_stream(cancel_token);
// inner_mut should allow mutation
event_stream.inner_mut().state.update_received_lsn(5000);
assert_eq!(event_stream.current_lsn(), 5000);
}
#[test]
fn test_event_stream_ref_accessors() {
use tokio_util::sync::CancellationToken;
let config = create_test_config();
let mut stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let event_stream_ref = stream.stream(cancel_token);
assert_eq!(event_stream_ref.current_lsn(), 0);
let _inner = event_stream_ref.inner();
}
#[test]
fn test_event_stream_ref_inner_mut() {
use tokio_util::sync::CancellationToken;
let config = create_test_config();
let mut stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let mut event_stream_ref = stream.stream(cancel_token);
event_stream_ref.inner_mut().state.update_received_lsn(7000);
assert_eq!(event_stream_ref.current_lsn(), 7000);
}
#[test]
fn test_process_wal_relation_then_insert() {
let mut stream = create_test_stream(create_test_config());
// Build a Relation payload manually: 'R' + relation_id(4) + namespace\0 + name\0 + replica_id(1) + ncols(2) + columns
let mut rel_payload = Vec::new();
rel_payload.push(b'R');
rel_payload.extend_from_slice(&100u32.to_be_bytes());
rel_payload.extend_from_slice(b"public\0");
rel_payload.extend_from_slice(b"items\0");
rel_payload.push(b'd'); // replica identity
rel_payload.extend_from_slice(&2u16.to_be_bytes()); // 2 columns
// Column 1: flags(1) + name\0 + type_oid(4) + type_modifier(4)
rel_payload.push(1u8); // is_key
rel_payload.extend_from_slice(b"id\0");
rel_payload.extend_from_slice(&23u32.to_be_bytes()); // int4
rel_payload.extend_from_slice(&(-1i32).to_be_bytes());
// Column 2
rel_payload.push(0u8); // not key
rel_payload.extend_from_slice(b"val\0");
rel_payload.extend_from_slice(&25u32.to_be_bytes()); // text
rel_payload.extend_from_slice(&(-1i32).to_be_bytes());
let wal = build_wal_message(0x1000, 0x1100, &rel_payload);
let result = stream.process_wal_message(wal).unwrap();
assert!(result.is_some()); // First sighting of a relation emits an event
assert!(stream.state.get_relation(100).is_some());
// Now build an Insert payload: 'I' + relation_id(4) + 'N' + ncols(2) + columns
let mut ins_payload = Vec::new();
ins_payload.push(b'I');
ins_payload.extend_from_slice(&100u32.to_be_bytes());
ins_payload.push(b'N'); // new tuple marker
ins_payload.extend_from_slice(&2u16.to_be_bytes()); // 2 columns
// Column 1: 't' + len(4) + data (text column)
ins_payload.push(b't');
ins_payload.extend_from_slice(&1u32.to_be_bytes());
ins_payload.push(b'5');
// Column 2: 't' + len(4) + data
ins_payload.push(b't');
ins_payload.extend_from_slice(&3u32.to_be_bytes());
ins_payload.extend_from_slice(b"abc");
let wal2 = build_wal_message(0x1100, 0x1200, &ins_payload);
let result2 = stream.process_wal_message(wal2).unwrap();
assert!(result2.is_some());
let event = result2.unwrap();
match event.event_type {
EventType::Insert {
schema,
table,
data,
..
} => {
assert_eq!(&*schema, "public");
assert_eq!(&*table, "items");
assert_eq!(data.get("id").unwrap(), "5");
assert_eq!(data.get("val").unwrap(), "abc");
}
_ => panic!("Expected Insert event"),
}
}
#[tokio::test]
async fn test_send_feedback_no_received_lsn() {
let mut stream = create_test_stream(create_test_config());
// send_feedback should return Ok(()) when last_received_lsn is 0
let result = stream.send_feedback().await;
assert!(result.is_ok());
}
#[test]
fn test_build_options_origin_any() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
4,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
)
.with_origin(Some(OriginFilter::Any));
let stream = create_test_stream(config);
let options = stream.build_replication_options().unwrap();
assert!(options.iter().any(|(k, v)| k == "origin" && v == "any"));
}
#[test]
fn test_build_options_no_streaming_mode_off() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
2,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream = create_test_stream(config);
let options = stream.build_replication_options().unwrap();
// StreamingMode::Off should NOT produce a "streaming" option
assert!(!options.iter().any(|(k, _)| k == "streaming"));
}
#[test]
fn test_validate_version_0() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
0,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream = create_test_stream(config);
let err = stream.validate_replication_options().unwrap_err();
assert!(err.to_string().contains("Unsupported protocol version: 0"));
}
#[test]
fn test_validate_version_4_with_all_features() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
4,
StreamingMode::Parallel,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
)
.with_messages(true)
.with_binary(true)
.with_two_phase(true)
.with_origin(Some(OriginFilter::None));
let stream = create_test_stream(config);
assert!(stream.validate_replication_options().is_ok());
}
#[tokio::test]
async fn test_maybe_send_feedback_with_changed_lsn() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
1,
StreamingMode::Off,
Duration::from_millis(1), // Very short feedback interval
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let mut stream = create_test_stream(config);
// Set received LSN (non-zero)
stream.state.update_received_lsn(5000);
// Update shared feedback to non-zero
stream.shared_lsn_feedback.update_flushed_lsn(3000);
stream.shared_lsn_feedback.update_applied_lsn(2000);
// Wait for feedback interval to pass
tokio::time::sleep(Duration::from_millis(5)).await;
// This should attempt to send feedback, but fail silently because of null conn
// Covers the maybe_send_feedback paths: get_feedback_lsn, flushed > 0, applied > 0, lsn_has_changed
stream.maybe_send_feedback().await;
// If no panic, the error was caught by the warn!() in maybe_send_feedback
}
#[tokio::test]
async fn test_maybe_send_feedback_lsn_not_changed() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
1,
StreamingMode::Off,
Duration::from_millis(1),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let mut stream = create_test_stream(config);
// Set received LSN
stream.state.update_received_lsn(5000);
stream.shared_lsn_feedback.update_flushed_lsn(3000);
stream.shared_lsn_feedback.update_applied_lsn(2000);
tokio::time::sleep(Duration::from_millis(5)).await;
stream.maybe_send_feedback().await; // First call - sends (or tries to)
// Mark the sent LSN values so the next call sees no change
stream.state.mark_feedback_sent_with_lsn(3000, 3000);
tokio::time::sleep(Duration::from_millis(5)).await;
// Same LSN values, should NOT try to send feedback
stream.maybe_send_feedback().await;
}
#[tokio::test]
async fn test_maybe_send_feedback_zero_shared_lsn() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
1,
StreamingMode::Off,
Duration::from_millis(1),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let mut stream = create_test_stream(config);
// Non-zero received but zero shared feedback
stream.state.update_received_lsn(5000);
// shared_lsn_feedback flushed and applied are 0 by default
tokio::time::sleep(Duration::from_millis(5)).await;
stream.maybe_send_feedback().await;
// This covers the f == 0 and a == 0 branches
}
#[tokio::test]
async fn test_send_feedback_with_shared_lsn_values() {
let config = create_test_config();
let mut stream = create_test_stream(config);
stream.state.update_received_lsn(10000);
stream.shared_lsn_feedback.update_flushed_lsn(8000);
stream.shared_lsn_feedback.update_applied_lsn(6000);
// send_feedback will fail on the null connection, but exercises the LSN logic
let result = stream.send_feedback().await;
// It should error because null connection can't send status update
assert!(result.is_err());
// But the local state should have been updated before the error
// Actually, send_feedback calls connection.send_standby_status_update first
// and updates state after. Since it errors on the BDcall, state won't update.
}
#[tokio::test]
async fn test_send_feedback_applied_greater_than_received() {
let config = create_test_config();
let mut stream = create_test_stream(config);
stream.state.update_received_lsn(5000);
// Consumer reports higher LSN than received (edge case, capped)
stream.shared_lsn_feedback.update_flushed_lsn(9000);
stream.shared_lsn_feedback.update_applied_lsn(9000);
// Will fail on null conn but exercises the min() capping logic
let result = stream.send_feedback().await;
assert!(result.is_err());
}
#[test]
fn test_convert_to_change_event_message_with_content() {
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Message {
flags: 1,
lsn: 0x6000,
prefix: "test_prefix".to_string(),
content: Bytes::from_static(b"hello world"),
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
// Message events are now surfaced with full content
assert!(result.is_some());
let event = result.unwrap();
assert_eq!(event.event_type_str(), "message");
if let EventType::Message {
flags,
prefix,
content,
..
} = &event.event_type
{
assert_eq!(*flags, 1);
assert_eq!(prefix.as_ref(), "test_prefix");
assert_eq!(content.as_ref(), b"hello world");
} else {
panic!("Expected Message event");
}
}
#[test]
fn test_convert_to_change_event_type_message() {
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Type {
type_id: 100,
namespace: "pg_catalog".to_string(),
type_name: "int4".to_string(),
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
// Type messages are now surfaced as events
assert!(result.is_some());
let event = result.unwrap();
assert_eq!(event.event_type_str(), "type");
}
#[test]
fn test_convert_update_no_old_tuple() {
use crate::protocol::{ColumnData, ColumnInfo, RelationInfo, TupleData};
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let relation = RelationInfo::new(
100,
"public",
"users",
b'd',
vec![
ColumnInfo::new(1, "id".to_string(), 23, -1),
ColumnInfo::new(0, "name".to_string(), 25, -1),
],
);
stream.state.add_relation(relation);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Update {
relation_id: 100,
old_tuple: None,
new_tuple: TupleData::new(vec![
ColumnData::text(b"1".to_vec()),
ColumnData::text(b"NewName".to_vec()),
]),
key_type: None,
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
assert!(result.is_some());
let event = result.unwrap();
match event.event_type {
EventType::Update {
old_data, new_data, ..
} => {
assert!(old_data.is_none());
assert_eq!(new_data.get("name").unwrap(), "NewName");
}
_ => panic!("Expected Update event"),
}
}
#[test]
fn test_convert_truncate_all_unknown() {
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Truncate {
relation_ids: vec![999],
flags: 0,
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
assert!(result.is_some());
match result.unwrap().event_type {
EventType::Truncate { tables, .. } => {
assert!(tables.is_empty());
}
_ => panic!("Expected Truncate event"),
}
}
#[tokio::test]
async fn test_maybe_send_feedback_flushed_capped() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
1,
StreamingMode::Off,
Duration::from_millis(1),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let mut stream = create_test_stream(config);
stream.state.update_received_lsn(0x1000);
stream.shared_lsn_feedback.update_flushed_lsn(0x9000);
stream.shared_lsn_feedback.update_applied_lsn(0x8000);
tokio::time::sleep(Duration::from_millis(5)).await;
stream.maybe_send_feedback().await;
}
// ---- Coverage: with_slot_options builder ----
#[test]
fn test_config_with_slot_options_builder() {
let options = ReplicationSlotOptions {
temporary: true,
snapshot: Some("export".to_string()),
two_phase: true,
reserve_wal: false,
failover: true,
};
let config = create_test_config().with_slot_options(options);
assert!(config.slot_options.temporary);
assert_eq!(config.slot_options.snapshot, Some("export".to_string()));
assert!(config.slot_options.two_phase);
assert!(!config.slot_options.reserve_wal);
assert!(config.slot_options.failover);
}
#[test]
fn test_config_with_slot_options_default() {
let config = create_test_config().with_slot_options(ReplicationSlotOptions::default());
assert!(!config.slot_options.temporary);
assert!(config.slot_options.snapshot.is_none());
assert!(!config.slot_options.two_phase);
assert!(!config.slot_options.reserve_wal);
assert!(!config.slot_options.failover);
}
// ---- Coverage: exported_snapshot_name / is_temporary_slot ----
#[test]
fn test_exported_snapshot_name_none() {
let stream = create_test_stream(create_test_config());
assert!(stream.exported_snapshot_name().is_none());
}
#[test]
fn test_exported_snapshot_name_some() {
let config = create_test_config();
let mut stream = create_test_stream(config);
stream.slot = SlotState::Ours {
consistent_point: None,
exported_snapshot: Some("00000001-00000001-1".to_string()),
};
assert_eq!(stream.exported_snapshot_name(), Some("00000001-00000001-1"));
}
#[test]
fn test_is_temporary_slot_false() {
let stream = create_test_stream(create_test_config());
assert!(!stream.is_temporary_slot());
}
#[test]
fn test_is_temporary_slot_true() {
let config = create_test_config().with_slot_options(ReplicationSlotOptions {
temporary: true,
..Default::default()
});
let stream = create_test_stream(config);
assert!(stream.is_temporary_slot());
}
// ---- Coverage: send_feedback with zero shared LSN (else branches) ----
#[tokio::test]
async fn test_send_feedback_zero_shared_lsn() {
let config = create_test_config();
let mut stream = create_test_stream(config);
stream.state.update_received_lsn(5000);
// shared_lsn_feedback flushed and applied are 0 by default
// This exercises the else branches (flushed_lsn = 0, applied_lsn = 0)
let result = stream.send_feedback().await;
// Errors because null connection can't send status, but the LSN logic is covered
assert!(result.is_err());
}
// ---- Coverage: full config builder chain with slot_options ----
#[test]
fn test_config_full_builder_chain_with_all_fields() {
let config = create_test_config()
.with_messages(true)
.with_binary(true)
.with_two_phase(true)
.with_origin(Some(OriginFilter::None))
.with_streaming_mode(StreamingMode::Parallel)
.with_slot_type(SlotType::Physical)
.with_slot_options(ReplicationSlotOptions {
temporary: true,
reserve_wal: true,
..Default::default()
});
assert!(config.messages);
assert!(config.binary);
assert!(config.two_phase);
assert_eq!(config.origin, Some(OriginFilter::None));
assert_eq!(config.streaming_mode, StreamingMode::Parallel);
assert!(matches!(config.slot_type, SlotType::Physical));
assert!(config.slot_options.temporary);
assert!(config.slot_options.reserve_wal);
}
// ---- Coverage: process_keepalive with reply_requested and valid received LSN ----
#[tokio::test]
async fn test_process_keepalive_reply_with_valid_lsn() {
let mut stream = create_test_stream(create_test_config());
// Set a non-zero received LSN so send_feedback doesn't short-circuit
stream.state.update_received_lsn(0x2000);
let data = build_keepalive_message(0x3000, true);
// Will attempt send_feedback, which errors on null conn — but keepalive should still succeed
// because process_keepalive calls send_feedback which returns Err, and that propagates
let result = stream.process_keepalive_message(&data).await;
// send_feedback on null connection will error
assert!(result.is_err());
}
// ---- Coverage: convert_to_change_event StreamAbort without optional fields ----
#[test]
fn test_convert_to_change_event_stream_abort_no_optionals() {
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::StreamAbort {
xid: 99,
subtransaction_xid: 100,
abort_lsn: None,
abort_timestamp: None,
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
assert!(result.is_some());
match result.unwrap().event_type {
EventType::StreamAbort {
transaction_id,
subtransaction_xid,
abort_lsn,
..
} => {
assert_eq!(transaction_id, 99);
assert_eq!(subtransaction_xid, 100);
assert!(abort_lsn.is_none());
}
_ => panic!("Expected StreamAbort"),
}
}
// ---- Coverage: convert_to_change_event insert with key_type in update ----
#[test]
fn test_convert_to_change_event_update_key_type_k() {
use crate::protocol::{ColumnData, ColumnInfo, RelationInfo, TupleData};
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let relation = RelationInfo::new(
100,
"public",
"users",
b'd',
vec![
ColumnInfo::new(1, "id".to_string(), 23, -1),
ColumnInfo::new(0, "name".to_string(), 25, -1),
],
);
stream.state.add_relation(relation);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Update {
relation_id: 100,
old_tuple: Some(TupleData::new(vec![ColumnData::text(b"1".to_vec())])),
new_tuple: TupleData::new(vec![
ColumnData::text(b"1".to_vec()),
ColumnData::text(b"Updated".to_vec()),
]),
key_type: Some('K'),
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x500).unwrap();
assert!(result.is_some());
let event = result.unwrap();
match event.event_type {
EventType::Update { key_columns, .. } => {
// 'K' => only flagged key columns
assert_eq!(key_columns, vec![Arc::from("id")]);
}
_ => panic!("Expected Update event"),
}
}
// ---- Coverage: EventStream/EventStreamRef exported_snapshot_name ----
#[test]
fn test_event_stream_exported_snapshot() {
use tokio_util::sync::CancellationToken;
let config = create_test_config();
let mut stream = create_test_stream(config);
stream.slot = SlotState::Ours {
consistent_point: None,
exported_snapshot: Some("snap-123".to_string()),
};
let cancel_token = CancellationToken::new();
let event_stream = stream.into_stream(cancel_token);
assert_eq!(
event_stream.inner().exported_snapshot_name(),
Some("snap-123")
);
}
#[test]
fn test_event_stream_is_temporary() {
use tokio_util::sync::CancellationToken;
let config = create_test_config().with_slot_options(ReplicationSlotOptions {
temporary: true,
..Default::default()
});
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let event_stream = stream.into_stream(cancel_token);
assert!(event_stream.inner().is_temporary_slot());
}
/// Test that `new()` with an invalid/unreachable connection string and a short
/// timeout produces a `Timeout` error (the retry loop takes too long).
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn test_new_with_unreachable_host_timeout() {
let config = ReplicationStreamConfig::new(
"test_slot".to_string(),
"test_pub".to_string(),
2,
StreamingMode::On,
Duration::from_secs(10),
Duration::from_millis(1), // extremely short timeout
Duration::from_secs(60),
RetryConfig {
max_attempts: 1,
initial_delay: Duration::from_millis(1),
max_delay: Duration::from_millis(10),
multiplier: 1.0,
max_duration: Duration::from_millis(10),
jitter: false,
},
);
let result = LogicalReplicationStream::new(
"postgresql://nobody:nothing@192.0.2.1:5432/nonexistent?replication=database&connect_timeout=1",
config,
)
.await;
// Must fail — either Timeout (from the outer timeout_or_error) or a connection error
assert!(result.is_err());
}
/// Test that `new()` with parameters that always fail on connect_with_retry
/// returns an error (non-timeout path).
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn test_new_with_invalid_connection_string_no_timeout() {
let config = ReplicationStreamConfig::new(
"test_slot".to_string(),
"test_pub".to_string(),
2,
StreamingMode::On,
Duration::from_secs(10),
Duration::from_secs(60), // generous timeout
Duration::from_secs(60),
RetryConfig {
max_attempts: 1,
initial_delay: Duration::from_millis(1),
max_delay: Duration::from_millis(10),
multiplier: 1.0,
max_duration: Duration::from_secs(30),
jitter: false,
},
);
let result = LogicalReplicationStream::new(
"host=invalid_host_that_does_not_exist port=99999 dbname=nope replication=database connect_timeout=1",
config,
)
.await;
assert!(
result.is_err(),
"Expected connection failure for invalid host"
);
}
/// Test that `new()` stores retry_config from the config into retry_handler,
/// and verify its fields are properly propagated through `create_test_stream`.
#[test]
fn test_new_retry_handler_config_propagated() {
let retry_config = RetryConfig {
max_attempts: 7,
initial_delay: Duration::from_millis(200),
max_delay: Duration::from_secs(10),
multiplier: 3.0,
max_duration: Duration::from_secs(120),
jitter: false,
};
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
2,
StreamingMode::On,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
retry_config,
);
let stream = create_test_stream(config.clone());
// Verify that the config fields are stored correctly
assert_eq!(stream.config.retry_config.max_attempts, 7);
assert_eq!(
stream.config.retry_config.initial_delay,
Duration::from_millis(200)
);
assert_eq!(stream.config.retry_config.multiplier, 3.0);
assert!(!stream.config.retry_config.jitter);
}
/// Verify that `new()` initializes all fields correctly via `create_test_stream`.
#[test]
fn test_new_initializes_all_fields() {
let config = create_test_config();
let stream = create_test_stream(config);
// slot_created must be false initially
assert!(!stream.slot.is_present());
// exported_snapshot_name must be None initially
assert!(stream.exported_snapshot_name().is_none());
// State should be default (wire position at 0)
assert_eq!(stream.state.last_received_lsn, 0);
// SharedLsnFeedback should be zeroed
let (f, a) = stream.shared_lsn_feedback.get_feedback_lsn();
assert_eq!(f, 0);
assert_eq!(a, 0);
}
/// Verify that `new()` creates a parser with the correct protocol version.
#[test]
fn test_new_parser_protocol_version() {
for version in 1..=4 {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
version,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream = create_test_stream(config);
assert_eq!(stream.config.protocol_version, version);
}
}
/// When `slot_created` is already `true`, `ensure_replication_slot()` should
/// short circuit and return `Ok(())` immediately without touching the connection.
#[tokio::test]
async fn test_ensure_replication_slot_already_created() {
let config = create_test_config();
let mut stream = create_test_stream(config);
stream.slot = SlotState::Ours {
consistent_point: None,
exported_snapshot: None,
};
// This should succeed immediately because slot_created is true,
// even though the connection is null.
let result = stream.ensure_replication_slot().await;
assert!(result.is_ok());
assert!(stream.slot.is_present());
}
/// Verify that calling `ensure_replication_slot()` twice when already created
/// is idempotent and always returns Ok.
#[tokio::test]
async fn test_ensure_replication_slot_idempotent() {
let config = create_test_config();
let mut stream = create_test_stream(config);
stream.slot = SlotState::Ours {
consistent_point: None,
exported_snapshot: None,
};
assert!(stream.ensure_replication_slot().await.is_ok());
assert!(stream.ensure_replication_slot().await.is_ok());
assert!(stream.slot.is_present());
}
/// Verify that the output plugin is set to "pgoutput" for Logical
/// and None for Physical slot types (tests the match branch).
#[test]
fn test_ensure_replication_slot_output_plugin_selection() {
// Logical slot → output_plugin = Some("pgoutput")
let logical_config = create_test_config().with_slot_type(SlotType::Logical);
assert!(matches!(logical_config.slot_type, SlotType::Logical));
// Physical slot → output_plugin = None
let physical_config = create_test_config().with_slot_type(SlotType::Physical);
assert!(matches!(physical_config.slot_type, SlotType::Physical));
}
/// Verify that exported_snapshot_name is None before slot creation.
#[test]
fn test_exported_snapshot_name_none_before_creation() {
let config = create_test_config();
let stream = create_test_stream(config);
assert!(stream.exported_snapshot_name().is_none());
}
/// Verify that we can set and retrieve exported_snapshot_name.
#[test]
fn test_exported_snapshot_name_set_and_get() {
let config = create_test_config();
let mut stream = create_test_stream(config);
stream.slot = SlotState::Ours {
consistent_point: None,
exported_snapshot: Some("00000003-00000028-1".to_string()),
};
assert_eq!(stream.exported_snapshot_name(), Some("00000003-00000028-1"));
}
// ========================================
// Tests for recover_connection() logic
// ========================================
/// When `slot_options.temporary` is true, `recover_connection()` should reset
/// `slot_created` to false so the slot can be recreated.
#[test]
fn test_recover_connection_resets_slot_created_for_temporary() {
let config = create_test_config().with_slot_options(ReplicationSlotOptions {
temporary: true,
..Default::default()
});
let mut stream = create_test_stream(config);
stream.slot = SlotState::Ours {
consistent_point: None,
exported_snapshot: None,
};
// Simulate the temporary-slot reset logic from recover_connection()
if stream.config.slot_options.temporary {
stream.slot = SlotState::Absent;
}
assert!(!stream.slot.is_present());
}
/// When `slot_options.temporary` is false, `recover_connection()` should NOT
/// reset `slot_created`.
#[test]
fn test_recover_connection_preserves_slot_created_for_persistent() {
let config = create_test_config().with_slot_options(ReplicationSlotOptions {
temporary: false,
..Default::default()
});
let mut stream = create_test_stream(config);
stream.slot = SlotState::Ours {
consistent_point: None,
exported_snapshot: None,
};
// Simulate the temporary-slot reset logic from recover_connection()
if stream.config.slot_options.temporary {
stream.slot = SlotState::Absent;
}
assert!(stream.slot.is_present());
}
/// Verify that recover_connection restarts replication from last_received_lsn.
#[test]
fn test_recover_connection_uses_last_received_lsn() {
let config = create_test_config();
let mut stream = create_test_stream(config);
stream.state.update_received_lsn(0xABCD_1234);
// The recovery logic reads last_received_lsn for restart position
let last_lsn = stream.state.last_received_lsn;
assert_eq!(last_lsn, 0xABCD_1234);
}
// ========================================
// Tests for timeout_or_error()
// ========================================
/// Test that `timeout_or_error()` returns inner error when it completes
/// before timeout.
#[tokio::test]
async fn test_timeout_or_error_propagates_inner_error_type() {
let result: Result<()> = timeout_or_error(Duration::from_secs(5), async {
Err(ReplicationError::authentication("bad password"))
})
.await;
assert!(result.is_err());
let err = result.unwrap_err();
assert!(matches!(err, ReplicationError::Authentication(_)));
assert!(err.to_string().contains("bad password"));
}
/// Test that `timeout_or_error()` with a transient error propagates correctly.
#[tokio::test]
async fn test_timeout_or_error_propagates_transient_error() {
let result: Result<()> = timeout_or_error(Duration::from_secs(5), async {
Err(ReplicationError::transient_connection("connection dropped"))
})
.await;
assert!(result.is_err());
let err = result.unwrap_err();
assert!(err.is_transient());
}
/// Test that `timeout_or_error()` returns Ok value when the future completes.
#[tokio::test]
async fn test_timeout_or_error_returns_value() {
let result = timeout_or_error(Duration::from_secs(5), async {
Ok::<String, ReplicationError>("hello".to_string())
})
.await;
assert!(result.is_ok());
assert_eq!(result.unwrap(), "hello");
}
/// Test that `timeout_or_error()` with zero-duration timeout triggers immediately.
#[tokio::test]
async fn test_timeout_or_error_zero_duration() {
let result = timeout_or_error(Duration::from_millis(0), async {
// This async block will never complete before a zero timeout
tokio::time::sleep(Duration::from_secs(10)).await;
Ok::<_, ReplicationError>(())
})
.await;
assert!(matches!(result, Err(ReplicationError::Timeout(_))));
}
/// Test that timeout error message contains the duration.
#[tokio::test]
async fn test_timeout_or_error_message_contains_duration() {
let result = timeout_or_error(Duration::from_millis(5), async {
tokio::time::sleep(Duration::from_secs(10)).await;
Ok::<_, ReplicationError>(())
})
.await;
if let Err(err) = result {
let msg = err.to_string();
assert!(msg.contains("timeout") || msg.contains("timed out"));
} else {
panic!("Expected timeout error");
}
}
// ========================================
// Tests for build_replication_options() edge cases
// ========================================
/// Test build_replication_options with only origin set, no streaming.
#[test]
fn test_build_replication_options_origin_only() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
2,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
)
.with_origin(Some(OriginFilter::Any));
let stream = create_test_stream(config);
let options = stream.build_replication_options().unwrap();
// Should have proto_version, publication_names, and origin
assert_eq!(options.len(), 3);
assert!(options.iter().any(|(k, v)| k == "origin" && v == "any"));
// Should NOT have streaming
assert!(!options.iter().any(|(k, _)| k == "streaming"));
}
/// Test that StreamingMode::Off does NOT add a streaming option.
#[test]
fn test_build_replication_options_no_streaming_off() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
1,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream = create_test_stream(config);
let options = stream.build_replication_options().unwrap();
assert!(!options.iter().any(|(k, _)| k == "streaming"));
}
/// Test build_replication_options with messages enabled but nothing else.
#[test]
fn test_build_replication_options_messages_only() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
1,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
)
.with_messages(true);
let stream = create_test_stream(config);
let options = stream.build_replication_options().unwrap();
assert!(options.iter().any(|(k, v)| k == "messages" && v == "on"));
assert!(!options.iter().any(|(k, _)| k == "binary"));
assert!(!options.iter().any(|(k, _)| k == "two_phase"));
}
/// Test build_replication_options with binary enabled.
#[test]
fn test_build_replication_options_binary_only() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
1,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
)
.with_binary(true);
let stream = create_test_stream(config);
let options = stream.build_replication_options().unwrap();
assert!(options.iter().any(|(k, v)| k == "binary" && v == "on"));
}
/// Test build_replication_options with two_phase enabled at v3.
#[test]
fn test_build_replication_options_two_phase() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
3,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
)
.with_two_phase(true);
let stream = create_test_stream(config);
let options = stream.build_replication_options().unwrap();
assert!(options.iter().any(|(k, v)| k == "two_phase" && v == "on"));
}
/// Test build_replication_options with parallel streaming at v4.
#[test]
fn test_build_replication_options_parallel_v4() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
4,
StreamingMode::Parallel,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream = create_test_stream(config);
let options = stream.build_replication_options().unwrap();
assert!(options
.iter()
.any(|(k, v)| k == "streaming" && v == "parallel"));
}
/// Test that quoted publication name is formatted correctly.
#[test]
fn test_build_replication_options_publication_name_quoting() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"my_publication".to_string(),
1,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream = create_test_stream(config);
let options = stream.build_replication_options().unwrap();
let pub_opt = options
.iter()
.find(|(k, _)| k == "publication_names")
.expect("publication_names option must exist");
assert_eq!(pub_opt.1, "\"my_publication\"");
}
// ========================================
// Tests for validate_replication_options() edge cases
// ========================================
/// Protocol version 0 should be rejected.
#[test]
fn test_validate_replication_options_version_zero() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
0,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream = create_test_stream(config);
let err = stream.validate_replication_options().unwrap_err();
assert!(err.to_string().contains("Unsupported protocol version: 0"));
}
/// All valid protocol versions (1-4) should pass validation with StreamingMode::Off.
#[test]
fn test_validate_replication_options_all_valid_versions() {
for v in 1..=4 {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
v,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream = create_test_stream(config);
assert!(
stream.validate_replication_options().is_ok(),
"Version {v} should be valid"
);
}
}
/// Very large protocol version should be rejected.
#[test]
fn test_validate_replication_options_large_version() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
u32::MAX,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream = create_test_stream(config);
assert!(stream.validate_replication_options().is_err());
}
/// two_phase at v3 should pass, at v2 should fail.
#[test]
fn test_validate_two_phase_boundary() {
// v2 with two_phase → error
let config_v2 = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
2,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
)
.with_two_phase(true);
let stream_v2 = create_test_stream(config_v2);
assert!(stream_v2.validate_replication_options().is_err());
// v3 with two_phase → ok
let config_v3 = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
3,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
)
.with_two_phase(true);
let stream_v3 = create_test_stream(config_v3);
assert!(stream_v3.validate_replication_options().is_ok());
}
/// streaming=parallel at v3 should fail, at v4 should pass.
#[test]
fn test_validate_parallel_boundary() {
let config_v3 = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
3,
StreamingMode::Parallel,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream_v3 = create_test_stream(config_v3);
assert!(stream_v3.validate_replication_options().is_err());
let config_v4 = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
4,
StreamingMode::Parallel,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream_v4 = create_test_stream(config_v4);
assert!(stream_v4.validate_replication_options().is_ok());
}
/// streaming=on at v1 should fail, at v2 should pass.
#[test]
fn test_validate_streaming_on_boundary() {
let config_v1 = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
1,
StreamingMode::On,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream_v1 = create_test_stream(config_v1);
assert!(stream_v1.validate_replication_options().is_err());
let config_v2 = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
2,
StreamingMode::On,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
let stream_v2 = create_test_stream(config_v2);
assert!(stream_v2.validate_replication_options().is_ok());
}
/// Validate that all options together at v4 passes.
#[test]
fn test_validate_all_options_v4() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
4,
StreamingMode::Parallel,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
)
.with_messages(true)
.with_binary(true)
.with_two_phase(true)
.with_origin(Some(OriginFilter::None));
let stream = create_test_stream(config);
assert!(stream.validate_replication_options().is_ok());
}
// ========================================
// Tests for initialize() path coverage
// ========================================
/// Verify that initialize() calls ensure_replication_slot(). Since we use
/// `create_test_stream` with a null connection, calling identify_system()
/// will fail, so we verify the error is from that call and not from
/// ensure_replication_slot().
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn test_initialize_fails_on_identify_system_with_null_conn() {
let config = create_test_config();
let mut stream = create_test_stream(config);
// initialize() calls identify_system() first, which will fail with null conn
let result = stream.initialize().await;
assert!(result.is_err());
}
/// When slot_created is true and identify_system is somehow skipped,
/// the ensure_replication_slot() short circuit still works.
#[tokio::test]
async fn test_ensure_replication_slot_after_manual_flag() {
let config = create_test_config();
let mut stream = create_test_stream(config);
// Manually mark slot as created
stream.slot = SlotState::Ours {
consistent_point: None,
exported_snapshot: None,
};
// ensure_replication_slot should succeed immediately
assert!(stream.ensure_replication_slot().await.is_ok());
}
// ========================================
// Tests for check_connection_health()
// ========================================
/// Health check should be skipped when the interval hasn't elapsed yet.
#[tokio::test]
async fn test_check_connection_health_skipped_when_too_soon() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
2,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(3600), // 1 hour health check interval
RetryConfig::default(),
);
let mut stream = create_test_stream(config);
// Since last_health_check was just set, this should skip the check
// and return Ok even with a null connection.
let result = stream.check_connection_health().await;
assert!(result.is_ok());
}
/// Health check detects dead connection when interval has elapsed.
/// With a null connection, is_alive() returns false, triggering recovery
/// which will fail since there's no real DB.
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn test_check_connection_health_detects_dead_connection() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
2,
StreamingMode::Off,
Duration::from_secs(10),
Duration::from_millis(1), // very short timeout
Duration::from_millis(0), // immediate health check
RetryConfig {
max_attempts: 1,
initial_delay: Duration::from_millis(1),
max_delay: Duration::from_millis(1),
multiplier: 1.0,
max_duration: Duration::from_millis(10),
jitter: false,
},
);
let mut stream = create_test_stream(config);
// Force health check interval to have elapsed
stream.last_health_check = Instant::now() - Duration::from_secs(3600);
// is_alive() will return false for a null connection,
// triggering recovery which will fail
let result = stream.check_connection_health().await;
assert!(result.is_err());
}
/// Test the full builder chain with all options.
#[test]
fn test_config_full_builder_chain() {
let config = ReplicationStreamConfig::new(
"my_slot".to_string(),
"my_pub".to_string(),
4,
StreamingMode::Parallel,
Duration::from_secs(5),
Duration::from_secs(15),
Duration::from_secs(120),
RetryConfig {
max_attempts: 10,
initial_delay: Duration::from_millis(500),
max_delay: Duration::from_secs(30),
multiplier: 2.5,
max_duration: Duration::from_secs(600),
jitter: true,
},
)
.with_messages(true)
.with_binary(true)
.with_two_phase(true)
.with_origin(Some(OriginFilter::Any))
.with_slot_options(ReplicationSlotOptions {
temporary: true,
snapshot: Some("export".to_string()),
two_phase: true,
failover: true,
..Default::default()
})
.with_slot_type(SlotType::Logical);
assert_eq!(config.slot_name, "my_slot");
assert_eq!(config.publication_name, "my_pub");
assert_eq!(config.protocol_version, 4);
assert_eq!(config.streaming_mode, StreamingMode::Parallel);
assert!(config.messages);
assert!(config.binary);
assert!(config.two_phase);
assert_eq!(config.origin, Some(OriginFilter::Any));
assert!(config.slot_options.temporary);
assert_eq!(config.slot_options.snapshot, Some("export".to_string()));
assert!(config.slot_options.two_phase);
assert!(config.slot_options.failover);
assert!(matches!(config.slot_type, SlotType::Logical));
assert_eq!(config.feedback_interval, Duration::from_secs(5));
assert_eq!(config.connection_timeout, Duration::from_secs(15));
assert_eq!(config.health_check_interval, Duration::from_secs(120));
assert_eq!(config.retry_config.max_attempts, 10);
}
/// Test config defaults for slot_options.
#[test]
fn test_config_default_slot_options() {
let config = create_test_config();
// Default snapshot is "nothing" per the constructor
assert_eq!(config.slot_options.snapshot, Some("nothing".to_string()));
assert!(!config.slot_options.temporary);
assert!(!config.slot_options.two_phase);
assert!(!config.slot_options.reserve_wal);
assert!(!config.slot_options.failover);
}
/// Test config with_slot_options replaces previous options entirely.
#[test]
fn test_config_with_slot_options_replaces_defaults() {
let config = create_test_config().with_slot_options(ReplicationSlotOptions {
temporary: true,
snapshot: Some("export".to_string()),
..Default::default()
});
assert!(config.slot_options.temporary);
assert_eq!(config.slot_options.snapshot, Some("export".to_string()));
// Fields not set should be Default
assert!(!config.slot_options.two_phase);
assert!(!config.slot_options.failover);
}
// ========================================
// Tests for shared_lsn_feedback in stream context
// ========================================
/// Verify that shared_lsn_feedback can be updated and read back from the stream.
#[test]
fn test_stream_shared_lsn_feedback_update_and_read() {
let config = create_test_config();
let stream = create_test_stream(config);
stream.shared_lsn_feedback.update_flushed_lsn(5000);
stream.shared_lsn_feedback.update_applied_lsn(4000);
let (f, a) = stream.shared_lsn_feedback.get_feedback_lsn();
assert_eq!(f, 5000);
assert_eq!(a, 4000);
}
/// Verify that shared_lsn_feedback does not allow backward movement.
#[test]
fn test_stream_shared_lsn_feedback_monotonic() {
let config = create_test_config();
let stream = create_test_stream(config);
stream.shared_lsn_feedback.update_applied_lsn(1000);
stream.shared_lsn_feedback.update_applied_lsn(500); // should be ignored
stream.shared_lsn_feedback.update_applied_lsn(2000);
stream.shared_lsn_feedback.update_applied_lsn(1500); // should be ignored
let (_, a) = stream.shared_lsn_feedback.get_feedback_lsn();
assert_eq!(a, 2000);
}
/// current_lsn returns 0 for a fresh stream.
#[test]
fn test_current_lsn_fresh_stream() {
let config = create_test_config();
let stream = create_test_stream(config);
assert_eq!(stream.current_lsn(), 0);
}
/// current_lsn reflects state.last_received_lsn.
#[test]
fn test_current_lsn_after_state_update() {
let config = create_test_config();
let mut stream = create_test_stream(config);
stream.state.update_received_lsn(0x1234_5678);
assert_eq!(stream.current_lsn(), 0x1234_5678);
}
/// is_temporary_slot returns the slot_options.temporary value.
#[test]
fn test_is_temporary_slot_default_false() {
let config = create_test_config();
let stream = create_test_stream(config);
assert!(!stream.is_temporary_slot());
}
#[test]
fn test_is_temporary_slot_explicitly_true() {
let config = create_test_config().with_slot_options(ReplicationSlotOptions {
temporary: true,
..Default::default()
});
let stream = create_test_stream(config);
assert!(stream.is_temporary_slot());
}
// ========================================
// Tests for EventStream wrapper methods
// ========================================
/// EventStream::current_lsn delegates to inner stream.
#[test]
fn test_event_stream_current_lsn() {
let config = create_test_config();
let mut stream = create_test_stream(config);
stream.state.update_received_lsn(42);
let cancel_token = CancellationToken::new();
let event_stream = stream.into_stream(cancel_token);
assert_eq!(event_stream.current_lsn(), 42);
}
/// EventStream inner/inner_mut provide access to the underlying stream.
#[test]
fn test_event_stream_inner_access() {
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let mut event_stream = stream.into_stream(cancel_token);
// inner() provides immutable access
assert_eq!(event_stream.inner().config.slot_name, "test_slot");
// inner_mut() provides mutable access
event_stream.inner_mut().state.update_received_lsn(999);
assert_eq!(event_stream.current_lsn(), 999);
}
/// EventStreamRef works similarly to EventStream.
#[test]
fn test_event_stream_ref_basic() {
let config = create_test_config();
let mut stream = create_test_stream(config);
stream.state.update_received_lsn(100);
let cancel_token = CancellationToken::new();
let event_stream_ref = stream.stream(cancel_token);
assert_eq!(event_stream_ref.current_lsn(), 100);
assert_eq!(event_stream_ref.inner().config.slot_name, "test_slot");
}
/// EventStreamRef inner_mut access.
#[test]
fn test_event_stream_ref_inner_mut_update() {
let config = create_test_config();
let mut stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let mut event_stream_ref = stream.stream(cancel_token);
event_stream_ref.inner_mut().state.update_received_lsn(777);
assert_eq!(event_stream_ref.current_lsn(), 777);
}
/// RetryConfig with max_attempts = 0 should still fail (no retries).
#[tokio::test]
async fn test_new_with_zero_retry_attempts() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
2,
StreamingMode::On,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig {
max_attempts: 0,
initial_delay: Duration::from_millis(1),
max_delay: Duration::from_millis(10),
multiplier: 1.0,
max_duration: Duration::from_secs(5),
jitter: false,
},
);
let result = LogicalReplicationStream::new(
"postgresql://nobody:nothing@127.0.0.1:1/nonexistent?replication=database&connect_timeout=1",
config,
)
.await;
// Should fail because max_attempts is 0 → no attempts made
assert!(result.is_err());
}
/// RetryConfig with very short max_duration should trigger duration exceeded.
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
async fn test_new_with_short_max_duration() {
let config = ReplicationStreamConfig::new(
"slot".to_string(),
"pub".to_string(),
2,
StreamingMode::On,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig {
max_attempts: 100,
initial_delay: Duration::from_millis(1),
max_delay: Duration::from_millis(10),
multiplier: 1.0,
max_duration: Duration::from_millis(1), // extremely short
jitter: false,
},
);
let result = LogicalReplicationStream::new(
"postgresql://nobody:nothing@127.0.0.1:1/nonexistent?replication=database&connect_timeout=1",
config,
)
.await;
assert!(result.is_err());
}
/// Verify that state.relations is empty by default.
#[test]
fn test_stream_state_relations_empty() {
let config = create_test_config();
let stream = create_test_stream(config);
assert!(stream.state.relations.is_empty());
}
/// Verify that last_health_check is recent after creation.
#[test]
fn test_stream_last_health_check_recent() {
let config = create_test_config();
let stream = create_test_stream(config);
// Should have been set within the last second
assert!(stream.last_health_check.elapsed() < Duration::from_secs(1));
}
/// Verify config is stored correctly.
#[test]
fn test_stream_config_stored() {
let config = ReplicationStreamConfig::new(
"custom_slot".to_string(),
"custom_pub".to_string(),
3,
StreamingMode::Off,
Duration::from_secs(15),
Duration::from_secs(45),
Duration::from_secs(90),
RetryConfig::default(),
);
let stream = create_test_stream(config);
assert_eq!(stream.config.slot_name, "custom_slot");
assert_eq!(stream.config.publication_name, "custom_pub");
assert_eq!(stream.config.protocol_version, 3);
assert_eq!(stream.config.streaming_mode, StreamingMode::Off);
assert_eq!(stream.config.feedback_interval, Duration::from_secs(15));
assert_eq!(stream.config.connection_timeout, Duration::from_secs(45));
assert_eq!(stream.config.health_check_interval, Duration::from_secs(90));
}
#[tokio::test]
async fn test_stop_returns_ok_with_default_state() {
let config = create_test_config();
let mut stream = create_test_stream(config);
let result = stream.stop().await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_stop_preserves_current_lsn() {
let config = create_test_config();
let mut stream = create_test_stream(config);
// Simulate having received some WAL data
stream.state.update_received_lsn(0x16B374D848);
let result = stream.stop().await;
assert!(result.is_ok());
// LSN should still be accessible after stop
assert_eq!(stream.current_lsn(), 0x16B374D848);
}
#[tokio::test]
async fn test_stop_with_shared_lsn_feedback_set() {
let config = create_test_config();
let mut stream = create_test_stream(config);
// Simulate consumer having updated feedback LSNs
stream.shared_lsn_feedback.update_flushed_lsn(5000);
stream.shared_lsn_feedback.update_applied_lsn(5000);
stream.state.update_received_lsn(10000);
let result = stream.stop().await;
assert!(result.is_ok());
// Feedback values should still be intact after stop
let (flushed, applied) = stream.shared_lsn_feedback.get_feedback_lsn();
assert_eq!(flushed, 5000);
assert_eq!(applied, 5000);
}
#[tokio::test]
async fn test_stop_can_be_called_multiple_times() {
let config = create_test_config();
let mut stream = create_test_stream(config);
stream.state.update_received_lsn(1000);
// Calling stop() multiple times should always succeed
assert!(stream.stop().await.is_ok());
assert!(stream.stop().await.is_ok());
assert!(stream.stop().await.is_ok());
}
#[tokio::test]
async fn test_stop_with_high_lsn_value() {
let config = create_test_config();
let mut stream = create_test_stream(config);
// Use a large LSN value near u64 max
let high_lsn: u64 = 0xFFFF_FFFF_FFFF_FFFE;
stream.state.update_received_lsn(high_lsn);
let result = stream.stop().await;
assert!(result.is_ok());
assert_eq!(stream.current_lsn(), high_lsn);
}
#[allow(dead_code)]
const _: () = {
fn assert_stream<T: futures_core::Stream>() {}
fn assert_fused_stream<T: futures_core::FusedStream>() {}
fn assert_unpin<T: Unpin>() {}
fn assert_send<T: Send>() {}
fn assertions() {
assert_stream::<EventStream>();
assert_fused_stream::<EventStream>();
assert_unpin::<EventStream>();
assert_send::<EventStream>();
}
};
#[test]
fn test_event_stream_is_terminated_initially_false() {
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let event_stream = stream.into_stream(cancel_token);
assert!(!event_stream.is_terminated());
}
#[tokio::test]
async fn test_event_stream_shutdown_sets_terminated() {
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let mut event_stream = stream.into_stream(cancel_token);
assert!(!event_stream.is_terminated());
let result = event_stream.shutdown().await;
assert!(result.is_ok());
assert!(event_stream.is_terminated());
}
#[tokio::test]
async fn test_event_stream_shutdown_cancels_token() {
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let cancel_clone = cancel_token.clone();
let mut event_stream = stream.into_stream(cancel_token);
assert!(!cancel_clone.is_cancelled());
event_stream.shutdown().await.ok();
assert!(cancel_clone.is_cancelled());
}
#[tokio::test]
async fn test_event_stream_shutdown_idempotent() {
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let mut event_stream = stream.into_stream(cancel_token);
// Calling shutdown multiple times should be safe
assert!(event_stream.shutdown().await.is_ok());
assert!(event_stream.shutdown().await.is_ok());
assert!(event_stream.is_terminated());
}
#[test]
fn test_event_stream_shared_feedback_always_available() {
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let event_stream = stream.into_stream(cancel_token);
// LSN feedback methods should work regardless of stream state
event_stream.update_flushed_lsn(1000);
event_stream.update_applied_lsn(2000);
let (flushed, applied) = event_stream.get_feedback_lsn();
assert_eq!(flushed, 2000); // applied also updates flushed
assert_eq!(applied, 2000);
}
#[test]
fn test_event_stream_inner_returns_ref() {
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let event_stream = stream.into_stream(cancel_token);
// inner() should return a reference when no poll is in flight
let inner = event_stream.inner();
assert_eq!(inner.current_lsn(), 0);
}
#[test]
fn test_event_stream_inner_mut_returns_ref() {
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let mut event_stream = stream.into_stream(cancel_token);
// inner_mut() should allow mutation when no poll is in flight
event_stream.inner_mut().state.update_received_lsn(9000);
assert_eq!(event_stream.current_lsn(), 9000);
}
#[test]
fn test_event_stream_next_event_name() {
// Verify that next_event() is the method name (not conflicting with StreamExt::next)
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let mut event_stream = stream.into_stream(cancel_token);
// This should compile - next_event is our native method
let _future = event_stream.next_event();
// Don't await it since we have a null connection
}
#[tokio::test]
async fn test_event_stream_fused_stream_after_shutdown() {
use futures_core::FusedStream;
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let mut event_stream = stream.into_stream(cancel_token);
assert!(!FusedStream::is_terminated(&event_stream));
event_stream.shutdown().await.ok();
assert!(FusedStream::is_terminated(&event_stream));
}
// ========================================
// process_wal_message with Bytes (zero-copy path coverage)
// ========================================
#[test]
fn test_process_wal_message_with_bytes_input() {
use bytes::Bytes;
let mut stream = create_test_stream(create_test_config());
let payload = build_begin_payload(0x2000, 42);
let data = build_wal_message(0x1000, 0x1500, &payload);
let bytes_data = Bytes::from(data);
let result = stream.process_wal_message(bytes_data).unwrap();
assert!(result.is_some());
let event = result.unwrap();
match event.event_type {
EventType::Begin { transaction_id, .. } => {
assert_eq!(transaction_id, 42);
}
_ => panic!("Expected Begin event"),
}
assert_eq!(stream.state.last_received_lsn, 0x1500);
}
#[test]
fn raw_xlogdata_struct_fields() {
use bytes::Bytes;
let raw = RawXLogData {
wal_start: Lsn::new(0x1000),
wal_end: Lsn::new(0x2000),
server_time: 42,
data: Bytes::from_static(b"payload"),
};
assert_eq!(raw.wal_start, Lsn::new(0x1000));
assert_eq!(raw.wal_end, Lsn::new(0x2000));
assert_eq!(raw.server_time, 42);
assert_eq!(&raw.data[..], b"payload");
}
#[test]
fn test_process_wal_message_bytes_commit() {
use bytes::Bytes;
let mut stream = create_test_stream(create_test_config());
let payload = build_commit_payload(0x2000, 0x2100);
let data = Bytes::from(build_wal_message(0x1000, 0x1500, &payload));
let result = stream.process_wal_message(data).unwrap();
assert!(result.is_some());
match result.unwrap().event_type {
EventType::Commit { .. } => {}
_ => panic!("Expected Commit event"),
}
}
#[test]
fn test_process_wal_message_bytes_too_short() {
use bytes::Bytes;
let mut stream = create_test_stream(create_test_config());
let data = Bytes::from(vec![b'w'; 10]);
let result = stream.process_wal_message(data);
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("too short"));
}
#[test]
fn test_process_wal_message_bytes_no_payload() {
use bytes::Bytes;
let mut stream = create_test_stream(create_test_config());
let data = Bytes::from(build_wal_message(0x1000, 0x1500, &[]));
let result = stream.process_wal_message(data).unwrap();
assert!(result.is_none());
}
#[test]
fn test_process_wal_message_bytes_updates_lsn() {
use bytes::Bytes;
let mut stream = create_test_stream(create_test_config());
assert_eq!(stream.state.last_received_lsn, 0);
let payload = build_begin_payload(0x2000, 1);
let data = Bytes::from(build_wal_message(0x1000, 0x5000, &payload));
let _ = stream.process_wal_message(data);
assert_eq!(stream.state.last_received_lsn, 0x5000);
}
#[test]
fn test_process_wal_message_bytes_zero_end_lsn() {
use bytes::Bytes;
let mut stream = create_test_stream(create_test_config());
stream.state.update_received_lsn(0x3000);
let payload = build_begin_payload(0x2000, 1);
let data = Bytes::from(build_wal_message(0x1000, 0, &payload));
let _ = stream.process_wal_message(data);
assert_eq!(stream.state.last_received_lsn, 0x3000); // unchanged
}
#[test]
fn test_process_wal_message_bytes_relation_then_insert() {
use bytes::Bytes;
let mut stream = create_test_stream(create_test_config());
// Build Relation payload
let mut rel_payload = Vec::new();
rel_payload.push(b'R');
rel_payload.extend_from_slice(&100u32.to_be_bytes());
rel_payload.extend_from_slice(b"public\0");
rel_payload.extend_from_slice(b"items\0");
rel_payload.push(b'd');
rel_payload.extend_from_slice(&1u16.to_be_bytes());
rel_payload.push(1u8);
rel_payload.extend_from_slice(b"id\0");
rel_payload.extend_from_slice(&23u32.to_be_bytes());
rel_payload.extend_from_slice(&(-1i32).to_be_bytes());
let wal = Bytes::from(build_wal_message(0x1000, 0x1500, &rel_payload));
let result = stream.process_wal_message(wal).unwrap();
assert!(result.is_some()); // First sighting of a relation emits an event
// Build Insert payload
let mut ins_payload = Vec::new();
ins_payload.push(b'I');
ins_payload.extend_from_slice(&100u32.to_be_bytes());
ins_payload.push(b'N');
ins_payload.extend_from_slice(&1u16.to_be_bytes());
ins_payload.push(b't');
ins_payload.extend_from_slice(&5u32.to_be_bytes());
ins_payload.extend_from_slice(b"hello");
let wal2 = Bytes::from(build_wal_message(0x2000, 0x2500, &ins_payload));
let result2 = stream.process_wal_message(wal2).unwrap();
assert!(result2.is_some());
match result2.unwrap().event_type {
EventType::Insert { schema, table, .. } => {
assert_eq!(&*schema, "public");
assert_eq!(&*table, "items");
}
_ => panic!("Expected Insert event"),
}
}
#[test]
fn decode_raw_frame_returns_payload_and_positions() {
let mut stream = create_test_stream(create_test_config());
// Build a 'w' frame with an explicit non-zero send_time to prove the
// header offsets are read correctly.
let mut frame = vec![b'w'];
frame.extend_from_slice(&0x1000u64.to_be_bytes()); // wal_start
frame.extend_from_slice(&0x2000u64.to_be_bytes()); // wal_end
frame.extend_from_slice(&0x0102_0304_0506_0708i64.to_be_bytes()); // send_time
frame.extend_from_slice(b"pgoutput-bytes"); // payload
let raw = stream.decode_raw_frame(bytes::Bytes::from(frame)).unwrap();
assert_eq!(raw.wal_start, Lsn::new(0x1000));
assert_eq!(raw.wal_end, Lsn::new(0x2000));
assert_eq!(raw.server_time, 0x0102_0304_0506_0708i64);
assert_eq!(&raw.data[..], b"pgoutput-bytes");
}
#[test]
fn decode_raw_frame_data_is_zero_copy_slice() {
let mut stream = create_test_stream(create_test_config());
let orig = bytes::Bytes::from(build_wal_message(1, 2, b"abcdef"));
let raw = stream.decode_raw_frame(orig.clone()).unwrap();
// slice(25..) must point into the same allocation (no memcpy).
assert_eq!(raw.data.as_ptr(), orig[25..].as_ptr());
}
#[test]
fn decode_raw_frame_header_only_has_empty_payload() {
let mut stream = create_test_stream(create_test_config());
// Empty payload → frame is exactly 25 bytes (header only). Decodes to a
// RawXLogData carrying the header positions and an empty `data` slice.
let frame = bytes::Bytes::from(build_wal_message(0x10, 0x20, &[]));
assert_eq!(frame.len(), 25);
let raw = stream.decode_raw_frame(frame).unwrap();
assert_eq!(raw.wal_start, Lsn::new(0x10));
assert_eq!(raw.wal_end, Lsn::new(0x20));
assert!(raw.data.is_empty());
}
#[test]
fn decode_raw_frame_too_short_errors() {
let mut stream = create_test_stream(create_test_config());
let data = bytes::Bytes::from(vec![b'w'; 10]);
let err = stream.decode_raw_frame(data).unwrap_err();
assert!(err.to_string().contains("too short"));
}
#[test]
fn decode_raw_frame_advances_received_lsn() {
let mut stream = create_test_stream(create_test_config());
assert_eq!(stream.state.last_received_lsn, 0);
let frame = bytes::Bytes::from(build_wal_message(0x1000, 0xABCD, b"x"));
stream.decode_raw_frame(frame).unwrap();
assert_eq!(stream.state.last_received_lsn, 0xABCD);
}
#[tokio::test]
async fn next_raw_event_arms_transport_stop_at_or_past_lsn() {
// Arming now lives in next_raw_event (decode_raw_frame is policy-free), so
// drive it through the pump with a seeded frame.
let token = CancellationToken::new();
// Below the stop LSN: terminal not armed.
let cfg = ReplicationStreamConfig::builder("slot", "pub").with_stop_at_lsn(0x2000u64);
let below = bytes::Bytes::from(build_wal_message(0x1000, 0x1FFF, b"x"));
let mut stream = create_test_stream_with_frames(cfg, vec![below]);
stream.next_raw_event(&token).await.unwrap();
assert_eq!(stream.replay_stop, ReplayStop::Running);
// At/past the stop LSN: armed with this message's wal_end.
let cfg2 = ReplicationStreamConfig::builder("slot", "pub").with_stop_at_lsn(0x2000u64);
let at = bytes::Bytes::from(build_wal_message(0x1000, 0x2000, b"x"));
let mut stream2 = create_test_stream_with_frames(cfg2, vec![at]);
stream2.next_raw_event(&token).await.unwrap();
assert_eq!(stream2.replay_stop, ReplayStop::Armed(Lsn::new(0x2000)));
}
// These drive the shared raw pump through a null connection pre-seeded with CopyData frames. Both backends now expose the frame-seeding seam (`null_for_testing_with_frames`), so these run on libpq and rustls-tls alike — the pump serves each frame from the pending queue before any FFI.
// The one exception below (`..._terminates_after_transport_stop`) drives CopyDone, a backend-specific FFI send that is unsafe on libpq's null-pointer conn, so it stays rustls-tls-only (integration-tested for libpq).
/// Build a test stream whose connection serves `frames` in order from its
/// `pending` queue, so the connection-driven raw path is unit-testable.
fn create_test_stream_with_frames(
config: ReplicationStreamConfig,
frames: Vec<bytes::Bytes>,
) -> LogicalReplicationStream {
let mut stream = create_test_stream(config);
stream.connection = PgReplicationConnection::null_for_testing_with_frames(frames);
stream
}
#[tokio::test]
async fn next_raw_event_returns_seeded_wal_payload() {
let payload = build_begin_payload(0x2000, 7);
let frame = bytes::Bytes::from(build_wal_message(0x1000, 0x2000, &payload));
let mut stream = create_test_stream_with_frames(create_test_config(), vec![frame]);
let token = CancellationToken::new();
let raw = stream.next_raw_event(&token).await.unwrap();
assert_eq!(raw.wal_start, Lsn::new(0x1000));
assert_eq!(raw.wal_end, Lsn::new(0x2000));
assert_eq!(&raw.data[..], payload.as_slice());
// The 'w' header's end_lsn advanced the received LSN.
assert_eq!(stream.state.last_received_lsn, 0x2000);
}
#[tokio::test]
async fn next_raw_event_returns_header_only_frame() {
// A header-only 'w' frame (len == 25) is returned directly with an empty
// payload — the pump does not skip it (a conformant server never sends
// one; a malformed frame surfaces as an empty-`data` event).
let header_only = bytes::Bytes::from(build_wal_message(0x1000, 0x1500, &[]));
let mut stream = create_test_stream_with_frames(create_test_config(), vec![header_only]);
let token = CancellationToken::new();
let raw = stream.next_raw_event(&token).await.unwrap();
assert_eq!(raw.wal_end, Lsn::new(0x1500));
assert!(raw.data.is_empty());
}
#[tokio::test]
async fn next_raw_event_handles_keepalive_then_wal() {
// A keepalive ('k', no reply requested) is consumed internally and
// advances the received LSN; the following 'w' frame is returned.
let mut ka = vec![b'k'];
ka.extend_from_slice(&0x1500u64.to_be_bytes()); // wal_end
ka.extend_from_slice(&0i64.to_be_bytes()); // timestamp
ka.push(0u8); // reply_requested = false
let payload = build_begin_payload(0x2000, 1);
let wal = bytes::Bytes::from(build_wal_message(0x1000, 0x2000, &payload));
let mut stream =
create_test_stream_with_frames(create_test_config(), vec![bytes::Bytes::from(ka), wal]);
let token = CancellationToken::new();
let raw = stream.next_raw_event(&token).await.unwrap();
assert_eq!(raw.wal_end, Lsn::new(0x2000));
assert_eq!(stream.state.last_received_lsn, 0x2000);
}
#[tokio::test]
async fn next_raw_event_skips_empty_and_unknown_frames() {
// An empty CopyData payload is skipped; an unknown message type is logged
// and skipped; the following 'w' frame is returned.
let payload = build_begin_payload(0x2000, 1);
let wal = bytes::Bytes::from(build_wal_message(0x1000, 0x2000, &payload));
let mut stream = create_test_stream_with_frames(
create_test_config(),
vec![
bytes::Bytes::new(),
bytes::Bytes::from(vec![b'z', 1, 2, 3]),
wal,
],
);
let token = CancellationToken::new();
let raw = stream.next_raw_event(&token).await.unwrap();
assert_eq!(raw.wal_end, Lsn::new(0x2000));
}
#[tokio::test]
async fn next_raw_event_cancelled_returns_error() {
// A pre-cancelled token short-circuits the pump before any read.
let mut stream = create_test_stream_with_frames(create_test_config(), vec![]);
let token = CancellationToken::new();
token.cancel();
let err = stream.next_raw_event(&token).await.unwrap_err();
assert!(matches!(err, ReplicationError::Cancelled(_)));
}
#[tokio::test]
async fn next_raw_event_runs_periodic_feedback_check() {
// Pre-arm the throttle counter so the next pumped frame lands on the
// feedback-check tick, exercising the maybe_send_feedback branch.
let payload = build_begin_payload(0x2000, 1);
let wal = bytes::Bytes::from(build_wal_message(0x1000, 0x2000, &payload));
let mut stream = create_test_stream_with_frames(create_test_config(), vec![wal]);
stream.feedback_check_counter = FEEDBACK_CHECK_EVENT_INTERVAL - 1;
let token = CancellationToken::new();
let raw = stream.next_raw_event(&token).await.unwrap();
assert_eq!(raw.wal_end, Lsn::new(0x2000));
}
#[cfg(feature = "rustls-tls")]
#[tokio::test]
async fn next_raw_event_terminates_after_transport_stop() {
// A frame whose wal_end reaches stop_at_lsn is delivered and arms the
// transport terminal; the next call sends CopyDone once (a no-op on the
// null connection) and reports StreamStopped — end-to-end through the
// raw entry point, not just the shared pump.
let cfg = ReplicationStreamConfig::builder("slot", "pub").with_stop_at_lsn(0x2000u64);
let payload = build_begin_payload(0x2000, 1);
let frame = bytes::Bytes::from(build_wal_message(0x1000, 0x2000, &payload));
let mut stream = create_test_stream_with_frames(cfg, vec![frame]);
let token = CancellationToken::new();
// First call delivers the crossing frame and arms the terminal.
let raw = stream.next_raw_event(&token).await.unwrap();
assert_eq!(raw.wal_end, Lsn::new(0x2000));
assert_eq!(stream.replay_stop, ReplayStop::Armed(Lsn::new(0x2000)));
// Second call terminates cleanly (CopyDone sent once, StreamStopped).
let err = stream.next_raw_event(&token).await.unwrap_err();
assert!(matches!(err, ReplicationError::StreamStopped(l) if l == Lsn::new(0x2000)));
assert_eq!(stream.replay_stop, ReplayStop::Stopped(Lsn::new(0x2000)));
}
#[tokio::test]
async fn next_event_returns_decoded_event() {
// Drive the TYPED path end-to-end through the seeded pump: a Begin frame
// pumps → process_wal_message → convert → a ChangeEvent.
let payload = build_begin_payload(0x2000, 42);
let frame = bytes::Bytes::from(build_wal_message(0x1000, 0x2000, &payload));
let mut stream = create_test_stream_with_frames(create_test_config(), vec![frame]);
let token = CancellationToken::new();
let event = stream.next_event(&token).await.unwrap();
match event.event_type {
EventType::Begin { transaction_id, .. } => assert_eq!(transaction_id, 42),
other => panic!("expected Begin, got {other:?}"),
}
}
#[tokio::test]
async fn next_event_arms_commit_boundary_stop() {
// A Commit whose end_lsn reaches stop_at_lsn is delivered AND arms the
// commit-boundary terminal (the typed path's bounded-replay mechanism,
// distinct from the raw transport stop).
let cfg = ReplicationStreamConfig::builder("slot", "pub").with_stop_at_lsn(0x2000u64);
let payload = build_commit_payload(0x2000, 0x2100); // commit_lsn, end_lsn
let frame = bytes::Bytes::from(build_wal_message(0x1000, 0x2100, &payload));
let mut stream = create_test_stream_with_frames(cfg, vec![frame]);
let token = CancellationToken::new();
let event = stream.next_event(&token).await.unwrap();
assert!(matches!(event.event_type, EventType::Commit { .. }));
// reached_stop_lsn returns the commit's end_lsn (0x2100 >= 0x2000).
assert_eq!(stream.replay_stop, ReplayStop::Armed(Lsn::new(0x2100)));
}
// into_stream() and EventStream construction coverage
// ========================================
#[test]
fn test_into_stream_creates_valid_event_stream() {
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let event_stream = stream.into_stream(cancel_token.clone());
// Verify all fields are properly initialized
assert!(event_stream.state.stream().is_some());
assert!(!event_stream.is_terminated());
assert!(!event_stream.state.is_polling());
assert!(!event_stream.is_terminated());
// Shared feedback should work
event_stream.update_flushed_lsn(100);
let (flushed, _) = event_stream.get_feedback_lsn();
assert_eq!(flushed, 100);
}
#[test]
fn test_into_stream_shared_feedback_is_cloned() {
let config = create_test_config();
let stream = create_test_stream(config);
let original_feedback = Arc::clone(&stream.shared_lsn_feedback);
let cancel_token = CancellationToken::new();
let event_stream = stream.into_stream(cancel_token);
// Updates via event_stream should be visible via the original Arc
event_stream.update_applied_lsn(5000);
let (flushed, applied) = original_feedback.get_feedback_lsn();
assert_eq!(applied, 5000);
assert_eq!(flushed, 5000);
}
#[test]
fn test_into_stream_inner_accessible() {
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let event_stream = stream.into_stream(cancel_token);
// inner() should return the stream
let inner = event_stream.inner();
assert_eq!(inner.current_lsn(), 0);
}
#[test]
fn test_into_stream_inner_mut_accessible() {
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let mut event_stream = stream.into_stream(cancel_token);
// Mutate via inner_mut
event_stream.inner_mut().state.update_received_lsn(9999);
assert_eq!(event_stream.current_lsn(), 9999);
}
// ========================================
// Stream::poll_next coverage
// ========================================
#[tokio::test]
async fn test_event_stream_poll_next_terminated_returns_none() {
use futures::StreamExt;
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let mut event_stream = stream.into_stream(cancel_token);
// Manually set terminated
event_stream.state.mark_terminated();
// poll_next should return None immediately
let result = event_stream.next().await;
assert!(result.is_none());
}
#[tokio::test]
async fn test_event_stream_poll_next_after_shutdown_returns_none() {
use futures::StreamExt;
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let mut event_stream = stream.into_stream(cancel_token);
event_stream.shutdown().await.ok();
assert!(event_stream.is_terminated());
// After shutdown, poll_next should return None
let result = event_stream.next().await;
assert!(result.is_none());
}
#[tokio::test]
async fn test_event_stream_poll_next_cancel_terminates() {
use futures::StreamExt;
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let cancel_clone = cancel_token.clone();
let mut event_stream = stream.into_stream(cancel_token);
// Cancel immediately
cancel_clone.cancel();
// poll_next should detect cancellation and return None
let result = event_stream.next().await;
assert!(result.is_none());
assert!(event_stream.is_terminated());
}
#[tokio::test]
async fn test_event_stream_poll_next_fused_after_cancel() {
use futures::StreamExt;
use futures_core::FusedStream;
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
cancel_token.cancel(); // Pre-cancel
let mut event_stream = stream.into_stream(cancel_token);
assert!(!event_stream.is_terminated());
// First poll should cancel
let _ = event_stream.next().await;
assert!(FusedStream::is_terminated(&event_stream));
// Subsequent polls should return None
let result = event_stream.next().await;
assert!(result.is_none());
}
#[tokio::test]
async fn test_event_stream_poll_next_restores_inner_after_cancel() {
use futures::StreamExt;
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
cancel_token.cancel();
let mut event_stream = stream.into_stream(cancel_token);
// Poll to trigger the take/restore cycle
let _ = event_stream.next().await;
// After poll completes, inner should be restored (even though terminated)
// The inner is restored on Ready, so it should be Some
assert!(event_stream.state.stream().is_some());
assert!(!event_stream.state.is_polling());
}
// ========================================
// shutdown() additional coverage
// ========================================
#[tokio::test]
async fn test_event_stream_shutdown_clears_inflight() {
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let mut event_stream = stream.into_stream(cancel_token);
event_stream.shutdown().await.ok();
assert!(!event_stream.state.is_polling());
assert!(event_stream.is_terminated());
assert!(event_stream.state.stream().is_some()); // inner remains available
}
/// Replaces `test_event_stream_shutdown_with_inner_none`, which built
/// `(inner: None, inflight: None)` by hand and asserted that `shutdown`
/// "handles None gracefully". That state is exactly where the old `shutdown`
/// landed when a poll was in flight — dropping the future that owned the
/// connection and skipping the final feedback — so the test pinned the bug
/// as correct. `StreamState` makes the state unrepresentable; what is left
/// to check is that a parked stream gets reclaimed, not abandoned.
#[tokio::test]
async fn test_event_stream_shutdown_reclaims_a_parked_stream() {
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let mut event_stream = stream.into_stream(cancel_token);
event_stream
.state
.park_with_result(Err(ReplicationError::Protocol("parked".to_string())));
let result = event_stream.shutdown().await;
assert!(result.is_ok(), "{result:?}");
assert!(event_stream.is_terminated());
assert!(
event_stream.state.stream().is_some(),
"the stream must be reclaimed from the parked future, not dropped"
);
}
// ========================================
// EventStream next_event() coverage
// ========================================
#[test]
fn test_event_stream_next_event_returns_future() {
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let mut event_stream = stream.into_stream(cancel_token);
// Verify that next_event returns a future (compile-time check)
let _fut = event_stream.next_event();
// Don't await — we have a null connection
}
// ========================================
// process_wal_message impl Into<Bytes> coverage
// ========================================
#[test]
fn test_process_wal_message_accepts_vec() {
let mut stream = create_test_stream(create_test_config());
let payload = build_begin_payload(0x2000, 10);
let data: Vec<u8> = build_wal_message(0x1000, 0x1500, &payload);
// Vec<u8> implements Into<Bytes>
let result = stream.process_wal_message(data).unwrap();
assert!(result.is_some());
}
#[test]
fn test_process_wal_message_accepts_bytes() {
use bytes::Bytes;
let mut stream = create_test_stream(create_test_config());
let payload = build_begin_payload(0x2000, 10);
let data = Bytes::from(build_wal_message(0x1000, 0x1500, &payload));
// Bytes implements Into<Bytes> (identity)
let result = stream.process_wal_message(data).unwrap();
assert!(result.is_some());
}
#[test]
fn test_process_wal_message_uses_zero_copy_reader() {
use bytes::Bytes;
let mut stream = create_test_stream(create_test_config());
// Build a relation + insert to test zero-copy Bytes slicing through the pipeline
let mut rel_payload = Vec::new();
rel_payload.push(b'R');
rel_payload.extend_from_slice(&50u32.to_be_bytes());
rel_payload.extend_from_slice(b"test_schema\0");
rel_payload.extend_from_slice(b"test_table\0");
rel_payload.push(b'd');
rel_payload.extend_from_slice(&2u16.to_be_bytes());
// col 1
rel_payload.push(1u8);
rel_payload.extend_from_slice(b"col_a\0");
rel_payload.extend_from_slice(&23u32.to_be_bytes());
rel_payload.extend_from_slice(&(-1i32).to_be_bytes());
// col 2
rel_payload.push(0u8);
rel_payload.extend_from_slice(b"col_b\0");
rel_payload.extend_from_slice(&25u32.to_be_bytes());
rel_payload.extend_from_slice(&(-1i32).to_be_bytes());
let wal = Bytes::from(build_wal_message(0x100, 0x200, &rel_payload));
let _ = stream.process_wal_message(wal);
// Insert with two columns
let mut ins = Vec::new();
ins.push(b'I');
ins.extend_from_slice(&50u32.to_be_bytes());
ins.push(b'N');
ins.extend_from_slice(&2u16.to_be_bytes());
ins.push(b't');
ins.extend_from_slice(&2u32.to_be_bytes());
ins.extend_from_slice(b"42");
ins.push(b't');
ins.extend_from_slice(&5u32.to_be_bytes());
ins.extend_from_slice(b"world");
let wal2 = Bytes::from(build_wal_message(0x300, 0x400, &ins));
let result = stream.process_wal_message(wal2).unwrap();
assert!(result.is_some());
let event = result.unwrap();
match event.event_type {
EventType::Insert {
schema,
table,
data,
..
} => {
assert_eq!(&*schema, "test_schema");
assert_eq!(&*table, "test_table");
assert_eq!(data.len(), 2);
}
_ => panic!("Expected Insert"),
}
}
// ========================================
// poll_next branch coverage via injected inflight futures
// ========================================
/// Helper: create an EventStream and inject a pre-built inflight future
/// so that poll_next exercises the result-handling branches without
/// needing a real PostgreSQL connection.
fn create_event_stream_with_inflight(result: Result<ChangeEvent>) -> EventStream {
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let mut event_stream = stream.into_stream(cancel_token);
event_stream.state.park_with_result(result);
event_stream
}
/// A `poll_next` that returned `Pending` parks the stream in `inflight`;
/// dropping the driving future (`tokio::time::timeout`, a losing `select!`
/// arm) does not clear it. `next_event` must resume that future — it used to
/// `.expect()` against a state ordinary use can reach, and panicked.
#[tokio::test]
async fn next_event_resumes_an_abandoned_poll_instead_of_panicking() {
let event = ChangeEvent {
event_type: EventType::Begin {
final_lsn: crate::types::Lsn::new(0x1000),
commit_timestamp: chrono::Utc::now(),
transaction_id: 7,
},
lsn: crate::types::Lsn::new(0x1000),
metadata: None,
};
let mut es = create_event_stream_with_inflight(Ok(event));
let got = es
.next_event()
.await
.expect("the parked event is still owed");
assert_eq!(got.lsn, crate::types::Lsn::new(0x1000));
assert!(es.state.stream().is_some(), "the stream must be reclaimed");
assert!(!es.state.is_polling(), "the parked future must be cleared");
// Reclaiming heals the state, so the inner-only accessors work again.
let _ = es.current_lsn();
}
/// `shutdown`'s documented step 2 is "Sends a final LSN feedback". With a
/// poll parked, the stream lives inside `inflight`, so the old code took the
/// `else` arm, skipped the feedback, returned `Ok(())`, and then dropped the
/// future that owned the connection — tearing down the COPY with no CopyDone.
#[tokio::test]
async fn shutdown_reclaims_the_stream_from_an_abandoned_poll() {
let mut es = create_event_stream_with_inflight(Err(ReplicationError::protocol(
"poll was abandoned".to_string(),
)));
es.shutdown().await.ok();
assert!(
es.state.stream().is_some(),
"shutdown must reclaim the stream so the final feedback can be sent"
);
assert!(es.is_terminated());
}
/// Two wire options, one user intent. The server accepts `two_phase 'on'` on
/// a slot created without `TWO_PHASE` and then silently never decodes
/// PREPARE, so setting only one of the pair fails quietly.
#[test]
fn with_two_phase_sets_both_wire_options() {
let config = ReplicationStreamConfig::builder("s", "p").with_two_phase(true);
assert!(config.two_phase, "START_REPLICATION option");
assert!(
config.slot_options.two_phase,
"CREATE_REPLICATION_SLOT option — without this the slot is not \
two-phase capable and PREPARE is never decoded"
);
let off = ReplicationStreamConfig::builder("s", "p").with_two_phase(false);
assert!(!off.two_phase);
assert!(!off.slot_options.two_phase);
}
#[tokio::test]
async fn test_poll_next_ok_event_branch() {
use futures::StreamExt;
let event = ChangeEvent {
event_type: EventType::Begin {
final_lsn: crate::types::Lsn::new(0x1000),
commit_timestamp: chrono::Utc::now(),
transaction_id: 42,
},
lsn: crate::types::Lsn::new(0x1000),
metadata: None,
};
let mut event_stream = create_event_stream_with_inflight(Ok(event));
// poll_next should return Ready(Some(Ok(event)))
let result = event_stream.next().await;
assert!(result.is_some());
let item = result.unwrap();
assert!(item.is_ok());
let ev = item.unwrap();
match ev.event_type {
EventType::Begin { transaction_id, .. } => assert_eq!(transaction_id, 42),
_ => panic!("Expected Begin event"),
}
// Stream should NOT be terminated after a successful event
assert!(!event_stream.is_terminated());
// Inner should be restored
assert!(event_stream.state.stream().is_some());
assert!(!event_stream.state.is_polling());
}
#[tokio::test]
async fn test_poll_next_permanent_error_branch() {
use futures::StreamExt;
let err = ReplicationError::PermanentConnection("fatal error".to_string());
let mut event_stream = create_event_stream_with_inflight(Err(err));
// poll_next should return Ready(Some(Err(...))) and set terminated
let result = event_stream.next().await;
assert!(result.is_some());
let item = result.unwrap();
assert!(item.is_err());
let e = item.unwrap_err();
assert!(e.is_permanent());
assert!(e.to_string().contains("fatal error"));
// Stream SHOULD be terminated after a permanent error
assert!(event_stream.is_terminated());
// Inner should be restored
assert!(event_stream.state.stream().is_some());
// Subsequent poll should return None (terminated)
let result2 = event_stream.next().await;
assert!(result2.is_none());
}
#[tokio::test]
async fn test_poll_next_transient_error_branch() {
use futures::StreamExt;
// Protocol error is neither permanent, nor cancelled, nor transient (by enum match) —
// it falls through to the catch-all Err(e) arm
let err = ReplicationError::Protocol("not in replication mode".to_string());
let mut event_stream = create_event_stream_with_inflight(Err(err));
// poll_next should return Ready(Some(Err(...))) but NOT terminate
let result = event_stream.next().await;
assert!(result.is_some());
let item = result.unwrap();
assert!(item.is_err());
let e = item.unwrap_err();
assert!(!e.is_permanent());
assert!(!e.is_cancelled());
// Stream should NOT be terminated — it's a transient error
assert!(!event_stream.is_terminated());
// Inner should be restored
assert!(event_stream.state.stream().is_some());
}
#[tokio::test]
async fn test_poll_next_cancelled_error_branch() {
use futures::StreamExt;
let err = ReplicationError::Cancelled("user cancelled".to_string());
let mut event_stream = create_event_stream_with_inflight(Err(err));
// poll_next should return Ready(None) and set terminated
let result = event_stream.next().await;
assert!(result.is_none());
assert!(event_stream.is_terminated());
assert!(event_stream.state.stream().is_some());
}
#[tokio::test]
async fn test_poll_next_authentication_error_is_permanent() {
use futures::StreamExt;
let err = ReplicationError::Authentication("bad credentials".to_string());
let mut event_stream = create_event_stream_with_inflight(Err(err));
let result = event_stream.next().await;
assert!(result.is_some());
let item = result.unwrap();
assert!(item.is_err());
assert!(item.unwrap_err().is_permanent());
assert!(event_stream.is_terminated());
}
#[tokio::test]
async fn test_poll_next_slot_error_is_permanent() {
use futures::StreamExt;
let err = ReplicationError::ReplicationSlot("slot does not exist".to_string());
let mut event_stream = create_event_stream_with_inflight(Err(err));
let result = event_stream.next().await;
assert!(result.is_some());
assert!(result.unwrap().unwrap_err().is_permanent());
assert!(event_stream.is_terminated());
}
#[tokio::test]
async fn test_poll_next_buffer_error_is_transient_fallthrough() {
use futures::StreamExt;
let err = ReplicationError::Buffer("buffer overflow".to_string());
let mut event_stream = create_event_stream_with_inflight(Err(err));
let result = event_stream.next().await;
assert!(result.is_some());
let item = result.unwrap();
assert!(item.is_err());
// Buffer errors are neither permanent nor cancelled — stream stays alive
assert!(!event_stream.is_terminated());
}
#[tokio::test]
async fn test_poll_next_multiple_transient_errors_keep_stream_alive() {
use futures::StreamExt;
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let mut event_stream = stream.into_stream(cancel_token);
// Inject first transient error
event_stream
.state
.park_with_result(Err(ReplicationError::Protocol("error 1".to_string())));
let result1 = event_stream.next().await;
assert!(result1.is_some());
assert!(result1.unwrap().is_err());
assert!(!event_stream.is_terminated());
// Inject second transient error
event_stream
.state
.park_with_result(Err(ReplicationError::Buffer("error 2".to_string())));
let result2 = event_stream.next().await;
assert!(result2.is_some());
assert!(result2.unwrap().is_err());
assert!(!event_stream.is_terminated()); // Still alive after multiple transient errors
// Now inject a permanent error
event_stream
.state
.park_with_result(Err(ReplicationError::PermanentConnection(
"done".to_string(),
)));
let result3 = event_stream.next().await;
assert!(result3.is_some());
assert!(result3.unwrap().is_err());
assert!(event_stream.is_terminated()); // Now terminated
}
#[tokio::test]
async fn test_poll_next_success_then_cancel() {
use futures::StreamExt;
let config = create_test_config();
let stream = create_test_stream(config);
let cancel_token = CancellationToken::new();
let mut event_stream = stream.into_stream(cancel_token);
// Inject a successful event
let event = ChangeEvent {
event_type: EventType::Commit {
commit_lsn: crate::types::Lsn::new(0x2000),
end_lsn: crate::types::Lsn::new(0x2100),
commit_timestamp: chrono::Utc::now(),
},
lsn: crate::types::Lsn::new(0x2000),
metadata: None,
};
event_stream.state.park_with_result(Ok(event));
let result1 = event_stream.next().await;
assert!(result1.is_some());
assert!(result1.unwrap().is_ok());
assert!(!event_stream.is_terminated());
// Now inject a cancellation
event_stream
.state
.park_with_result(Err(ReplicationError::Cancelled(
"shutting down".to_string(),
)));
let result2 = event_stream.next().await;
assert!(result2.is_none());
assert!(event_stream.is_terminated());
}
// ================================================================
// Two-phase commit convert_to_change_event tests
// ================================================================
#[test]
fn test_convert_to_change_event_begin_prepare() {
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::BeginPrepare {
prepare_lsn: 0x5000,
end_lsn: 0x5100,
timestamp: 756_864_000_000_000, // some PG timestamp
xid: 42,
gid: "gid_test_bp".to_string(),
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x4900).unwrap();
assert!(result.is_some());
let event = result.unwrap();
assert_eq!(event.event_type_str(), "begin_prepare");
assert_eq!(event.lsn.value(), 0x4900);
if let EventType::BeginPrepare {
transaction_id,
prepare_lsn,
end_lsn,
gid,
..
} = &event.event_type
{
assert_eq!(*transaction_id, 42);
assert_eq!(prepare_lsn.value(), 0x5000);
assert_eq!(end_lsn.value(), 0x5100);
assert_eq!(gid.as_ref(), "gid_test_bp");
} else {
panic!("Expected BeginPrepare event");
}
}
#[test]
fn test_convert_to_change_event_prepare() {
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Prepare {
flags: 0,
prepare_lsn: 0x6000,
end_lsn: 0x6100,
timestamp: 756_864_000_000_000,
xid: 43,
gid: "gid_test_p".to_string(),
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x6000).unwrap();
assert!(result.is_some());
let event = result.unwrap();
assert_eq!(event.event_type_str(), "prepare");
assert_eq!(event.lsn.value(), 0x6000);
if let EventType::Prepare {
flags,
transaction_id,
prepare_lsn,
end_lsn,
gid,
..
} = &event.event_type
{
assert_eq!(*flags, 0);
assert_eq!(*transaction_id, 43);
assert_eq!(prepare_lsn.value(), 0x6000);
assert_eq!(end_lsn.value(), 0x6100);
assert_eq!(gid.as_ref(), "gid_test_p");
} else {
panic!("Expected Prepare event");
}
}
#[test]
fn test_convert_to_change_event_commit_prepared() {
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::CommitPrepared {
flags: 0,
commit_lsn: 0x7000,
end_lsn: 0x7100,
timestamp: 756_864_000_000_000,
xid: 44,
gid: "gid_test_cp".to_string(),
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x7000).unwrap();
assert!(result.is_some());
let event = result.unwrap();
assert_eq!(event.event_type_str(), "commit_prepared");
assert_eq!(event.lsn.value(), 0x7000);
if let EventType::CommitPrepared {
flags,
transaction_id,
commit_lsn,
end_lsn,
gid,
..
} = &event.event_type
{
assert_eq!(*flags, 0);
assert_eq!(*transaction_id, 44);
assert_eq!(commit_lsn.value(), 0x7000);
assert_eq!(end_lsn.value(), 0x7100);
assert_eq!(gid.as_ref(), "gid_test_cp");
} else {
panic!("Expected CommitPrepared event");
}
}
#[test]
fn test_convert_to_change_event_rollback_prepared() {
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::RollbackPrepared {
flags: 0,
prepare_end_lsn: 0x8000,
rollback_end_lsn: 0x8100,
prepare_timestamp: 756_864_000_000_000,
rollback_timestamp: 756_864_060_000_000,
xid: 45,
gid: "gid_test_rp".to_string(),
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x8000).unwrap();
assert!(result.is_some());
let event = result.unwrap();
assert_eq!(event.event_type_str(), "rollback_prepared");
assert_eq!(event.lsn.value(), 0x8000);
if let EventType::RollbackPrepared {
flags,
transaction_id,
prepare_end_lsn,
rollback_end_lsn,
gid,
..
} = &event.event_type
{
assert_eq!(*flags, 0);
assert_eq!(*transaction_id, 45);
assert_eq!(prepare_end_lsn.value(), 0x8000);
assert_eq!(rollback_end_lsn.value(), 0x8100);
assert_eq!(gid.as_ref(), "gid_test_rp");
} else {
panic!("Expected RollbackPrepared event");
}
}
#[test]
fn test_convert_to_change_event_stream_prepare() {
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::StreamPrepare {
flags: 0,
prepare_lsn: 0x9000,
end_lsn: 0x9100,
timestamp: 756_864_000_000_000,
xid: 46,
gid: "gid_test_sp".to_string(),
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0x9000).unwrap();
assert!(result.is_some());
let event = result.unwrap();
assert_eq!(event.event_type_str(), "stream_prepare");
assert_eq!(event.lsn.value(), 0x9000);
if let EventType::StreamPrepare {
flags,
transaction_id,
prepare_lsn,
end_lsn,
gid,
..
} = &event.event_type
{
assert_eq!(*flags, 0);
assert_eq!(*transaction_id, 46);
assert_eq!(prepare_lsn.value(), 0x9000);
assert_eq!(end_lsn.value(), 0x9100);
assert_eq!(gid.as_ref(), "gid_test_sp");
} else {
panic!("Expected StreamPrepare event");
}
}
// ================================================================
// Origin event with field validation
// ================================================================
#[test]
fn test_convert_to_change_event_origin_fields() {
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Origin {
origin_lsn: 0xA000,
origin_name: "upstream_dc".to_string(),
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0xA100).unwrap();
assert!(result.is_some());
let event = result.unwrap();
assert_eq!(event.event_type_str(), "origin");
assert_eq!(event.lsn.value(), 0xA100);
if let EventType::Origin {
origin_lsn,
origin_name,
} = &event.event_type
{
assert_eq!(origin_lsn.value(), 0xA000);
assert_eq!(origin_name.as_ref(), "upstream_dc");
} else {
panic!("Expected Origin event");
}
}
// ================================================================
// Relation schema-change detection tests
// ================================================================
#[test]
fn test_convert_to_change_event_relation_first_time_emits_then_dedups() {
use crate::protocol::ColumnInfo;
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
let mk = || {
StreamingReplicationMessage::new(LogicalReplicationMessage::Relation {
relation_id: 200,
namespace: Arc::from("public"),
relation_name: Arc::from("orders"),
replica_identity: b'd',
columns: vec![
ColumnInfo::new(1, "id".to_string(), 23, -1),
ColumnInfo::new(0, "total".to_string(), 1700, -1),
],
})
};
// First sighting: emit, and cache the schema.
let first =
crate::protocol::message_to_change_event(&mut stream.state, mk(), 0xB000).unwrap();
assert!(first.is_some(), "First-time relation should emit an event");
assert!(stream.state.get_relation(200).is_some());
// pgoutput repeats Relation at the start of every transaction; an
// unchanged repeat must stay silent.
let repeat =
crate::protocol::message_to_change_event(&mut stream.state, mk(), 0xB100).unwrap();
assert!(repeat.is_none(), "Unchanged repeat should return None");
}
#[test]
fn test_convert_to_change_event_relation_same_schema_returns_none() {
use crate::protocol::{ColumnInfo, RelationInfo};
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
// Pre-populate cache
let relation = RelationInfo::new(
300,
"public",
"items",
b'd',
vec![
ColumnInfo::new(1, "id".to_string(), 23, -1),
ColumnInfo::new(0, "name".to_string(), 25, -1),
],
);
stream.state.add_relation(relation);
// Send identical relation message — no schema change
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Relation {
relation_id: 300,
namespace: Arc::from("public"),
relation_name: Arc::from("items"),
replica_identity: b'd',
columns: vec![
ColumnInfo::new(1, "id".to_string(), 23, -1),
ColumnInfo::new(0, "name".to_string(), 25, -1),
],
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0xC000).unwrap();
assert!(
result.is_none(),
"Same schema should not emit a Relation event"
);
}
#[test]
fn test_convert_to_change_event_relation_schema_change_new_column() {
use crate::protocol::{ColumnInfo, RelationInfo};
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
// Pre-populate cache with 2 columns
let relation = RelationInfo::new(
400,
"public",
"accounts",
b'd',
vec![
ColumnInfo::new(1, "id".to_string(), 23, -1),
ColumnInfo::new(0, "email".to_string(), 25, -1),
],
);
stream.state.add_relation(relation);
// Send relation message with an extra column — schema change!
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Relation {
relation_id: 400,
namespace: Arc::from("public"),
relation_name: Arc::from("accounts"),
replica_identity: b'd',
columns: vec![
ColumnInfo::new(1, "id".to_string(), 23, -1),
ColumnInfo::new(0, "email".to_string(), 25, -1),
ColumnInfo::new(0, "phone".to_string(), 25, -1),
],
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0xD000).unwrap();
assert!(
result.is_some(),
"Schema change should emit a Relation event"
);
let event = result.unwrap();
assert_eq!(event.event_type_str(), "relation");
assert_eq!(event.lsn.value(), 0xD000);
if let EventType::Relation {
relation_id,
namespace,
relation_name,
columns,
..
} = &event.event_type
{
assert_eq!(*relation_id, 400);
assert_eq!(namespace.as_ref(), "public");
assert_eq!(relation_name.as_ref(), "accounts");
assert_eq!(columns.len(), 3);
assert_eq!(columns[2].name.as_ref(), "phone");
} else {
panic!("Expected Relation event");
}
}
#[test]
fn test_convert_to_change_event_relation_schema_change_type_change() {
use crate::protocol::{ColumnInfo, RelationInfo};
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
// Pre-populate with type_id=25 (text)
let relation = RelationInfo::new(
500,
"public",
"products",
b'd',
vec![
ColumnInfo::new(1, "id".to_string(), 23, -1),
ColumnInfo::new(0, "price".to_string(), 25, -1),
],
);
stream.state.add_relation(relation);
// Send relation with price changed to type_id=1700 (numeric)
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Relation {
relation_id: 500,
namespace: Arc::from("public"),
relation_name: Arc::from("products"),
replica_identity: b'd',
columns: vec![
ColumnInfo::new(1, "id".to_string(), 23, -1),
ColumnInfo::new(0, "price".to_string(), 1700, -1),
],
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0xE000).unwrap();
assert!(
result.is_some(),
"Column type change should emit a Relation event"
);
let event = result.unwrap();
assert_eq!(event.event_type_str(), "relation");
}
#[test]
fn test_convert_to_change_event_relation_schema_change_replica_identity() {
use crate::protocol::{ColumnInfo, RelationInfo};
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
// Pre-populate with replica_identity = 'd' (Default)
let relation = RelationInfo::new(
600,
"public",
"logs",
b'd',
vec![ColumnInfo::new(1, "id".to_string(), 23, -1)],
);
stream.state.add_relation(relation);
// Send relation with replica_identity changed to 'f' (Full)
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Relation {
relation_id: 600,
namespace: Arc::from("public"),
relation_name: Arc::from("logs"),
replica_identity: b'f',
columns: vec![ColumnInfo::new(1, "id".to_string(), 23, -1)],
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0xF000).unwrap();
assert!(
result.is_some(),
"Replica identity change should emit a Relation event"
);
let event = result.unwrap();
if let EventType::Relation {
replica_identity, ..
} = &event.event_type
{
assert_eq!(*replica_identity, ReplicaIdentity::Full);
} else {
panic!("Expected Relation event");
}
}
#[test]
fn test_convert_to_change_event_relation_schema_change_column_rename() {
use crate::protocol::{ColumnInfo, RelationInfo};
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
// Pre-populate
let relation = RelationInfo::new(
700,
"public",
"events",
b'd',
vec![
ColumnInfo::new(1, "id".to_string(), 23, -1),
ColumnInfo::new(0, "name".to_string(), 25, -1),
],
);
stream.state.add_relation(relation);
// Column renamed: "name" -> "title"
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Relation {
relation_id: 700,
namespace: Arc::from("public"),
relation_name: Arc::from("events"),
replica_identity: b'd',
columns: vec![
ColumnInfo::new(1, "id".to_string(), 23, -1),
ColumnInfo::new(0, "title".to_string(), 25, -1),
],
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0xF100).unwrap();
assert!(
result.is_some(),
"Column rename should emit a Relation event"
);
let event = result.unwrap();
if let EventType::Relation { columns, .. } = &event.event_type {
assert_eq!(columns[1].name.as_ref(), "title");
} else {
panic!("Expected Relation event");
}
}
#[test]
fn test_convert_to_change_event_relation_schema_change_namespace() {
use crate::protocol::{ColumnInfo, RelationInfo};
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
// Pre-populate in "public" schema
let relation = RelationInfo::new(
800,
"public",
"data",
b'd',
vec![ColumnInfo::new(1, "id".to_string(), 23, -1)],
);
stream.state.add_relation(relation);
// Same relation_id but moved to "archive" schema
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Relation {
relation_id: 800,
namespace: Arc::from("archive"),
relation_name: Arc::from("data"),
replica_identity: b'd',
columns: vec![ColumnInfo::new(1, "id".to_string(), 23, -1)],
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0xF200).unwrap();
assert!(
result.is_some(),
"Namespace change should emit a Relation event"
);
let event = result.unwrap();
if let EventType::Relation { namespace, .. } = &event.event_type {
assert_eq!(namespace.as_ref(), "archive");
} else {
panic!("Expected Relation event");
}
}
#[test]
fn test_convert_to_change_event_relation_schema_change_key_flag() {
use crate::protocol::{ColumnInfo, RelationInfo};
use crate::{LogicalReplicationMessage, StreamingReplicationMessage};
let config = create_test_config();
let mut stream = create_test_stream(config);
// Pre-populate: "email" is NOT a key column (flags=0)
let relation = RelationInfo::new(
900,
"public",
"users",
b'd',
vec![
ColumnInfo::new(1, "id".to_string(), 23, -1),
ColumnInfo::new(0, "email".to_string(), 25, -1),
],
);
stream.state.add_relation(relation);
// Now "email" becomes a key column (flags=1)
let msg = StreamingReplicationMessage::new(LogicalReplicationMessage::Relation {
relation_id: 900,
namespace: Arc::from("public"),
relation_name: Arc::from("users"),
replica_identity: b'd',
columns: vec![
ColumnInfo::new(1, "id".to_string(), 23, -1),
ColumnInfo::new(1, "email".to_string(), 25, -1),
],
});
let result =
crate::protocol::message_to_change_event(&mut stream.state, msg, 0xF300).unwrap();
assert!(
result.is_some(),
"Key flag change should emit a Relation event"
);
let event = result.unwrap();
if let EventType::Relation { columns, .. } = &event.event_type {
assert!(columns[1].is_key, "email should now be a key column");
} else {
panic!("Expected Relation event");
}
}
}
#[cfg(test)]
mod ergonomics_config_tests {
use super::*;
use std::time::Duration;
// builder() must match the exact defaults new() produces.
#[test]
fn builder_matches_new_defaults() {
let built = ReplicationStreamConfig::builder("slot_a", "pub_a")
.with_protocol_version(2)
.with_streaming_mode(StreamingMode::On)
.with_feedback_interval(Duration::from_secs(10))
.with_connection_timeout(Duration::from_secs(30))
.with_health_check_interval(Duration::from_secs(60));
let made = ReplicationStreamConfig::new(
"slot_a".to_string(),
"pub_a".to_string(),
2,
StreamingMode::On,
Duration::from_secs(10),
Duration::from_secs(30),
Duration::from_secs(60),
RetryConfig::default(),
);
assert_eq!(built.slot_name, made.slot_name);
assert_eq!(built.publication_name, made.publication_name);
assert_eq!(built.protocol_version, made.protocol_version);
assert_eq!(built.streaming_mode, made.streaming_mode);
assert_eq!(built.feedback_interval, made.feedback_interval);
assert_eq!(built.connection_timeout, made.connection_timeout);
assert_eq!(built.health_check_interval, made.health_check_interval);
// Critical: snapshot default must be Some("nothing").
assert_eq!(built.slot_options.snapshot, Some("nothing".to_string()));
assert_eq!(built.slot_options.snapshot, made.slot_options.snapshot);
}
#[test]
fn default_has_nothing_snapshot() {
let d = ReplicationStreamConfig::default();
assert_eq!(d.slot_options.snapshot, Some("nothing".to_string()));
assert_eq!(d.protocol_version, 2);
}
}
#[cfg(test)]
mod for_each_event_tests {
use super::*;
use crate::{ChangeEvent, ColumnValue, Lsn, RowData};
use std::collections::VecDeque;
use std::sync::atomic::{AtomicU64, Ordering};
struct MockSource {
events: VecDeque<Result<ChangeEvent>>,
last_ack: AtomicU64,
}
impl MockSource {
fn new(events: Vec<Result<ChangeEvent>>) -> Self {
Self {
events: events.into(),
last_ack: AtomicU64::new(0),
}
}
}
impl EventSource for MockSource {
async fn recv(&mut self) -> Result<ChangeEvent> {
self.events
.pop_front()
.unwrap_or_else(|| Err(ReplicationError::Cancelled("drained".into())))
}
fn ack(&self, lsn: XLogRecPtr) {
self.last_ack.store(lsn, Ordering::SeqCst);
}
}
fn ins(lsn: u64) -> ChangeEvent {
let data = RowData::from_pairs(vec![("id", ColumnValue::text("1"))]);
ChangeEvent::insert("public", "t", 1, data, Lsn::new(lsn))
}
#[tokio::test]
async fn acks_after_each_ok_handler() {
let mut src = MockSource::new(vec![Ok(ins(10)), Ok(ins(20))]);
let mut seen = Vec::new();
let out = for_each_event_impl(&mut src, |ev| {
seen.push(ev.lsn.value());
async { Ok(()) }
})
.await;
assert!(out.is_ok());
assert_eq!(seen, vec![10, 20]);
assert_eq!(src.last_ack.load(Ordering::SeqCst), 20);
}
#[tokio::test]
async fn handler_err_stops_without_acking_that_event() {
let mut src = MockSource::new(vec![Ok(ins(10)), Ok(ins(20))]);
let mut n = 0;
let out = for_each_event_impl(&mut src, |_ev| {
n += 1;
async move {
if n == 2 {
Err(ReplicationError::deserialize("boom"))
} else {
Ok(())
}
}
})
.await;
assert!(out.is_err());
// First event (10) acked, second (20) failed before ack.
assert_eq!(src.last_ack.load(Ordering::SeqCst), 10);
}
#[tokio::test]
async fn cancelled_recv_exits_ok() {
let mut src = MockSource::new(vec![Err(ReplicationError::Cancelled("stop".into()))]);
let out = for_each_event_impl(&mut src, |_ev| async { Ok(()) }).await;
assert!(out.is_ok());
}
#[tokio::test]
async fn non_cancelled_recv_err_propagates() {
let mut src = MockSource::new(vec![Err(ReplicationError::deserialize("io"))]);
let out = for_each_event_impl(&mut src, |_ev| async { Ok(()) }).await;
assert!(out.is_err());
}
#[tokio::test]
async fn for_each_event_stops_gracefully_on_stream_stopped() {
let mut src = MockSource::new(vec![
Ok(ins(10)),
Ok(ins(20)),
Err(ReplicationError::stream_stopped(Lsn::new(25))),
]);
let mut count = 0u32;
let res = for_each_event_impl(&mut src, |_ev| {
count += 1;
async { Ok(()) }
})
.await;
assert!(res.is_ok(), "StreamStopped must end the loop gracefully");
assert_eq!(count, 2, "both events before the stop are handled");
assert_eq!(src.last_ack.load(std::sync::atomic::Ordering::SeqCst), 20);
}
}
#[cfg(test)]
mod stop_at_lsn_config_tests {
use super::*;
#[test]
fn default_config_has_no_stop_lsn() {
let cfg = ReplicationStreamConfig::default();
assert_eq!(cfg.stop_at_lsn, None);
let cfg2 = ReplicationStreamConfig::builder("slot", "pub");
assert_eq!(cfg2.stop_at_lsn, None);
}
#[test]
fn with_stop_at_lsn_sets_target() {
let cfg = ReplicationStreamConfig::builder("slot", "pub").with_stop_at_lsn(0x2000u64);
assert_eq!(cfg.stop_at_lsn, Some(Lsn::new(0x2000)));
}
}