proposal: context: add context.Reader/context.Writer

[Manbeardo]

Proposal Details

The stdlib has a fair amount of readers/writers that block until more content is available. For example:

• io.Pipe()
• os.Pipe()
• net.TCPConn

It can be tricky to handle context cancellation promptly when a goroutine is waiting for a blocking read/write.

Existing workarounds

Let's look at some workarounds to implement a version of io.Copy that accepts a context and closes the writer when it's done.

Putting cleanup code in context.AfterFunc()

func CopyAndClose(ctx context.Context, r io.Reader, w io.WriteCloser) error {
    stop := context.AfterFunc(ctx, func() {
        _ = w.Close()
    })
    defer func() {
        if stop() {
            _ = w.Close()
        }
    }()
    buf := make([]byte, 1024)
    for {
        if err := context.Cause(ctx); err != nil {
            return err
        }

        n, err := r.Read(buf)
        if errors.Is(err, io.EOF) {
            return nil
        }

        if err := context.Cause(ctx); err != nil {
            return err
        }

        _, err = w.Write(buf[:n])
        if err != nil {
            return err
        }
    }
}

Problems:

• Doesn't return immediately if it's blocked on a read/write call
• Requires even more synchronization than in the example in order to guarantee the function will never attempt to call w.Write after w has been closed.

Using additional goroutines and channels to unblock function return

func CopyAndClose(ctx context.Context, r io.Reader, w io.WriteCloser) error {
    eg := &errgroup.Group{}
    c := make(chan []byte)

    eg.Go(func() error {
        defer close(c)
        buf := make([]byte, 1024)
        buf2 := make([]byte, 1024)
        for {
            n, err := r.Read(buf)
            if errors.Is(err, io.EOF) {
                return nil
            }
            select {
            case <-ctx.Done():
                return context.Cause(ctx)
            case c <- buf[:n]:
            }
            buf, buf2 = buf2, buf
        }
    })
    eg.Go(func() error {
        defer w.Close()
        for {
            select {
            case <-ctx.Done():
                return context.Cause(ctx)
            case buf, ok := <-c:
                if !ok {
                    return nil
                }
                _, err := w.Write(buf)
                if err != nil {
                    return err
                }
            }
        }
    })

    done := make(chan error)
    go func() {
        done <- eg.Wait()
        close(done)
    }()

    select {
    case err := <-done:
        return err
    case <-ctx.Done():
        go func() {
            <-done
        }()
        return context.Cause(ctx)
    }
}

Problems:

• Doesn't guarantee that w has been closed when the function returns
• Leaves several dangling goroutines that won't finish until the blocking Read/Write calls finish

Proposal

Add something like this to context:

type Reader interface {
    io.Reader
    ContextRead(ctx Context, p []byte) (n int, err error)
}

func Read(ctx Context, r io.Reader, p []byte) (n int, err error) {
    if r, ok := r.(Reader); ok {
        return r.ContextRead(ctx, p)
    }
    if err := Cause(ctx); err != nil {
        return 0, err
    }
    return r.Read(p)
}

type Writer interface {
    io.Writer
    ContextWrite(ctx Context, p []byte) (n int, err error)
}

func Write(ctx Context, w io.Writer, p []byte) (n int, err error) {
    if w, ok := w.(Writer); ok {
        return w.ContextWrite(ctx, p)
    }
    if err := Cause(ctx); err != nil {
        return 0, err
    }
    return w.Write(p)
}

This would allow every reader/writer to implement their own ContextRead/ContextWrite methods as possible/necessary in order to unblock their calls on context cancellation. Having it in the stdlib would be especially beneficial because many of the readers/writers that folks would want to use this with are in the stdlib.

With these proposed changes, CopyAndClose would look like this:

func CopyAndClose(ctx context.Context, r io.Reader, w io.WriteCloser) error {
    defer w.Close()

    buf := make([]byte, 1024)
    for {
        n, err := context.Read(ctx, r, buf)
        if errors.Is(err, io.EOF) {
            return nil
        } else if err != nil {
            return err
        }

        _, err = context.Write(ctx, w, buf[:n])
        if err != nil {
            return err
        }
    }
}

[ianlancetaylor]

There's a lot of overlap with #20280 here.

[Manbeardo]

Wow, GitHub search has completely fallen on its face for me. My search keywords were "context" and "reader", yet that proposal didn't show up on the first page of results.

[ianlancetaylor]

On the other hand it's the top result when I do a Google search for "site:github.com golang issue context reader".

I also find GitHub issue search to be suboptimal.

[apparentlymart]

Two initial reactions:

1. It feels more reasonable for package io to import package context than the other way around, if one has to import the other. Also these operations feel more "I/O-related" than "Context-related". Therefore I'd suggest placing these new declarations in package io.
2. Elsewhere in standard library a late-added variant of a function that takes a context has been named with Context as a suffix, rather than as a prefix. For example, signal.Notify was augmented as signal.NotifyContext to support the additional context argument. Therefore I'd suggest naming the interface methods ReadContext and WriteContext, and (continuing the previous point) the functions io.ReadContext and io.WriteContext.

Aside from those two minor (bikesheddy) points, this does seem like a plausible way to retrofit the io interfaces with optional context support.

[Manbeardo]

It feels more reasonable for package io to import package context than the other way around, if one has to import the other. Also these operations feel more "I/O-related" than "Context-related". Therefore I'd suggest placing these new declarations in package io.

I put io.Reader and io.Writer in the example for expediency, but it'd probably make sense for the actual implementation to use its own definition of those interfaces to avoid importing io. Since those interface definitions are the only dependencies, there's no hard requirement to import io when defining them in the context package.

Elsewhere in standard library a late-added variant of a function that takes a context has been named with Context as a suffix, rather than as a prefix. For example, signal.Notify was augmented as signal.NotifyContext to support the additional context argument. Therefore I'd suggest naming the interface methods ReadContext and WriteContext, and (continuing the previous point) the functions io.ReadContext and io.WriteContext.

Yeah, that makes sense. I originally deviated from the "context suffix" naming convention because my first concept put it in io and my choices (when observing the "-er" interface naming convention) were between:

type ContextReader interface {
    Reader
    ContextRead(ctx context.Context, p []byte) (n int, err error)
}

or

type ReadContexter interface {
    Reader
    ReadContext(ctx context.Context, p []byte) (n int, err error)
}

If the interfaces are defined as context.Reader/context.Writer, that concern is no longer relevant.

[earthboundkid]

A problem with io.ReaderFrom and WriterTo is there is no way to compose the objects that wrap something that may or may not have the method. context.Reader/Writer should preempt this mistake by using errors.ErrUnsupported.

func Read(ctx Context, r io.Reader, p []byte) (n int, err error) {
    if r, ok := r.(Reader); ok {
        if n, err := r.ContextRead(ctx, p); err != errors.ErrUnsupported {
            return n, err
        } // else fallthrough
    }
    if err := Cause(ctx); err != nil {
        return 0, err
    }
    return r.Read(p)
}

[Manbeardo]

@earthboundkid I'm not sure I understand how that problem would apply here? If a context.Reader wraps an io.Reader, wouldn't you get the most intuitive result by having the wrapper invoke context.Read on the io.Reader?

For example:

type InverseByteReader struct {
    r io.Reader
}

func (ibr *InverseByteReader) ReadContext(ctx context.Context, p []byte) (int, error) {
    n, err := context.Read(ctx, ibr.r, p)
    if err != nil {
        return n, err
    }
    for i, b := range p[:n] {
        p[i] = b^0b11111111
    }
    return n, err
}