feat(wal): stream real per-block server acks for faster PostgreSQL feedback

core/internal/client/klioclient/grpcclient/common_test.go

@@ -29,12 +29,6 @@ import (
 
 const fakeWalContent = "deadbeef"
 
-// WALClient is the interface that wraps the backend WAL storage.
-type WALClient interface {
-	// StoreWAL upload a WAL file to a remote store
-	StoreWAL(ctx context.Context, name string, content []byte, sendToTier2 bool) error
-}
-
 type testingRepository struct {
 	prefilledSnapshots int
 	client             *TemporaryConnection
@@ -111,10 +105,10 @@ func runSnapshotLookupBenchmark(b *testing.B, repo *testingRepository) {
 	}
 }
 
-func addFakeWals(ctx context.Context, repo WALClient, start int, count int) error {
+func addFakeWals(ctx context.Context, repo *TemporaryConnection, start int, count int) error {
 	for i := range count {
 		walName := fmt.Sprintf("%024X", start+i)
-		err := repo.StoreWAL(ctx, walName, []byte(fakeWalContent), false)
+		err := repo.UploadFile(ctx, walName, []byte(fakeWalContent), false)
 		if err != nil {
 			return fmt.Errorf("while generating fake wal: %w", err)
 		}

core/internal/client/klioclient/grpcclient/connection.go

@@ -30,43 +30,48 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
 
+	"github.com/cloudnative-pg/klio/core/internal/client/klioclient"
 	klioGRPC "github.com/cloudnative-pg/klio/core/internal/grpc"
 	"github.com/cloudnative-pg/klio/core/internal/wal"
 	"github.com/cloudnative-pg/klio/core/pkg/config"
 )
 
-type grpcWALStream struct {
-	innerStream klioGRPC.WAL_PutClient
-	segmentSize uint64
-	clusterName string
-	sentBytes   uint64
-	walName     string
-	sendToTier2 bool
+// Connection represents a connection to a Klio server.
+type Connection struct {
+	klioGRPC.WALClient
+
+	clientConfig   *config.ClientConfig
+	grpcConnection *grpc.ClientConn
 }
 
-// Close implements common.WALStream.
-func (g *grpcWALStream) Close(_ context.Context) error {
-	result, err := g.innerStream.CloseAndRecv()
+// StoreWALStreaming implements the WAL streaming service.
+//
+//nolint:ireturn
+func (c *Connection) StoreWALStreaming(
+	ctx context.Context,
+	name string,
+	segmentSize uint64,
+	sendToTier2 bool,
+	walStartLSN uint64,
+	feedbackChannel chan<- wal.Feedback,
+) (klioclient.WALUploader, error) {
+	stream, err := c.Put(ctx)
 	if err != nil {
-		return fmt.Errorf("while flushing WAL file: %w", err)
+		return nil, fmt.Errorf("while starting uploading a WAL file: %w", err)
 	}
 
-	if result.GetWrittenSize() != g.sentBytes {
-		return &IncompleteWALFileError{
-			uploadedSize: result.GetWrittenSize(),
-			expectedSize: g.sentBytes,
-		}
+	g := &grpcWALStream{
+		innerStream:     stream,
+		segmentSize:     segmentSize,
+		clusterName:     c.clientConfig.ClusterName,
+		walName:         name,
+		sendToTier2:     sendToTier2,
+		walStartLSN:     walStartLSN,
+		feedbackChannel: feedbackChannel,
 	}
+	g.startFeedbackReader()
 
-	return nil
-}
-
-// Connection represents a connection to a Klio server.
-type Connection struct {
-	klioGRPC.WALClient
-
-	clientConfig   *config.ClientConfig
-	grpcConnection *grpc.ClientConn
+	return g, nil
 }
 
 // Connect opens a connection to a Klio server.

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

core/internal/client/klioclient/grpcclient/walclient.go

@@ -29,21 +29,21 @@ import (
 	klioGRPC "github.com/cloudnative-pg/klio/core/internal/grpc"
 )
 
-// StoreWAL uploads a WAL in the WAL server
-// Important: this function uploads a full WAL file.
-func (c *Connection) StoreWAL(ctx context.Context, name string, content []byte, sendToTier2 bool) error {
+// UploadFile uploads a file to the Klio server.
+// No feedback is expected from the server.
+func (c *Connection) UploadFile(ctx context.Context, name string, content []byte, sendToTier2 bool) error {
 	stream, err := c.Put(ctx)
 	if err != nil {
-		return fmt.Errorf("while starting uploading a WAL file: %w", err)
+		return fmt.Errorf("while starting uploading a file: %w", err)
 	}
 
 	walReader := bytes.NewBuffer(content)
-
 	buffer := make([]byte, 4096)
+
 	for {
 		readBytes, readError := walReader.Read(buffer)
 		if readError != nil && !errors.Is(readError, io.EOF) {
-			return fmt.Errorf("error while reading WAL (reading from buffer): %w", readError)
+			return fmt.Errorf("error while reading file (reading from buffer): %w", readError)
 		}
 
 		if err := stream.Send(&klioGRPC.PutRequest{
@@ -52,31 +52,30 @@ func (c *Connection) StoreWAL(ctx context.Context, name string, content []byte,
 			SegmentSize: uint64(len(content)),
 			WalBlock:    buffer[:readBytes],
 			SendToTier2: sendToTier2,
+
+			// We're not interested to the feedback from the Klio server,
+			// so we just start this WAL file from LSN zero.
+			WalStartLsn: 0,
 		}); err != nil {
-			return fmt.Errorf("error while sending WAL block (sending via GRPC): %w", err)
+			return fmt.Errorf("error while sending file 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
 }
-
-// StoreHistoryFile uses the underlying GRPC connection to store a history file.
-func (c *Connection) StoreHistoryFile(ctx context.Context, name string, content []byte, sendToTier2 bool) error {
-	return c.StoreWAL(ctx, name, content, sendToTier2)
-}

core/internal/client/klioclient/grpcclient/walstreamer.go → core/internal/client/klioclient/grpcclient/waldownloader.go

@@ -34,27 +34,6 @@ import (
 	klioGRPC "github.com/cloudnative-pg/klio/core/internal/grpc"
 )
 
-// StoreWALStreaming implements the WAL streaming service.
-func (c *Connection) StoreWALStreaming(
-	ctx context.Context,
-	name string,
-	segmentSize uint64,
-	sendToTier2 bool,
-) (*klioclient.WALUploader, 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{
-		innerStream: stream,
-		segmentSize: segmentSize,
-		clusterName: c.clientConfig.ClusterName,
-		walName:     name,
-		sendToTier2: sendToTier2,
-	}), nil
-}
-
 // GetWALStreaming get a WAL from a remote connection.
 func (c *Connection) GetWALStreaming(ctx context.Context, walName string, out io.Writer) error { //nolint:cyclop
 	client, err := c.Get(ctx, &klioGRPC.GetRequest{

core/internal/client/klioclient/grpcclient/waluploader.go

@@ -21,25 +21,67 @@ package grpcclient
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"io"
 	"time"
 
 	klioGRPC "github.com/cloudnative-pg/klio/core/internal/grpc"
 	"github.com/cloudnative-pg/klio/core/internal/opentelemetry"
+	"github.com/cloudnative-pg/klio/core/internal/wal"
 )
 
-// SendBlock implements common.WALUploaderImpl.
+// grpcWALStream uploads a WAL file to a Klio server over a single
+// bidirectional gRPC stream.
+//
+// A background goroutine reads the acknowledgments and streams them
+// to feedbackChannel.
+type grpcWALStream struct {
+	innerStream klioGRPC.WAL_PutClient
+	segmentSize uint64
+	clusterName string
+	sentBytes   uint64
+	walName     string
+	walStartLSN uint64
+	sendToTier2 bool
+
+	feedbackChannel chan<- wal.Feedback
+
+	// ackErr and ackDone follow the same pattern as messageReceiver.err in
+	// nonblocking_receive.go: ackErr is written at most once, by the
+	// background reader, always before it closes ackDone. The Go memory
+	// model guarantees that observing ackDone closed happens after that
+	// write, so reading ackErr once ackDone is (or is seen to be) closed is
+	// safe without any extra synchronization.
+	ackErr  error
+	ackDone chan struct{}
+}
+
+// SendBlock implements klioclient.WALUploader. It pipelines: the block is
+// handed to gRPC and SendBlock returns without waiting for the server's
+// per-block acknowledgment. The background goroutine started by
+// startFeedbackReader pushes each acknowledgment to feedbackChannel as it
+// arrives, so the caller learns how much of what has been sent is actually
+// confirmed durable from there.
 func (g *grpcWALStream) SendBlock(ctx context.Context, block []byte) error {
+	if err := g.ackedErr(); err != nil {
+		return err
+	}
+
 	sendStart := time.Now()
+
 	err := g.innerStream.Send(&klioGRPC.PutRequest{
 		ClusterName: g.clusterName,
 		WalName:     g.walName,
 		SegmentSize: g.segmentSize,
 		WalBlock:    block,
 		SendToTier2: g.sendToTier2,
+		WalStartLsn: g.walStartLSN,
 	})
 
-	opentelemetry.RecordDuration(ctx, opentelemetry.ClientWal.BlockDuration, time.Since(sendStart), err,
+	sendDuration := time.Since(sendStart)
+
+	opentelemetry.RecordDuration(ctx, opentelemetry.ClientWal.BlockDuration, sendDuration, err,
 		opentelemetry.AttributeKeyClusterName.Of(g.clusterName),
 		opentelemetry.PathPut.Attribute(), opentelemetry.StageSend.Attribute())
 
@@ -51,3 +93,56 @@ func (g *grpcWALStream) SendBlock(ctx context.Context, block []byte) error {
 
 	return nil
 }
+
+// Close stops sending on the stream and waits for the background ack
+// reader to observe the end of the stream, surfacing any error it recorded.
+func (g *grpcWALStream) Close(_ context.Context) error {
+	if err := g.innerStream.CloseSend(); err != nil {
+		return err //nolint:wrapcheck
+	}
+
+	<-g.ackDone
+
+	return g.ackErr
+}
+
+// startFeedbackReader starts the background goroutine that drains the server's
+// per-block acknowledgments.
+func (g *grpcWALStream) startFeedbackReader() {
+	g.ackDone = make(chan struct{})
+
+	go func() {
+		defer close(g.ackDone)
+
+		for {
+			result, err := g.innerStream.Recv()
+			if err != nil {
+				if !errors.Is(err, io.EOF) {
+					g.ackErr = fmt.Errorf("error while receiving WAL block ack: %w", err)
+				}
+
+				return
+			}
+
+			if g.feedbackChannel != nil {
+				g.feedbackChannel <- wal.Feedback{
+					FlushLSN:  result.GetFlushLsn(),
+					WriteLSN:  result.GetWriteLsn(),
+					ReplayLSN: result.GetFlushLsn(),
+				}
+			}
+		}
+	}()
+}
+
+// ackedErr returns the error recorded by the background ack reader once it
+// has observed the end of the stream, so SendBlock can fail fast instead of
+// sending into a stream already known to be dead.
+func (g *grpcWALStream) ackedErr() error {
+	select {
+	case <-g.ackDone:
+		return g.ackErr
+	default:
+		return nil
+	}
+}

core/internal/client/klioclient/interfaces.go

@@ -97,34 +97,12 @@ type Client interface {
 	Close(ctx context.Context)
 }
 
-// WALUploaderImpl is the underlying implementation of a WAL
+// WALUploader is the underlying implementation of a WAL
 // uploader.
-type WALUploaderImpl interface {
-	// SendBlock sends a WAL Block
+type WALUploader interface {
+	// SendBlock sends a WAL Block.
 	SendBlock(ctx context.Context, block []byte) 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
-}