[#277] rpc + zfs: drop zfs.StreamCopier, use io.ReadCloser instead

util/bytecounter/bytecounter_readcloser.go

@@ -0,0 +1,39 @@
+package bytecounter
+
+import (
+	"io"
+	"sync/atomic"
+)
+
+// ReadCloser wraps an io.ReadCloser, reimplementing
+// its interface and counting the bytes written to during copying.
+type ReadCloser interface {
+	io.ReadCloser
+	Count() int64
+}
+
+// NewReadCloser wraps rc.
+func NewReadCloser(rc io.ReadCloser) ReadCloser {
+	return &readCloser{rc, 0}
+}
+
+type readCloser struct {
+	rc    io.ReadCloser
+	count int64
+}
+
+func (r *readCloser) Count() int64 {
+	return atomic.LoadInt64(&r.count)
+}
+
+var _ io.ReadCloser = &readCloser{}
+
+func (r *readCloser) Close() error {
+	return r.rc.Close()
+}
+
+func (r *readCloser) Read(p []byte) (int, error) {
+	n, err := r.rc.Read(p)
+	atomic.AddInt64(&r.count, int64(n))
+	return n, err
+}

util/bytecounter/bytecounter_reader.go

@@ -1,49 +0,0 @@
-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)
-}

util/bytecounter/bytecounter_streamcopier.go

@@ -1,71 +0,0 @@
-package bytecounter
-
-import (
-	"io"
-	"sync/atomic"
-
-	"github.com/zrepl/zrepl/zfs"
-)
-
-// StreamCopier wraps a zfs.StreamCopier, reimplementing
-// its interface and counting the bytes written to during copying.
-type StreamCopier interface {
-	zfs.StreamCopier
-	Count() int64
-}
-
-// NewStreamCopier wraps sc into a StreamCopier.
-// If sc is io.Reader, it is guaranteed that the returned StreamCopier
-// implements that interface, too.
-func NewStreamCopier(sc zfs.StreamCopier) StreamCopier {
-	bsc := &streamCopier{sc, 0}
-	if scr, ok := sc.(io.Reader); ok {
-		return streamCopierAndReader{bsc, scr}
-	} else {
-		return bsc
-	}
-}
-
-type streamCopier struct {
-	sc    zfs.StreamCopier
-	count int64
-}
-
-// proxy writer used by streamCopier
-type streamCopierWriter struct {
-	parent *streamCopier
-	w      io.Writer
-}
-
-func (w streamCopierWriter) Write(p []byte) (n int, err error) {
-	n, err = w.w.Write(p)
-	atomic.AddInt64(&w.parent.count, int64(n))
-	return
-}
-
-func (s *streamCopier) Count() int64 {
-	return atomic.LoadInt64(&s.count)
-}
-
-var _ zfs.StreamCopier = &streamCopier{}
-
-func (s streamCopier) Close() error {
-	return s.sc.Close()
-}
-
-func (s *streamCopier) WriteStreamTo(w io.Writer) zfs.StreamCopierError {
-	ww := streamCopierWriter{s, w}
-	return s.sc.WriteStreamTo(ww)
-}
-
-// a streamCopier whose underlying sc is an io.Reader
-type streamCopierAndReader struct {
-	*streamCopier
-	asReader io.Reader
-}
-
-func (scr streamCopierAndReader) Read(p []byte) (int, error) {
-	n, err := scr.asReader.Read(p)
-	atomic.AddInt64(&scr.streamCopier.count, int64(n))
-	return n, err
-}

util/bytecounter/bytecounter_streamcopier_test.go

@@ -1,39 +0,0 @@
-package bytecounter
-
-import (
-	"io"
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-
-	"github.com/zrepl/zrepl/zfs"
-)
-
-type mockStreamCopierAndReader struct {
-	zfs.StreamCopier // to satisfy interface
-	reads            int
-}
-
-func (r *mockStreamCopierAndReader) Read(p []byte) (int, error) {
-	r.reads++
-	return len(p), nil
-}
-
-var _ io.Reader = &mockStreamCopierAndReader{}
-
-func TestNewStreamCopierReexportsReader(t *testing.T) {
-	mock := &mockStreamCopierAndReader{}
-	x := NewStreamCopier(mock)
-
-	r, ok := x.(io.Reader)
-	if !ok {
-		t.Fatalf("%T does not implement io.Reader, hence reader cannout have been wrapped", x)
-	}
-
-	var buf [23]byte
-	n, err := r.Read(buf[:])
-	assert.True(t, mock.reads == 1)
-	assert.True(t, n == len(buf))
-	assert.NoError(t, err)
-	assert.True(t, x.Count() == 23)
-}