retry.go

  1/*
  2 *
  3 * Copyright 2023 Google LLC
  4 *
  5 * Licensed under the Apache License, Version 2.0 (the "License");
  6 * you may not use this file except in compliance with the License.
  7 * You may obtain a copy of the License at
  8 *
  9 *     https://www.apache.org/licenses/LICENSE-2.0
 10 *
 11 * Unless required by applicable law or agreed to in writing, software
 12 * distributed under the License is distributed on an "AS IS" BASIS,
 13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 14 * See the License for the specific language governing permissions and
 15 * limitations under the License.
 16 *
 17 */
 18
 19// Package retry provides a retry helper for talking to S2A gRPC server.
 20// The implementation is modeled after
 21// https://github.com/googleapis/google-cloud-go/blob/main/compute/metadata/retry.go
 22package retry
 23
 24import (
 25	"context"
 26	"math/rand"
 27	"time"
 28
 29	"google.golang.org/grpc/grpclog"
 30)
 31
 32const (
 33	maxRetryAttempts = 5
 34	maxRetryForLoops = 10
 35)
 36
 37type defaultBackoff struct {
 38	max time.Duration
 39	mul float64
 40	cur time.Duration
 41}
 42
 43// Pause returns a duration, which is used as the backoff wait time
 44// before the next retry.
 45func (b *defaultBackoff) Pause() time.Duration {
 46	d := time.Duration(1 + rand.Int63n(int64(b.cur)))
 47	b.cur = time.Duration(float64(b.cur) * b.mul)
 48	if b.cur > b.max {
 49		b.cur = b.max
 50	}
 51	return d
 52}
 53
 54// Sleep will wait for the specified duration or return on context
 55// expiration.
 56func Sleep(ctx context.Context, d time.Duration) error {
 57	t := time.NewTimer(d)
 58	select {
 59	case <-ctx.Done():
 60		t.Stop()
 61		return ctx.Err()
 62	case <-t.C:
 63		return nil
 64	}
 65}
 66
 67// NewRetryer creates an instance of S2ARetryer using the defaultBackoff
 68// implementation.
 69var NewRetryer = func() *S2ARetryer {
 70	return &S2ARetryer{bo: &defaultBackoff{
 71		cur: 100 * time.Millisecond,
 72		max: 30 * time.Second,
 73		mul: 2,
 74	}}
 75}
 76
 77type backoff interface {
 78	Pause() time.Duration
 79}
 80
 81// S2ARetryer implements a retry helper for talking to S2A gRPC server.
 82type S2ARetryer struct {
 83	bo       backoff
 84	attempts int
 85}
 86
 87// Attempts return the number of retries attempted.
 88func (r *S2ARetryer) Attempts() int {
 89	return r.attempts
 90}
 91
 92// Retry returns a boolean indicating whether retry should be performed
 93// and the backoff duration.
 94func (r *S2ARetryer) Retry(err error) (time.Duration, bool) {
 95	if err == nil {
 96		return 0, false
 97	}
 98	if r.attempts >= maxRetryAttempts {
 99		return 0, false
100	}
101	r.attempts++
102	return r.bo.Pause(), true
103}
104
105// Run uses S2ARetryer to execute the function passed in, until success or reaching
106// max number of retry attempts.
107func Run(ctx context.Context, f func() error) {
108	retryer := NewRetryer()
109	forLoopCnt := 0
110	var err error
111	for {
112		err = f()
113		if bo, shouldRetry := retryer.Retry(err); shouldRetry {
114			if grpclog.V(1) {
115				grpclog.Infof("will attempt retry: %v", err)
116			}
117			if ctx.Err() != nil {
118				if grpclog.V(1) {
119					grpclog.Infof("exit retry loop due to context error: %v", ctx.Err())
120				}
121				break
122			}
123			if errSleep := Sleep(ctx, bo); errSleep != nil {
124				if grpclog.V(1) {
125					grpclog.Infof("exit retry loop due to sleep error: %v", errSleep)
126				}
127				break
128			}
129			// This shouldn't happen, just make sure we are not stuck in the for loops.
130			forLoopCnt++
131			if forLoopCnt > maxRetryForLoops {
132				if grpclog.V(1) {
133					grpclog.Infof("exit the for loop after too many retries")
134				}
135				break
136			}
137			continue
138		}
139		if grpclog.V(1) {
140			grpclog.Infof("retry conditions not met, exit the loop")
141		}
142		break
143	}
144}