contextutil: add CancelWithReason

[jordanlewis]

contextutil.WithCancel now returns a new context type that holds its
reason for cancellation when canceled with CancelWithReason instead of
its normal cancelfunc.

Alternative / new approach to #29318.

Release note: None

[cockroach-teamcity]

This change is 

[petermattis]

Reviewable status: complete! 0 of 0 LGTMs obtained

---

pkg/util/contextutil/context.go, line 51 at r1 (raw file):

// cancellation given in err if the context was created with
// contextutil.WithCancel.
func CancelWithReason(ctx context.Context, cancelFunc context.CancelFunc, err error) {

The interface here feels a bit fragile. Specifically, it looks easy to call contextutil.WithCancel and then forget to call CancelWithReason and merely invoke the context.CancelFunc directly. Perhaps we should define type CancelWithReasonFunc func(error) and return that from contextutil.WithCancel. That would force all callers to pass in a reason or nil.

Another drawback to this approach is that you're only capturing explicitly cancellation, not cancellation due to deadline exceeded. Perhaps that's not a big deal because we can already get context.DeadlineExceeded.

A more significant drawback (flaw?) is that I'm not sure reasonCancelCtx is kosher because the context implementation wants to be able to type assert to specific implementations in order to find the parent cancellation context.

// parentCancelCtx follows a chain of parent references until it finds a
// *cancelCtx. This function understands how each of the concrete types in this
// package represents its parent.
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
	for {
		switch c := parent.(type) {
		case *cancelCtx:
			return c, true
		case *timerCtx:
			return &c.cancelCtx, true
		case *valueCtx:
			parent = c.Context
		default:
			return nil, false
		}
	}
}

By wrapping a context with reasonCancelCtx you're preventing that function from working which would break the cancellation chain. I think. Easy enough to test.

[dt]

As discussed offline, I also don't think we want to return our fancy error from .Err() since some callers, including in libraries we use, compare .Err()'s return to the exported sentinel context.Canceled. I think we need a new method that returns the fancy error, and then our callsites that want helpful error (but don't plan to compare it to the stdlib's sentinel) need to be aware of that method and call it instead, maybe via a type-sniffing helper like contextutil.Err(ctx) or something.

[petermattis]

Here's a concrete idea that addresses both my concern and David's: create our base contexts with a special WithValue context that has a pointer to a reason. We populate the reason when the cancel method is invoked and we can retrieve it later with a GetReason method.

type reasonKey struct{}

type CancelWithReasonFunc func(string)

func WithReason(ctx context.Context) (context.Context, CancelWithReasonFunc) {
  ptr := new(unsafe.Pointer)
  ctx = context.WithValue(ctx, reasonKey{}, ptr)
  ctx, cancel := context.WithCancel(ctx)
  return ctx, func(reason string) {
    atomic.StorePointer(ptr, &reason)
    cancel()
  }
}

func GetReason(ctx context.Context) string {
  i := ctx.Value(reasonKey{})
  switch t := i.(type) {
    case *unsafe.Pointer:
      return *(string*)(atomic.LoadPointer(t))
  }
  return ""
}

[jordanlewis]

Thanks for the suggestion @petermattis, I'll try that ouit

[dt]

I lost another couple hours this week to debugging a mystery "context canceled" so my interest in this is revived.

I'm thinking I should pursue the approach @petermattis sketched out above, though I'm thinking maybe an error rather than a string reason so we can e.g. use errors.Wrap and capture stacks too if we want. What kills me though is that the stdlib cacelContext already stores an error when it is cancelled that it returns when you call .Err() -- -- it is just always the sentinel context.Canceled error. I wish they'd let you pass a reason error instead and just used the sentinel on nil (or even just allowed a nil error cancellation).

[jordanlewis]

@dt @petermattis I've necromanced this thread. PTAL.

[petermattis]

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @andreimatei, @dt, @jordanlewis, and @mjibson)

---

pkg/util/contextutil/context.go, line 38 at r2 (raw file):

type CancelWithReasonFunc func(error)

