readers: add NoCloserNotify to find out when a request body has been closed

NoCloserNotify hides the Close method of the reader passed in like
NoCloser, but calls a notify function (once only) when the returned
body is closed. This lets callers of http.NewRequest find out when the
transport has finished with a request body, as it is documented to
possibly close it in a different goroutine after the request has
finished.

lib/readers/noclose.go

@@ -1,6 +1,9 @@
 package readers
 
-import "io"
+import (
+	"io"
+	"sync"
+)
 
 // noClose is used to wrap an io.Reader to stop it being upgraded
 type noClose struct {
@@ -43,3 +46,54 @@ func NoCloser(in io.Reader) io.Reader {
 	}
 	return noClose{in: in}
 }
+
+// noCloseNotify is used to wrap an io.Reader replacing its Close method
+type noCloseNotify struct {
+	in     io.Reader
+	once   sync.Once
+	notify func()
+}
+
+// Read implements io.Reader by passing it straight on
+func (nc *noCloseNotify) Read(p []byte) (n int, err error) {
+	return nc.in.Read(p)
+}
+
+// Close calls notify the first time it is called. It never closes the
+// underlying reader.
+func (nc *noCloseNotify) Close() error {
+	nc.once.Do(nc.notify)
+	return nil
+}
+
+// noCloseNotifyWriterTo is a noCloseNotify which also forwards io.WriterTo
+type noCloseNotifyWriterTo struct {
+	noCloseNotify
+}
+
+// WriteTo implements io.WriterTo by passing it straight on
+func (nc *noCloseNotifyWriterTo) WriteTo(w io.Writer) (n int64, err error) {
+	return nc.in.(io.WriterTo).WriteTo(w)
+}
+
+// NoCloserNotify makes sure the io.Reader passed in can't be upgraded
+// to an io.Closer, like NoCloser, but returns an io.ReadCloser whose
+// Close calls notify (once only) instead of closing in.
+//
+// This is for use with http.NewRequest to find out when the transport
+// has finished with the request body. The transport always closes the
+// body, but is documented to possibly do so in a different goroutine
+// after the request has finished, so notify signals when the
+// transport can no longer be reading the body.
+//
+// If in implements io.WriterTo then the returned reader does too so
+// that io.Copy can still use the more efficient path.
+func NoCloserNotify(in io.Reader, notify func()) io.ReadCloser {
+	if in == nil {
+		return nil
+	}
+	if _, canWriteTo := in.(io.WriterTo); canWriteTo {
+		return &noCloseNotifyWriterTo{noCloseNotify{in: in, notify: notify}}
+	}
+	return &noCloseNotify{in: in, notify: notify}
+}

lib/readers/noclose_test.go

@@ -66,3 +66,39 @@ func TestNoCloser(t *testing.T) {
 	assert.Equal(t, int64(42), n)
 	assert.Equal(t, errRead, err)
 }
+
+func TestNoCloserNotify(t *testing.T) {
+	assert.Nil(t, NoCloserNotify(nil, func() {}))
+
+	notified := 0
+	notify := func() { notified++ }
+
+	// A reader without Close or WriteTo
+	nc := NoCloserNotify(readOnly{}, notify)
+	_, err := nc.Read(nil)
+	assert.Equal(t, io.EOF, err)
+	_, hasWriteTo := nc.(io.WriterTo)
+	assert.False(t, hasWriteTo)
+
+	// Close doesn't close the underlying reader (whose Close returns
+	// io.EOF) and notifies once only
+	notified = 0
+	nc = NoCloserNotify(readClose{}, notify)
+	_, err = nc.Read(nil)
+	assert.Equal(t, errRead, err)
+	assert.NoError(t, nc.Close())
+	assert.Equal(t, 1, notified)
+	assert.NoError(t, nc.Close())
+	assert.Equal(t, 1, notified)
+
+	// WriteTo is forwarded
+	notified = 0
+	ncw := NoCloserNotify(readCloseWriteTo{}, notify)
+	wt, hasWriteTo := ncw.(io.WriterTo)
+	assert.True(t, hasWriteTo)
+	n, err := wt.WriteTo(nil)
+	assert.Equal(t, int64(42), n)
+	assert.Equal(t, errRead, err)
+	assert.NoError(t, ncw.Close())
+	assert.Equal(t, 1, notified)
+}