reader and writer use some retry decision functions

Author: rekby

internal/topic/retriable_error.go

@@ -1,8 +1,10 @@
 package topic
 
 import (
+	"fmt"
 	"time"
 
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/backoff"
 	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
 )
 
@@ -15,3 +17,36 @@ func CheckResetReconnectionCounters(lastTry, now time.Time, connectionTimeout ti
 	const resetAttemptEmpiricalCoefficient = 10
 	return now.Sub(lastTry) > connectionTimeout*resetAttemptEmpiricalCoefficient
 }
+
+func CheckRetryMode(err error, isPreCheck bool, currentAttempt int, customCheckFunc PublicCheckRetryFunc) (
+	_ backoff.Backoff,
+	isRetriable bool,
+) {
+	isRetriable = true
+
+	decision := PublicRetryDecisionDefault
+	if customCheckFunc != nil {
+		decision = customCheckFunc(NewCheckRetryArgs(isPreCheck, currentAttempt, err))
+	}
+
+	switch decision {
+	case PublicRetryDecisionDefault:
+		isRetriable = IsRetryableError(err)
+	case PublicRetryDecisionRetry:
+		isRetriable = true
+	case PublicRetryDecisionStop:
+		isRetriable = false
+	default:
+		panic(fmt.Errorf("unexpected retry decision: %v", decision))
+	}
+
+	if !isRetriable {
+		return nil, false
+	}
+
+	if retry.Check(err).BackoffType() == backoff.TypeFast {
+		return backoff.Fast, true
+	}
+
+	return backoff.Slow, true
+}

internal/topic/topicreaderinternal/stream_reconnector.go

