From cdb2d5e16dd42690601ec436e816d24490ec00d3 Mon Sep 17 00:00:00 2001 From: Leonardo Cecchi Date: Tue, 25 Aug 2026 18:07:02 +0200 Subject: [PATCH 1/2] refactor(wal): stream WAL blocks to server over bidi gRPC, drop client buffering Change the Put RPC from client-streaming to bidi-streaming so the Klio server acknowledges each WAL block as it is written, instead of only returning a single result when the stream closes. The gRPC client and walserver upload handler are updated accordingly, and the WAL receiver now runs two goroutines: one pumping PostgreSQL WAL data to the server, the other reading per-block feedback and turning it into standby status updates for PostgreSQL. This makes the in-memory WAL buffer, its periodic flush timer and the related config (buffer_size, flush_timeout_ms, standby_message_timeout_seconds, SetDefaults) unnecessary, since blocks are now sent to the server as they arrive and confirmed via the server's own feedback rather than a client-side timer. Signed-off-by: Leonardo Cecchi --- core/cmd/backup/delete.go | 3 - core/cmd/backup/get-metadata.go | 3 - core/cmd/backup/list.go | 3 - core/cmd/backup/run.go | 3 - core/cmd/backup/verify.go | 3 - core/cmd/get-metadata.go | 3 - core/cmd/get-wal.go | 3 - core/cmd/reset-lsn.go | 3 - core/cmd/restore.go | 3 - core/cmd/retention/get.go | 3 - core/cmd/retention/set.go | 3 - core/cmd/send-wal.go | 3 - core/cmd/walplayer/play.go | 3 - .../klioclient/grpcclient/connection.go | 13 +- .../client/klioclient/grpcclient/errors.go | 4 + .../client/klioclient/grpcclient/walclient.go | 20 +-- .../klioclient/grpcclient/walstreamer.go | 6 +- .../klioclient/grpcclient/waluploader.go | 11 ++ core/internal/client/klioclient/interfaces.go | 26 +-- core/internal/client/sendwal/buffer/buffer.go | 68 +------- core/internal/client/sendwal/buffer/grpc.go | 94 ++++++++++- .../internal/client/sendwal/buffer/handler.go | 11 +- core/internal/client/sendwal/buffer/memory.go | 101 ------------ core/internal/client/sendwal/receiver.go | 152 +++++++++--------- core/internal/grpc/klio_wal.pb.go | 6 +- core/internal/grpc/klio_wal_grpc.pb.go | 13 +- core/internal/server/walserver/upload.go | 10 +- core/pkg/config/client.go | 40 ----- core/pkg/config/client_validate.go | 9 -- core/pkg/config/client_validate_test.go | 48 +----- core/pkg/config/decode_test.go | 12 +- core/proto/klio_wal.proto | 2 +- documentation/web/docs/developer/_protocol.md | 2 +- operator/internal/klioconfig/config.go | 4 - operator/pkg/config/client.go | 40 ----- 35 files changed, 241 insertions(+), 490 deletions(-) delete mode 100644 core/internal/client/sendwal/buffer/memory.go diff --git a/core/cmd/backup/delete.go b/core/cmd/backup/delete.go index 6ab86859..2d0db005 100644 --- a/core/cmd/backup/delete.go +++ b/core/cmd/backup/delete.go @@ -47,9 +47,6 @@ var deleteCmd = &cobra.Command{ return fmt.Errorf("could not unmarshal configuration: %w", err) } - // Sets the default values, to be overridden by the user configuration - configuration.SetDefaults() - if configuration.Client == (config.ClientConfig{}) { return cli.ErrClientSectionIsRequired } diff --git a/core/cmd/backup/get-metadata.go b/core/cmd/backup/get-metadata.go index 546398e7..f68df954 100644 --- a/core/cmd/backup/get-metadata.go +++ b/core/cmd/backup/get-metadata.go @@ -53,9 +53,6 @@ func runGetMetadata(cmd *cobra.Command, args []string) error { return fmt.Errorf("could not unmarshal configuration: %w", err) } - // Sets the default values, to be overridden by the user configuration - configuration.SetDefaults() - if configuration.Client == (config.ClientConfig{}) { return cli.ErrClientSectionIsRequired } diff --git a/core/cmd/backup/list.go b/core/cmd/backup/list.go index 06c7fc87..07071bed 100644 --- a/core/cmd/backup/list.go +++ b/core/cmd/backup/list.go @@ -47,9 +47,6 @@ var listCmd = &cobra.Command{ return fmt.Errorf("could not unmarshal configuration: %w", err) } - // Sets the default values, to be overridden by the user configuration - configuration.SetDefaults() - if configuration.Client == (config.ClientConfig{}) { return cli.ErrClientSectionIsRequired } diff --git a/core/cmd/backup/run.go b/core/cmd/backup/run.go index d12916e3..d0fc1c01 100644 --- a/core/cmd/backup/run.go +++ b/core/cmd/backup/run.go @@ -61,9 +61,6 @@ func runBackup(cmd *cobra.Command, _ []string) error { return fmt.Errorf("could not unmarshal configuration: %w", err) } - // Sets the default values, to be overridden by the user configuration - configuration.SetDefaults() - if configuration.Client == (config.ClientConfig{}) { return cli.ErrClientSectionIsRequired } diff --git a/core/cmd/backup/verify.go b/core/cmd/backup/verify.go index 23fb8a3d..efccda21 100644 --- a/core/cmd/backup/verify.go +++ b/core/cmd/backup/verify.go @@ -146,9 +146,6 @@ func loadAndValidateConfig() (*config.Data, error) { return nil, fmt.Errorf("could not unmarshal configuration: %w", err) } - // Sets the default values, to be overridden by the user configuration - configuration.SetDefaults() - if configuration.Client == (config.ClientConfig{}) { return nil, cli.ErrClientSectionIsRequired } diff --git a/core/cmd/get-metadata.go b/core/cmd/get-metadata.go index d24a3b4d..6fcb9481 100644 --- a/core/cmd/get-metadata.go +++ b/core/cmd/get-metadata.go @@ -49,9 +49,6 @@ var getMetadataCmd = &cobra.Command{ return fmt.Errorf("could not unmarshal configuration: %w", err) } - // Sets the defaults values, to be overridden by the user configuration - configuration.SetDefaults() - if configuration.Client == (config.ClientConfig{}) { return cli.ErrClientSectionIsRequired } diff --git a/core/cmd/get-wal.go b/core/cmd/get-wal.go index b4f3cb33..381a2c9c 100644 --- a/core/cmd/get-wal.go +++ b/core/cmd/get-wal.go @@ -56,9 +56,6 @@ var getWalCmd = &cobra.Command{ return fmt.Errorf("could not unmarshal configuration: %w", err) } - // Sets the defaults values, to be overridden by the user configuration - configuration.SetDefaults() - if configuration.Client == (config.ClientConfig{}) { return cli.ErrClientSectionIsRequired } diff --git a/core/cmd/reset-lsn.go b/core/cmd/reset-lsn.go index 50bd60fb..d66d9b19 100644 --- a/core/cmd/reset-lsn.go +++ b/core/cmd/reset-lsn.go @@ -49,9 +49,6 @@ var resetLSNCommand = &cobra.Command{ return fmt.Errorf("could not unmarshal configuration: %w", err) } - // Sets the defaults values, to be overridden by the user configuration - configuration.SetDefaults() - if configuration.Source == (config.SourceConfig{}) { return cli.ErrSourceSectionIsRequired } diff --git a/core/cmd/restore.go b/core/cmd/restore.go index 306e4c33..1196059d 100644 --- a/core/cmd/restore.go +++ b/core/cmd/restore.go @@ -62,9 +62,6 @@ var restoreCmd = &cobra.Command{ return fmt.Errorf("could not unmarshal configuration: %w", err) } - // Sets the default values, to be overridden by the user configuration - configuration.SetDefaults() - if configuration.Client == (config.ClientConfig{}) { return cli.ErrClientSectionIsRequired } diff --git a/core/cmd/retention/get.go b/core/cmd/retention/get.go index 0022ff22..fe3cf4ba 100644 --- a/core/cmd/retention/get.go +++ b/core/cmd/retention/get.go @@ -47,9 +47,6 @@ var getCmd = &cobra.Command{ return fmt.Errorf("could not unmarshal configuration: %w", err) } - // Sets the default values, to be overridden by the user configuration - configuration.SetDefaults() - if configuration.Client == (config.ClientConfig{}) { return cli.ErrClientSectionIsRequired } diff --git a/core/cmd/retention/set.go b/core/cmd/retention/set.go index 00c41d6d..1a983c64 100644 --- a/core/cmd/retention/set.go +++ b/core/cmd/retention/set.go @@ -47,9 +47,6 @@ var setCmd = &cobra.Command{ return fmt.Errorf("could not unmarshal configuration: %w", err) } - // Sets the defaults values, to be overridden by the user configuration - configuration.SetDefaults() - if configuration.Client == (config.ClientConfig{}) { return cli.ErrClientSectionIsRequired } diff --git a/core/cmd/send-wal.go b/core/cmd/send-wal.go index d7127e1e..e1bad6f9 100644 --- a/core/cmd/send-wal.go +++ b/core/cmd/send-wal.go @@ -67,9 +67,6 @@ var sendWalCmd = &cobra.Command{ return fmt.Errorf("could not unmarshal configuration: %w", err) } - // Sets the defaults values, to be overridden by the user configuration - configuration.SetDefaults() - if configuration.Source == (config.SourceConfig{}) { return cli.ErrSourceSectionIsRequired } diff --git a/core/cmd/walplayer/play.go b/core/cmd/walplayer/play.go index 1d178b3c..e077867c 100644 --- a/core/cmd/walplayer/play.go +++ b/core/cmd/walplayer/play.go @@ -54,9 +54,6 @@ var playCmd = &cobra.Command{ return fmt.Errorf("could not unmarshal configuration: %w", err) } - // Sets the defaults values, to be overridden by the user configuration - configuration.SetDefaults() - if configuration.Client == (config.ClientConfig{}) { return cli.ErrClientSectionIsRequired } diff --git a/core/internal/client/klioclient/grpcclient/connection.go b/core/internal/client/klioclient/grpcclient/connection.go index 25064d94..f41d2938 100644 --- a/core/internal/client/klioclient/grpcclient/connection.go +++ b/core/internal/client/klioclient/grpcclient/connection.go @@ -46,18 +46,7 @@ type grpcWALStream struct { // Close implements common.WALStream. func (g *grpcWALStream) Close(_ context.Context) error { - result, err := g.innerStream.CloseAndRecv() - if err != nil { - return fmt.Errorf("while flushing WAL file: %w", err) - } - - if result.GetWrittenSize() != g.sentBytes { - return &IncompleteWALFileError{ - uploadedSize: result.GetWrittenSize(), - expectedSize: g.sentBytes, - } - } - + g.innerStream.CloseSend() return nil } diff --git a/core/internal/client/klioclient/grpcclient/errors.go b/core/internal/client/klioclient/grpcclient/errors.go index ae0fa9eb..9184c399 100644 --- a/core/internal/client/klioclient/grpcclient/errors.go +++ b/core/internal/client/klioclient/grpcclient/errors.go @@ -27,6 +27,10 @@ import ( // ErrInconsistentCertificate is raised when the server certificate cannot be parsed. var ErrInconsistentCertificate = errors.New("inconsistent server certificate (parsing)") +// ErrNoResultReceived is raised when the server closes the WAL upload stream +// without sending a result. +var ErrNoResultReceived = errors.New("server closed stream without sending a result") + // IncompleteWALFileError is raised when a WAL file has been uploaded incompletely. type IncompleteWALFileError struct { uploadedSize uint64 diff --git a/core/internal/client/klioclient/grpcclient/walclient.go b/core/internal/client/klioclient/grpcclient/walclient.go index 388280fc..b92309ee 100644 --- a/core/internal/client/klioclient/grpcclient/walclient.go +++ b/core/internal/client/klioclient/grpcclient/walclient.go @@ -38,8 +38,8 @@ func (c *Connection) StoreWAL(ctx context.Context, name string, content []byte, } walReader := bytes.NewBuffer(content) - buffer := make([]byte, 4096) + for { readBytes, readError := walReader.Read(buffer) if readError != nil && !errors.Is(readError, io.EOF) { @@ -56,23 +56,23 @@ func (c *Connection) StoreWAL(ctx context.Context, name string, content []byte, return fmt.Errorf("error while sending WAL block (sending via GRPC): %w", err) } + _, err := stream.Recv() + if errors.Is(err, io.EOF) { + return ErrNoResultReceived + } + if err != nil { + return fmt.Errorf("while flushing WAL file: %w", err) + } + if errors.Is(readError, io.EOF) { break } } - result, err := stream.CloseAndRecv() - if err != nil { + if err := stream.CloseSend(); err != nil { return fmt.Errorf("while flushing WAL file: %w", err) } - if result.GetWrittenSize() != uint64(len(content)) { - return &IncompleteWALFileError{ - uploadedSize: result.GetWrittenSize(), - expectedSize: uint64(len(content)), - } - } - return nil } diff --git a/core/internal/client/klioclient/grpcclient/walstreamer.go b/core/internal/client/klioclient/grpcclient/walstreamer.go index 327afb23..ee8680f3 100644 --- a/core/internal/client/klioclient/grpcclient/walstreamer.go +++ b/core/internal/client/klioclient/grpcclient/walstreamer.go @@ -40,19 +40,19 @@ func (c *Connection) StoreWALStreaming( name string, segmentSize uint64, sendToTier2 bool, -) (*klioclient.WALUploader, error) { +) (klioclient.WALUploaderImpl, error) { stream, err := c.Put(ctx) if err != nil { return nil, fmt.Errorf("while starting uploading a WAL file: %w", err) } - return klioclient.NewWALUploader(&grpcWALStream{ + return &grpcWALStream{ innerStream: stream, segmentSize: segmentSize, clusterName: c.clientConfig.ClusterName, walName: name, sendToTier2: sendToTier2, - }), nil + }, nil } // GetWALStreaming get a WAL from a remote connection. diff --git a/core/internal/client/klioclient/grpcclient/waluploader.go b/core/internal/client/klioclient/grpcclient/waluploader.go index 478b9bb7..0183a5d8 100644 --- a/core/internal/client/klioclient/grpcclient/waluploader.go +++ b/core/internal/client/klioclient/grpcclient/waluploader.go @@ -51,3 +51,14 @@ func (g *grpcWALStream) SendBlock(ctx context.Context, block []byte) error { return nil } + +// ReceiveFeedback receives the feedback from the server, that is telling us how +// many bytes it flushed. +func (g *grpcWALStream) ReceiveFeedback(ctx context.Context) (uint64, error) { + result, err := g.innerStream.Recv() + if err != nil { + return 0, err + } + + return result.GetWrittenSize(), nil +} diff --git a/core/internal/client/klioclient/interfaces.go b/core/internal/client/klioclient/interfaces.go index 6c4642f4..adaa7ca9 100644 --- a/core/internal/client/klioclient/interfaces.go +++ b/core/internal/client/klioclient/interfaces.go @@ -103,28 +103,10 @@ type WALUploaderImpl interface { // SendBlock sends a WAL Block SendBlock(ctx context.Context, block []byte) error + // ReceiveFeedback receives a feedback from the server that is telling + // us how many bytes it has flushed. + ReceiveFeedback(ctx context.Context) (uint64, error) + // Close closes the WAL streaming session Close(ctx context.Context) error } - -// WALUploader allows the user the upload a WAL file to a remote store, block by block. -type WALUploader struct { - impl WALUploaderImpl -} - -// NewWALUploader creates a WAL uploader given the underlying implementation. -func NewWALUploader(impl WALUploaderImpl) *WALUploader { - return &WALUploader{ - impl: impl, - } -} - -// SendBlock sends a WAL Block. -func (u *WALUploader) SendBlock(ctx context.Context, block []byte) error { - return u.impl.SendBlock(ctx, block) //nolint:wrapcheck -} - -// Close closes the WAL streaming session. -func (u *WALUploader) Close(ctx context.Context) error { - return u.impl.Close(ctx) //nolint:wrapcheck -} diff --git a/core/internal/client/sendwal/buffer/buffer.go b/core/internal/client/sendwal/buffer/buffer.go index 5be1513e..038ef341 100644 --- a/core/internal/client/sendwal/buffer/buffer.go +++ b/core/internal/client/sendwal/buffer/buffer.go @@ -20,7 +20,6 @@ SPDX-License-Identifier: Apache-2.0 package buffer import ( - "bytes" "context" "fmt" @@ -28,11 +27,6 @@ import ( "github.com/cloudnative-pg/machinery/pkg/types" ) -// maximumBufferSizeFactor allows configuring the higher limit of memory -// allocation of the WAL buffer. It is multiplied to the configured -// buffer size to get the limit. -const maximumBufferSizeFactor = 2 - // Data is the implementation of the WAL buffer. type Data struct { segmentSize uint64 @@ -40,23 +34,17 @@ type Data struct { handler Handler - writeLSN uint64 - flushLSN uint64 - buffer *bytes.Buffer - bufferSize int + writeLSN uint64 } // New creates a new WAL buffer. -func New(tli int, walSegmentSize uint64, handler Handler, bufferSize int) *Data { +func New(tli int, walSegmentSize uint64, handler Handler) *Data { result := &Data{ segmentSize: walSegmentSize, tli: tli, handler: handler, - bufferSize: bufferSize, } - result.buffer = result.newBuffer() - return result } @@ -138,25 +126,11 @@ func (wal *Data) ProcessWALData(ctx context.Context, data []byte, startWAL types return nil } -// FlushLSN gets the latest LSN that was flushed down to the Klio server. -func (wal *Data) FlushLSN() uint64 { - return wal.flushLSN -} - // WriteLSN gets the latest LSN that was written into the memory. func (wal *Data) WriteLSN() uint64 { return wal.writeLSN } -// Flush flushes the buffer to the Klio server connection. -func (wal *Data) Flush(ctx context.Context) error { - return wal.flushInternal(ctx) -} - -func (wal *Data) newBuffer() *bytes.Buffer { - return bytes.NewBuffer(make([]byte, 0, wal.bufferSize)) -} - func (wal *Data) openWALPos(ctx context.Context, blockpos uint64) error { contextLogger := log.FromContext(ctx) contextLogger.Info("Opening WAL file", "blockpos", types.Int64ToLSN(blockpos)) @@ -166,48 +140,26 @@ func (wal *Data) openWALPos(ctx context.Context, blockpos uint64) error { } wal.writeLSN = blockpos - wal.flushLSN = blockpos return nil } func (wal *Data) writeToWALFile(ctx context.Context, data []byte) error { - if _, err := wal.buffer.Write(data); err != nil { - return fmt.Errorf("while writing to buffer: %w", err) - } - - wal.writeLSN += uint64(len(data)) - - if wal.buffer.Len() >= wal.bufferSize { - return wal.Flush(ctx) - } - - return nil -} - -func (wal *Data) flushInternal(ctx context.Context) error { contextLogger := log.FromContext(ctx) - if wal.handler == nil || !wal.handler.HasWALFileOpened() || wal.buffer.Len() == 0 { + if wal.handler == nil || !wal.handler.HasWALFileOpened() || len(data) == 0 { return nil } contextLogger.Debug("Writing block", - "blockpos", types.Int64ToLSN(wal.writeLSN), "blocksize", wal.buffer.Len()) - _, err := wal.handler.Write(ctx, wal.buffer.Bytes()) - if err != nil { + "blockpos", types.Int64ToLSN(wal.writeLSN), + "blocksize", len(data), + ) + if err := wal.handler.Write(ctx, data); err != nil { return fmt.Errorf("while writing to WAL handler: %w", err) } - // Clear content but keeps the slice capacity - wal.buffer.Reset() - - // Prevent memory bloat in long-running processes. - if wal.buffer.Cap() > wal.bufferSize*maximumBufferSizeFactor { - wal.buffer = wal.newBuffer() - } - - wal.flushLSN = wal.writeLSN + wal.writeLSN += uint64(len(data)) return nil } @@ -216,10 +168,6 @@ func (wal *Data) closeCurrentWAL(ctx context.Context) error { contextLogger := log.FromContext(ctx) contextLogger.Debug("Closing WAL file") - if err := wal.Flush(ctx); err != nil { - return fmt.Errorf("while flushing WAL handler: %w", err) - } - if err := wal.handler.CloseWAL(ctx); err != nil { return fmt.Errorf("while closing current WAL file: %w", err) } diff --git a/core/internal/client/sendwal/buffer/grpc.go b/core/internal/client/sendwal/buffer/grpc.go index 6cf88658..b24bf57e 100644 --- a/core/internal/client/sendwal/buffer/grpc.go +++ b/core/internal/client/sendwal/buffer/grpc.go @@ -22,6 +22,7 @@ package buffer import ( "context" "fmt" + "sync" "github.com/cloudnative-pg/machinery/pkg/log" "github.com/cloudnative-pg/machinery/pkg/types" @@ -36,13 +37,22 @@ type KlioClientStreamingHandler struct { conn *grpcclient.Connection stream klioclient.WALUploaderImpl - offset uint64 + + walStartLSN uint64 + offset uint64 + serverOffset uint64 sendToTier2 bool tli int segmentSize uint64 currentWALFile string + + mu sync.RWMutex + // streamReady is closed whenever stream transitions from nil to + // non-nil, waking up any Write/Receive call blocked in waitForStream. + // It is replaced with a fresh channel whenever stream is cleared. + streamReady chan struct{} } // NewKlioClientHandler creates a new klio handler. @@ -58,17 +68,22 @@ func NewKlioClientHandler( segmentSize: segmentSize, stream: nil, sendToTier2: sendToTier2, + streamReady: make(chan struct{}), } } // OpenWAL implements the Handler interface. func (wal *KlioClientStreamingHandler) OpenWAL(ctx context.Context, blockpos uint64) error { + wal.mu.Lock() + defer wal.mu.Unlock() + currentWALFile, err := types.Int64ToLSN(blockpos).WALFileName(wal.tli, wal.segmentSize) if err != nil { return fmt.Errorf("while creating WAL file name (pos %v): %w", blockpos, err) } wal.offset = 0 + wal.walStartLSN = blockpos wal.currentWALFile = currentWALFile stream, err := wal.conn.StoreWALStreaming(ctx, wal.currentWALFile, wal.segmentSize, wal.sendToTier2) @@ -77,17 +92,24 @@ func (wal *KlioClientStreamingHandler) OpenWAL(ctx context.Context, blockpos uin } wal.stream = stream + close(wal.streamReady) return nil } // HasWALFileOpened implements the Handler interface. func (wal *KlioClientStreamingHandler) HasWALFileOpened() bool { + wal.mu.RLock() + defer wal.mu.RUnlock() + return wal.currentWALFile != "" } // CloseWAL implements the Handler interface. func (wal *KlioClientStreamingHandler) CloseWAL(ctx context.Context) error { + wal.mu.Lock() + defer wal.mu.Unlock() + contextLogger := log.FromContext(ctx) contextLogger.Debug("Closing WAL File", "walFileName", wal.currentWALFile) @@ -98,23 +120,85 @@ func (wal *KlioClientStreamingHandler) CloseWAL(ctx context.Context) error { wal.currentWALFile = "" wal.stream = nil + wal.streamReady = make(chan struct{}) return nil } // CurrentOffset implements the Handler interface. func (wal *KlioClientStreamingHandler) CurrentOffset() (uint64, error) { + wal.mu.RLock() + defer wal.mu.RUnlock() + return wal.offset, nil } // Write implements the Handler interface. -func (wal *KlioClientStreamingHandler) Write(ctx context.Context, block []byte) (int, error) { - err := wal.stream.SendBlock(ctx, block) +func (wal *KlioClientStreamingHandler) Write(ctx context.Context, block []byte) error { + stream, err := wal.waitForStream(ctx) if err != nil { - return 0, err //nolint:wrapcheck + return err + } + + wal.mu.RLock() + defer wal.mu.RUnlock() + + if err := stream.SendBlock(ctx, block); err != nil { + return err //nolint:wrapcheck } wal.offset += uint64(len(block)) - return len(block), nil + return nil +} + +// ReceiveFeedback implements the Handler interface. +func (wal *KlioClientStreamingHandler) Receive(ctx context.Context) error { + stream, err := wal.waitForStream(ctx) + if err != nil { + return err + } + + wal.mu.RLock() + defer wal.mu.RUnlock() + + serverOffset, err := stream.ReceiveFeedback(ctx) + if err != nil { + return err + } + + wal.serverOffset = serverOffset + + return nil +} + +func (wal *KlioClientStreamingHandler) LatestWrittenLSN() uint64 { + return wal.walStartLSN + wal.offset +} + +func (wal *KlioClientStreamingHandler) LatestConfirmedLSN() uint64 { + return wal.walStartLSN + wal.serverOffset +} + +// waitForStream blocks until a WAL stream has been opened, or the context +// is done. +// +//nolint:ireturn // stream is already stored as an interface field +func (wal *KlioClientStreamingHandler) waitForStream(ctx context.Context) (klioclient.WALUploaderImpl, error) { + for { + wal.mu.RLock() + stream := wal.stream + ready := wal.streamReady + wal.mu.RUnlock() + + if stream != nil { + return stream, nil + } + + select { + case <-ready: + case <-ctx.Done(): + return nil, ctx.Err() //nolint:wrapcheck + } + } } diff --git a/core/internal/client/sendwal/buffer/handler.go b/core/internal/client/sendwal/buffer/handler.go index d17859a4..6af15a1a 100644 --- a/core/internal/client/sendwal/buffer/handler.go +++ b/core/internal/client/sendwal/buffer/handler.go @@ -37,6 +37,13 @@ type Handler interface { // CurrentOffset returns the current offset in the WAL file CurrentOffset() (uint64, error) - // Write writes data in the current WAL file - Write(ctx context.Context, p []byte) (n int, err error) + // Write blocks until we sent the packet to the socket + Write(ctx context.Context, p []byte) error + + // Receive blocks until we have confirmation message from the server + Receive(ctx context.Context) error + + LatestWrittenLSN() uint64 + + LatestConfirmedLSN() uint64 } diff --git a/core/internal/client/sendwal/buffer/memory.go b/core/internal/client/sendwal/buffer/memory.go deleted file mode 100644 index 3021af04..00000000 --- a/core/internal/client/sendwal/buffer/memory.go +++ /dev/null @@ -1,101 +0,0 @@ -/* -Copyright © contributors to CloudNativePG, established as -CloudNativePG a Series of LF Projects, LLC. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. - -SPDX-License-Identifier: Apache-2.0 -*/ - -package buffer - -import ( - "bytes" - "context" - "fmt" - - "github.com/ccoveille/go-safecast/v2" - "github.com/cloudnative-pg/machinery/pkg/log" - "github.com/cloudnative-pg/machinery/pkg/types" -) - -// Flusher is the type of functions that are called -// to write a WAL file. -type Flusher func(walName string, data []byte) error - -// MemBufferHandler is the handler of WAL files that writes in a memory buffer -// and, when the WAL is completed, flushes it via a Flusher function. -type MemBufferHandler struct { - currentWALFile string - buffer bytes.Buffer - logger log.Logger - flusher Flusher - - tli int - segmentSize uint64 -} - -// NewMemBufferHandler creates a new memory buffer handler. -func NewMemBufferHandler(logger log.Logger, tli int, segmentSize uint64, flusher Flusher) *MemBufferHandler { - return &MemBufferHandler{ - currentWALFile: "", - buffer: *bytes.NewBuffer(make([]byte, 0, segmentSize)), - logger: logger, - flusher: flusher, - tli: tli, - segmentSize: segmentSize, - } -} - -// HasWALFileOpened implements the Handler interface. -func (wal *MemBufferHandler) HasWALFileOpened() bool { - return wal.currentWALFile != "" -} - -// OpenWAL implements the Handler interface. -func (wal *MemBufferHandler) OpenWAL(_ context.Context, blockpos uint64) error { - var err error - - wal.currentWALFile, err = types.Int64ToLSN(blockpos).WALFileName(wal.tli, wal.segmentSize) - if err != nil { - return fmt.Errorf("while creating WAL file name (pos %v): %w", blockpos, err) - } - wal.buffer.Reset() - - wal.logger.Debug("Opening WAL File", "walFileName", wal.currentWALFile) - - return nil -} - -// CloseWAL implements the Handler interface. -func (wal *MemBufferHandler) CloseWAL(_ context.Context) error { - wal.logger.Debug("Closing WAL File", "walFileName", wal.currentWALFile) - if err := wal.flusher(wal.currentWALFile, wal.buffer.Bytes()); err != nil { - return err - } - - wal.currentWALFile = "" - wal.buffer.Reset() - - return nil -} - -// CurrentOffset implements the Handler interface. -func (wal *MemBufferHandler) CurrentOffset() (uint64, error) { - return safecast.Convert[uint64](wal.buffer.Len()) -} - -// Write implements the Handler interface. -func (wal *MemBufferHandler) Write(_ context.Context, p []byte) (int, error) { - return wal.buffer.Write(p) //nolint:wrapcheck -} diff --git a/core/internal/client/sendwal/receiver.go b/core/internal/client/sendwal/receiver.go index 245d5a68..0e9d0fde 100644 --- a/core/internal/client/sendwal/receiver.go +++ b/core/internal/client/sendwal/receiver.go @@ -23,9 +23,10 @@ import ( "context" "errors" "fmt" + "io" "path" "strings" - "time" + "sync" "github.com/cloudnative-pg/cloudnative-pg/pkg/postgres" "github.com/cloudnative-pg/machinery/pkg/log" @@ -36,6 +37,7 @@ import ( "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/metric" "go.opentelemetry.io/otel/trace" + "k8s.io/apimachinery/pkg/util/wait" "github.com/cloudnative-pg/klio/core/internal/client/klioclient/grpcclient" "github.com/cloudnative-pg/klio/core/internal/client/sendwal/buffer" @@ -51,6 +53,8 @@ type Process struct { infrastructure *infrastructure.Postgres client *grpcclient.Connection sendToTier2 bool + + feedbackMux sync.Mutex } // New creates a new receiver. @@ -444,21 +448,36 @@ func (s *Process) startReplication( int(timeline), walSegmentSize, klioHandler, - s.config.Source.BufferSize, ) - copyDoneResult, err := s.manageWALStream(ctx, conn, walBuffer) - if err != nil { - return err - } + var copyDoneResult *pglogrepl.CopyDoneResult + var endOfPGConnectionError error + + var wg wait.Group + wg.StartWithContext(ctx, func(ctx context.Context) { + copyDoneResult, endOfPGConnectionError = s.fromPostgresToWALServer(ctx, conn, walBuffer, klioHandler) + if endOfPGConnectionError != nil { + return + } - if klioHandler.HasWALFileOpened() { - // If the transmission terminated but there is still a WAL file in progress, - // we close it. - // This happens when PG is shut down. - if err := klioHandler.CloseWAL(ctx); err != nil { - return fmt.Errorf("while closing the WAL file: %w", err) + if klioHandler.HasWALFileOpened() { + // If the transmission terminated but there is still a WAL file in progress, + // we close it. + // This happens when PG is shut down. + if err := klioHandler.CloseWAL(ctx); err != nil { + endOfPGConnectionError = fmt.Errorf("while closing the WAL file: %w", err) + } + } + }) + wg.StartWithContext(ctx, func(ctx context.Context) { + if err := s.fromWALServerToPostgres(ctx, conn, klioHandler); err != nil { + contextLogger.Error(err, "Error while receiving feedback from the Klio server") } + }) + wg.Wait() + + if endOfPGConnectionError != nil { + return endOfPGConnectionError } // Check if the timeline has changed and restart replication if needed @@ -486,67 +505,52 @@ func (s *Process) startReplication( return nil } -//nolint:gocognit,cyclop -func (s *Process) manageWALStream( +func (s *Process) fromWALServerToPostgres( ctx context.Context, conn *pgconn.PgConn, - buffer *buffer.Data, -) (*pglogrepl.CopyDoneResult, error) { + klioHandler buffer.Handler, +) error { contextLogger := log.FromContext(ctx) - flushDeadline := s.config.Source.FlushTimeout() - nextFlushDeadline := time.Now().Add(flushDeadline) - - feedbackDeadline := s.config.Source.StandbyMessageTimeout() - nextFeedbackDeadline := time.Now().Add(feedbackDeadline) - -loop: for { - if time.Now().After(nextFlushDeadline) { - flushedLSN := buffer.FlushLSN() - - if err := buffer.Flush(ctx); err != nil { - contextLogger.Error(err, "Failed flush WAL data") - return nil, fmt.Errorf("while flushing WAL data: %w", err) - } - - // When flush really written something down to the Klio server, - // the FlushedLSN will be different. In that case, we want to immediately - // give feedback to the PostgreSQL server. This ultimately - // will result in updated data in pg_stat_replication. - if flushedLSN != buffer.FlushLSN() { - nextFeedbackDeadline = time.Time{} + if err := klioHandler.Receive(ctx); err != nil { + if errors.Is(err, io.EOF) { + break } - nextFlushDeadline = time.Now().Add(flushDeadline) + return err } - if time.Now().After(nextFeedbackDeadline) { - // We communicate back to PostgreSQL the feedback when: - // - // 1. the feedback deadline exceeded - // 2. we received something from streaming replication - s.sendFeedback(ctx, conn, buffer) - nextFeedbackDeadline = time.Now().Add(feedbackDeadline) + if err := s.sendFeedback(ctx, conn, klioHandler); err != nil { + contextLogger.Error(err, "Error while sending feedback to PostgreSQL") } + } + + return nil +} - standbyMessageDeadlineContext, cancel := context.WithDeadline(ctx, nextFlushDeadline) - msg, err := conn.ReceiveMessage(standbyMessageDeadlineContext) - cancel() +//nolint:gocognit,cyclop +func (s *Process) fromPostgresToWALServer( + ctx context.Context, + conn *pgconn.PgConn, + buffer *buffer.Data, + handler buffer.Handler, +) (*pglogrepl.CopyDoneResult, error) { + contextLogger := log.FromContext(ctx) +loop: + for { + msg, err := conn.ReceiveMessage(ctx) if err != nil { - if pgconn.Timeout(err) { - continue - } if errors.Is(err, context.Canceled) { break } - contextLogger.Error(err, "receive message failed") + contextLogger.Error(err, "receive message failed") break } - log.FromContext(ctx).Trace( + contextLogger.Trace( "Received message", "msgType", fmt.Sprintf("%T", msg)) @@ -567,7 +571,12 @@ loop: ) if pkm.ReplyRequested { - s.sendFeedback(ctx, conn, buffer) + if err := s.sendFeedback(ctx, conn, handler); err != nil { + contextLogger.Error( + err, + "Error while sending requested feedback to PostgreSQL. We may be disconnected.", + ) + } } case pglogrepl.XLogDataByteID: @@ -584,10 +593,6 @@ loop: return nil, fmt.Errorf("could not process WAL data at %s: %w", xld.WALStart, err) } - // Force the code to communicate back to PostgreSQL the current status without waiting for - // a flush - nextFeedbackDeadline = time.Time{} - default: contextLogger.Info("Received unexpected copydata message", "msg", msg) return nil, NewUnexpectedCopydataMessageError(msg.Data) @@ -622,24 +627,19 @@ loop: return copyDoneResult, nil } -func (s *Process) sendFeedback(ctx context.Context, conn *pgconn.PgConn, buffer *buffer.Data) { - contextLogger := log.FromContext(ctx) +func (s *Process) sendFeedback( + ctx context.Context, + conn *pgconn.PgConn, + handler buffer.Handler, +) error { + s.feedbackMux.Lock() + defer s.feedbackMux.Unlock() - err := pglogrepl.SendStandbyStatusUpdate( - ctx, - conn, - pglogrepl.StandbyStatusUpdate{ - WALWritePosition: pglogrepl.LSN(buffer.WriteLSN()), - WALFlushPosition: pglogrepl.LSN(buffer.FlushLSN()), - WALApplyPosition: pglogrepl.LSN(buffer.FlushLSN()), - }, - ) - if err != nil { - contextLogger.Error(err, "Failed to send standby status update, skipping") - } else { - contextLogger.Debug( - "Sent Standby status message", - "write_lsn", types.Int64ToLSN(buffer.WriteLSN()), - "flush_lsn", types.Int64ToLSN(buffer.FlushLSN())) + msg := pglogrepl.StandbyStatusUpdate{ + WALWritePosition: pglogrepl.LSN(handler.LatestWrittenLSN()), + WALFlushPosition: pglogrepl.LSN(handler.LatestConfirmedLSN()), + WALApplyPosition: pglogrepl.LSN(handler.LatestConfirmedLSN()), } + + return pglogrepl.SendStandbyStatusUpdate(ctx, conn, msg) } diff --git a/core/internal/grpc/klio_wal.pb.go b/core/internal/grpc/klio_wal.pb.go index 207d602c..613c386b 100644 --- a/core/internal/grpc/klio_wal.pb.go +++ b/core/internal/grpc/klio_wal.pb.go @@ -934,9 +934,9 @@ const file_proto_klio_wal_proto_rawDesc = "" + "\x16tier2_retention_policy\x18\t \x01(\tR\x14tier2RetentionPolicy\"f\n" + "\x11CloseBackupResult\x12%\n" + "\x0etier2_schedule\x18\x01 \x01(\bR\rtier2Schedule\x12*\n" + - "\x11missing_wal_files\x18\x02 \x03(\tR\x0fmissingWalFiles2\xd8\x03\n" + - "\x03WAL\x12:\n" + - "\x03Put\x12\x17.klio.wal.v1.PutRequest\x1a\x16.klio.wal.v1.PutResult\"\x00(\x01\x12:\n" + + "\x11missing_wal_files\x18\x02 \x03(\tR\x0fmissingWalFiles2\xda\x03\n" + + "\x03WAL\x12<\n" + + "\x03Put\x12\x17.klio.wal.v1.PutRequest\x1a\x16.klio.wal.v1.PutResult\"\x00(\x010\x01\x12:\n" + "\x03Get\x12\x17.klio.wal.v1.GetRequest\x1a\x16.klio.wal.v1.GetResult\"\x000\x01\x12N\n" + "\vGetMetadata\x12\x1f.klio.wal.v1.GetMetadataRequest\x1a\x1c.klio.wal.v1.ClusterMetadata\"\x00\x12\\\n" + "\x0fRequestWALStart\x12#.klio.wal.v1.RequestWALStartRequest\x1a\".klio.wal.v1.RequestWALStartResult\"\x00\x12Y\n" + diff --git a/core/internal/grpc/klio_wal_grpc.pb.go b/core/internal/grpc/klio_wal_grpc.pb.go index c78a82b5..f8d46c9b 100644 --- a/core/internal/grpc/klio_wal_grpc.pb.go +++ b/core/internal/grpc/klio_wal_grpc.pb.go @@ -50,7 +50,7 @@ const ( // // For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream. type WALClient interface { - Put(ctx context.Context, opts ...grpc.CallOption) (grpc.ClientStreamingClient[PutRequest, PutResult], error) + Put(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[PutRequest, PutResult], error) Get(ctx context.Context, in *GetRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[GetResult], error) GetMetadata(ctx context.Context, in *GetMetadataRequest, opts ...grpc.CallOption) (*ClusterMetadata, error) RequestWALStart(ctx context.Context, in *RequestWALStartRequest, opts ...grpc.CallOption) (*RequestWALStartResult, error) @@ -66,7 +66,7 @@ func NewWALClient(cc grpc.ClientConnInterface) WALClient { return &wALClient{cc} } -func (c *wALClient) Put(ctx context.Context, opts ...grpc.CallOption) (grpc.ClientStreamingClient[PutRequest, PutResult], error) { +func (c *wALClient) Put(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[PutRequest, PutResult], error) { cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) stream, err := c.cc.NewStream(ctx, &WAL_ServiceDesc.Streams[0], WAL_Put_FullMethodName, cOpts...) if err != nil { @@ -77,7 +77,7 @@ func (c *wALClient) Put(ctx context.Context, opts ...grpc.CallOption) (grpc.Clie } // This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. -type WAL_PutClient = grpc.ClientStreamingClient[PutRequest, PutResult] +type WAL_PutClient = grpc.BidiStreamingClient[PutRequest, PutResult] func (c *wALClient) Get(ctx context.Context, in *GetRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[GetResult], error) { cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) @@ -142,7 +142,7 @@ func (c *wALClient) CloseBackup(ctx context.Context, in *CloseBackupRequest, opt // All implementations must embed UnimplementedWALServer // for forward compatibility. type WALServer interface { - Put(grpc.ClientStreamingServer[PutRequest, PutResult]) error + Put(grpc.BidiStreamingServer[PutRequest, PutResult]) error Get(*GetRequest, grpc.ServerStreamingServer[GetResult]) error GetMetadata(context.Context, *GetMetadataRequest) (*ClusterMetadata, error) RequestWALStart(context.Context, *RequestWALStartRequest) (*RequestWALStartResult, error) @@ -158,7 +158,7 @@ type WALServer interface { // pointer dereference when methods are called. type UnimplementedWALServer struct{} -func (UnimplementedWALServer) Put(grpc.ClientStreamingServer[PutRequest, PutResult]) error { +func (UnimplementedWALServer) Put(grpc.BidiStreamingServer[PutRequest, PutResult]) error { return status.Error(codes.Unimplemented, "method Put not implemented") } func (UnimplementedWALServer) Get(*GetRequest, grpc.ServerStreamingServer[GetResult]) error { @@ -202,7 +202,7 @@ func _WAL_Put_Handler(srv interface{}, stream grpc.ServerStream) error { } // This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. -type WAL_PutServer = grpc.ClientStreamingServer[PutRequest, PutResult] +type WAL_PutServer = grpc.BidiStreamingServer[PutRequest, PutResult] func _WAL_Get_Handler(srv interface{}, stream grpc.ServerStream) error { m := new(GetRequest) @@ -315,6 +315,7 @@ var WAL_ServiceDesc = grpc.ServiceDesc{ { StreamName: "Put", Handler: _WAL_Put_Handler, + ServerStreams: true, ClientStreams: true, }, { diff --git a/core/internal/server/walserver/upload.go b/core/internal/server/walserver/upload.go index bb37dac5..ec8cdaf6 100644 --- a/core/internal/server/walserver/upload.go +++ b/core/internal/server/walserver/upload.go @@ -222,6 +222,12 @@ func (h *putHandler) processBlock(ctx context.Context, request *grpc.PutRequest) h.writtenSize += uint64(len(request.GetWalBlock())) h.recordLatestWrittenLSN(ctx) + if err := h.req.Send(&grpc.PutResult{ + WrittenSize: h.writtenSize, + }); err != nil { + h.logger.Error(err, "Error while sending feedback to the client, skipping") + } + return nil } @@ -417,7 +423,7 @@ func (h *putHandler) finalize(ctx context.Context) error { return err } - if err := h.req.SendAndClose(&grpc.PutResult{ + if err := h.req.Send(&grpc.PutResult{ WrittenSize: h.writtenSize, }); err != nil { h.logger.Warning( @@ -436,7 +442,7 @@ func (h *putHandler) finalize(ctx context.Context) error { // closeEmpty reports an empty result when no WAL block was ever received. func (h *putHandler) closeEmpty() error { - if err := h.req.SendAndClose(&grpc.PutResult{ + if err := h.req.Send(&grpc.PutResult{ WrittenSize: 0, }); err != nil { h.logger.Error(err, "Error while closing empty WAL file") diff --git a/core/pkg/config/client.go b/core/pkg/config/client.go index a6726d59..c59e3493 100644 --- a/core/pkg/config/client.go +++ b/core/pkg/config/client.go @@ -19,8 +19,6 @@ SPDX-License-Identifier: Apache-2.0 package config -import "time" - // Data is the configuration. // // This struct is used to generate a secret in the Kubernetes cluster, so its serialization must be stable. @@ -60,13 +58,6 @@ type Data struct { WALPrefetch WALPrefetchConfig `json:"wal_prefetch" mapstructure:"wal_prefetch"` } -// SetDefaults sets the default values of the configuration. -func (d *Data) SetDefaults() { - if d.Source != (SourceConfig{}) { - d.Source.SetDefaults() - } -} - // SourceConfig is the configuration of the WAL receiver. type SourceConfig struct { // DSN is the database service we should get the WALs from @@ -78,18 +69,6 @@ type SourceConfig struct { // Slot is the name of the replication slot to be used Slot string `json:"slot" mapstructure:"slot"` - - // StandbyMessageTimeoutSeconds is the timeout after which the WAL - // receiver will send a status update - StandbyMessageTimeoutSeconds int `json:"standby_message_timeout_seconds" mapstructure:"standby_message_timeout_seconds"` //nolint:lll - - // FlushTimeoutMilliseconds is the timeout in milliseconds after which buffered - // WAL data is automatically flushed to the Klio server - FlushTimeoutMilliseconds int `json:"flush_timeout_ms" mapstructure:"flush_timeout_ms"` - - // BufferSize is the maximum size in bytes of the in-memory WAL buffer before - // triggering an automatic flush - BufferSize int `json:"buffer_size" mapstructure:"buffer_size"` } // ClientConfig is the configuration of the Klio client. @@ -152,22 +131,3 @@ type WALPrefetchConfig struct { // MaxConcurrentDownloads is the maximum number of concurrent WAL downloads. MaxConcurrentDownloads int `json:"max_concurrent_downloads" mapstructure:"max_concurrent_downloads"` } - -// SetDefaults sets the default values of the configuration. -func (s *SourceConfig) SetDefaults() { - s.StandbyMessageTimeoutSeconds = 10 - s.FlushTimeoutMilliseconds = 200 - s.BufferSize = 2 * 1024 * 1024 // 2 MB -} - -// StandbyMessageTimeout returns the stanby message timeout in a -// time.Duration. -func (s *SourceConfig) StandbyMessageTimeout() time.Duration { - return time.Second * time.Duration(s.StandbyMessageTimeoutSeconds) -} - -// FlushTimeout returns the timeout after which the WALs are -// flushed. -func (s *SourceConfig) FlushTimeout() time.Duration { - return time.Millisecond * time.Duration(s.FlushTimeoutMilliseconds) -} diff --git a/core/pkg/config/client_validate.go b/core/pkg/config/client_validate.go index 1690d9e1..23b7bd3f 100644 --- a/core/pkg/config/client_validate.go +++ b/core/pkg/config/client_validate.go @@ -58,15 +58,6 @@ func (s *SourceConfig) Validate() error { errs = errors.Join(errs, errors.New( "invalid source config: slot name can only contain lower-case letters, numbers, and underscores")) } - if s.StandbyMessageTimeoutSeconds < 1 { - errs = errors.Join(errs, errors.New("invalid source config: standby_message_timeout_seconds must be at least 1")) - } - if s.FlushTimeoutMilliseconds < 1 { - errs = errors.Join(errs, errors.New("invalid source config: flush_timeout_ms must be at least 1")) - } - if s.BufferSize < 1 { - errs = errors.Join(errs, errors.New("invalid source config: buffer_size must be at least 1")) - } return errs } diff --git a/core/pkg/config/client_validate_test.go b/core/pkg/config/client_validate_test.go index 50652a21..dbdaa21a 100644 --- a/core/pkg/config/client_validate_test.go +++ b/core/pkg/config/client_validate_test.go @@ -34,12 +34,9 @@ func TestSourceConfigValidate(t *testing.T) { { name: "Valid config", config: SourceConfig{ - DSN: "postgres://...", - StandardDSN: "postgres://...", - Slot: "my_slot_123", - StandbyMessageTimeoutSeconds: 10, - FlushTimeoutMilliseconds: 100, - BufferSize: 1024, + DSN: "postgres://...", + StandardDSN: "postgres://...", + Slot: "my_slot_123", }, wantErr: false, }, @@ -61,45 +58,6 @@ func TestSourceConfigValidate(t *testing.T) { wantErr: true, substr: "slot name can only contain lower-case letters", }, - { - name: "Standby Message Timeout too low", - config: SourceConfig{ - DSN: "valid", - StandardDSN: "valid", - Slot: "valid", - StandbyMessageTimeoutSeconds: 0, - FlushTimeoutMilliseconds: 100, - BufferSize: 1024, - }, - wantErr: true, - substr: "must be at least 1", - }, - { - name: "Flush Timeout too low", - config: SourceConfig{ - DSN: "valid", - StandardDSN: "valid", - Slot: "valid", - StandbyMessageTimeoutSeconds: 100, - FlushTimeoutMilliseconds: 0, - BufferSize: 1024, - }, - wantErr: true, - substr: "must be at least 1", - }, - { - name: "Buffer Size too low", - config: SourceConfig{ - DSN: "valid", - StandardDSN: "valid", - Slot: "valid", - StandbyMessageTimeoutSeconds: 100, - FlushTimeoutMilliseconds: 200, - BufferSize: 0, - }, - wantErr: true, - substr: "must be at least 1", - }, } for _, tt := range tests { diff --git a/core/pkg/config/decode_test.go b/core/pkg/config/decode_test.go index d1339b12..aca2bf3d 100644 --- a/core/pkg/config/decode_test.go +++ b/core/pkg/config/decode_test.go @@ -78,18 +78,12 @@ source: dsn: "postgres://localhost:5432/mydb" standard_dsn: "postgres://localhost:5432/mydb" slot: my_slot - standby_message_timeout_seconds: 15 - flush_timeout_ms: 300 - buffer_size: 4096 `, want: Data{ Source: SourceConfig{ - DSN: "postgres://localhost:5432/mydb", - StandardDSN: "postgres://localhost:5432/mydb", - Slot: "my_slot", - StandbyMessageTimeoutSeconds: 15, - FlushTimeoutMilliseconds: 300, - BufferSize: 4096, + DSN: "postgres://localhost:5432/mydb", + StandardDSN: "postgres://localhost:5432/mydb", + Slot: "my_slot", }, }, }, diff --git a/core/proto/klio_wal.proto b/core/proto/klio_wal.proto index afbc3c15..0d2baf52 100644 --- a/core/proto/klio_wal.proto +++ b/core/proto/klio_wal.proto @@ -25,7 +25,7 @@ import "google/protobuf/timestamp.proto"; option go_package = "github.com/cloudnative-pg/klio/core/internal/grpc"; service WAL { - rpc Put(stream PutRequest) returns (PutResult) {} + rpc Put(stream PutRequest) returns (stream PutResult) {} rpc Get(GetRequest) returns (stream GetResult) {} rpc GetMetadata(GetMetadataRequest) returns (ClusterMetadata) {} diff --git a/documentation/web/docs/developer/_protocol.md b/documentation/web/docs/developer/_protocol.md index 4a1665e2..2547a0c9 100644 --- a/documentation/web/docs/developer/_protocol.md +++ b/documentation/web/docs/developer/_protocol.md @@ -539,7 +539,7 @@ feature. | Method Name | Request Type | Response Type | Description | | ----------- | ------------ | ------------- | ------------| -| Put | [PutRequest](#klio-wal-v1-PutRequest) stream | [PutResult](#klio-wal-v1-PutResult) | | +| Put | [PutRequest](#klio-wal-v1-PutRequest) stream | [PutResult](#klio-wal-v1-PutResult) stream | | | Get | [GetRequest](#klio-wal-v1-GetRequest) | [GetResult](#klio-wal-v1-GetResult) stream | | | GetMetadata | [GetMetadataRequest](#klio-wal-v1-GetMetadataRequest) | [ClusterMetadata](#klio-wal-v1-ClusterMetadata) | | | RequestWALStart | [RequestWALStartRequest](#klio-wal-v1-RequestWALStartRequest) | [RequestWALStartResult](#klio-wal-v1-RequestWALStartResult) | | diff --git a/operator/internal/klioconfig/config.go b/operator/internal/klioconfig/config.go index 89b43d0d..0823cb53 100644 --- a/operator/internal/klioconfig/config.go +++ b/operator/internal/klioconfig/config.go @@ -103,10 +103,6 @@ func GenerateConfig( DSN: "user=postgres replication=yes application_name=klio", StandardDSN: "user=postgres application_name=klio", Slot: "klio", - // The following parameters are not used by the plugin, but here with their default for completeness - StandbyMessageTimeoutSeconds: 0, - FlushTimeoutMilliseconds: 0, - BufferSize: 0, }, Client: config.ClientConfig{ ClusterName: spec.ClusterName, diff --git a/operator/pkg/config/client.go b/operator/pkg/config/client.go index a6726d59..c59e3493 100644 --- a/operator/pkg/config/client.go +++ b/operator/pkg/config/client.go @@ -19,8 +19,6 @@ SPDX-License-Identifier: Apache-2.0 package config -import "time" - // Data is the configuration. // // This struct is used to generate a secret in the Kubernetes cluster, so its serialization must be stable. @@ -60,13 +58,6 @@ type Data struct { WALPrefetch WALPrefetchConfig `json:"wal_prefetch" mapstructure:"wal_prefetch"` } -// SetDefaults sets the default values of the configuration. -func (d *Data) SetDefaults() { - if d.Source != (SourceConfig{}) { - d.Source.SetDefaults() - } -} - // SourceConfig is the configuration of the WAL receiver. type SourceConfig struct { // DSN is the database service we should get the WALs from @@ -78,18 +69,6 @@ type SourceConfig struct { // Slot is the name of the replication slot to be used Slot string `json:"slot" mapstructure:"slot"` - - // StandbyMessageTimeoutSeconds is the timeout after which the WAL - // receiver will send a status update - StandbyMessageTimeoutSeconds int `json:"standby_message_timeout_seconds" mapstructure:"standby_message_timeout_seconds"` //nolint:lll - - // FlushTimeoutMilliseconds is the timeout in milliseconds after which buffered - // WAL data is automatically flushed to the Klio server - FlushTimeoutMilliseconds int `json:"flush_timeout_ms" mapstructure:"flush_timeout_ms"` - - // BufferSize is the maximum size in bytes of the in-memory WAL buffer before - // triggering an automatic flush - BufferSize int `json:"buffer_size" mapstructure:"buffer_size"` } // ClientConfig is the configuration of the Klio client. @@ -152,22 +131,3 @@ type WALPrefetchConfig struct { // MaxConcurrentDownloads is the maximum number of concurrent WAL downloads. MaxConcurrentDownloads int `json:"max_concurrent_downloads" mapstructure:"max_concurrent_downloads"` } - -// SetDefaults sets the default values of the configuration. -func (s *SourceConfig) SetDefaults() { - s.StandbyMessageTimeoutSeconds = 10 - s.FlushTimeoutMilliseconds = 200 - s.BufferSize = 2 * 1024 * 1024 // 2 MB -} - -// StandbyMessageTimeout returns the stanby message timeout in a -// time.Duration. -func (s *SourceConfig) StandbyMessageTimeout() time.Duration { - return time.Second * time.Duration(s.StandbyMessageTimeoutSeconds) -} - -// FlushTimeout returns the timeout after which the WALs are -// flushed. -func (s *SourceConfig) FlushTimeout() time.Duration { - return time.Millisecond * time.Duration(s.FlushTimeoutMilliseconds) -} From a855882dfce8c479c59935f62389e2dba8602ccb Mon Sep 17 00:00:00 2001 From: Leonardo Cecchi Date: Thu, 27 Aug 2026 11:53:28 +0200 Subject: [PATCH 2/2] wip --- core/internal/client/sendwal/buffer/buffer.go | 15 ++++- core/internal/client/sendwal/buffer/grpc.go | 53 ++++++++++++++-- .../internal/client/sendwal/buffer/handler.go | 6 +- .../client/sendwal/nonblocking_receive.go | 62 +++++++++++++++++++ core/internal/client/sendwal/receiver.go | 17 ++++- 5 files changed, 146 insertions(+), 7 deletions(-) create mode 100644 core/internal/client/sendwal/nonblocking_receive.go diff --git a/core/internal/client/sendwal/buffer/buffer.go b/core/internal/client/sendwal/buffer/buffer.go index 038ef341..0e3e2750 100644 --- a/core/internal/client/sendwal/buffer/buffer.go +++ b/core/internal/client/sendwal/buffer/buffer.go @@ -35,14 +35,21 @@ type Data struct { handler Handler writeLSN uint64 + + // onWALOpened, if set, is invoked right after a new WAL file has been + // successfully opened, so the caller can report progress upstream + // (e.g. sending PostgreSQL a status update at WAL file boundaries, + // like src/bin/pg_basebackup/receivelog.c does). + onWALOpened func(ctx context.Context) error } // New creates a new WAL buffer. -func New(tli int, walSegmentSize uint64, handler Handler) *Data { +func New(tli int, walSegmentSize uint64, handler Handler, onWALOpened func(ctx context.Context) error) *Data { result := &Data{ segmentSize: walSegmentSize, tli: tli, handler: handler, + onWALOpened: onWALOpened, } return result @@ -141,6 +148,12 @@ func (wal *Data) openWALPos(ctx context.Context, blockpos uint64) error { wal.writeLSN = blockpos + if wal.onWALOpened != nil { + if err := wal.onWALOpened(ctx); err != nil { + return fmt.Errorf("while reporting progress after opening WAL file: %w", err) + } + } + return nil } diff --git a/core/internal/client/sendwal/buffer/grpc.go b/core/internal/client/sendwal/buffer/grpc.go index b24bf57e..4b76474c 100644 --- a/core/internal/client/sendwal/buffer/grpc.go +++ b/core/internal/client/sendwal/buffer/grpc.go @@ -20,6 +20,7 @@ SPDX-License-Identifier: Apache-2.0 package buffer import ( + "bytes" "context" "fmt" "sync" @@ -31,6 +32,12 @@ import ( "github.com/cloudnative-pg/klio/core/internal/client/klioclient/grpcclient" ) +// maxPendingBytes caps how much WAL data Write buffers before forcing a +// Flush, so a client catching up on a large backlog cannot grow the +// buffer unboundedly while PostgreSQL keeps streaming without ever +// blocking the reader. +const maxPendingBytes = 4 * 1024 * 1024 // 4 MiB + // KlioClientStreamingHandler is a handler that streams directly to a // Klio server. type KlioClientStreamingHandler struct { @@ -42,6 +49,10 @@ type KlioClientStreamingHandler struct { offset uint64 serverOffset uint64 + // pending buffers WAL bytes written by Write that have not yet been + // sent to the server. It is drained by Flush. + pending bytes.Buffer + sendToTier2 bool tli int @@ -114,6 +125,12 @@ func (wal *KlioClientStreamingHandler) CloseWAL(ctx context.Context) error { contextLogger.Debug("Closing WAL File", "walFileName", wal.currentWALFile) + // Flush directly instead of calling Flush, which would try to take + // wal.mu itself and deadlock: we already hold the exclusive Lock here. + if err := wal.flushLocked(ctx); err != nil { + return err + } + if err := wal.stream.Close(ctx); err != nil { return err //nolint:wrapcheck } @@ -135,19 +152,47 @@ func (wal *KlioClientStreamingHandler) CurrentOffset() (uint64, error) { // Write implements the Handler interface. func (wal *KlioClientStreamingHandler) Write(ctx context.Context, block []byte) error { - stream, err := wal.waitForStream(ctx) - if err != nil { + if _, err := wal.waitForStream(ctx); err != nil { return err } wal.mu.RLock() defer wal.mu.RUnlock() - if err := stream.SendBlock(ctx, block); err != nil { + wal.pending.Write(block) + wal.offset += uint64(len(block)) + + if wal.pending.Len() >= maxPendingBytes { + return wal.flushLocked(ctx) + } + + return nil +} + +// Flush implements the Handler interface. +func (wal *KlioClientStreamingHandler) Flush(ctx context.Context) error { + // Unlike Write and Receive, it does not wait for a stream to be open: + // if nothing has been buffered yet, there is nothing to flush, and a + // stream may not exist yet (e.g. before the first Write of a WAL file). + wal.mu.RLock() + defer wal.mu.RUnlock() + + return wal.flushLocked(ctx) +} + +// flushLocked sends any buffered data to the server. The caller must +// already hold wal.mu, either via the RLock taken by Write/Flush or the +// exclusive Lock held by CloseWAL. +func (wal *KlioClientStreamingHandler) flushLocked(ctx context.Context) error { + if wal.pending.Len() == 0 { + return nil + } + + if err := wal.stream.SendBlock(ctx, wal.pending.Bytes()); err != nil { return err //nolint:wrapcheck } - wal.offset += uint64(len(block)) + wal.pending.Reset() return nil } diff --git a/core/internal/client/sendwal/buffer/handler.go b/core/internal/client/sendwal/buffer/handler.go index 6af15a1a..7b17f397 100644 --- a/core/internal/client/sendwal/buffer/handler.go +++ b/core/internal/client/sendwal/buffer/handler.go @@ -37,9 +37,13 @@ type Handler interface { // CurrentOffset returns the current offset in the WAL file CurrentOffset() (uint64, error) - // Write blocks until we sent the packet to the socket + // Write buffers p to be sent to the server. It does not guarantee that + // p has reached the socket: call Flush for that. Write(ctx context.Context, p []byte) error + // Flush sends any data buffered by Write to the server. + Flush(ctx context.Context) error + // Receive blocks until we have confirmation message from the server Receive(ctx context.Context) error diff --git a/core/internal/client/sendwal/nonblocking_receive.go b/core/internal/client/sendwal/nonblocking_receive.go new file mode 100644 index 00000000..5064e288 --- /dev/null +++ b/core/internal/client/sendwal/nonblocking_receive.go @@ -0,0 +1,62 @@ +/* +Copyright © contributors to CloudNativePG, established as +CloudNativePG a Series of LF Projects, LLC. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +SPDX-License-Identifier: Apache-2.0 +*/ + +package sendwal + +import ( + "context" + "errors" + "fmt" + "net" + "time" + + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/pgproto3" +) + +// tryReceiveMessage attempts to receive the next message from PostgreSQL +// without blocking. It returns (nil, nil) if no message is currently +// available on the wire, rather than waiting for one to arrive. +func tryReceiveMessage(conn *pgconn.PgConn) (pgproto3.BackendMessage, error) { + if err := conn.Conn().SetReadDeadline(time.Now()); err != nil { + return nil, fmt.Errorf("failed to set read deadline: %w", err) + } + defer func() { + _ = conn.Conn().SetReadDeadline(time.Time{}) + }() + + // The peek is implemented with a deadline set to the current time on the + // underlying connection: this is safe because a deliberate read timeout + // does not poison a *pgconn.PgConn (unlike other read errors, it does not + // trigger the connection's async-close path). context.Background() is + // used for the underlying ReceiveMessage call so that pgconn's own + // context-cancellation watcher, which also manipulates the connection's + // read deadline, cannot race with the deadline set here. + msg, err := conn.ReceiveMessage(context.Background()) + if err != nil { + var netErr net.Error + if errors.As(err, &netErr) && netErr.Timeout() { + return nil, nil + } + + return nil, err + } + + return msg, nil +} diff --git a/core/internal/client/sendwal/receiver.go b/core/internal/client/sendwal/receiver.go index 0e9d0fde..b766845c 100644 --- a/core/internal/client/sendwal/receiver.go +++ b/core/internal/client/sendwal/receiver.go @@ -448,6 +448,9 @@ func (s *Process) startReplication( int(timeline), walSegmentSize, klioHandler, + func(ctx context.Context) error { + return s.sendFeedback(ctx, conn, klioHandler) + }, ) var copyDoneResult *pglogrepl.CopyDoneResult @@ -540,7 +543,19 @@ func (s *Process) fromPostgresToWALServer( loop: for { - msg, err := conn.ReceiveMessage(ctx) + // Check whether PostgreSQL already has a message waiting for us + // without blocking. If not, flush any WAL data buffered so far + // before waiting for the next one: this batches consecutive WAL + // blocks into fewer, larger writes to the WAL server instead of + // sending (and flushing) one per block. + msg, err := tryReceiveMessage(conn) + if err == nil && msg == nil { + if flushErr := handler.Flush(ctx); flushErr != nil { + contextLogger.Error(flushErr, "Error while flushing buffered WAL data") + } + + msg, err = conn.ReceiveMessage(ctx) + } if err != nil { if errors.Is(err, context.Canceled) { break