zrepl/util/bytecounter/bytecounter_reader.go

package bytecounter

import (
	"io"
	"sync/atomic"
	"time"
)

type ByteCounterReader struct {
	reader io.ReadCloser

	// called & accessed synchronously during Read, no external access
	cb       func(full int64)
	cbEvery  time.Duration
	lastCbAt time.Time

	// set atomically because it may be read by multiple threads
	bytes int64
}

func NewByteCounterReader(reader io.ReadCloser) *ByteCounterReader {
	return &ByteCounterReader{
		reader: reader,
	}
}

func (b *ByteCounterReader) SetCallback(every time.Duration, cb func(full int64)) {
	b.cbEvery = every
	b.cb = cb
}

func (b *ByteCounterReader) Close() error {
	return b.reader.Close()
}

func (b *ByteCounterReader) Read(p []byte) (n int, err error) {
	n, err = b.reader.Read(p)
	full := atomic.AddInt64(&b.bytes, int64(n))
	now := time.Now()
	if b.cb != nil && now.Sub(b.lastCbAt) > b.cbEvery {
		b.cb(full)
		b.lastCbAt = now
	}
	return n, err
}

func (b *ByteCounterReader) Bytes() int64 {
	return atomic.LoadInt64(&b.bytes)
}