Refactor the multi-host salt install page.

[arvados.git] / services / keepstore / count.go

diff --git a/services/keepstore/count.go b/services/keepstore/count.go
index a9f74368b4b4965e83796aef32478856ed1bed4a..700ca19dec958cbf978bc875e2b1f8e71f24bf49 100644 (file)
--- a/services/keepstore/count.go
+++ b/services/keepstore/count.go
@@ -1,4 +1,8 @@
-package main
+// Copyright (C) The Arvados Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
+package keepstore
 
 import (
        "io"
@@ -18,6 +22,10 @@ func NewCountingReader(r io.Reader, f func(uint64)) io.ReadCloser {
        }
 }
 
+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
+}