context: add APIs for setting a cancelation cause when deadline or timer expires

Fixes #56661

Change-Id: I1c23ebc52e6b7ae6ee956614e1a0a45d6ecbd5b4
Reviewed-on: https://go-review.googlesource.com/c/go/+/449318
Run-TryBot: Sameer Ajmani <sameer@golang.org>
Reviewed-by: Damien Neil <dneil@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>

1 files changed, 16 insertions, 2 deletions

diff --git a/src/context/context.go b/src/context/context.go
index f3fe1a474e..6bf6ec8dcc 100644
--- a/src/context/context.go
+++ b/src/context/context.go
@@ -492,6 +492,13 @@ func (c *cancelCtx) cancel(removeFromParent bool, err, cause error) {
 // Canceling this context releases resources associated with it, so code should
 // call cancel as soon as the operations running in this Context complete.
 func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
+	return WithDeadlineCause(parent, d, nil)
+}
+
+// WithDeadlineCause behaves like WithDeadline but also sets the cause of the
+// returned Context when the deadline is exceeded. The returned CancelFunc does
+// not set the cause.
+func WithDeadlineCause(parent Context, d time.Time, cause error) (Context, CancelFunc) {
 	if parent == nil {
 		panic("cannot create context from nil parent")
 	}
@@ -506,14 +513,14 @@ func WithDeadline(parent Context, d time.Time) (Context, CancelFunc) {
 	propagateCancel(parent, c)
 	dur := time.Until(d)
 	if dur <= 0 {
-		c.cancel(true, DeadlineExceeded, nil) // deadline has already passed
+		c.cancel(true, DeadlineExceeded, cause) // deadline has already passed
 		return c, func() { c.cancel(false, Canceled, nil) }
 	}
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	if c.err == nil {
 		c.timer = time.AfterFunc(dur, func() {
-			c.cancel(true, DeadlineExceeded, nil)
+			c.cancel(true, DeadlineExceeded, cause)
 		})
 	}
 	return c, func() { c.cancel(true, Canceled, nil) }
@@ -567,6 +574,13 @@ func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
 	return WithDeadline(parent, time.Now().Add(timeout))
 }
 
+// WithTimeoutCause behaves like WithTimeout but also sets the cause of the
+// returned Context when the timout expires. The returned CancelFunc does
+// not set the cause.
+func WithTimeoutCause(parent Context, timeout time.Duration, cause error) (Context, CancelFunc) {
+	return WithDeadlineCause(parent, time.Now().Add(timeout), cause)
+}
+
 // WithValue returns a copy of parent in which the value associated with key is
 // val.
 //