readers: make NoCloser pass through WriteTo so io.Copy keeps its fast path
NoCloser hides the Close method of an io.Reader but in doing so it also hid io.WriterTo if the underlying reader implemented it, forcing io.Copy to fall back to a buffered Read loop. Following io.NopCloser, return a variant which forwards WriteTo when the wrapped reader supports it.
This commit is contained in:
+17
-1
@@ -7,16 +7,29 @@ type noClose struct {
|
||||
in io.Reader
|
||||
}
|
||||
|
||||
// Read implements io.Closer by passing it straight on
|
||||
// Read implements io.Reader by passing it straight on
|
||||
func (nc noClose) Read(p []byte) (n int, err error) {
|
||||
return nc.in.Read(p)
|
||||
}
|
||||
|
||||
// noCloseWriterTo is a noClose which also forwards io.WriterTo
|
||||
type noCloseWriterTo struct {
|
||||
noClose
|
||||
}
|
||||
|
||||
// WriteTo implements io.WriterTo by passing it straight on
|
||||
func (nc noCloseWriterTo) WriteTo(w io.Writer) (n int64, err error) {
|
||||
return nc.in.(io.WriterTo).WriteTo(w)
|
||||
}
|
||||
|
||||
// NoCloser makes sure that the io.Reader passed in can't upgraded to
|
||||
// an io.Closer.
|
||||
//
|
||||
// This is for use with http.NewRequest to make sure the body doesn't
|
||||
// get upgraded to an io.Closer and the body closed unexpectedly.
|
||||
//
|
||||
// If in implements io.WriterTo then the returned reader does too so
|
||||
// that io.Copy can still use the more efficient path.
|
||||
func NoCloser(in io.Reader) io.Reader {
|
||||
if in == nil {
|
||||
return in
|
||||
@@ -25,5 +38,8 @@ func NoCloser(in io.Reader) io.Reader {
|
||||
if _, canClose := in.(io.Closer); !canClose {
|
||||
return in
|
||||
}
|
||||
if _, canWriteTo := in.(io.WriterTo); canWriteTo {
|
||||
return noCloseWriterTo{noClose{in: in}}
|
||||
}
|
||||
return noClose{in: in}
|
||||
}
|
||||
|
||||
@@ -26,6 +26,14 @@ func (readClose) Close() (err error) {
|
||||
return io.EOF
|
||||
}
|
||||
|
||||
type readCloseWriteTo struct {
|
||||
readClose
|
||||
}
|
||||
|
||||
func (readCloseWriteTo) WriteTo(w io.Writer) (n int64, err error) {
|
||||
return 42, errRead
|
||||
}
|
||||
|
||||
func TestNoCloser(t *testing.T) {
|
||||
assert.Equal(t, nil, NoCloser(nil))
|
||||
|
||||
@@ -41,4 +49,20 @@ func TestNoCloser(t *testing.T) {
|
||||
|
||||
_, err := nc.Read(nil)
|
||||
assert.Equal(t, errRead, err)
|
||||
|
||||
_, hasWriteTo := nc.(io.WriterTo)
|
||||
assert.False(t, hasWriteTo)
|
||||
|
||||
rcw := readCloseWriteTo{}
|
||||
ncw := NoCloser(rcw)
|
||||
assert.NotEqual(t, ncw, rcw)
|
||||
|
||||
_, hasClose = ncw.(io.Closer)
|
||||
assert.False(t, hasClose)
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user