lib/pool: add DelayAccounting() to fix accounting when reading hashes

lib/pool/reader_writer.go

@@ -19,6 +19,8 @@ type RW struct {
 	out        int       // offset we are reading from
 	lastOffset int       // size in last page
 	account    RWAccount // account for a read
+	reads      int       // count how many times the data has been read
+	accountOn  int       // only account on or after this read
 }
 
 var (
@@ -50,10 +52,40 @@ func (rw *RW) SetAccounting(account RWAccount) *RW {
 	return rw
 }
 
+// DelayAccountinger enables an accounting delay
+type DelayAccountinger interface {
+	// DelayAccounting makes sure the accounting function only
+	// gets called on the i-th or later read of the data from this
+	// point (counting from 1).
+	//
+	// This is useful so that we don't account initial reads of
+	// the data e.g. when calculating hashes.
+	//
+	// Set this to 0 to account everything.
+	DelayAccounting(i int)
+}
+
+// DelayAccounting makes sure the accounting function only gets called
+// on the i-th or later read of the data from this point (counting
+// from 1).
+//
+// This is useful so that we don't account initial reads of the data
+// e.g. when calculating hashes.
+//
+// Set this to 0 to account everything.
+func (rw *RW) DelayAccounting(i int) {
+	rw.accountOn = i
+	rw.reads = 0
+}
+
 // Returns the page and offset of i for reading.
 //
 // Ensure there are pages before calling this.
 func (rw *RW) readPage(i int) (page []byte) {
+	// Count a read of the data if we read the first page
+	if i == 0 {
+		rw.reads++
+	}
 	pageNumber := i / rw.pool.bufferSize
 	offset := i % rw.pool.bufferSize
 	page = rw.pages[pageNumber]
@@ -69,7 +101,14 @@ func (rw *RW) accountRead(n int) error {
 	if rw.account == nil {
 		return nil
 	}
-	return rw.account(n)
+	// Don't start accounting until we've reached this many reads
+	//
+	// rw.reads will be 1 the first time this is called
+	// rw.accountOn 2 means start accounting on the 2nd read through
+	if rw.reads >= rw.accountOn {
+		return rw.account(n)
+	}
+	return nil
 }
 
 // Read reads up to len(p) bytes into p. It returns the number of
@@ -227,10 +266,11 @@ func (rw *RW) Size() int64 {
 
 // Check interfaces
 var (
-	_ io.Reader     = (*RW)(nil)
+	_ io.Reader         = (*RW)(nil)
-	_ io.ReaderFrom = (*RW)(nil)
+	_ io.ReaderFrom     = (*RW)(nil)
-	_ io.Writer     = (*RW)(nil)
+	_ io.Writer         = (*RW)(nil)
-	_ io.WriterTo   = (*RW)(nil)
+	_ io.WriterTo       = (*RW)(nil)
-	_ io.Seeker     = (*RW)(nil)
+	_ io.Seeker         = (*RW)(nil)
-	_ io.Closer     = (*RW)(nil)
+	_ io.Closer         = (*RW)(nil)
+	_ DelayAccountinger = (*RW)(nil)
 )

lib/pool/reader_writer_test.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/rclone/rclone/lib/random"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 const blockSize = 4096
@@ -178,71 +179,164 @@ func TestRW(t *testing.T) {
 		assert.Equal(t, testData[7:10], dst)
 	})
 
-	errBoom := errors.New("accounting error")
+	t.Run("Account", func(t *testing.T) {
+		errBoom := errors.New("accounting error")
 
-	t.Run("AccountRead", func(t *testing.T) {
+		t.Run("Read", func(t *testing.T) {
-		// Test accounting errors
+			rw := newRW()
-		rw := newRW()
+			defer close(rw)
-		defer close(rw)
 
-		var total int
+			var total int
-		rw.SetAccounting(func(n int) error {
+			rw.SetAccounting(func(n int) error {
-			total += n
+				total += n
-			return nil
+				return nil
+			})
+
+			dst = make([]byte, 3)
+			n, err = rw.Read(dst)
+			assert.Equal(t, 3, n)
+			assert.NoError(t, err)
+			assert.Equal(t, 3, total)
 		})
 
-		dst = make([]byte, 3)
+		t.Run("WriteTo", func(t *testing.T) {
-		n, err = rw.Read(dst)
+			rw := newRW()
-		assert.Equal(t, 3, n)
+			defer close(rw)
-		assert.NoError(t, err)
+			var b bytes.Buffer
-		assert.Equal(t, 3, total)
-	})
 
-	t.Run("AccountWriteTo", func(t *testing.T) {
+			var total int
-		rw := newRW()
+			rw.SetAccounting(func(n int) error {
-		defer close(rw)
+				total += n
-		var b bytes.Buffer
+				return nil
+			})
 
-		var total int
+			n, err := rw.WriteTo(&b)
-		rw.SetAccounting(func(n int) error {
+			assert.NoError(t, err)
-			total += n
+			assert.Equal(t, 10, total)
-			return nil
+			assert.Equal(t, int64(10), n)
+			assert.Equal(t, testData, b.Bytes())
 		})
 
-		n, err := rw.WriteTo(&b)
+		t.Run("ReadDelay", func(t *testing.T) {
-		assert.NoError(t, err)
+			rw := newRW()
-		assert.Equal(t, 10, total)
+			defer close(rw)
-		assert.Equal(t, int64(10), n)
-		assert.Equal(t, testData, b.Bytes())
-	})
 
-	t.Run("AccountReadError", func(t *testing.T) {
+			var total int
-		// Test accounting errors
+			rw.SetAccounting(func(n int) error {
-		rw := newRW()
+				total += n
-		defer close(rw)
+				return nil
+			})
 
-		rw.SetAccounting(func(n int) error {
+			rewind := func() {
-			return errBoom
+				_, err := rw.Seek(0, io.SeekStart)
+				require.NoError(t, err)
+			}
+
+			rw.DelayAccounting(3)
+
+			dst = make([]byte, 16)
+
+			n, err = rw.Read(dst)
+			assert.Equal(t, 10, n)
+			assert.Equal(t, io.EOF, err)
+			assert.Equal(t, 0, total)
+			rewind()
+
+			n, err = rw.Read(dst)
+			assert.Equal(t, 10, n)
+			assert.Equal(t, io.EOF, err)
+			assert.Equal(t, 0, total)
+			rewind()
+
+			n, err = rw.Read(dst)
+			assert.Equal(t, 10, n)
+			assert.Equal(t, io.EOF, err)
+			assert.Equal(t, 10, total)
+			rewind()
+
+			n, err = rw.Read(dst)
+			assert.Equal(t, 10, n)
+			assert.Equal(t, io.EOF, err)
+			assert.Equal(t, 20, total)
+			rewind()
 		})
 
-		dst = make([]byte, 3)
+		t.Run("WriteToDelay", func(t *testing.T) {
-		n, err = rw.Read(dst)
+			rw := newRW()
-		assert.Equal(t, 3, n)
+			defer close(rw)
-		assert.Equal(t, errBoom, err)
+			var b bytes.Buffer
-	})
 
-	t.Run("AccountWriteToError", func(t *testing.T) {
+			var total int
-		rw := newRW()
+			rw.SetAccounting(func(n int) error {
-		defer close(rw)
+				total += n
-		rw.SetAccounting(func(n int) error {
+				return nil
-			return errBoom
+			})
+
+			rw.DelayAccounting(3)
+
+			rewind := func() {
+				_, err := rw.Seek(0, io.SeekStart)
+				require.NoError(t, err)
+				b.Reset()
+			}
+
+			n, err := rw.WriteTo(&b)
+			assert.NoError(t, err)
+			assert.Equal(t, 0, total)
+			assert.Equal(t, int64(10), n)
+			assert.Equal(t, testData, b.Bytes())
+			rewind()
+
+			n, err = rw.WriteTo(&b)
+			assert.NoError(t, err)
+			assert.Equal(t, 0, total)
+			assert.Equal(t, int64(10), n)
+			assert.Equal(t, testData, b.Bytes())
+			rewind()
+
+			n, err = rw.WriteTo(&b)
+			assert.NoError(t, err)
+			assert.Equal(t, 10, total)
+			assert.Equal(t, int64(10), n)
+			assert.Equal(t, testData, b.Bytes())
+			rewind()
+
+			n, err = rw.WriteTo(&b)
+			assert.NoError(t, err)
+			assert.Equal(t, 20, total)
+			assert.Equal(t, int64(10), n)
+			assert.Equal(t, testData, b.Bytes())
+			rewind()
 		})
-		var b bytes.Buffer
 
-		n, err := rw.WriteTo(&b)
+		t.Run("ReadError", func(t *testing.T) {
-		assert.Equal(t, errBoom, err)
+			// Test accounting errors
-		assert.Equal(t, int64(10), n)
+			rw := newRW()
-		assert.Equal(t, testData, b.Bytes())
+			defer close(rw)
+
+			rw.SetAccounting(func(n int) error {
+				return errBoom
+			})
+
+			dst = make([]byte, 3)
+			n, err = rw.Read(dst)
+			assert.Equal(t, 3, n)
+			assert.Equal(t, errBoom, err)
+		})
+
+		t.Run("WriteToError", func(t *testing.T) {
+			rw := newRW()
+			defer close(rw)
+			rw.SetAccounting(func(n int) error {
+				return errBoom
+			})
+			var b bytes.Buffer
+
+			n, err := rw.WriteTo(&b)
+			assert.Equal(t, errBoom, err)
+			assert.Equal(t, int64(10), n)
+			assert.Equal(t, testData, b.Bytes())
+		})
 	})
 
 }
@@ -363,26 +457,35 @@ func TestRWBoundaryConditions(t *testing.T) {
 		assert.Equal(t, int64(len(data)), nn)
 	}
 
+	type test struct {
+		name string
+		fn   func(*RW, []byte, int)
+	}
+
 	// Read and Write the data with a range of block sizes and functions
-	for _, writeFn := range []func(*RW, []byte, int){write, readFrom} {
+	for _, write := range []test{{"Write", write}, {"ReadFrom", readFrom}} {
-		for _, readFn := range []func(*RW, []byte, int){read, writeTo} {
+		t.Run(write.name, func(t *testing.T) {
-			for _, size := range sizes {
+			for _, read := range []test{{"Read", read}, {"WriteTo", writeTo}} {
-				data := buf[:size]
+				t.Run(read.name, func(t *testing.T) {
-				for _, chunkSize := range sizes {
+					for _, size := range sizes {
-					//t.Logf("Testing size=%d chunkSize=%d", useWrite, size, chunkSize)
+						data := buf[:size]
-					rw := NewRW(rwPool)
+						for _, chunkSize := range sizes {
-					assert.Equal(t, int64(0), rw.Size())
+							//t.Logf("Testing size=%d chunkSize=%d", useWrite, size, chunkSize)
-					accounted = 0
+							rw := NewRW(rwPool)
-					rw.SetAccounting(account)
+							assert.Equal(t, int64(0), rw.Size())
-					assert.Equal(t, 0, accounted)
+							accounted = 0
-					writeFn(rw, data, chunkSize)
+							rw.SetAccounting(account)
-					assert.Equal(t, int64(size), rw.Size())
+							assert.Equal(t, 0, accounted)
-					assert.Equal(t, 0, accounted)
+							write.fn(rw, data, chunkSize)
-					readFn(rw, data, chunkSize)
+							assert.Equal(t, int64(size), rw.Size())
-					assert.NoError(t, rw.Close())
+							assert.Equal(t, 0, accounted)
-					assert.Equal(t, size, accounted)
+							read.fn(rw, data, chunkSize)
-				}
+							assert.NoError(t, rw.Close())
+							assert.Equal(t, size, accounted)
+						}
+					}
+				})
 			}
-		}
+		})
 	}
 }