Merge pull request #864 Introduce topicwriter Writer.Flush() method from zveznicht/issues/851/flush

@zveznicht thanks for the pr :)

Author: rekby

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Added Flush method for topic writer
+
 ## v3.66.0
 * Added experimental package `retry/budget` for limit second and subsequent retry attempts 
 * Refactored internals for enabling `containedctx` linter

internal/topic/topicwriterinternal/queue.go

@@ -320,6 +320,16 @@ func (q *messageQueue) Wait(ctx context.Context, waiter MessageQueueAckWaiter) e
 	}
 }
 
+// WaitLastWritten waits for last written message gets ack.
+func (q *messageQueue) WaitLastWritten(ctx context.Context) error {
+	var lastIndex int
+	q.m.WithRLock(func() {
+		lastIndex = q.lastWrittenIndex
+	})
+
+	return q.Wait(ctx, MessageQueueAckWaiter{sequenseNumbers: []int{lastIndex}})
+}
+
 type MessageQueueAckWaiter struct {
 	sequenseNumbers []int
 }

internal/topic/topicwriterinternal/writer.go

@@ -55,3 +55,7 @@ func (w *Writer) WaitInit(ctx context.Context) (info InitialInfo, err error) {
 func (w *Writer) Close(ctx context.Context) error {
 	return w.streamWriter.Close(ctx)
 }
+
+func (w *Writer) Flush(ctx context.Context) error {
+	return w.streamWriter.Flush(ctx)
+}

internal/topic/topicwriterinternal/writer_reconnector.go

@@ -325,6 +325,10 @@ func (w *WriterReconnector) createMessagesWithContent(messages []PublicMessage)
 	return res, nil
 }
 
+func (w *WriterReconnector) Flush(ctx context.Context) error {
+	return w.queue.WaitLastWritten(ctx)
+}
+
 func (w *WriterReconnector) Close(ctx context.Context) error {
 	return w.close(ctx, xerrors.WithStackTrace(errStopWriterReconnector))
 }

internal/topic/topicwriterinternal/writer_reconnector_test.go

@@ -530,6 +530,70 @@ func TestWriterImpl_Reconnect(t *testing.T) {
 	})
 }
 
+func TestWriterImpl_CloseWithFlush(t *testing.T) {
+	e := newTestEnv(t, nil)
+
+	messageTime := time.Date(2023, 9, 7, 11, 34, 0, 0, time.UTC)
+	messageData := []byte("123")
+
+	const seqNo = 36
+
+	writeCompleted := make(empty.Chan)
+	e.stream.EXPECT().Send(&rawtopicwriter.WriteRequest{
+		Messages: []rawtopicwriter.MessageData{
+			{
+				SeqNo:            seqNo,
+				CreatedAt:        messageTime,
+				UncompressedSize: int64(len(messageData)),
+				Partitioning:     rawtopicwriter.Partitioning{},
+				Data:             messageData,
+			},
+		},
+		Codec: rawtopiccommon.CodecRaw,
+	}).Return(nil)
+
+	closeCompleted := make(empty.Chan)
+	go func() {
+		err := e.writer.Write(e.ctx, []PublicMessage{{
+			SeqNo:     seqNo,
+			CreatedAt: messageTime,
+			Data:      bytes.NewReader(messageData),
+		}})
+		close(writeCompleted)
+		require.NoError(t, err)
+	}()
+
+	<-writeCompleted
+
+	go func() {
+		require.NoError(t, e.writer.Flush(e.ctx))
+		require.NoError(t, e.writer.Close(e.ctx))
+		close(closeCompleted)
+	}()
+
+	select {
+	case <-closeCompleted:
+		t.Fatal("flush and close must complete only after message is acked")
+	case <-time.After(100 * time.Millisecond):
+		// pass
+	}
+
+	e.sendFromServer(&rawtopicwriter.WriteResult{
+		Acks: []rawtopicwriter.WriteAck{
+			{
+				SeqNo: seqNo,
+				MessageWriteStatus: rawtopicwriter.MessageWriteStatus{
+					Type:          rawtopicwriter.WriteStatusTypeWritten,
+					WrittenOffset: 4,
+				},
+			},
+		},
+		PartitionID: e.partitionID,
+	})
+
+	xtest.WaitChannelClosed(t, closeCompleted)
+}
+
 func TestAllMessagesHasSameBufCodec(t *testing.T) {
 	t.Run("Empty", func(t *testing.T) {
 		require.True(t, allMessagesHasSameBufCodec(nil))
@@ -891,7 +955,7 @@ func newTestEnv(t testing.TB, options *testEnvOptions) *testEnv {
 	require.NoError(t, res.writer.waitFirstInitResponse(res.ctx))
 
 	t.Cleanup(func() {
-		res.writer.close(context.Background(), errors.New("stop writer test environment"))
+		_ = res.writer.close(context.Background(), errors.New("stop writer test environment"))
 		close(res.stopReadEvents)
 		<-streamClosed
 	})

internal/topic/topicwriterinternal/writer_stream_interface.go

@@ -9,6 +9,7 @@ type StreamWriter interface {
 	Write(ctx context.Context, messages []PublicMessage) error
 	WaitInit(ctx context.Context) (info InitialInfo, err error)
 	Close(ctx context.Context) error
+	Flush(ctx context.Context) error
 }
 
 type InitialInfo struct {

internal/topic/topicwriterinternal/writer_stream_interface_mock_test.go

@@ -69,3 +69,8 @@ func (w *Writer) WaitInitInfo(ctx context.Context) (info PublicInitialInfo, err
 func (w *Writer) Close(ctx context.Context) error {
 	return w.inner.Close(ctx)
 }
+
+// Flush waits till all in-flight messages are acknowledged.
+func (w *Writer) Flush(ctx context.Context) error {
+	return w.inner.Flush(ctx)
+}