From 2e81c5f740198a2d63ecb68c6ed2cf0bc4bb2d9d Mon Sep 17 00:00:00 2001
From: Dmytro Kozlov <kozlovdmitriyy@gmail.com>
Date: Tue, 13 Jun 2023 14:54:24 +0300
Subject: [PATCH] vmctl: finish retries if context canceled (#4442)

vmctl: interrupt backoff retries if import context is cancelled

Co-authored-by: Roman Khavronenko <roman@victoriametrics.com>
---
 app/vmctl/backoff/backoff.go      | 15 +++++++++++++--
 app/vmctl/backoff/backoff_test.go | 25 ++++++++++++++++++++++++-
 docs/CHANGELOG.md                 |  1 +
 3 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/app/vmctl/backoff/backoff.go b/app/vmctl/backoff/backoff.go
index 431d6a0c4..9dc9637e8 100644
--- a/app/vmctl/backoff/backoff.go
+++ b/app/vmctl/backoff/backoff.go
@@ -42,7 +42,6 @@ func New() *Backoff {
 func (b *Backoff) Retry(ctx context.Context, cb retryableFunc) (uint64, error) {
 	var attempt uint64
 	for i := 0; i < b.retries; i++ {
-		// @TODO we should use context to cancel retries
 		err := cb()
 		if err == nil {
 			return attempt, nil
@@ -55,7 +54,19 @@ func (b *Backoff) Retry(ctx context.Context, cb retryableFunc) (uint64, error) {
 		backoff := float64(b.minDuration) * math.Pow(b.factor, float64(i))
 		dur := time.Duration(backoff)
 		logger.Errorf("got error: %s on attempt: %d; will retry in %v", err, attempt, dur)
-		time.Sleep(time.Duration(backoff))
+
+		t := time.NewTimer(dur)
+		select {
+		case <-t.C:
+			// duration elapsed, loop
+		case <-ctx.Done():
+			// context cancelled, kill the timer if it hasn't fired, and return
+			// the last error we got
+			if !t.Stop() {
+				<-t.C
+			}
+			return attempt, err
+		}
 	}
 	return attempt, fmt.Errorf("execution failed after %d retry attempts", b.retries)
 }
diff --git a/app/vmctl/backoff/backoff_test.go b/app/vmctl/backoff/backoff_test.go
index e9972a27e..46556d2d2 100644
--- a/app/vmctl/backoff/backoff_test.go
+++ b/app/vmctl/backoff/backoff_test.go
@@ -16,7 +16,7 @@ func TestRetry_Do(t *testing.T) {
 		backoffMinDuration time.Duration
 		retryableFunc      retryableFunc
 		ctx                context.Context
-		withCancel         bool
+		cancelTimeout      time.Duration
 		want               uint64
 		wantErr            bool
 	}{
@@ -79,10 +79,33 @@ func TestRetry_Do(t *testing.T) {
 			want:    5,
 			wantErr: true,
 		},
+		{
+			name:               "cancel context",
+			backoffRetries:     5,
+			backoffFactor:      0.1,
+			backoffMinDuration: time.Millisecond * 10,
+			retryableFunc: func() error {
+				t := time.NewTicker(time.Millisecond * 5)
+				defer t.Stop()
+				for range t.C {
+					return fmt.Errorf("got some error")
+				}
+				return nil
+			},
+			ctx:           context.Background(),
+			cancelTimeout: time.Second * 5,
+			want:          3,
+			wantErr:       true,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			r := New()
+			if tt.cancelTimeout != 0 {
+				newCtx, cancelFn := context.WithTimeout(tt.ctx, tt.cancelTimeout)
+				tt.ctx = newCtx
+				defer cancelFn()
+			}
 			got, err := r.Retry(tt.ctx, tt.retryableFunc)
 			if (err != nil) != tt.wantErr {
 				t.Errorf("Retry() error = %v, wantErr %v", err, tt.wantErr)
diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
index d39dd5ea8..06bfe5364 100644
--- a/docs/CHANGELOG.md
+++ b/docs/CHANGELOG.md
@@ -27,6 +27,7 @@ The following tip changes can be tested by building VictoriaMetrics components f
 * SECURITY: upgrade Go builder from Go1.20.4 to Go1.20.5. See [the list of issues addressed in Go1.20.5](https://github.com/golang/go/issues?q=milestone%3AGo1.20.5+label%3ACherryPickApproved).
 
 * FEATURE: [vmctl](https://docs.victoriametrics.com/vmctl.html): add verbose output for docker installations or when TTY isn't available. See [this issue](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/4081).
+* FEATURE: [vmctl](https://docs.victoriametrics.com/vmctl.html): interrupt backoff retries when import process is cancelled. The change makes vmctl more responsive in case of errors during the import. See [this pull request](https://github.com/VictoriaMetrics/VictoriaMetrics/pull/4442).
 * FEATURE: vmstorage: suppress "broken pipe" errors for search queries on vmstorage side. See [this commit](https://github.com/VictoriaMetrics/VictoriaMetrics/pull/4418/commits/a6a7795b9e1f210d614a2c5f9a3016b97ded4792).
 * FEATURE: [Official Grafana dashboards for VictoriaMetrics](https://grafana.com/orgs/victoriametrics): add panel for tracking rate of syscalls while writing or reading from disk via `process_io_(read|write)_syscalls_total` metrics.
 * FEATURE: add ability to fine-tune Graphite API limits via the following command-line flags: