Skip to content

Extract runReplicationStream — deduplicate startReplication / startPh… #428

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
fox0430 merged 2 commits into from
Jul 5, 2026
+104 −219
Merged

Extract runReplicationStream — deduplicate startReplication / startPh… #428

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a6aecc2
Extract runReplicationStream — deduplicate startReplication / startPh…
fox0430 Jul 5, 2026
d747909
nph
fox0430 Jul 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
323 changes: 104 additions & 219 deletions async_postgres/pg_replication.nim
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1015,6 +1015,104 @@ proc invalidateAbandonedStream(conn: PgConnection) =
if conn.state == csReplicating:
conn.state = csClosed

proc runReplicationStream(
conn: PgConnection,
startLsn: Lsn,
autoKeepaliveReply: bool,
statusInterval: async_backend.Duration,
callback: ReplicationCallback,
context: string,
): Future[void] {.async.} =
## Shared replication stream body. Caller must have already sent the
## ``START_REPLICATION`` query. ``context`` appears in error messages
## (e.g. ``"replication"`` / ``"physical replication"``).
var queryError: ref PgQueryError

block waitCopyBoth:
while true:
while (let opt = conn.nextMessage(); opt.isSome):
let msg = opt.get
case msg.kind
of bmkCopyBothResponse:
conn.state = csReplicating
break waitCopyBoth
of bmkErrorResponse:
queryError = newPgQueryError(msg.errorFields)
of bmkReadyForQuery:
conn.txStatus = msg.txStatus
conn.state = csReady
if queryError != nil:
raise queryError
raise newException(
PgConnectionError,
"START_REPLICATION " & context & " ended without CopyBothResponse",
)
else:
discard
await conn.fillRecvBuf()

conn.resetReplLsnTracking(startLsn)

defer:
conn.invalidateAbandonedStream()

var lastStatusSent = Moment.now()
var pendingRead: Future[void] = nil
when hasChronos:
defer:
if pendingRead != nil and not pendingRead.finished:
pendingRead.cancelSoon()

block recvLoop:
while true:
while (let opt = conn.nextMessage(); opt.isSome):
let msg = opt.get
case msg.kind
of bmkCopyData:
lastStatusSent = await conn.handleReplicationData(
msg.copyData, autoKeepaliveReply, callback, lastStatusSent
)
of bmkCopyDone:
await conn.sendMsg(@copyDoneMsg)
break recvLoop
of bmkErrorResponse:
queryError = newPgQueryError(msg.errorFields)
of bmkReadyForQuery:
conn.txStatus = msg.txStatus
conn.state = csReady
if queryError != nil:
raise queryError
return
else:
discard
lastStatusSent = await conn.maybeSendPeriodicStatus(
autoKeepaliveReply, statusInterval, lastStatusSent
)
if conn.state == csClosed:
raise newException(PgConnectionError, "Connection closed during " & context)
pendingRead =
await conn.replFillRecvBuf(statusInterval, lastStatusSent, pendingRead)
lastStatusSent = await conn.maybeSendPeriodicStatus(
autoKeepaliveReply, statusInterval, lastStatusSent
)

block drainLoop:
while true:
while (let opt = conn.nextMessage(); opt.isSome):
let msg = opt.get
case msg.kind
of bmkErrorResponse:
queryError = newPgQueryError(msg.errorFields)
of bmkReadyForQuery:
conn.txStatus = msg.txStatus
conn.state = csReady
if queryError != nil:
raise queryError
break drainLoop
else:
discard
await conn.fillRecvBuf()

