X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/16f704326f44fd1e5e5e60b936c9b5895d6a6ff8..2c6557f613fcf6cdcebb08c321a5d061aeb780c6:/services/keepstore/count.go diff --git a/services/keepstore/count.go b/services/keepstore/count.go index ccba4f1355..51434a803e 100644 --- a/services/keepstore/count.go +++ b/services/keepstore/count.go @@ -2,26 +2,30 @@ // // SPDX-License-Identifier: AGPL-3.0 -package main +package keepstore import ( "io" ) -func NewCountingWriter(w io.Writer, f func(uint64)) io.WriteCloser { +func newCountingWriter(w io.Writer, f func(uint64)) io.WriteCloser { return &countingReadWriter{ writer: w, counter: f, } } -func NewCountingReader(r io.Reader, f func(uint64)) io.ReadCloser { +func newCountingReader(r io.Reader, f func(uint64)) io.ReadCloser { return &countingReadWriter{ reader: r, counter: f, } } +func newCountingReaderAtSeeker(r readerAtSeeker, f func(uint64)) *countingReaderAtSeeker { + return &countingReaderAtSeeker{readerAtSeeker: r, counter: f} +} + type countingReadWriter struct { reader io.Reader writer io.Writer @@ -46,3 +50,25 @@ func (crw *countingReadWriter) Close() error { } return nil } + +type readerAtSeeker interface { + io.ReadSeeker + io.ReaderAt +} + +type countingReaderAtSeeker struct { + readerAtSeeker + counter func(uint64) +} + +func (crw *countingReaderAtSeeker) Read(buf []byte) (int, error) { + n, err := crw.readerAtSeeker.Read(buf) + crw.counter(uint64(n)) + return n, err +} + +func (crw *countingReaderAtSeeker) ReadAt(buf []byte, off int64) (int, error) { + n, err := crw.readerAtSeeker.ReadAt(buf, off) + crw.counter(uint64(n)) + return n, err +}