Merge pull request #803 added WaitInit to Writer and Reader

Author: rekby

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Added `WaitInit` and `WaitInitInfo` method to the topic reader and writer
+
 ## v3.52.2
 * Removed support of placeholder "_" for ignoring columns in `database/sql` result sets
 

internal/topic/topicreaderinternal/batched_stream_reader_interface.go

@@ -7,6 +7,7 @@ import (
 //go:generate mockgen -source batched_stream_reader_interface.go -destination batched_stream_reader_mock_test.go -package topicreaderinternal -write_package_comment=false
 
 type batchedStreamReader interface {
+	WaitInit(ctx context.Context) error
 	ReadMessageBatch(ctx context.Context, opts ReadMessageBatchOptions) (*PublicBatch, error)
 	Commit(ctx context.Context, commitRange commitRange) error
 	CloseWithError(ctx context.Context, err error) error

internal/topic/topicreaderinternal/batched_stream_reader_mock_test.go

@@ -119,6 +119,10 @@ func NewReader(
 	return res
 }
 
+func (r *Reader) WaitInit(ctx context.Context) error {
+	return r.reader.WaitInit(ctx)
+}
+
 func (r *Reader) ID() int64 {
 	return r.readerID
 }

internal/topic/topicreaderinternal/reader.go

@@ -7,6 +7,7 @@ import (
 	"testing"
 
 	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/empty"
@@ -153,3 +154,19 @@ func TestReader_Commit(t *testing.T) {
 		require.ErrorIs(t, err, errCommitSessionFromOtherReader)
 	})
 }
+
+func TestReader_WaitInit(t *testing.T) {
+	mc := gomock.NewController(t)
+	defer mc.Finish()
+
+	readerID := nextReaderID()
+	baseReader := NewMockbatchedStreamReader(mc)
+	reader := &Reader{
+		reader:   baseReader,
+		readerID: readerID,
+	}
+
+	baseReader.EXPECT().WaitInit(gomock.Any())
+	err := reader.WaitInit(context.Background())
+	assert.NoError(t, err)
+}

internal/topic/topicreaderinternal/reader_test.go

@@ -163,6 +163,13 @@ func newTopicStreamReaderStopped(
 	return res
 }
 
+func (r *topicStreamReaderImpl) WaitInit(_ context.Context) error {
+	if !r.started {
+		return errors.New("not started: can be started only after initialize from constructor")
+	}
+	return nil
+}
+
 func (r *topicStreamReaderImpl) ReadMessageBatch(
 	ctx context.Context,
 	opts ReadMessageBatchOptions,

internal/topic/topicreaderinternal/stream_reader_impl.go

@@ -340,6 +340,21 @@ func TestTopicStreamReaderImpl_Create(t *testing.T) {
 	})
 }
 
+func TestTopicStreamReaderImpl_WaitInit(t *testing.T) {
+	t.Run("OK", func(t *testing.T) {
+		e := newTopicReaderTestEnv(t)
+		e.Start()
+		err := e.reader.WaitInit(context.Background())
+		require.NoError(t, err)
+	})
+
+	t.Run("not started", func(t *testing.T) {
+		e := newTopicReaderTestEnv(t)
+		err := e.reader.WaitInit(context.Background())
+		require.Error(t, err)
+	})
+}
+
 func TestStreamReaderImpl_OnPartitionCloseHandle(t *testing.T) {
 	xtest.TestManyTimesWithName(t, "GracefulFalseCancelPartitionContext", func(t testing.TB) {
 		e := newTopicReaderTestEnv(t)

internal/topic/topicreaderinternal/stream_reader_impl_test.go

@@ -49,6 +49,10 @@ type readerReconnector struct {
 	streamVal                  batchedStreamReader
 	streamErr                  error
 	closedErr                  error
+
+	initErr    error
+	initDone   bool
+	initDoneCh empty.Chan
 }
 
 //nolint:revive
@@ -70,6 +74,7 @@ func newReaderReconnector(
 		baseContext:    baseContext,
 		retrySettings:  retrySettings,
 	}
+
 	if res.connectTimeout == 0 {
 		res.connectTimeout = value.InfiniteDuration
 	}
@@ -146,6 +151,13 @@ func (r *readerReconnector) CloseWithError(ctx context.Context, err error) error
 				closeErr = streamCloseErr
 			}
 		}
+
+		r.m.WithLock(func() {
+			if !r.initDone {
+				r.initErr = closeErr
+				close(r.initDoneCh)
+			}
+		})
 	})
 	return closeErr
 }
@@ -163,6 +175,7 @@ func (r *readerReconnector) initChannelsAndClock() {
 	}
 	r.reconnectFromBadStream = make(chan reconnectRequest, 1)
 	r.streamConnectionInProgress = make(empty.Chan)
+	r.initDoneCh = make(empty.Chan)
 	close(r.streamConnectionInProgress) // no progress at start
 }
 
@@ -261,6 +274,10 @@ func (r *readerReconnector) reconnect(ctx context.Context, reason error, oldRead
 		r.streamErr = err
 		if err == nil {
 			r.streamVal = newStream
+			if !r.initDone {
+				r.initDone = true
+				close(r.initDoneCh)
+			}
 		}
 	})
 	return err
@@ -312,10 +329,22 @@ func (r *readerReconnector) connectWithTimeout() (_ batchedStreamReader, err err
 	if res.err == nil {
 		return res.stream, nil
 	}
-
 	return nil, res.err
 }
 
