xio: Add Pipe - io.Pipe analog that supports cancellation

I need this to support sysread(/head/watch) cancellation in WCFS
filesystem [1,2,3].

[1] wendelin.core@b17aeb8c
[2] wendelin.core@f05271b1
[3] wendelin.core@5ba816da

xio/pipe.go

+// Copyright 2009 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE-go file.
+
+// Pipe adapter to connect code expecting an xio.Reader
+// with code expecting an xio.Writer.
+
+package xio
+
+import (
+	"context"
+	"io"
+	"sync"
+)
+
+// onceError is an object that will only store an error once.
+type onceError struct {
+	sync.Mutex // guards following
+	err        error
+}
+
+func (a *onceError) Store(err error) {
+	a.Lock()
+	defer a.Unlock()
+	if a.err != nil {
+		return
+	}
+	a.err = err
+}
+func (a *onceError) Load() error {
+	a.Lock()
+	defer a.Unlock()
+	return a.err
+}
+
+// A pipe is the shared pipe structure underlying PipeReader and PipeWriter.
+type pipe struct {
+	wrMu sync.Mutex // Serializes Write operations
+	wrCh chan []byte
+	rdCh chan int
+
+	once sync.Once // Protects closing done
+	done chan struct{}
+	rerr onceError
+	werr onceError
+}
+
+func (p *pipe) Read(ctx context.Context, b []byte) (n int, err error) {
+	select {
+	case <-p.done:
+		return 0, p.readCloseError()
+	case <-ctx.Done():
+		return 0, ctx.Err()
+	default:
+	}
+
+	select {
+	case bw := <-p.wrCh:
+		nr := copy(b, bw)
+		p.rdCh <- nr
+		return nr, nil
+	case <-p.done:
+		return 0, p.readCloseError()
+	case <-ctx.Done():
+		return 0, ctx.Err()
+	}
+}
+
+func (p *pipe) readCloseError() error {
+	rerr := p.rerr.Load()
+	if werr := p.werr.Load(); rerr == nil && werr != nil {
+		return werr
+	}
+	return io.ErrClosedPipe
+}
+
+func (p *pipe) CloseRead(err error) error {
+	if err == nil {
+		err = io.ErrClosedPipe
+	}
+	p.rerr.Store(err)
+	p.once.Do(func() { close(p.done) })
+	return nil
+}
+
+func (p *pipe) Write(ctx context.Context, b []byte) (n int, err error) {
+	select {
+	case <-p.done:
+		return 0, p.writeCloseError()
+	case <-ctx.Done():
+		return 0, ctx.Err()
+	default:
+		p.wrMu.Lock()
+		defer p.wrMu.Unlock()
+	}
+
+	for once := true; once || len(b) > 0; once = false {
+		select {
+		case p.wrCh <- b:
+			nw := <-p.rdCh
+			b = b[nw:]
+			n += nw
+		case <-p.done:
+			return n, p.writeCloseError()
+		case <-ctx.Done():
+			return n, ctx.Err()
+		}
+	}
+	return n, nil
+}
+
+func (p *pipe) writeCloseError() error {
+	werr := p.werr.Load()
+	if rerr := p.rerr.Load(); werr == nil && rerr != nil {
+		return rerr
+	}
+	return io.ErrClosedPipe
+}
+
+func (p *pipe) CloseWrite(err error) error {
+	if err == nil {
+		err = io.EOF
+	}
+	p.werr.Store(err)
+	p.once.Do(func() { close(p.done) })
+	return nil
+}
+
+// A PipeReader is the read half of a pipe.
+//
+// It is similar to io.PipeReader, but additionally provides cancellation support for Read.
+type PipeReader struct {
+	p *pipe
+}
+
+// Read implements xio.Reader interface:
+// it reads data from the pipe, blocking until a writer
+// arrives or the write end is closed.
+// If the write end is closed with an error, that error is
+// returned as err; otherwise err is EOF.
+func (r *PipeReader) Read(ctx context.Context, data []byte) (n int, err error) {
+	return r.p.Read(ctx, data)
+}
+
+// Close closes the reader; subsequent writes to the
+// write half of the pipe will return the error io.ErrClosedPipe.
+func (r *PipeReader) Close() error {
+	return r.CloseWithError(nil)
+}
+
+// CloseWithError closes the reader; subsequent writes
+// to the write half of the pipe will return the error err.
+//
+// CloseWithError never overwrites the previous error if it exists
+// and always returns nil.
+func (r *PipeReader) CloseWithError(err error) error {
+	return r.p.CloseRead(err)
+}
+
+// A PipeWriter is the write half of a pipe.
+//
+// It is similar to io.PipeWriter, but additionally provides cancellation support for Write.
+type PipeWriter struct {
+	p *pipe
+}
+
+// Write implements xio.Writer interface:
+// it writes data to the pipe, blocking until one or more readers
+// have consumed all the data or the read end is closed.
+// If the read end is closed with an error, that err is
+// returned as err; otherwise err is io.ErrClosedPipe.
+func (w *PipeWriter) Write(ctx context.Context, data []byte) (n int, err error) {
+	return w.p.Write(ctx, data)
+}
+
+// Close closes the writer; subsequent reads from the
+// read half of the pipe will return no bytes and EOF.
+func (w *PipeWriter) Close() error {
+	return w.CloseWithError(nil)
+}
+
+// CloseWithError closes the writer; subsequent reads from the
+// read half of the pipe will return no bytes and the error err,
+// or EOF if err is nil.
+//
+// CloseWithError never overwrites the previous error if it exists
+// and always returns nil.
+func (w *PipeWriter) CloseWithError(err error) error {
+	return w.p.CloseWrite(err)
+}
+
+// Pipe creates a synchronous in-memory pipe.
+// It can be used to connect code expecting a xio.Reader
+// with code expecting a xio.Writer.
+//
+// Reads and Writes on the pipe are matched one to one
+// except when multiple Reads are needed to consume a single Write.
+// That is, each Write to the PipeWriter blocks until it has satisfied
+// one or more Reads from the PipeReader that fully consume
+// the written data.
+// The data is copied directly from the Write to the corresponding
+// Read (or Reads); there is no internal buffering.
+//
+// It is safe to call Read and Write in parallel with each other or with Close.
+// Parallel calls to Read and parallel calls to Write are also safe:
+// the individual calls will be gated sequentially.
+//
+// Pipe is similar to io.Pipe but additionally provides cancellation support
+// for Read and Write.
+func Pipe() (*PipeReader, *PipeWriter) {
+	p := &pipe{
+		wrCh: make(chan []byte),
+		rdCh: make(chan int),
+		done: make(chan struct{}),
+	}
+	return &PipeReader{p}, &PipeWriter{p}
+}

