1// Package ctxio provides io.Reader and io.Writer wrappers that
  2// respect context.Contexts. Use these at the interface between
  3// your context code and your io.
  4//
  5// WARNING: read the code. see how writes and reads will continue
  6// until you cancel the io. Maybe this package should provide
  7// versions of io.ReadCloser and io.WriteCloser that automatically
  8// call .Close when the context expires. But for now -- since in my
  9// use cases I have long-lived connections with ephemeral io wrappers
 10// -- this has yet to be a need.
 11package ctxio
 12
 13import (
 14	"io"
 15
 16	context "golang.org/x/net/context"
 17)
 18
 19type ioret struct {
 20	n   int
 21	err error
 22}
 23
 24type Writer interface {
 25	io.Writer
 26}
 27
 28type ctxWriter struct {
 29	w   io.Writer
 30	ctx context.Context
 31}
 32
 33// NewWriter wraps a writer to make it respect given Context.
 34// If there is a blocking write, the returned Writer will return
 35// whenever the context is cancelled (the return values are n=0
 36// and err=ctx.Err().)
 37//
 38// Note well: this wrapper DOES NOT ACTUALLY cancel the underlying
 39// write-- there is no way to do that with the standard go io
 40// interface. So the read and write _will_ happen or hang. So, use
 41// this sparingly, make sure to cancel the read or write as necesary
 42// (e.g. closing a connection whose context is up, etc.)
 43//
 44// Furthermore, in order to protect your memory from being read
 45// _after_ you've cancelled the context, this io.Writer will
 46// first make a **copy** of the buffer.
 47func NewWriter(ctx context.Context, w io.Writer) *ctxWriter {
 48	if ctx == nil {
 49		ctx = context.Background()
 50	}
 51	return &ctxWriter{ctx: ctx, w: w}
 52}
 53
 54func (w *ctxWriter) Write(buf []byte) (int, error) {
 55	buf2 := make([]byte, len(buf))
 56	copy(buf2, buf)
 57
 58	c := make(chan ioret, 1)
 59
 60	go func() {
 61		n, err := w.w.Write(buf2)
 62		c <- ioret{n, err}
 63		close(c)
 64	}()
 65
 66	select {
 67	case r := <-c:
 68		return r.n, r.err
 69	case <-w.ctx.Done():
 70		return 0, w.ctx.Err()
 71	}
 72}
 73
 74type Reader interface {
 75	io.Reader
 76}
 77
 78type ctxReader struct {
 79	r   io.Reader
 80	ctx context.Context
 81}
 82
 83// NewReader wraps a reader to make it respect given Context.
 84// If there is a blocking read, the returned Reader will return
 85// whenever the context is cancelled (the return values are n=0
 86// and err=ctx.Err().)
 87//
 88// Note well: this wrapper DOES NOT ACTUALLY cancel the underlying
 89// write-- there is no way to do that with the standard go io
 90// interface. So the read and write _will_ happen or hang. So, use
 91// this sparingly, make sure to cancel the read or write as necesary
 92// (e.g. closing a connection whose context is up, etc.)
 93//
 94// Furthermore, in order to protect your memory from being read
 95// _before_ you've cancelled the context, this io.Reader will
 96// allocate a buffer of the same size, and **copy** into the client's
 97// if the read succeeds in time.
 98func NewReader(ctx context.Context, r io.Reader) *ctxReader {
 99	return &ctxReader{ctx: ctx, r: r}
100}
101
102func (r *ctxReader) Read(buf []byte) (int, error) {
103	buf2 := make([]byte, len(buf))
104
105	c := make(chan ioret, 1)
106
107	go func() {
108		n, err := r.r.Read(buf2)
109		c <- ioret{n, err}
110		close(c)
111	}()
112
113	select {
114	case ret := <-c:
115		copy(buf, buf2)
116		return ret.n, ret.err
117	case <-r.ctx.Done():
118		return 0, r.ctx.Err()
119	}
120}