services/keepstore/count.go

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