func WithCancelReason(ctx context.Context) (context.Context, CancelWithReasonFunc) {

Worth documenting what this does and why I might want to use it. I believe this adds 3 additional allocations to creating a new cancellation context: one for the unsafe.Pointer, one for the `context.WithValue, and one for the closure. I might be mistaken about that last one. Are these extra allocations going to be problematic? Feel free to argue that the extra allocations are a small price for better error messages during cancellation.

[jordanlewis]

Worth documenting what this does and why I might want to use it.

We must have crossed wires - I pushed a patchset that includes documentation already.

Good point about the allocations. I'll add a caveat to the function documentation. I will take you up on your offer to argue that the allocations are a small price to pay for the improved error messages.

[petermattis]

I'd love to see some uses of this and demonstration of the value it provides.

We must have crossed wires - I pushed a patchset that includes documentation already.

Yep, or I failed with reviewable. Regardless, it looks good.

Reviewable status: complete! 1 of 0 LGTMs obtained (waiting on @andreimatei, @dt, @jordanlewis, and @mjibson)

---

pkg/util/contextutil/context.go, line 44 at r3 (raw file):

// WithCancelReason adds a CancelFunc to this context, returning a new
// cancellable context and a CancelWithReasonFunc, which is like
// context.CancelFunc, except it also takes an "reason" error. The context that

Nit: s/an/a/g

---

pkg/util/contextutil/context.go, line 52 at r3 (raw file):

	ctx, cancel := wrap(context.WithCancel(ctx))
	return ctx, func(reason error) {
		atomic.StorePointer(ptr, unsafe.Pointer(&reason))

I think this pushes reason to the heap. If ptr was an atomic.Value I think that allocation would be avoided. Cancellation might be rare enough not to care about this.

[jordanlewis]

bors r+

[jordanlewis]

bors r-

[craig]

Canceled

[jordanlewis]

Reviewable status: complete! 0 of 0 LGTMs obtained (and 1 stale) (waiting on @andreimatei, @dt, @mjibson, and @petermattis)

---

pkg/util/contextutil/context.go, line 38 at r2 (raw file):

Previously, petermattis (Peter Mattis) wrote…

Worth documenting what this does and why I might want to use it. I believe this adds 3 additional allocations to creating a new cancellation context: one for the unsafe.Pointer, one for the `context.WithValue, and one for the closure. I might be mistaken about that last one. Are these extra allocations going to be problematic? Feel free to argue that the extra allocations are a small price for better error messages during cancellation.

Done.

---

pkg/util/contextutil/context.go, line 52 at r3 (raw file):

Previously, petermattis (Peter Mattis) wrote…

I think this pushes reason to the heap. If ptr was an atomic.Value I think that allocation would be avoided. Cancellation might be rare enough not to care about this.

Done.

[petermattis]

Reviewable status: complete! 0 of 0 LGTMs obtained (and 1 stale) (waiting on @andreimatei, @dt, and @mjibson)

[jordanlewis]

TFTRs!

@dt where were the main spots we needed to add this again? :)

[jordanlewis]

bors r+

[craig]

Build failed

• GitHub CI (Cockroach)

[andreimatei]

Reviewable status: complete! 0 of 0 LGTMs obtained (and 1 stale) (waiting on @andreimatei, @dt, @jordanlewis, and @mjibson)

---

pkg/util/contextutil/context.go, line 38 at r5 (raw file):

type reasonKey struct{}

// CancelWithReasonFunc is the

incomplete comment.
And I'd say pass a message here, not an error. That's what I'm doing on various Close(msg) methods that I'm adding and I found it to work better than an error. You can also have two flavors of the func.

---

pkg/util/contextutil/context.go, line 45 at r5 (raw file):

// context.CancelFunc, except it also takes a "reason" error. The context that
// is canceled with this CancelWithReasonFunc will immediately be updated to
// contain this "reason". The reason can be retrieved with GetCancelReason.

Please add some words about how this all interacts (or doesn't interact) with deadlines.

[jordanlewis]

TFTRs!

bors r+

Reviewable status: complete! 0 of 0 LGTMs obtained (and 1 stale) (waiting on @andreimatei, @dt, and @mjibson)

---

pkg/util/contextutil/context.go, line 38 at r5 (raw file):

Previously, andreimatei (Andrei Matei) wrote…

incomplete comment.
And I'd say pass a message here, not an error. That's what I'm doing on various Close(msg) methods that I'm adding and I found it to work better than an error. You can also have two flavors of the func.

Done.

---

pkg/util/contextutil/context.go, line 45 at r5 (raw file):

Previously, andreimatei (Andrei Matei) wrote…

Please add some words about how this all interacts (or doesn't interact) with deadlines.

Done.

[craig]

Build succeeded

• GitHub CI (Cockroach)