Snippet of RPC receive written with contexts

ideal_ioctx.go

type Conn struct {
	wc     io.WriteCloser
	cancel context.CancelFunc
	tasks  sync.WaitGroup
	// ... other internal state ...
}

func NewConn(rwc io.ReadWriteCloser) *Conn {
	ctx, cancel := context.WithCancel(context.Background())
	c := &Conn{
		wc:     rwc,
		cancel: cancel,
	}
	c.tasks.Add(1)
	go func() {
		defer c.tasks.Done()
		c.receive(ctx, rwc)
	}()
	return c
}

func (c *Conn) receive(ctx context.Context, r io.Reader) {
	for {
		msg, err := decode(ctx, r)
		if err != nil {
			return
		}
		response, err := process(ctx, msg)
		if err != nil {
			return
		}
		// Write can be canceled with a Context.
		if _, err := c.rwc.Write(ctx, response); err != nil {
			return
		}
	}
}

func decode(ctx context.Context, r io.Reader) (*capnp.Message, error) {
	var hdr [8]byte
	// Importantly, I'd like to be able to plumb Context through
	// functions like io.ReadFull that make many Read calls (like
	// decode does).
	if _, err := io.ReadFull(ctx, r, hdr[:]); err != nil {
		return nil, err
	}
	// ... parse header, read rest of body ...
}

func (c *Conn) Close() error {
	c.cancel()
	c.tasks.Wait()
	err := c.wc.Close()
	return err
}