add batch flusher

Author: rekby

internal/topic/topicreaderinternal/batcher.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xslices"
 	"sync/atomic"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/empty"
@@ -26,6 +27,7 @@ type batcher struct {
 	closed                                        bool
 	closeChan                                     empty.Chan
 	messages                                      batcherMessagesMap
+	sessionsForFlush                              []*topicreadercommon.PartitionSession
 }
 
 func newBatcher() *batcher {
@@ -81,6 +83,20 @@ func (b *batcher) PushRawMessage(session *topicreadercommon.PartitionSession, m
 	return b.addNeedLock(session, newBatcherItemRawMessage(m))
 }
 
+func (b *batcher) FlushPartitionSession(session *topicreadercommon.PartitionSession) {
+	b.m.WithLock(func() {
+		if b.closed {
+			return
+		}
+
+		if _, ok := b.messages[session]; !ok {
+			return
+		}
+
+		b.sessionsForFlush = append(b.sessionsForFlush, session)
+	})
+}
+
 func (b *batcher) addNeedLock(session *topicreadercommon.PartitionSession, item batcherMessageOrderItem) error {
 	var currentItems batcherMessageOrderItems
 	var ok bool
@@ -250,31 +266,59 @@ func (b *batcher) findNeedLock(filter batcherGetOptions) batcherResultCandidate
 		return batcherResultCandidate{}
 	}
 
-	rawMessageOpts := batcherGetOptions{rawMessagesOnly: true}
-
 	var batchResult batcherResultCandidate
 	needBatchResult := true
 
-	for k, items := range b.messages {
-		head, rest, ok := rawMessageOpts.cutBatchItemsHead(items)
-		if ok {
-			return newBatcherResultCandidate(k, head, rest, true)
+	for _, session := range b.sessionsForFlush {
+		items := b.messages[session]
+		candidate := b.findNeedLockProcessItem(filter, session, items, true)
+		if !candidate.Result.IsEmpty() {
+			return candidate
 		}
+	}
 
-		if needBatchResult {
-			head, rest, ok = b.applyForceFlagToOptions(filter).cutBatchItemsHead(items)
-			if !ok {
-				continue
-			}
+	for session, items := range b.messages {
+		candidate := b.findNeedLockProcessItem(filter, session, items, needBatchResult)
+
+		if candidate.Result.IsRawMessage() {
+			return candidate
+		}
 
+		if candidate.Result.IsBatch() {
+			batchResult = candidate
 			needBatchResult = false
-			batchResult = newBatcherResultCandidate(k, head, rest, true)
 		}
 	}
 
 	return batchResult
 }
 
+func (b *batcher) findNeedLockProcessItem(
+	filter batcherGetOptions,
+	k *topicreadercommon.PartitionSession,
+	items batcherMessageOrderItems,
+	needBatchResult bool,
+) (
+	result batcherResultCandidate,
+) {
+	rawMessageOpts := batcherGetOptions{rawMessagesOnly: true}
+	head, rest, ok := rawMessageOpts.cutBatchItemsHead(items)
+	if ok {
+		return newBatcherResultCandidate(k, head, rest, true)
+	}
+
+	if needBatchResult {
+		head, rest, ok = b.applyForceFlagToOptions(filter).cutBatchItemsHead(items)
+		if !ok {
+			return batcherResultCandidate{}
+		}
+
+		return newBatcherResultCandidate(k, head, rest, true)
+	}
+
+	return batcherResultCandidate{}
+}
+
 func (b *batcher) applyForceFlagToOptions(options batcherGetOptions) batcherGetOptions {
 	if !b.forceIgnoreMinRestrictionsOnNextMessagesBatch {
 		return options
@@ -289,6 +333,9 @@ func (b *batcher) applyForceFlagToOptions(options batcherGetOptions) batcherGetO
 func (b *batcher) applyNeedLock(res *batcherResultCandidate) {
 	if res.Rest.IsEmpty() && res.WaiterIndex >= 0 {
 		delete(b.messages, res.Key)
+		if len(b.sessionsForFlush) > 0 && b.sessionsForFlush[0] == res.Key {
+			b.sessionsForFlush = xslices.Delete(b.sessionsForFlush, 0, 1)
+		}
 	} else {
 		b.messages[res.Key] = res.Rest
 	}

internal/topic/topicreaderinternal/batcher_test.go

@@ -316,6 +316,62 @@ func TestBatcher_PopMinIgnored(t *testing.T) {
 	})
 }
 
+func TestBatcher_PopFlushed(t *testing.T) {
+	ctx := context.Background()
+	//rnd := rand.New(rand.NewSource(0))
+
+	s1 := topicreadercommon.NewPartitionSession(ctx,
+		"test",
+		1,
+		-1,
+		"test",
+		partitionSessionID(1),
+		1,
+		0,
+	)
+	s2 := topicreadercommon.NewPartitionSession(ctx,
+		"test",
+		2,
+		-1,
+		"test",
+		partitionSessionID(1),
+		1,
+		0,
+	)
+
+	sessions := []*topicreadercommon.PartitionSession{s1, s2}
+	for flushIndex := range sessions {
+		b := newBatcher()
+		err := b.PushBatches(
+			xtest.Must(
+				topicreadercommon.NewBatch(
+					s1,
+					[]*topicreadercommon.PublicMessage{
+						topicreadercommon.NewPublicMessageBuilder().PartitionSession(s1).Build(),
+					},
+				),
+			),
+			xtest.Must(
+				topicreadercommon.NewBatch(
+					s2,
+					[]*topicreadercommon.PublicMessage{
+						topicreadercommon.NewPublicMessageBuilder().PartitionSession(s2).Build(),
+					},
+				),
+			),
+		)
+		require.NoError(t, err)
+
+		s := sessions[flushIndex]
+		b.FlushPartitionSession(s)
+
+		res, err := b.Pop(ctx, batcherGetOptions{})
+		require.NoError(t, err)
+
+		require.Same(t, s, topicreadercommon.BatchGetPartitionSession(res.Batch))
+	}
+}
+
 func TestBatcherConcurency(t *testing.T) {
 	xtest.TestManyTimesWithName(t, "OneBatch", func(tb testing.TB) {
 		b := newBatcher()

internal/xslices/delete.go

@@ -0,0 +1,14 @@
+//go:build go1.22
+
+package xslices
+
+import "slices"
+
+// Delete removes the elements s[i:j] from s, returning the modified slice.
+// Delete panics if j > len(s) or s[i:j] is not a valid slice of s.
+// Delete is O(len(s)-i), so if many items must be deleted, it is better to
+// make a single call deleting them all together than to delete one at a time.
+// Delete zeroes the elements s[len(s)-(j-i):len(s)].
+func Delete[S ~[]E, E any](s S, i, j int) S {
+	return slices.Delete(s, i, j)
+}

internal/xslices/delete_old.go

@@ -0,0 +1,21 @@
+//go:build !go1.22
+
+package xslices
+
+// Delete removes the elements s[i:j] from s, returning the modified slice.
+// Delete panics if j > len(s) or s[i:j] is not a valid slice of s.
+// Delete is O(len(s)-i), so if many items must be deleted, it is better to
+// make a single call deleting them all together than to delete one at a time.
+// Delete zeroes the elements s[len(s)-(j-i):len(s)].
+func Delete[S ~[]E, E any](s S, i, j int) S {
+	_ = s[i:j:len(s)] // bounds check
+
+	if i == j {
+		return s
+	}
+
+	oldlen := len(s)
+	s = append(s[:i], s[j:]...)
+	clear(s[len(s):oldlen]) // zero/nil out the obsolete elements, for GC
+	return s
+}