X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/6f461f4d0a996da85140982846a5d5c10ccfaae4..2c6557f613fcf6cdcebb08c321a5d061aeb780c6:/services/keepstore/count.go?ds=sidebyside

diff --git a/services/keepstore/count.go b/services/keepstore/count.go
index a9f74368b4..51434a803e 100644
--- a/services/keepstore/count.go
+++ b/services/keepstore/count.go
@@ -1,23 +1,31 @@
-package main
+// Copyright (C) The Arvados Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
+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
@@ -42,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
+}