proc startReplication*(
conn: PgConnection,
slotName: string,
Expand Down Expand Up @@ -1148,120 +1246,9 @@ proc startReplication*(
sql.add(")")

await conn.sendMsg(encodeQuery(sql))

# Wait for CopyBothResponse. The block exits only via the CopyBothResponse
# break or by raising — fillRecvBuf re-raises connection errors, and a
# ReadyForQuery before CopyBothResponse always raises.
var queryError: ref PgQueryError

block waitCopyBoth:
while true:
while (let opt = conn.nextMessage(); opt.isSome):
let msg = opt.get
case msg.kind
of bmkCopyBothResponse:
conn.state = csReplicating
break waitCopyBoth
of bmkErrorResponse:
queryError = newPgQueryError(msg.errorFields)
of bmkReadyForQuery:
conn.txStatus = msg.txStatus
conn.state = csReady
if queryError != nil:
raise queryError
raise newException(
PgConnectionError, "START_REPLICATION ended without CopyBothResponse"
)
else:
discard
await conn.fillRecvBuf()

# Track the highest end LSN of WAL data actually received from the wire —
# computed as XLogData.startLsn + data.len, *not* XLogData.walEnd. ``walEnd``
# is the server's current WAL position at the time the message was sent and can
# be ahead of the bytes this message carries; acknowledging ``walEnd`` would
# falsely advance ``confirmed_flush_lsn`` past unprocessed WAL and cause data
# loss on slot restart. The position lives on the connection (the single source
# of truth, read by confirmFlushed and the auto-reply); reset it and the
# confirmed-flush position to the resume point so a reused connection does not
# inherit a stale value.
conn.resetReplLsnTracking(startLsn)

# If anything unwinds the streaming loop while the connection is still
# mid-CopyBoth — most often the user callback raising — the protocol stream is
# left out of sync. Poison the connection on the way out so it is not stranded
# in csReplicating (where every later call raises a misleading PgStateError).
# The clean stop (CopyDone -> ReadyForQuery) returns it to csReady first, so
# this is a no-op on the normal exit. See invalidateAbandonedStream.
defer:
conn.invalidateAbandonedStream()

var lastStatusSent = Moment.now()
var pendingRead: Future[void] = nil # in-flight timed read, threaded across waits
when hasChronos:
# If an error unwinds the loop while a timed read is still in flight, abandon
# it so it is not left dangling on the torn-down connection. Normal exits
# (CopyDone / ReadyForQuery) only happen after a read has been consumed, so
# `pendingRead` is nil then and this is a no-op — never a cancel-then-reread.
defer:
if pendingRead != nil and not pendingRead.finished:
pendingRead.cancelSoon()

# Streaming loop
block recvLoop:
while true:
while (let opt = conn.nextMessage(); opt.isSome):
let msg = opt.get
case msg.kind
of bmkCopyData:
lastStatusSent = await conn.handleReplicationData(
msg.copyData, autoKeepaliveReply, callback, lastStatusSent
)
of bmkCopyDone:
# Server ended the stream; reply with CopyDone before draining
await conn.sendMsg(@copyDoneMsg)
break recvLoop
of bmkErrorResponse:
queryError = newPgQueryError(msg.errorFields)
of bmkReadyForQuery:
conn.txStatus = msg.txStatus
conn.state = csReady
if queryError != nil:
raise queryError
return
else:
discard
# Send a proactive status update if one is due after draining this batch
# (covers a busy stream on both backends).
lastStatusSent = await conn.maybeSendPeriodicStatus(
autoKeepaliveReply, statusInterval, lastStatusSent
)
if conn.state == csClosed:
raise newException(PgConnectionError, "Connection closed during replication")
pendingRead =
await conn.replFillRecvBuf(statusInterval, lastStatusSent, pendingRead)
# Send again after a timed wake with no new data (idle coverage on chronos).
lastStatusSent = await conn.maybeSendPeriodicStatus(
autoKeepaliveReply, statusInterval, lastStatusSent
)

# After CopyDone, drain to ReadyForQuery
block drainLoop:
while true:
while (let opt = conn.nextMessage(); opt.isSome):
let msg = opt.get
case msg.kind
of bmkErrorResponse:
queryError = newPgQueryError(msg.errorFields)
of bmkReadyForQuery:
conn.txStatus = msg.txStatus
conn.state = csReady
if queryError != nil:
raise queryError
break drainLoop
else:
discard
await conn.fillRecvBuf()
await runReplicationStream(
conn, startLsn, autoKeepaliveReply, statusInterval, callback, "replication"
)

proc stopReplication*(conn: PgConnection): Future[void] {.async.} =
## Gracefully terminate the replication stream.
Expand Down Expand Up @@ -1352,108 +1339,6 @@ proc startPhysicalReplication*(
sql.add(" TIMELINE " & $timeline)

await conn.sendMsg(encodeQuery(sql))

# See startReplication for the invariants of this block — it can only exit
# via the CopyBothResponse break or by raising.
var queryError: ref PgQueryError

block waitCopyBoth:
while true:
while (let opt = conn.nextMessage(); opt.isSome):
let msg = opt.get
case msg.kind
of bmkCopyBothResponse:
conn.state = csReplicating
break waitCopyBoth
of bmkErrorResponse:
queryError = newPgQueryError(msg.errorFields)
of bmkReadyForQuery:
conn.txStatus = msg.txStatus
conn.state = csReady
if queryError != nil:
raise queryError
raise newException(
PgConnectionError,
"START_REPLICATION PHYSICAL ended without CopyBothResponse",
)
else:
discard
await conn.fillRecvBuf()

# Highest end LSN of WAL data actually received — see startReplication for
# why we track ``startLsn + data.len`` rather than ``walEnd``, and why the
# position lives on the connection. Reset it and the confirmed-flush position
# to the resume point so a reused connection does not inherit a stale value.
conn.resetReplLsnTracking(startLsn)

# See startReplication: poison the connection if the streaming loop unwinds
# while still mid-CopyBoth (e.g. the callback raised), so it is not stranded in
# csReplicating. No-op on the clean CopyDone -> ReadyForQuery exit.
defer:
conn.invalidateAbandonedStream()

var lastStatusSent = Moment.now()
var pendingRead: Future[void] = nil # in-flight timed read, threaded across waits
when hasChronos:
# See startReplication: drop a still-in-flight timed read if an error unwinds
# the loop; nil (and so a no-op) on the normal CopyDone / ReadyForQuery exits.
defer:
if pendingRead != nil and not pendingRead.finished:
pendingRead.cancelSoon()

block recvLoop:
while true:
while (let opt = conn.nextMessage(); opt.isSome):
let msg = opt.get
case msg.kind
of bmkCopyData:
lastStatusSent = await conn.handleReplicationData(
msg.copyData, autoKeepaliveReply, callback, lastStatusSent
)
of bmkCopyDone:
await conn.sendMsg(@copyDoneMsg)
break recvLoop
of bmkErrorResponse:
queryError = newPgQueryError(msg.errorFields)
of bmkReadyForQuery:
conn.txStatus = msg.txStatus
conn.state = csReady
if queryError != nil:
raise queryError
return
else:
discard
# See startReplication: send a proactive status update when one is due,
# both after draining a batch and after a timed idle wake.
lastStatusSent = await conn.maybeSendPeriodicStatus(
autoKeepaliveReply, statusInterval, lastStatusSent
)
if conn.state == csClosed:
raise newException(
PgConnectionError, "Connection closed during physical replication"
)
pendingRead =
await conn.replFillRecvBuf(statusInterval, lastStatusSent, pendingRead)
lastStatusSent = await conn.maybeSendPeriodicStatus(
autoKeepaliveReply, statusInterval, lastStatusSent
)

# After CopyDone, drain to ReadyForQuery. Accept the optional timeline-switch
# result set (RowDescription + DataRow + CommandComplete) the server emits
# when the stream stopped because the standby reached the end of a timeline.
block drainLoop:
while true:
while (let opt = conn.nextMessage(); opt.isSome):
let msg = opt.get
case msg.kind
of bmkErrorResponse:
queryError = newPgQueryError(msg.errorFields)
of bmkReadyForQuery:
conn.txStatus = msg.txStatus
conn.state = csReady
if queryError != nil:
raise queryError
break drainLoop
else:
discard
await conn.fillRecvBuf()
await runReplicationStream(
conn, startLsn, autoKeepaliveReply, statusInterval, callback, "physical replication"
)