Fixes #33372
(cherry picked from commit 889d4683a1)
Co-authored-by: Marcus Efraimsson <marcus.efraimsson@gmail.com>
This commit is contained in:
committed by
GitHub
parent
908a951fc0
commit
a8ecdca826
@@ -0,0 +1,39 @@
|
||||
package httpclient
|
||||
|
||||
import (
|
||||
"io"
|
||||
)
|
||||
|
||||
type CloseCallbackFunc func(bytesRead int64)
|
||||
|
||||
// CountBytesReader counts the total amount of bytes read from the underlying reader.
|
||||
//
|
||||
// The provided callback func will be called before the underlying reader is closed.
|
||||
func CountBytesReader(reader io.ReadCloser, callback CloseCallbackFunc) io.ReadCloser {
|
||||
if reader == nil {
|
||||
panic("reader cannot be nil")
|
||||
}
|
||||
|
||||
if callback == nil {
|
||||
panic("callback cannot be nil")
|
||||
}
|
||||
|
||||
return &countBytesReader{reader: reader, callback: callback}
|
||||
}
|
||||
|
||||
type countBytesReader struct {
|
||||
reader io.ReadCloser
|
||||
callback CloseCallbackFunc
|
||||
counter int64
|
||||
}
|
||||
|
||||
func (r *countBytesReader) Read(p []byte) (int, error) {
|
||||
n, err := r.reader.Read(p)
|
||||
r.counter += int64(n)
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (r *countBytesReader) Close() error {
|
||||
r.callback(r.counter)
|
||||
return r.reader.Close()
|
||||
}
|
||||
Reference in New Issue
Block a user