services/keepstore/count.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package main
   6
   7 import (
   8         "io"
   9 )
  10
  11 func NewCountingWriter(w io.Writer, f func(uint64)) io.WriteCloser {
  12         return &countingReadWriter{
  13                 writer:  w,
  14                 counter: f,
  15         }
  16 }
  17
  18 func NewCountingReader(r io.Reader, f func(uint64)) io.ReadCloser {
  19         return &countingReadWriter{
  20                 reader:  r,
  21                 counter: f,
  22         }
  23 }
  24
  25 type countingReadWriter struct {
  26         reader  io.Reader
  27         writer  io.Writer
  28         counter func(uint64)
  29 }
  30
  31 func (crw *countingReadWriter) Read(buf []byte) (int, error) {
  32         n, err := crw.reader.Read(buf)
  33         crw.counter(uint64(n))
  34         return n, err
  35 }
  36
  37 func (crw *countingReadWriter) Write(buf []byte) (int, error) {
  38         n, err := crw.writer.Write(buf)
  39         crw.counter(uint64(n))
  40         return n, err
  41 }
  42
  43 func (crw *countingReadWriter) Close() error {
  44         if c, ok := crw.writer.(io.Closer); ok {
  45                 return c.Close()
  46         }
  47         return nil
  48 }