retry.go

Documentation: github.com/google/s2a-go/retry

     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
    22  package retry
    23  
    24  import (
    25  	"context"
    26  	"math/rand"
    27  	"time"
    28  
    29  	"google.golang.org/grpc/grpclog"
    30  )
    31  
    32  const (
    33  	maxRetryAttempts = 5
    34  	maxRetryForLoops = 10
    35  )
    36  
    37  type 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.
    45  func (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.
    56  func 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.
    69  var NewRetryer = func() *S2ARetryer {
    70  	return &S2ARetryer{bo: &defaultBackoff{
    71  		cur: 100 * time.Millisecond,
    72  		max: 30 * time.Second,
    73  		mul: 2,
    74  	}}
    75  }
    76  
    77  type backoff interface {
    78  	Pause() time.Duration
    79  }
    80  
    81  // S2ARetryer implements a retry helper for talking to S2A gRPC server.
    82  type S2ARetryer struct {
    83  	bo       backoff
    84  	attempts int
    85  }
    86  
    87  // Attempts return the number of retries attempted.
    88  func (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.
    94  func (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.
   107  func 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  }
   145
View as plain text