Datasources: Set response size metric based on actual bytes read (#40303) (#40867)

Fixes #33372

(cherry picked from commit 889d4683a1)

Co-authored-by: Marcus Efraimsson <marcus.efraimsson@gmail.com>
This commit is contained in:
Grot (@grafanabot)
2021-10-25 05:01:16 -06:00
committed by GitHub
parent 908a951fc0
commit a8ecdca826
3 changed files with 87 additions and 8 deletions
@@ -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()
}