Resetting streams based on contexts

[Stebalien]

At the moment, there isn't a great way to "cancel" a request if a given context is canceled.

The solution I'm using in my project is:

// Resets the stream on cancel, or when the returned function is called.
func ResetOnCancel(ctx context.Context, s network.Stream) func() error {
	errCh := make(chan error, 1)
	cancel := context.AfterFunc(ctx, func() {
		errCh <- s.Reset()
		close(errCh)
	})
	return func() error {
		if cancel() {
			return s.Reset()
		} else {
			return <-errCh
		}
	}
}

This way, I open streams as follows...

stream, _ := host.NewStream(ctx, ...)
defer ResetOnCancel(ctx, stream)()

...and have them reset on either return or when the context is canceled.

But it would be nice to have an official way to do this. Either upstream this function or add something to the interface like:

type Stream interface {
    SetContext(ctx context.Context)
}

Which would behave like SetDeadline but with contexts. Now that go has context.AfterFunc, this is actually possible without having to spawn a new goroutine. This can also be done at the network stream layer instead of at the transport stream layer, which should make it a minimally breaking change.

Thoughts?

[MarcoPolo]

I wonder if we should have used the context passed into .NewStream? That feels more natural to me than a SetContext method, but it would be a breaking change. The benefits probably don't outweigh the subtleties of introducing that breaking change.

the .SetContext interface seems like the next best thing (We certainly can't do a .NewStreamWithContext 🙈). And maybe clearing the context is useful in cases where the caller is reusing streams (for whatever reason, even if ideally they shouldn't need to).

If anyone else thinks they would use this please let me know.

[Stebalien]

I wonder if we should have used the context passed into .NewStream? That feels more natural to me than a SetContext method, but it would be a breaking change. The benefits probably don't outweigh the subtleties of introducing that breaking change.

The problem is that NewStream can block as well and the rule was always "don't store the context, only use it for the duration of the call".

But... now there's https://pkg.go.dev/net/http#NewRequestWithContext.

Although NewStream is a bit different as streams can be long-lived and often outlive the current "request". Resetting the stream when the context is canceled could lead to hard to track down bugs.

Ideally Read/Write/Close would take a context but... there's a lengthy debate on that here: golang/go#20280.

The remaining question is: should SetContext cause the stream to be reset, or should it behave more like SetDeadline (canceling reads/writes). TBH, it should probably behave more like the latter even though that's a bit less convenient.