feat(v2): add WithTimeout option (#259)

v2/call_option.go

@@ -218,6 +218,14 @@ func (p pathOpt) Resolve(s *CallSettings) {
 	s.Path = p.p
 }
 
+type timeoutOpt struct {
+	t time.Duration
+}
+
+func (t timeoutOpt) Resolve(s *CallSettings) {
+	s.timeout = t.t
+}
+
 // WithPath applies a Path override to the HTTP-based APICall.
 //
 // This is for internal use only.
@@ -230,6 +238,15 @@ func WithGRPCOptions(opt ...grpc.CallOption) CallOption {
 	return grpcOpt(append([]grpc.CallOption(nil), opt...))
 }
 
+// WithTimeout is a convenience option for setting a context.WithTimeout on the
+// singular context.Context used for **all** APICall attempts. Calculated from
+// the start of the first APICall attempt.
+// If the context.Context provided to Invoke already has a Deadline set, that
+// will always be respected over the deadline calculated using this option.
+func WithTimeout(t time.Duration) CallOption {
+	return &timeoutOpt{t: t}
+}
+
 // CallSettings allow fine-grained control over how calls are made.
 type CallSettings struct {
 	// Retry returns a Retryer to be used to control retry logic of a method call.
@@ -241,4 +258,8 @@ type CallSettings struct {
 
 	// Path is an HTTP override for an APICall.
 	Path string
+
+	// Timeout defines the amount of time that Invoke has to complete.
+	// Unexported so it cannot be changed by the code in an APICall.
+	timeout time.Duration
 }

v2/call_option_test.go

@@ -129,3 +129,14 @@ func TestOnHTTPCodes(t *testing.T) {
 		}
 	}
 }
+
+func TestWithTimeout(t *testing.T) {
+	settings := CallSettings{}
+	to := 10 * time.Second
+
+	WithTimeout(to).Resolve(&settings)
+
+	if settings.timeout != to {
+		t.Errorf("got %v, want %v", settings.timeout, to)
+	}
+}

v2/invoke.go

@@ -68,6 +68,16 @@ type sleeper func(ctx context.Context, d time.Duration) error
 // invoke implements Invoke, taking an additional sleeper argument for testing.
 func invoke(ctx context.Context, call APICall, settings CallSettings, sp sleeper) error {
 	var retryer Retryer
+
+	// Only use the value provided via WithTimeout if the context doesn't
+	// already have a deadline. This is important for backwards compatibility if
+	// the user already set a deadline on the context given to Invoke.
+	if _, ok := ctx.Deadline(); !ok && settings.timeout != 0 {
+		c, cc := context.WithTimeout(ctx, settings.timeout)
+		defer cc()
+		ctx = c
+	}
+
 	for {
 		err := call(ctx, settings)
 		if err == nil {

v2/invoke_test.go

@@ -202,3 +202,65 @@ func TestInvokeRetryTimeout(t *testing.T) {
 		t.Errorf("found error %s, want %s", err, context.Canceled)
 	}
 }
+
+func TestInvokeWithTimeout(t *testing.T) {
+	// Dummy APICall that sleeps for the given amount of time. This simulates an
+	// APICall executing, allowing us to verify which deadline was respected,
+	// that which is already set on the Context, or the one calculated using the
+	// WithTimeout option's value.
+	sleepingCall := func(sleep time.Duration) APICall {
+		return func(ctx context.Context, _ CallSettings) error {
+			time.Sleep(sleep)
+			return ctx.Err()
+		}
+	}
+
+	bg := context.Background()
+	preset, pcc := context.WithTimeout(bg, 10*time.Millisecond)
+	defer pcc()
+
+	for _, tst := range []struct {
+		name    string
+		timeout time.Duration
+		sleep   time.Duration
+		ctx     context.Context
+		want    error
+	}{
+		{
+			name:    "success",
+			timeout: 10 * time.Millisecond,
+			sleep:   1 * time.Millisecond,
+			ctx:     bg,
+			want:    nil,
+		},
+		{
+			name:    "respect_context_deadline",
+			timeout: 1 * time.Millisecond,
+			sleep:   3 * time.Millisecond,
+			ctx:     preset,
+			want:    nil,
+		},
+		{
+			name:    "with_timeout_deadline_exceeded",
+			timeout: 1 * time.Millisecond,
+			sleep:   3 * time.Millisecond,
+			ctx:     bg,
+			want:    context.DeadlineExceeded,
+		},
+	} {
+		t.Run(tst.name, func(t *testing.T) {
+			// Recording sleep isn't really necessary since there is
+			// no retry here, but we need a sleeper so might as well.
+			var sp recordSleeper
+			var settings CallSettings
+
+			WithTimeout(tst.timeout).Resolve(&settings)
+
+			err := invoke(tst.ctx, sleepingCall(tst.sleep), settings, sp.sleep)
+
+			if err != tst.want {
+				t.Errorf("found error %v, want %v", err, tst.want)
+			}
+		})
+	}
+}