github.com/grafana/pyroscope@v1.18.0/pkg/util/retry/hedged_test.go (about)

     1  package retry
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"testing"
     7  	"time"
     8  )
     9  
    10  func Test_Hedging(t *testing.T) {
    11  	delay := time.Millisecond * 100 // Might be flaky.
    12  	ctx := context.Background()
    13  	e1 := errors.New("e1")
    14  	e2 := errors.New("e2")
    15  
    16  	type attempt struct {
    17  		duration time.Duration
    18  		err      error
    19  	}
    20  
    21  	type testCase struct {
    22  		description string
    23  		failFast    bool
    24  		attempts    [2]attempt
    25  		expectRetry bool
    26  		expectError error
    27  	}
    28  
    29  	createCall := func(c testCase) func(context.Context, bool) (bool, error) {
    30  		return func(ctx context.Context, isRetry bool) (bool, error) {
    31  			d := c.attempts[0].duration
    32  			e := c.attempts[0].err
    33  			if isRetry {
    34  				d = c.attempts[1].duration
    35  				e = c.attempts[1].err
    36  			}
    37  			select {
    38  			case <-time.After(d):
    39  			case <-ctx.Done():
    40  			}
    41  			return isRetry, e
    42  		}
    43  	}
    44  
    45  	testCases := []testCase{
    46  		{
    47  			description: "Attempt fails before retry and FailFast",
    48  			failFast:    true,
    49  			attempts: [2]attempt{
    50  				{delay / 2, e1},
    51  				{delay / 2, nil},
    52  			},
    53  			expectRetry: false,
    54  			expectError: e1,
    55  		},
    56  		{
    57  			description: "Attempt fails before retry and not FailFast",
    58  			failFast:    false,
    59  			attempts: [2]attempt{
    60  				{delay / 2, e1},
    61  				{delay / 2, nil},
    62  			},
    63  			expectRetry: true,
    64  			expectError: nil,
    65  		},
    66  
    67  		{
    68  			description: "Attempt fails before retry and retry fails and FailFast",
    69  			failFast:    true,
    70  			attempts: [2]attempt{
    71  				{delay / 2, e1},
    72  				{delay / 2, e2},
    73  			},
    74  			expectRetry: false,
    75  			expectError: e1,
    76  		},
    77  		{
    78  			description: "Attempt fails before retry and retry fails and not FailFast",
    79  			failFast:    false,
    80  			attempts: [2]attempt{
    81  				{delay / 2, e1},
    82  				{delay / 2, e2},
    83  			},
    84  			expectRetry: true,
    85  			expectError: e2,
    86  		},
    87  
    88  		{
    89  			description: "Attempt fails after retry and FailFast",
    90  			failFast:    true,
    91  			attempts: [2]attempt{
    92  				{delay * 2, e1},
    93  				{delay * 2, nil},
    94  			},
    95  			expectRetry: false,
    96  			expectError: e1,
    97  		},
    98  		{
    99  			description: "Attempt fails after retry and not FailFast",
   100  			failFast:    false,
   101  			attempts: [2]attempt{
   102  				{delay * 2, e1},
   103  				{delay * 2, nil},
   104  			},
   105  			expectRetry: true,
   106  			expectError: nil,
   107  		},
   108  
   109  		{
   110  			description: "Attempt fails after retry and retry fails and FailFast",
   111  			failFast:    true,
   112  			attempts: [2]attempt{
   113  				{delay * 2, e1},
   114  				{delay * 2, e2},
   115  			},
   116  			expectRetry: false,
   117  			expectError: e1,
   118  		},
   119  		{
   120  			description: "Attempt fails after retry and retry fails and not FailFast",
   121  			failFast:    false,
   122  			attempts: [2]attempt{
   123  				{delay * 2, e1},
   124  				{delay * 2, e2},
   125  			},
   126  			expectRetry: true,
   127  			expectError: e2,
   128  		},
   129  	}
   130  
   131  	for _, c := range testCases {
   132  		c := c
   133  		t.Run(c.description, func(t *testing.T) {
   134  			t.Parallel()
   135  			// With FailFast, it's expected the call returns
   136  			// immediately after the first response. Increasing
   137  			// the delay to make sure the test doesn't flake.
   138  			d := delay
   139  			if c.failFast {
   140  				d = time.Second * 10
   141  			}
   142  			a := Hedged[bool]{
   143  				Call:     createCall(c),
   144  				Trigger:  time.After(d),
   145  				FailFast: c.failFast,
   146  			}
   147  			r, err := a.Do(ctx)
   148  			if r != c.expectRetry {
   149  				t.Fatal("expected retry")
   150  			}
   151  			if !errors.Is(err, c.expectError) {
   152  				t.Fatal("expected error", c.expectError)
   153  			}
   154  		})
   155  	}
   156  }