xio/pipe_test.go

+// Copyright 2009 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE-go file.
+
+package xio_test
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+	. "lab.nexedi.com/kirr/go123/xio"
+	"sort"
+	"strings"
+	"testing"
+	"time"
+)
+
+var bg = context.Background()
+
+func checkWrite(t *testing.T, w Writer, data []byte, c chan int) {
+	n, err := w.Write(bg, data)
+	if err != nil {
+		t.Errorf("write: %v", err)
+	}
+	if n != len(data) {
+		t.Errorf("short write: %d != %d", n, len(data))
+	}
+	c <- 0
+}
+
+// Test a single read/write pair.
+func TestPipe1(t *testing.T) {
+	c := make(chan int)
+	r, w := Pipe()
+	var buf = make([]byte, 64)
+	go checkWrite(t, w, []byte("hello, world"), c)
+	n, err := r.Read(bg, buf)
+	if err != nil {
+		t.Errorf("read: %v", err)
+	} else if n != 12 || string(buf[0:12]) != "hello, world" {
+		t.Errorf("bad read: got %q", buf[0:n])
+	}
+	<-c
+	r.Close()
+	w.Close()
+}
+
+func reader(t *testing.T, r Reader, c chan int) {
+	var buf = make([]byte, 64)
+	for {
+		n, err := r.Read(bg, buf)
+		if err == io.EOF {
+			c <- 0
+			break
+		}
+		if err != nil {
+			t.Errorf("read: %v", err)
+		}
+		c <- n
+	}
+}
+
+// Test a sequence of read/write pairs.
+func TestPipe2(t *testing.T) {
+	c := make(chan int)
+	r, w := Pipe()
+	go reader(t, r, c)
+	var buf = make([]byte, 64)
+	for i := 0; i < 5; i++ {
+		p := buf[0 : 5+i*10]
+		n, err := w.Write(bg, p)
+		if n != len(p) {
+			t.Errorf("wrote %d, got %d", len(p), n)
+		}
+		if err != nil {
+			t.Errorf("write: %v", err)
+		}
+		nn := <-c
+		if nn != n {
+			t.Errorf("wrote %d, read got %d", n, nn)
+		}
+	}
+	w.Close()
+	nn := <-c
+	if nn != 0 {
+		t.Errorf("final read got %d", nn)
+	}
+}
+
+type pipeReturn struct {
+	n   int
+	err error
+}
+
+// Test a large write that requires multiple reads to satisfy.
+func writer(w io.WriteCloser, buf []byte, c chan pipeReturn) {
+	n, err := w.Write(buf)
+	w.Close()
+	c <- pipeReturn{n, err}
+}
+
+func TestPipe3(t *testing.T) {
+	c := make(chan pipeReturn)
+	r, w := Pipe()
+	var wdat = make([]byte, 128)
+	for i := 0; i < len(wdat); i++ {
+		wdat[i] = byte(i)
+	}
+	go writer(BindCtxWC(w, bg), wdat, c)
+	var rdat = make([]byte, 1024)
+	tot := 0
+	for n := 1; n <= 256; n *= 2 {
+		nn, err := r.Read(bg, rdat[tot : tot+n])
+		if err != nil && err != io.EOF {
+			t.Fatalf("read: %v", err)
+		}
+
+		// only final two reads should be short - 1 byte, then 0
+		expect := n
+		if n == 128 {
+			expect = 1
+		} else if n == 256 {
+			expect = 0
+			if err != io.EOF {
+				t.Fatalf("read at end: %v", err)
+			}
+		}
+		if nn != expect {
+			t.Fatalf("read %d, expected %d, got %d", n, expect, nn)
+		}
+		tot += nn
+	}
+	pr := <-c
+	if pr.n != 128 || pr.err != nil {
+		t.Fatalf("write 128: %d, %v", pr.n, pr.err)
+	}
+	if tot != 128 {
+		t.Fatalf("total read %d != 128", tot)
+	}
+	for i := 0; i < 128; i++ {
+		if rdat[i] != byte(i) {
+			t.Fatalf("rdat[%d] = %d", i, rdat[i])
+		}
+	}
+}
+
+// Test read after/before writer close.
+
+type closer interface {
+	CloseWithError(error) error
+	Close() error
+}
+
+type pipeTest struct {
+	async          bool
+	err            error
+	closeWithError bool
+}
+
+func (p pipeTest) String() string {
+	return fmt.Sprintf("async=%v err=%v closeWithError=%v", p.async, p.err, p.closeWithError)
+}
+
+var pipeTests = []pipeTest{
+	{true, nil, false},
+	{true, nil, true},
+	{true, io.ErrShortWrite, true},
+	{false, nil, false},
+	{false, nil, true},
+	{false, io.ErrShortWrite, true},
+}
+
+func delayClose(t *testing.T, cl closer, ch chan int, tt pipeTest) {
+	time.Sleep(1 * time.Millisecond)
+	var err error
+	if tt.closeWithError {
+		err = cl.CloseWithError(tt.err)
+	} else {
+		err = cl.Close()
+	}
+	if err != nil {
+		t.Errorf("delayClose: %v", err)
+	}
+	ch <- 0
+}
+
+func TestPipeReadClose(t *testing.T) {
+	for _, tt := range pipeTests {
+		c := make(chan int, 1)
+		r, w := Pipe()
+		if tt.async {
+			go delayClose(t, w, c, tt)
+		} else {
+			delayClose(t, w, c, tt)
+		}
+		var buf = make([]byte, 64)
+		n, err := r.Read(bg, buf)
+		<-c
+		want := tt.err
+		if want == nil {
+			want = io.EOF
+		}
+		if err != want {
+			t.Errorf("read from closed pipe: %v want %v", err, want)
+		}
+		if n != 0 {
+			t.Errorf("read on closed pipe returned %d", n)
+		}
+		if err = r.Close(); err != nil {
+			t.Errorf("r.Close: %v", err)
+		}
+	}
+}
+
+// Test close on Read side during Read.
+func TestPipeReadClose2(t *testing.T) {
+	c := make(chan int, 1)
+	r, _ := Pipe()
+	go delayClose(t, r, c, pipeTest{})
+	n, err := r.Read(bg, make([]byte, 64))
+	<-c
+	if n != 0 || err != io.ErrClosedPipe {
+		t.Errorf("read from closed pipe: %v, %v want %v, %v", n, err, 0, io.ErrClosedPipe)
+	}
+}
+
+// Test write after/before reader close.
+
+func TestPipeWriteClose(t *testing.T) {
+	for _, tt := range pipeTests {
+		c := make(chan int, 1)
+		r, w := Pipe()
+		if tt.async {
+			go delayClose(t, r, c, tt)
+		} else {
+			delayClose(t, r, c, tt)
+		}
+		n, err := io.WriteString(BindCtxW(w, bg), "hello, world")
+		<-c
+		expect := tt.err
+		if expect == nil {
+			expect = io.ErrClosedPipe
+		}
+		if err != expect {
+			t.Errorf("write on closed pipe: %v want %v", err, expect)
+		}
+		if n != 0 {
+			t.Errorf("write on closed pipe returned %d", n)
+		}
+		if err = w.Close(); err != nil {
+			t.Errorf("w.Close: %v", err)
+		}
+	}
+}
+
+// Test close on Write side during Write.
+func TestPipeWriteClose2(t *testing.T) {
+	c := make(chan int, 1)
+	_, w := Pipe()
+	go delayClose(t, w, c, pipeTest{})
+	n, err := w.Write(bg, make([]byte, 64))
+	<-c
+	if n != 0 || err != io.ErrClosedPipe {
+		t.Errorf("write to closed pipe: %v, %v want %v, %v", n, err, 0, io.ErrClosedPipe)
+	}
+}
+
+func TestWriteEmpty(t *testing.T) {
+	r, w := Pipe()
+	go func() {
+		w.Write(bg, []byte{})
+		w.Close()
+	}()
+	var b [2]byte
+	io.ReadFull(BindCtxR(r, bg), b[0:2])
+	r.Close()
+}
+
+func TestWriteNil(t *testing.T) {
+	r, w := Pipe()
+	go func() {
+		w.Write(bg, nil)
+		w.Close()
+	}()
+	var b [2]byte
+	io.ReadFull(BindCtxR(r, bg), b[0:2])
+	r.Close()
+}
+
+func TestWriteAfterWriterClose(t *testing.T) {
+	r, w := Pipe()
+
+	done := make(chan bool)
+	var writeErr error
+	go func() {
+		_, err := w.Write(bg, []byte("hello"))
+		if err != nil {
+			t.Errorf("got error: %q; expected none", err)
+		}
+		w.Close()
+		_, writeErr = w.Write(bg, []byte("world"))
+		done <- true
+	}()
+
+	buf := make([]byte, 100)
+	var result string
+	n, err := io.ReadFull(BindCtxR(r, bg), buf)
+	if err != nil && err != io.ErrUnexpectedEOF {
+		t.Fatalf("got: %q; want: %q", err, io.ErrUnexpectedEOF)
+	}
+	result = string(buf[0:n])
+	<-done
+
+	if result != "hello" {
+		t.Errorf("got: %q; want: %q", result, "hello")
+	}
+	if writeErr != io.ErrClosedPipe {
+		t.Errorf("got: %q; want: %q", writeErr, io.ErrClosedPipe)
+	}
+}
+
+func TestPipeCloseError(t *testing.T) {
+	type testError1 struct{ error }
+	type testError2 struct{ error }
+
+	r, w := Pipe()
+	r.CloseWithError(testError1{})
+	if _, err := w.Write(bg, nil); err != (testError1{}) {
+		t.Errorf("Write error: got %T, want testError1", err)
+	}
+	r.CloseWithError(testError2{})
+	if _, err := w.Write(bg, nil); err != (testError1{}) {
+		t.Errorf("Write error: got %T, want testError1", err)
+	}
+
+	r, w = Pipe()
+	w.CloseWithError(testError1{})
+	if _, err := r.Read(bg, nil); err != (testError1{}) {
+		t.Errorf("Read error: got %T, want testError1", err)
+	}
+	w.CloseWithError(testError2{})
+	if _, err := r.Read(bg, nil); err != (testError1{}) {
+		t.Errorf("Read error: got %T, want testError1", err)
+	}
+}
+
+func TestPipeConcurrent(t *testing.T) {
+	const (
+		input    = "0123456789abcdef"
+		count    = 8
+		readSize = 2
+	)
+
+	t.Run("Write", func(t *testing.T) {
+		r, w := Pipe()
+
+		for i := 0; i < count; i++ {
+			go func() {
+				time.Sleep(time.Millisecond) // Increase probability of race
+				if n, err := w.Write(bg, []byte(input)); n != len(input) || err != nil {
+					t.Errorf("Write() = (%d, %v); want (%d, nil)", n, err, len(input))
+				}
+			}()
+		}
+
+		buf := make([]byte, count*len(input))
+		for i := 0; i < len(buf); i += readSize {
+			if n, err := r.Read(bg, buf[i : i+readSize]); n != readSize || err != nil {
+				t.Errorf("Read() = (%d, %v); want (%d, nil)", n, err, readSize)
+			}
+		}
+
+		// Since each Write is fully gated, if multiple Read calls were needed,
+		// the contents of Write should still appear together in the output.
+		got := string(buf)
+		want := strings.Repeat(input, count)
+		if got != want {
+			t.Errorf("got: %q; want: %q", got, want)
+		}
+	})
+
+	t.Run("Read", func(t *testing.T) {
+		r, w := Pipe()
+
+		c := make(chan []byte, count*len(input)/readSize)
+		for i := 0; i < cap(c); i++ {
+			go func() {
+				time.Sleep(time.Millisecond) // Increase probability of race
+				buf := make([]byte, readSize)
+				if n, err := r.Read(bg, buf); n != readSize || err != nil {
+					t.Errorf("Read() = (%d, %v); want (%d, nil)", n, err, readSize)
+				}
+				c <- buf
+			}()
+		}
+
+		for i := 0; i < count; i++ {
+			if n, err := w.Write(bg, []byte(input)); n != len(input) || err != nil {
+				t.Errorf("Write() = (%d, %v); want (%d, nil)", n, err, len(input))
+			}
+		}
+
+		// Since each read is independent, the only guarantee about the output
+		// is that it is a permutation of the input in readSized groups.
+		got := make([]byte, 0, count*len(input))
+		for i := 0; i < cap(c); i++ {
+			got = append(got, (<-c)...)
+		}
+		got = sortBytesInGroups(got, readSize)
+		want := bytes.Repeat([]byte(input), count)
+		want = sortBytesInGroups(want, readSize)
+		if string(got) != string(want) {
+			t.Errorf("got: %q; want: %q", got, want)
+		}
+	})
+}
+
+func sortBytesInGroups(b []byte, n int) []byte {
+	var groups [][]byte
+	for len(b) > 0 {
+		groups = append(groups, b[:n])
+		b = b[n:]
+	}
+	sort.Slice(groups, func(i, j int) bool { return bytes.Compare(groups[i], groups[j]) < 0 })
+	return bytes.Join(groups, nil)
+}
+
+
+// Verify that .Read and .Write handle cancellation.
+func TestPipeCancel(t *testing.T) {
+	buf := make([]byte, 64)
+	r, _ := Pipe()
+	ctx, cancel := context.WithCancel(bg)
+
+	go func() {
+		time.Sleep(1*time.Millisecond)
+		cancel()
+	}()
+
+	n, err := r.Read(ctx, buf)
+	if eok := context.Canceled; !(n == 0 && err == eok) {
+		t.Errorf("read: got (%v, %v)  ; want (%v, %v)", n, err, 0, eok)
+	}
+
+
+	_, w := Pipe()
+	ctx, cancel = context.WithTimeout(bg, 1*time.Millisecond)
+
+	n, err = w.Write(ctx, buf)
+	if eok := context.DeadlineExceeded; !(n == 0 && err == eok) {
+		t.Errorf("write: got (%v, %v)  ; want (%v, %v)", n, err, 0, eok)
+	}
+
+}

xio/xio.go

@@ -26,6 +26,8 @@
 //	  It is the opposite operation for BindCtx, but for arbitrary io.X
 //	  returned xio.X handles context only on best-effort basis. In
 //	  particular IO cancellation is not reliably handled for os.File .
+//	- Pipe amends io.Pipe and creates synchronous in-memory pipe that
+//	  supports IO cancellation.
 //
 // Miscellaneous utilities:
 //