10484: Report volume IO stats for S3 volumes.

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

diff --git a/services/keepstore/count.go b/services/keepstore/count.go
new file mode 100644 (file)
index 0000000..a9f7436
--- /dev/null
+++ b/services/keepstore/count.go
@@ -0,0 +1,44 @@
+package main
+
+import (
+       "io"
+)
+
+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 {
+       return &countingReadWriter{
+               reader:  r,
+               counter: f,
+       }
+}
+
+type countingReadWriter struct {
+       reader  io.Reader
+       writer  io.Writer
+       counter func(uint64)
+}
+
+func (crw *countingReadWriter) Read(buf []byte) (int, error) {
+       n, err := crw.reader.Read(buf)
+       crw.counter(uint64(n))
+       return n, err
+}
+
+func (crw *countingReadWriter) Write(buf []byte) (int, error) {
+       n, err := crw.writer.Write(buf)
+       crw.counter(uint64(n))
+       return n, err
+}
+
+func (crw *countingReadWriter) Close() error {
+       if c, ok := crw.writer.(io.Closer); ok {
+               return c.Close()
+       }
+       return nil
+}