fix(wal): gate flush feedback on durable server acknowledgements

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

@@ -23,8 +23,11 @@ import (
 	"context"
 	"crypto/tls"
 	"crypto/x509"
+	"errors"
 	"fmt"
+	"io"
 	"os"
+	"sync/atomic"
 
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
 	"google.golang.org/grpc"
@@ -42,25 +45,102 @@ type grpcWALStream struct {
 	sentBytes   uint64
 	walName     string
 	sendToTier2 bool
+
+	// syncedBytes is the highest cumulative durable size acknowledged by the
+	// server for this WAL file. It is written by the ack reader goroutine and
+	// read by the WAL processing goroutine, so it must be accessed atomically.
+	syncedBytes atomic.Uint64
+
+	// ackDone is closed when the ack reader goroutine has consumed every
+	// acknowledgement up to the end of the stream. ackErr, set before ackDone
+	// is closed, carries any non-EOF error observed while reading.
+	ackDone chan struct{}
+	ackErr  error
+}
+
+// newGRPCWALStream wraps a freshly opened Put stream and starts the background
+// goroutine that consumes the server durability acknowledgements.
+func newGRPCWALStream(
+	stream klioGRPC.WAL_PutClient,
+	name string,
+	segmentSize uint64,
+	clusterName string,
+	sendToTier2 bool,
+) *grpcWALStream {
+	g := &grpcWALStream{
+		innerStream: stream,
+		segmentSize: segmentSize,
+		clusterName: clusterName,
+		walName:     name,
+		sendToTier2: sendToTier2,
+		ackDone:     make(chan struct{}),
+	}
+
+	go g.readAcks()
+
+	return g
+}
+
+// SyncedOffset implements common.WALStream. It returns the highest durable size
+// acknowledged so far for this WAL file.
+func (g *grpcWALStream) SyncedOffset() (uint64, error) {
+	// If the reader has already terminated with a non-EOF error, the
+	// acknowledgements can no longer be trusted.
+	select {
+	case <-g.ackDone:
+		if g.ackErr != nil {
+			return 0, fmt.Errorf("while receiving WAL acknowledgements: %w", g.ackErr)
+		}
+	default:
+	}
+
+	return g.syncedBytes.Load(), nil
 }
 
-// Close implements common.WALStream.
+// Close implements common.WALStream. It stops sending, waits for every
+// acknowledgement to be drained and verifies that the whole file is durable.
 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 err := g.innerStream.CloseSend(); err != nil {
+		return fmt.Errorf("while closing the WAL upload stream: %w", err)
+	}
+
+	// Wait for the reader to drain all acknowledgements up to the end of the
+	// stream, so syncedBytes reflects the final durable size.
+	<-g.ackDone
+	if g.ackErr != nil {
+		return fmt.Errorf("while receiving WAL acknowledgements: %w", g.ackErr)
 	}
 
-	if result.GetWrittenSize() != g.sentBytes {
+	if synced := g.syncedBytes.Load(); synced != g.sentBytes {
 		return &IncompleteWALFileError{
-			uploadedSize: result.GetWrittenSize(),
+			uploadedSize: synced,
 			expectedSize: g.sentBytes,
 		}
 	}
 
 	return nil
 }
 
+// readAcks consumes the server acknowledgements until the stream ends, tracking
+// the latest durable size. gRPC allows a single concurrent reader alongside the
+// single writer used by SendBlock.
+func (g *grpcWALStream) readAcks() {
+	defer close(g.ackDone)
+
+	for {
+		result, err := g.innerStream.Recv()
+		if err != nil {
+			if !errors.Is(err, io.EOF) {
+				g.ackErr = err
+			}
+
+			return
+		}
+
+		g.syncedBytes.Store(result.GetWrittenSize())
+	}
+}
+
 // Connection represents a connection to a Klio server.
 type Connection struct {
 	klioGRPC.WALClient

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

@@ -61,21 +61,42 @@ func (c *Connection) StoreWAL(ctx context.Context, name string, content []byte,
 		}
 	}
 
-	result, err := stream.CloseAndRecv()
+	if err := stream.CloseSend(); err != nil {
+		return fmt.Errorf("while closing the WAL upload stream: %w", err)
+	}
+
+	writtenSize, err := drainPutAcks(stream)
 	if err != nil {
-		return fmt.Errorf("while flushing WAL file: %w", err)
+		return err
 	}
 
-	if result.GetWrittenSize() != uint64(len(content)) {
+	if writtenSize != uint64(len(content)) {
 		return &IncompleteWALFileError{
-			uploadedSize: result.GetWrittenSize(),
+			uploadedSize: writtenSize,
 			expectedSize: uint64(len(content)),
 		}
 	}
 
 	return nil
 }
 
+// drainPutAcks consumes the durability acknowledgements of a Put stream until
+// it ends, returning the last (cumulative) durable size reported by the server.
+func drainPutAcks(stream klioGRPC.WAL_PutClient) (uint64, error) {
+	var writtenSize uint64
+	for {
+		result, err := stream.Recv()
+		if errors.Is(err, io.EOF) {
+			return writtenSize, nil
+		}
+		if err != nil {
+			return 0, fmt.Errorf("while flushing WAL file: %w", err)
+		}
+
+		writtenSize = result.GetWrittenSize()
+	}
+}
+
 // 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

@@ -46,13 +46,13 @@ func (c *Connection) StoreWALStreaming(
 		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
+	return klioclient.NewWALUploader(newGRPCWALStream(
+		stream,
+		name,
+		segmentSize,
+		c.clientConfig.ClusterName,
+		sendToTier2,
+	)), nil
 }
 
 // GetWALStreaming get a WAL from a remote connection.

core/internal/client/klioclient/interfaces.go

@@ -103,6 +103,10 @@ type WALUploaderImpl interface {
 	// SendBlock sends a WAL Block
 	SendBlock(ctx context.Context, block []byte) error
 
+	// SyncedOffset returns the number of bytes the server has acknowledged as
+	// durably persisted for the current WAL file.
+	SyncedOffset() (uint64, error)
+
 	// Close closes the WAL streaming session
 	Close(ctx context.Context) error
 }
@@ -124,6 +128,12 @@ func (u *WALUploader) SendBlock(ctx context.Context, block []byte) error {
 	return u.impl.SendBlock(ctx, block) //nolint:wrapcheck
 }
 
+// SyncedOffset returns the number of bytes the server has acknowledged as
+// durably persisted for the current WAL file.
+func (u *WALUploader) SyncedOffset() (uint64, error) {
+	return u.impl.SyncedOffset() //nolint:wrapcheck
+}
+
 // Close closes the WAL streaming session.
 func (u *WALUploader) Close(ctx context.Context) error {
 	return u.impl.Close(ctx) //nolint:wrapcheck

core/internal/client/sendwal/buffer/buffer.go

@@ -40,19 +40,27 @@ type Data struct {
 
 	handler Handler
 
-	writeLSN   uint64
-	flushLSN   uint64
-	buffer     *bytes.Buffer
-	bufferSize int
+	writeLSN     uint64
+	flushLSN     uint64
+	fileStartLSN uint64
+	buffer       *bytes.Buffer
+	bufferSize   int
+
+	// requireDurableAck gates flushLSN advancement on the destination
+	// acknowledging that the data has been durably persisted, instead of
+	// advancing it as soon as the data has been handed to the send buffer.
+	requireDurableAck bool
 }
 
-// New creates a new WAL buffer.
-func New(tli int, walSegmentSize uint64, handler Handler, bufferSize int) *Data {
+// New creates a new WAL buffer. When requireDurableAck is true, the flush LSN
+// only advances up to data the handler reports as durably persisted.
+func New(tli int, walSegmentSize uint64, handler Handler, bufferSize int, requireDurableAck bool) *Data {
 	result := &Data{
-		segmentSize: walSegmentSize,
-		tli:         tli,
-		handler:     handler,
-		bufferSize:  bufferSize,
+		segmentSize:       walSegmentSize,
+		tli:               tli,
+		handler:           handler,
+		bufferSize:        bufferSize,
+		requireDurableAck: requireDurableAck,
 	}
 
 	result.buffer = result.newBuffer()
@@ -167,6 +175,7 @@ func (wal *Data) openWALPos(ctx context.Context, blockpos uint64) error {
 
 	wal.writeLSN = blockpos
 	wal.flushLSN = blockpos
+	wal.fileStartLSN = blockpos
 
 	return nil
 }
@@ -188,26 +197,50 @@ func (wal *Data) writeToWALFile(ctx context.Context, data []byte) error {
 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() {
 		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 {
-		return fmt.Errorf("while writing to WAL handler: %w", err)
+	if wal.buffer.Len() > 0 {
+		contextLogger.Debug("Writing block",
+			"blockpos", types.Int64ToLSN(wal.writeLSN), "blocksize", wal.buffer.Len())
+		_, err := wal.handler.Write(ctx, wal.buffer.Bytes())
+		if 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()
+		}
 	}
 
-	// Clear content but keeps the slice capacity
-	wal.buffer.Reset()
+	return wal.advanceFlushLSN()
+}
 
-	// Prevent memory bloat in long-running processes.
-	if wal.buffer.Cap() > wal.bufferSize*maximumBufferSizeFactor {
-		wal.buffer = wal.newBuffer()
+// advanceFlushLSN moves the flush position forward. In durable-ack mode it only
+// advances up to data the handler reports as durably persisted by the server,
+// which may lag the written position; otherwise it tracks the written position.
+// It is also called on idle flush ticks so that acknowledgements arriving after
+// the last write are still reflected in the flush position.
+func (wal *Data) advanceFlushLSN() error {
+	if !wal.requireDurableAck {
+		wal.flushLSN = wal.writeLSN
+		return nil
 	}
 
-	wal.flushLSN = wal.writeLSN
+	synced, err := wal.handler.SyncedOffset()
+	if err != nil {
+		return fmt.Errorf("while reading the durable offset: %w", err)
+	}
+
+	durableLSN := wal.fileStartLSN + synced
+	if durableLSN > wal.flushLSN {
+		wal.flushLSN = durableLSN
+	}
 
 	return nil
 }
@@ -224,5 +257,10 @@ func (wal *Data) closeCurrentWAL(ctx context.Context) error {
 		return fmt.Errorf("while closing current WAL file: %w", err)
 	}
 
+	// Closing the WAL file drains and verifies every outstanding
+	// acknowledgement, so the whole segment is now durable: advance the flush
+	// position to the segment boundary regardless of the mode.
+	wal.flushLSN = wal.writeLSN
+
 	return nil
 }