fixed race on write err in discovery.go + refactored repeater API

Author: asmyasnikov

CHANGELOG.md

@@ -8,7 +8,8 @@
 * Removed `internal/lazy.Discovery` (discovery client always initialized)
 * Fixed `trace.Table` event structs
 * Refactored grpc options for define dns-balancing configuration
-* Refactored `retry.Retry` signature (added `retry.WithID`, `retry.WithTrace` and `retry.WithIdempotent` opt-in args, required param `isIdempotentOperation` removed) 
+* Refactored `retry.Retry` signature (added `retry.WithID`, `retry.WithTrace` and `retry.WithIdempotent` opt-in args, required param `isIdempotentOperation` removed)
+* Refactored package `internal/repeater`
 
 ## 3.9.4
 * Fixed data race on closing session pool

internal/discovery/discovery.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"net"
 	"strconv"
-	"time"
 
 	"google.golang.org/protobuf/proto"
 
@@ -57,15 +56,14 @@ func New(
 		repeater.NewRepeater(
 			deadline.ContextWithoutDeadline(ctx),
 			c.config.Interval(),
-			func(ctx context.Context) {
+			func(ctx context.Context) (err error) {
 				next, err = c.Discover(ctx)
-				// if nothing endpoint - re-discover after one second
-				// and use old endpoint list
-				if err != nil || len(next) == 0 {
-					go func() {
-						time.Sleep(time.Second)
-						crudExplorer.Force()
-					}()
+				if err != nil {
+					return err
+				}
+
+				// if nothing endpoint - use old endpoint list
+				if len(next) == 0 {
 					return
 				}
 
@@ -94,7 +92,10 @@ func New(
 						)
 					},
 				)
+
 				curr = next
+
+				return nil
 			},
 		),
 	)

internal/repeater/repeat.go

@@ -20,7 +20,7 @@ type repeater struct {
 	interval time.Duration
 
 	// Task is a function that must be executed periodically.
-	task func(context.Context)
+	task func(context.Context) error
 
 	timer    timeutil.Timer
 	stopOnce sync.Once
@@ -35,7 +35,7 @@ type repeater struct {
 func NewRepeater(
 	ctx context.Context,
 	interval time.Duration,
-	task func(ctx context.Context),
+	task func(ctx context.Context) (err error),
 ) Repeater {
 	if interval <= 0 {
 		return nil
@@ -72,22 +72,30 @@ func (r *repeater) Force() {
 	}
 }
 
+func (r *repeater) singleTask() {
+	if err := r.task(r.ctx); err != nil {
+		r.timer.Reset(time.Second)
+	} else {
+		r.timer.Reset(r.interval)
+	}
+}
+
 func (r *repeater) worker() {
 	defer func() {
 		close(r.done)
 	}()
-	r.task(r.ctx)
+	r.singleTask()
 	for {
 		select {
 		case <-r.stop:
 			return
 		case <-r.timer.C():
-			r.task(r.ctx)
+			r.singleTask()
 		case <-r.force:
 			if !r.timer.Stop() {
 				<-r.timer.C()
 			}
+			r.timer.Reset(r.interval)
 		}
-		r.timer.Reset(r.interval)
 	}
 }

internal/repeater/test/repeat_test.go

@@ -28,8 +28,9 @@ func TestRepeater(t *testing.T) {
 
 	exec := make(chan struct{})
 	r := repeater.NewRepeater(ctx, 42*time.Second,
-		func(_ context.Context) {
+		func(_ context.Context) error {
 			exec <- struct{}{}
+			return nil
 		})
 
 	timerC <- time.Now()
@@ -59,10 +60,11 @@ func TestRepeaterCancellation(t *testing.T) {
 	defer cancel()
 
 	r := repeater.NewRepeater(ctx, 42*time.Second,
-		func(ctx context.Context) {
+		func(ctx context.Context) error {
 			enter <- struct{}{}
 			<-ctx.Done()
 			exit <- struct{}{}
+			return nil
 		})
 
 	// Run callback in a separate goroutine to avoid deadlock.
@@ -104,12 +106,14 @@ func noRecv(ch interface{}, timeout time.Duration) error {
 }
 
 func assertRecv(t *testing.T, timeout time.Duration, ch interface{}) {
+	t.Helper()
 	if err := recv(ch, timeout); err != nil {
 		t.Fatalf("%s: %v", testutil.FileLine(2), err)
 	}
 }
 
 func assertNoRecv(t *testing.T, timeout time.Duration, ch interface{}) {
+	t.Helper()
 	if err := noRecv(ch, timeout); err != nil {
 		t.Fatalf("%s: %v", testutil.FileLine(2), err)
 	}