diff --git a/async_postgres/pg_replication.nim b/async_postgres/pg_replication.nim index 3a0b457..1c814f0 100644 --- a/async_postgres/pg_replication.nim +++ b/async_postgres/pg_replication.nim @@ -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, @@ -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. @@ -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" + )