+func (r *readerReconnector) WaitInit(ctx context.Context) error {
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-r.initDoneCh:
+		return r.initErr
+	}
+}
+
 func (r *readerReconnector) fireReconnectOnRetryableError(stream batchedStreamReader, err error) {
 	if !r.isRetriableError(err) {
 		return

internal/topic/topicreaderinternal/stream_reconnector.go

@@ -130,7 +130,11 @@ func TestTopicReaderReconnectorReadMessageBatch(t *testing.T) {
 	})
 
 	xtest.TestManyTimesWithName(t, "OnClose", func(t testing.TB) {
-		reconnector := &readerReconnector{tracer: &trace.Topic{}}
+		reconnector := &readerReconnector{
+			tracer:    &trace.Topic{},
+			streamErr: errUnconnected,
+		}
+		reconnector.initChannelsAndClock()
 		testErr := errors.New("test'")
 
 		go func() {
@@ -256,7 +260,10 @@ func TestTopicReaderReconnectorConnectionLoop(t *testing.T) {
 	t.Run("StartWithCancelledContext", func(t *testing.T) {
 		ctx, cancel := xcontext.WithCancel(context.Background())
 		cancel()
-		reconnector := &readerReconnector{tracer: &trace.Topic{}}
+		reconnector := &readerReconnector{
+			tracer: &trace.Topic{},
+		}
+		reconnector.initChannelsAndClock()
 		reconnector.reconnectionLoop(ctx) // must return
 	})
 }
@@ -295,6 +302,84 @@ func TestTopicReaderReconnectorStart(t *testing.T) {
 	_ = reconnector.CloseWithError(ctx, nil)
 }
 
+func TestTopicReaderReconnectorWaitInit(t *testing.T) {
+	t.Run("OK", func(t *testing.T) {
+		mc := gomock.NewController(t)
+		defer mc.Finish()
+
+		reconnector := &readerReconnector{
+			tracer: &trace.Topic{},
+		}
+		reconnector.initChannelsAndClock()
+
+		stream := NewMockbatchedStreamReader(mc)
+
+		reconnector.readerConnect = readerConnectFuncMock(readerConnectFuncAnswer{
+			callback: func(ctx context.Context) (batchedStreamReader, error) {
+				return stream, nil
+			},
+		})
+
+		reconnector.start()
+
+		err := reconnector.WaitInit(context.Background())
+		require.NoError(t, err)
+
+		// one more run is needed to check idempotency
+		err = reconnector.WaitInit(context.Background())
+		require.NoError(t, err)
+	})
+
+	t.Run("contextDeadlineInProgress", func(t *testing.T) {
+		mc := gomock.NewController(t)
+		defer mc.Finish()
+
+		reconnector := &readerReconnector{
+			tracer: &trace.Topic{},
+		}
+		reconnector.initChannelsAndClock()
+
+		stream := NewMockbatchedStreamReader(mc)
+
+		ctx, cancel := context.WithCancel(context.Background())
+		reconnector.readerConnect = readerConnectFuncMock(readerConnectFuncAnswer{
+			callback: func(ctx context.Context) (batchedStreamReader, error) {
+				cancel()
+				return stream, nil
+			},
+		})
+		reconnector.start()
+
+		err := reconnector.WaitInit(ctx)
+		require.ErrorIs(t, err, ctx.Err())
+	})
+
+	t.Run("contextDeadlineBeforeStart", func(t *testing.T) {
+		mc := gomock.NewController(t)
+		defer mc.Finish()
+
+		reconnector := &readerReconnector{
+			tracer: &trace.Topic{},
+		}
+		reconnector.initDoneCh = make(empty.Chan, 1)
+		reconnector.initChannelsAndClock()
+
+		stream := NewMockbatchedStreamReader(mc)
+
+		reconnector.readerConnect = readerConnectFuncMock(readerConnectFuncAnswer{
+			callback: func(ctx context.Context) (batchedStreamReader, error) {
+				return stream, nil
+			},
+		})
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		err := reconnector.WaitInit(ctx)
+
+		require.ErrorIs(t, err, ctx.Err())
+	})
+}
+
 func TestTopicReaderReconnectorFireReconnectOnRetryableError(t *testing.T) {
 	t.Run("Ok", func(t *testing.T) {
 		mc := gomock.NewController(t)
@@ -332,7 +417,9 @@ func TestTopicReaderReconnectorFireReconnectOnRetryableError(t *testing.T) {
 		mc := gomock.NewController(t)
 		defer mc.Finish()
 
-		reconnector := &readerReconnector{tracer: &trace.Topic{}}
+		reconnector := &readerReconnector{
+			tracer: &trace.Topic{},
+		}
 		stream := NewMockbatchedStreamReader(mc)
 		reconnector.initChannelsAndClock()
 

internal/topic/topicreaderinternal/stream_reconnector_test.go

@@ -48,6 +48,10 @@ func (w *Writer) Write(ctx context.Context, messages ...Message) error {
 	return w.streamWriter.Write(ctx, messages)
 }
 
+func (w *Writer) WaitInit(ctx context.Context) (info InitialInfo, err error) {
+	return w.streamWriter.WaitInit(ctx)
+}
+
 func (w *Writer) Close(ctx context.Context) error {
 	return w.streamWriter.Close(ctx)
 }