@@ -18,7 +18,6 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsync"
-	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
@@ -254,46 +253,15 @@ func (r *readerReconnector) reconnect(ctx context.Context, oldReader batchedStre
 }
 
 func (r *readerReconnector) isRetriableError(err error) bool {
-	_, res := r.checkErrRetryModeLogic(true, 0, err)
+	_, res := topic.CheckRetryMode(err, true, 0, r.retrySettings.CheckError)
 	return res
 }
 
 func (r *readerReconnector) checkErrRetryMode(attempts int, err error) (
 	backoffType backoff.Backoff,
 	isRetriableErr bool,
 ) {
-	return r.checkErrRetryModeLogic(false, attempts, err)
-}
-
-func (r *readerReconnector) checkErrRetryModeLogic(preCheck bool, attempts int, err error) (
-	backoffType backoff.Backoff,
-	isRetriableErr bool,
-) {
-	isRetriableErr = true
-	if r.retrySettings.CheckError != nil {
-		switch decision := r.retrySettings.CheckError(topic.NewCheckRetryArgs(preCheck, attempts, err)); decision {
-		case topic.PublicRetryDecisionDefault:
-			isRetriableErr = topic.IsRetryableError(err)
-		case topic.PublicRetryDecisionRetry:
-			isRetriableErr = true
-		case topic.PublicRetryDecisionStop:
-			isRetriableErr = false
-		default:
-			panic(fmt.Errorf("unexpected retry decision: %v", decision))
-		}
-	}
-	if !isRetriableErr {
-		return nil, false
-	}
-	if preCheck {
-		return nil, true
-	}
-
-	if retry.Check(err).BackoffType() == backoff.TypeFast {
-		return backoff.Fast, true
-	}
-
-	return backoff.Slow, true
+	return topic.CheckRetryMode(err, false, attempts, r.retrySettings.CheckError)
 }
 
 func (r *readerReconnector) connectWithTimeout() (_ batchedStreamReader, err error) {

internal/topic/topicwriterinternal/writer_reconnector.go

@@ -15,7 +15,6 @@ import (
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/credentials"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/background"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/backoff"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/config"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/empty"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/grpcwrapper/rawtopic/rawtopiccommon"
@@ -104,7 +103,8 @@ func newWriterReconnectorConfig(options ...PublicWriterOption) WriterReconnector
 }
 
 type WriterReconnector struct {
-	cfg WriterReconnectorConfig
+	cfg           WriterReconnectorConfig
+	retrySettings topic.RetrySettings
 
 	semaphore                      *semaphore.Weighted
 	queue                          messageQueue
@@ -337,8 +337,8 @@ func (w *WriterReconnector) connectionLoop(ctx context.Context) {
 		streamCtxCancel(xerrors.WithStackTrace(errCloseWriterReconnectorConnectionLoop))
 	}()
 
-	var prevConnectionError error
-	var prevConnectionTime time.Time
+	var reconnectReason error
+	var prevAttemptTime time.Time
 
 	for {
 		if ctx.Err() != nil {
@@ -348,54 +348,61 @@ func (w *WriterReconnector) connectionLoop(ctx context.Context) {
 		streamCtxCancel(xerrors.WithStackTrace(errCloseWriterReconnectorReconnect))
 		streamCtx, streamCtxCancel = createStreamContext()
 
-		attempt++
-
-		// delay if reconnect
-		if attempt > 1 {
-			delay := backoff.Fast.Delay(attempt - 2)
-			select {
-			case <-doneCtx:
-				return
-			case <-w.clock.After(delay):
-				// pass
-			}
-		}
-
-		traceOnDone := trace.TopicOnWriterReconnect(
-			w.cfg.tracer,
-			w.writerInstanceID,
-			w.cfg.topic,
-			w.cfg.producerID,
-			attempt,
-		)
-
-		stream, trackedErr := w.connectWithTimeout(streamCtx)
-		traceOnDone(trackedErr)
-
-		var writer *SingleStreamWriter
-		if trackedErr == nil {
+		now := time.Now()
+		if topic.CheckResetReconnectionCounters(prevAttemptTime, now, w.cfg.connectTimeout) {
 			attempt = 0
-			w.queue.ResetSentProgress()
-			writer, trackedErr = NewSingleStreamWriter(ctx, w.createWriterStreamConfig(stream))
+		} else {
+			attempt++
 		}
-		if trackedErr == nil {
-			w.onWriterChange(writer)
-			trackedErr = writer.CloseWait(ctx)
+		prevAttemptTime = now
+
+		if reconnectReason != nil {
+			if backoff, retry := topic.CheckRetryMode(reconnectReason, false, attempt, w.retrySettings.CheckError); retry {
+				delay := backoff.Delay(attempt)
+				select {
+				case <-doneCtx:
+					return
+				case <-w.clock.After(delay):
+					// pass
+				}
+			} else {
+				_ = w.close(ctx, reconnectReason)
+				return
+			}
 		}
-		w.onWriterChange(nil)
 
-		if !w.isRetriableErr(trackedErr) {
-			closeCtx, cancel := context.WithCancel(ctx)
-			cancel()
-			_ = w.close(closeCtx, trackedErr)
-			return
+		writer, err := w.startWriteStream(ctx, streamCtx, attempt)
+		w.onWriterChange(writer)
+		if err == nil {
+			reconnectReason = writer.WaitClose(ctx)
+		} else {
+			reconnectReason = err
 		}
-		// next iteration
 	}
 }
 
-func (w *WriterReconnector) isRetriableErr(err error) bool {
-	return topic.IsRetryableError(err)
+func (w *WriterReconnector) startWriteStream(ctx context.Context, streamCtx context.Context, attempt int) (
+	writer *SingleStreamWriter,
+	err error,
+) {
+	traceOnDone := trace.TopicOnWriterReconnect(
+		w.cfg.tracer,
+		w.writerInstanceID,
+		w.cfg.topic,
+		w.cfg.producerID,
+		attempt,
+	)
+	defer func() {
+		traceOnDone(err)
+	}()
+
+	stream, err := w.connectWithTimeout(streamCtx)
+	if err != nil {
+		return nil, err
+	}
+
+	w.queue.ResetSentProgress()
+	return NewSingleStreamWriter(ctx, w.createWriterStreamConfig(stream))
 }
 
 func (w *WriterReconnector) needReceiveLastSeqNo() bool {

internal/topic/topicwriterinternal/writer_reconnector_test.go

@@ -774,6 +774,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"))
 		close(res.stopReadEvents)
 		<-streamClosed
 	})

internal/topic/topicwriterinternal/writer_single_stream.go

@@ -103,7 +103,7 @@ func (w *SingleStreamWriter) close(ctx context.Context, reason error) error {
 	return resErr
 }
 
-func (w *SingleStreamWriter) CloseWait(ctx context.Context) error {
+func (w *SingleStreamWriter) WaitClose(ctx context.Context) error {
 	select {
 	case <-ctx.Done():
 		return ctx.Err()