go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/common/retry/limited.go (about)

     1  // Copyright 2016 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package retry
    16  
    17  import (
    18  	"context"
    19  	"time"
    20  
    21  	"go.chromium.org/luci/common/clock"
    22  )
    23  
    24  // Limited is an Iterator implementation that may be limited by a maximum number
    25  // of retries and/or time.
    26  type Limited struct {
    27  	// Delay is the next generated delay.
    28  	Delay time.Duration
    29  
    30  	// Retries, if >= 0, is the number of remaining retries. If <0, no retry
    31  	// count will be applied.
    32  	Retries int
    33  
    34  	// MaxTotal is the maximum total elapsed time. If <= 0, no maximum will be
    35  	// enforced.
    36  	MaxTotal time.Duration
    37  
    38  	// The time when the generator initially started.
    39  	startTime time.Time
    40  }
    41  
    42  var _ Iterator = (*Limited)(nil)
    43  
    44  // Next implements the Iterator interface.
    45  func (i *Limited) Next(ctx context.Context, _ error) time.Duration {
    46  	switch {
    47  	case i.Retries == 0:
    48  		return Stop
    49  	case i.Retries > 0:
    50  		i.Retries--
    51  	}
    52  
    53  	// If there is a maximum total time, enforce it.
    54  	if i.MaxTotal > 0 {
    55  		now := clock.Now(ctx)
    56  		if i.startTime.IsZero() {
    57  			i.startTime = now
    58  		}
    59  
    60  		var elapsed time.Duration
    61  		if now.After(i.startTime) {
    62  			elapsed = now.Sub(i.startTime)
    63  		}
    64  
    65  		// Remaining time is the difference between total allowed time and elapsed
    66  		// time.
    67  		remaining := i.MaxTotal - elapsed
    68  		if remaining <= 0 {
    69  			// No more time!
    70  			i.Retries = 0
    71  			return Stop
    72  		}
    73  	}
    74  
    75  	return i.Delay
    76  }