Merge pull request #1363 Fixed listener stop partition

Author: rekby

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Fixed experimental topic listener handle stop partition event
+
 ## v3.76.1
 * Fixed `query.WithCommit()` flag behaviour for `tx.Execute` in query service
 * OAuth 2.0 token exchange: allowed multiple resource parameters in according to https://www.rfc-editor.org/rfc/rfc8693

examples/topic/topicreader/topicreader_handler.go

@@ -55,7 +55,7 @@ func (h *TopicEventsHandler) OnReadMessages(
 
 func (h *TopicEventsHandler) OnStartPartitionSessionRequest(
 	ctx context.Context,
-	event *topiclistener.StartPartitionSessionRequest,
+	event *topiclistener.EventStartPartitionSession,
 ) error {
 	lockID, offset, err := lockPartition(ctx, event.PartitionSession.TopicPath, event.PartitionSession.PartitionID)
 
@@ -75,7 +75,7 @@ func (h *TopicEventsHandler) OnStartPartitionSessionRequest(
 
 func (h *TopicEventsHandler) OnStopPartitionSessionRequest(
 	ctx context.Context,
-	event *topiclistener.StopPartitionSessionRequest,
+	event *topiclistener.EventStopPartitionSession,
 ) error {
 	h.m.Lock()
 	lockID := h.locks[event.PartitionSession.PartitionSessionID]

internal/topic/topiclistenerinternal/event_handler.go

@@ -3,6 +3,7 @@ package topiclistenerinternal
 import (
 	"context"
 
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/empty"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/topic/topicreadercommon"
 	"github.com/ydb-platform/ydb-go-sdk/v3/topic/topicreader"
 )
@@ -16,7 +17,7 @@ type EventHandler interface {
 	// You can set topiclistener.StartPartitionSessionConfirm for change default settings.
 	//
 	// Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
-	OnStartPartitionSessionRequest(ctx context.Context, event *PublicStartPartitionSessionEvent) error
+	OnStartPartitionSessionRequest(ctx context.Context, event *PublicEventStartPartitionSession) error
 
 	// OnReadMessages called with batch of messages. Max count of messages limited by internal buffer size
 	//
@@ -31,7 +32,7 @@ type EventHandler interface {
 	// It is guaranteed about the method will be called least once.
 	//
 	// Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
-	OnStopPartitionSessionRequest(ctx context.Context, event *PublicStopPartitionSessionEvent) error
+	OnStopPartitionSessionRequest(ctx context.Context, event *PublicEventStopPartitionSession) error
 }
 
 // PublicReadMessages
@@ -42,25 +43,37 @@ type PublicReadMessages struct {
 	Batch            *topicreader.Batch
 }
 
-// PublicStartPartitionSessionEvent
+// PublicEventStartPartitionSession
 //
 // Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
-type PublicStartPartitionSessionEvent struct {
+type PublicEventStartPartitionSession struct {
 	PartitionSession topicreadercommon.PublicPartitionSession
 	CommittedOffset  int64
 	PartitionOffsets PublicOffsetsRange
-	respChan         chan PublicStartPartitionSessionConfirm
+	confirm          confirmStorage[PublicStartPartitionSessionConfirm]
+}
+
+func NewPublicStartPartitionSessionEvent(
+	session topicreadercommon.PublicPartitionSession,
+	committedOffset int64,
+	partitionOffsets PublicOffsetsRange,
+) *PublicEventStartPartitionSession {
+	return &PublicEventStartPartitionSession{
+		PartitionSession: session,
+		CommittedOffset:  committedOffset,
+		PartitionOffsets: partitionOffsets,
+	}
 }
 
 // Confirm
 //
 // Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
-func (e *PublicStartPartitionSessionEvent) Confirm() {
+func (e *PublicEventStartPartitionSession) Confirm() {
 	e.ConfirmWithParams(PublicStartPartitionSessionConfirm{})
 }
 
-func (e *PublicStartPartitionSessionEvent) ConfirmWithParams(p PublicStartPartitionSessionConfirm) {
-	e.respChan <- p
+func (e *PublicEventStartPartitionSession) ConfirmWithParams(p PublicStartPartitionSessionConfirm) {
+	e.confirm.Set(p)
 }
 
 // PublicStartPartitionSessionConfirm
@@ -97,10 +110,10 @@ type PublicOffsetsRange struct {
 	End   int64
 }
 
-// PublicStopPartitionSessionEvent
+// PublicEventStopPartitionSession
 //
 // Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
-type PublicStopPartitionSessionEvent struct {
+type PublicEventStopPartitionSession struct {
 	PartitionSession topicreadercommon.PublicPartitionSession
 
 	// Graceful mean about server is waiting for client finish work with the partition and confirm stop the work
@@ -109,17 +122,24 @@ type PublicStopPartitionSessionEvent struct {
 	Graceful        bool
 	CommittedOffset int64
 
-	resp chan PublicStopPartitionSessionConfirm
+	confirm confirmStorage[empty.Struct]
 }
 
-// Confirm
-//
-// Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
-func (e *PublicStopPartitionSessionEvent) Confirm() {
-	e.resp <- PublicStopPartitionSessionConfirm{}
+func NewPublicStopPartitionSessionEvent(
+	partitionSession topicreadercommon.PublicPartitionSession,
+	graceful bool,
+	committedOffset int64,
+) *PublicEventStopPartitionSession {
+	return &PublicEventStopPartitionSession{
+		PartitionSession: partitionSession,
+		Graceful:         graceful,
+		CommittedOffset:  committedOffset,
+	}
 }
 
-// PublicStopPartitionSessionConfirm
+// Confirm
 //
 // Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
-type PublicStopPartitionSessionConfirm struct{}
+func (e *PublicEventStopPartitionSession) Confirm() {
+	e.confirm.Set(empty.Struct{})
+}

internal/topic/topiclistenerinternal/event_handler_mock_test.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"reflect"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -268,15 +269,15 @@ func (l *streamListener) onStartPartitionRequest(
 		PartitionSessionID: m.PartitionSession.PartitionSessionID,
 	}
 
-	event := &PublicStartPartitionSessionEvent{
-		PartitionSession: session.ToPublic(),
-		CommittedOffset:  m.CommittedOffset.ToInt64(),
-		PartitionOffsets: PublicOffsetsRange{
+	event := NewPublicStartPartitionSessionEvent(
+		session.ToPublic(),
+		m.CommittedOffset.ToInt64(),
+		PublicOffsetsRange{
 			Start: m.PartitionOffsets.Start.ToInt64(),
 			End:   m.PartitionOffsets.End.ToInt64(),
 		},
-		respChan: make(chan PublicStartPartitionSessionConfirm, 1),
-	}
+	)
+
 	err := l.handler.OnStartPartitionSessionRequest(ctx, event)
 	if err != nil {
 		return err
@@ -286,8 +287,8 @@ func (l *streamListener) onStartPartitionRequest(
 	select {
 	case <-ctx.Done():
 		return ctx.Err()
-	case userResp = <-event.respChan:
-		// pass
+	case <-event.confirm.Done():
+		userResp, _ = event.confirm.Get()
 	}
 
 	if userResp.readOffset != nil {
@@ -329,12 +330,11 @@ func (l *streamListener) onStopPartitionRequest(
 		handlerCtx = session.Context()
 	}
 
-	event := &PublicStopPartitionSessionEvent{
-		PartitionSession: session.ToPublic(),
-		Graceful:         m.Graceful,
-		CommittedOffset:  m.CommittedOffset.ToInt64(),
-		resp:             make(chan PublicStopPartitionSessionConfirm, 1),
-	}
+	event := NewPublicStopPartitionSessionEvent(
+		session.ToPublic(),
+		m.Graceful,
+		m.CommittedOffset.ToInt64(),
+	)
 
 	if err = l.handler.OnStopPartitionSessionRequest(handlerCtx, event); err != nil {
 		return err
@@ -344,15 +344,15 @@ func (l *streamListener) onStopPartitionRequest(
 		// remove partition on the confirmation or on the listener closed
 		select {
 		case <-l.background.Done():
-		case <-event.resp:
+		case <-event.confirm.Done():
 		}
 		_, _ = l.sessions.Remove(m.PartitionSessionID)
 	}()
 
 	select {
 	case <-ctx.Done():
 		return ctx.Err()
-	case <-event.resp:
+	case <-event.confirm.Done():
 		// pass
 	}
 
@@ -396,3 +396,42 @@ func (l *streamListener) sendMessage(m rawtopicreader.ClientMessage) {
 	default:
 	}
 }
+
+type confirmStorage[T any] struct {
+	doneChan      empty.Chan
+	confirmed     atomic.Bool
+	val           T
+	confirmAction sync.Once
+	initAction    sync.Once
+}
+
+func (c *confirmStorage[T]) init() {
+	c.initAction.Do(func() {
+		c.doneChan = make(empty.Chan)
+	})
+}
+
+func (c *confirmStorage[T]) Set(val T) {
+	c.init()
+	c.confirmAction.Do(func() {
+		c.val = val
+		c.confirmed.Store(true)
+		close(c.doneChan)
+	})
+}
+
+func (c *confirmStorage[T]) Done() empty.ChanReadonly {
+	c.init()
+
+	return c.doneChan
+}
+
+func (c *confirmStorage[T]) Get() (val T, ok bool) {
+	c.init()
+
+	if c.confirmed.Load() {
+		return c.val, true
+	}
+
+	return val, false
+}

internal/topic/topiclistenerinternal/stream_listener.go

@@ -15,14 +15,15 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopicreader"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawydb"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/topic/topicreadercommon"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
 )
 
 func TestStreamListener_OnReceiveServerMessage(t *testing.T) {
 	const batchBytes = 100
 
 	const seqNo int64 = 4
 
-	t.Run("onReadResponse", func(t *testing.T) {
+	xtest.TestManyTimesWithName(t, "onReadResponse", func(t testing.TB) {
 		e := fixenv.New(t)
 		ctx := sf.Context(e)
 
@@ -69,7 +70,7 @@ func TestStreamListener_OnReceiveServerMessage(t *testing.T) {
 			},
 		})
 	})
-	t.Run("onStartPartitionSession", func(t *testing.T) {
+	xtest.TestManyTimesWithName(t, "onStartPartitionSession", func(t testing.TB) {
 		e := fixenv.New(t)
 
 		respReadOffset := int64(16)
@@ -78,7 +79,7 @@ func TestStreamListener_OnReceiveServerMessage(t *testing.T) {
 		EventHandlerMock(e).EXPECT().OnStartPartitionSessionRequest(
 			gomock.Any(),
 			gomock.Any(),
-		).DoAndReturn(func(ctx context.Context, event *PublicStartPartitionSessionEvent) error {
+		).DoAndReturn(func(ctx context.Context, event *PublicEventStartPartitionSession) error {
 			require.Equal(t, topicreadercommon.PublicPartitionSession{
 				PartitionSessionID: 1, // ClientPartitionSessionID
 				TopicPath:          "asd",
@@ -130,7 +131,7 @@ func TestStreamListener_OnReceiveServerMessage(t *testing.T) {
 		require.NoError(t, err)
 		require.NotNil(t, session)
 	})
-	t.Run("onStopPartitionRequest", func(t *testing.T) {
+	xtest.TestManyTimesWithName(t, "onStopPartitionRequest", func(t testing.TB) {
 		e := fixenv.New(t)
 		ctx := sf.Context(e)
 
@@ -139,7 +140,7 @@ func TestStreamListener_OnReceiveServerMessage(t *testing.T) {
 		EventHandlerMock(e).EXPECT().OnStopPartitionSessionRequest(
 			PartitionSession(e).Context(),
 			gomock.Any(),
-		).DoAndReturn(func(ctx context.Context, event *PublicStopPartitionSessionEvent) error {
+		).DoAndReturn(func(ctx context.Context, event *PublicEventStopPartitionSession) error {
 			require.Equal(t, PartitionSession(e).ClientPartitionSessionID, event.PartitionSession.PartitionSessionID)
 			require.True(t, event.Graceful)
 			require.Equal(t, int64(5), event.CommittedOffset)
@@ -173,7 +174,7 @@ func TestStreamListener_CloseSessionsOnCloseListener(t *testing.T) {
 	EventHandlerMock(e).EXPECT().OnStopPartitionSessionRequest(
 		PartitionSession(e).Context(),
 		gomock.Any(),
-	).Do(func(ctx context.Context, event *PublicStopPartitionSessionEvent) error {
+	).Do(func(ctx context.Context, event *PublicEventStopPartitionSession) error {
 		require.Equal(t, PartitionSession(e).ClientPartitionSessionID, event.PartitionSession.PartitionSessionID)
 		require.False(t, event.Graceful)
 		require.Equal(t, PartitionSession(e).CommittedOffset().ToInt64(), event.CommittedOffset)

internal/topic/topiclistenerinternal/stream_listener_test.go

@@ -51,8 +51,8 @@ type TestTopicListener_Handler struct {
 
 	listener         *topiclistener.TopicListener
 	readMessages     *topiclistener.ReadMessages
-	onPartitionStart *topiclistener.StartPartitionSessionRequest
-	onPartitionStop  *topiclistener.StopPartitionSessionRequest
+	onPartitionStart *topiclistener.EventStartPartitionSession
+	onPartitionStop  *topiclistener.EventStopPartitionSession
 	done             empty.Chan
 }
 
@@ -63,7 +63,7 @@ func (h *TestTopicListener_Handler) OnReaderCreated(event *topiclistener.ReaderR
 
 func (h *TestTopicListener_Handler) OnStartPartitionSessionRequest(
 	ctx context.Context,
-	event *topiclistener.StartPartitionSessionRequest,
+	event *topiclistener.EventStartPartitionSession,
 ) error {
 	h.onPartitionStart = event
 	event.Confirm()
@@ -72,7 +72,7 @@ func (h *TestTopicListener_Handler) OnStartPartitionSessionRequest(
 
 func (h *TestTopicListener_Handler) OnStopPartitionSessionRequest(
 	ctx context.Context,
-	event *topiclistener.StopPartitionSessionRequest,
+	event *topiclistener.EventStopPartitionSession,
 ) error {
 	h.onPartitionStop = event
 	event.Confirm()