services/keepstore/count.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package keepstore
   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 func NewCountingReaderAtSeeker(r readerAtSeeker, f func(uint64)) *countingReaderAtSeeker {
  26         return &countingReaderAtSeeker{readerAtSeeker: r, counter: f}
  27 }
  28
  29 type countingReadWriter struct {
  30         reader  io.Reader
  31         writer  io.Writer
  32         counter func(uint64)
  33 }
  34
  35 func (crw *countingReadWriter) Read(buf []byte) (int, error) {
  36         n, err := crw.reader.Read(buf)
  37         crw.counter(uint64(n))
  38         return n, err
  39 }
  40
  41 func (crw *countingReadWriter) Write(buf []byte) (int, error) {
  42         n, err := crw.writer.Write(buf)
  43         crw.counter(uint64(n))
  44         return n, err
  45 }
  46
  47 func (crw *countingReadWriter) Close() error {
  48         if c, ok := crw.writer.(io.Closer); ok {
  49                 return c.Close()
  50         }
  51         return nil
  52 }
  53
  54 type readerAtSeeker interface {
  55         io.ReadSeeker
  56         io.ReaderAt
  57 }
  58
  59 type countingReaderAtSeeker struct {
  60         readerAtSeeker
  61         counter func(uint64)
  62 }
  63
  64 func (crw *countingReaderAtSeeker) Read(buf []byte) (int, error) {
  65         n, err := crw.readerAtSeeker.Read(buf)
  66         crw.counter(uint64(n))
  67         return n, err
  68 }
  69
  70 func (crw *countingReaderAtSeeker) ReadAt(buf []byte, off int64) (int, error) {
  71         n, err := crw.readerAtSeeker.ReadAt(buf, off)
  72         crw.counter(uint64(n))
  73         return n, err
  74 }