From 7744f98960fdd88b44a112c5668f6e8e94ae7419 Mon Sep 17 00:00:00 2001 From: Nick Craig-Wood Date: Mon, 31 Aug 2026 17:43:27 +0100 Subject: [PATCH] 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 | 56 ++++++++++++++++++++++++++++++++++++- lib/readers/noclose_test.go | 36 ++++++++++++++++++++++++ 2 files changed, 91 insertions(+), 1 deletion(-) diff --git a/lib/readers/noclose.go b/lib/readers/noclose.go index 7516ef317..96fb42ec6 100644 --- a/lib/readers/noclose.go +++ b/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} +} diff --git a/lib/readers/noclose_test.go b/lib/readers/noclose_test.go index 6bfa43b34..7bfa19a41 100644 --- a/lib/readers/noclose_test.go +++ b/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) +}