chunker_test.go

package desync

import (
	"bytes"
	"crypto/sha512"
	"math/bits"
	"os"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

const (
	ChunkSizeAvgDefault uint64 = 64 * 1024
	ChunkSizeMinDefault        = ChunkSizeAvgDefault / 4
	ChunkSizeMaxDefault        = ChunkSizeAvgDefault * 4
)

func TestChunkerLargeFile(t *testing.T) {
	f, err := os.Open("testdata/chunker.input")
	require.NoError(t, err)
	defer f.Close()

	expected := []struct {
		Start uint64
		Size  uint64
		ID    string
	}{
		{Start: 0, Size: 81590, ID: "ad951d7f65c27828ce390f3c81c41d75f80e4527169ad072ad720b56220f5be4"},
		{Start: 81590, Size: 46796, ID: "ef6df312072ccefe965f07669b2819902f4e9889ebe7c35a38f1dc11ee99f212"},
		{Start: 128386, Size: 36543, ID: "a816e22f4105741972eb34909b6f8ffa569759a1c2cf82ab88394b3db9019f23"},
		{Start: 164929, Size: 83172, ID: "8b8e4a274f06dc3c92d49869a699a5a8255c0bf0b48a4d3c3689aaa3e9cff090"},
		{Start: 248101, Size: 76749, ID: "583d08fc16d8d191af362a1aaecea6af062cc8afab1b301786bb717aa1b425b4"},
		{Start: 324850, Size: 79550, ID: "aefa8c5a3c86896110565b6a3748c2f985892e8ab0073730cac390cb478a913a"},
		{Start: 404400, Size: 41484, ID: "8e39f02975c8d0596e46f643b90cd290b7c0386845132eee4d415c63317773a4"},
		{Start: 445884, Size: 20326, ID: "d689ca889f2f7ba26896681214f0f0f5f5177d5820d99b1f11ddb76b693bddee"},
		{Start: 466210, Size: 31652, ID: "259de367c7ef2f51133d04e744f05918ceb93bd4b9c2bb6621ffeae70501dd09"},
		{Start: 497862, Size: 19995, ID: "01ae987ec457cacc8b3528e3254bc9c93b3f0c0b2a51619e15be16e678ef016d"},
		{Start: 517857, Size: 103873, ID: "78618b2d0539ecf45c08c7334e1c61051725767a76ba9108ad5298c6fd7cde1b"},
		{Start: 621730, Size: 38087, ID: "f44e6992cccadb08d8e18174ba3d6dd6365bdfb9906a58a9f82621ace0461c0d"},
		{Start: 659817, Size: 38377, ID: "abbf9935aaa535538c5fbff069481c343c2770207d88b94584314ee33050ae4f"},
		{Start: 698194, Size: 23449, ID: "a6c737b95ab514d6538c6ef4c42ef2f08b201c3426a88b95e67e517510cd1fb9"},
		{Start: 721643, Size: 47321, ID: "51d44e2d355d5c5b846543d47ba9569f12bbc3d49970c91913a8e3efef45e47e"},
		{Start: 768964, Size: 86692, ID: "90f7e061ed2fb1ed9594297851f8528d3ac355c98457b5dce08ee7d88f801b26"},
		{Start: 855656, Size: 28268, ID: "2dea144e5d771420e90b6e96c1e97e9c6afeda2c37ae7c95ceaf3ee2550efa08"},
		{Start: 883924, Size: 65465, ID: "7a94e051c82ec7abba32883b2eee9a2832e8e9bcc3b3151743fef533e2d46e70"},
		{Start: 949389, Size: 33255, ID: "32edd2d382045ad64d5fbd1a574f8191b700b9e0a2406bd90d2eefcf77168846"},
		{Start: 982644, Size: 65932, ID: "a8bfdadaecbee1ed16ce23d8bf771d1b3fbca2e631fc71b5adb3846c1bb2d542"},
	}

	c, err := NewChunker(f, ChunkSizeMinDefault, ChunkSizeAvgDefault, ChunkSizeMaxDefault)
	require.NoError(t, err)

	for i, e := range expected {
		start, buf, err := c.Next()
		require.NoError(t, err)
		chunkID := ChunkID(sha512.Sum512_256(buf))
		require.Equal(t, e.ID, chunkID.String(), "chunk #%d hash", i+1)
		require.Equal(t, e.Start, start, "chunk #%d start", i+1)
		require.Equal(t, e.Size, uint64(len(buf)), "chunk #%d size", i+1)
	}
	// Should get a size of 0 at the end
	_, buf, err := c.Next()
	require.NoError(t, err)
	require.Empty(t, buf, "expected size 0 at the end")
}

func TestChunkerEmptyFile(t *testing.T) {
	r := bytes.NewReader([]byte{})
	c, err := NewChunker(r, ChunkSizeMinDefault, ChunkSizeAvgDefault, ChunkSizeMaxDefault)
	require.NoError(t, err)
	start, buf, err := c.Next()
	require.NoError(t, err)
	require.Empty(t, buf)
	require.Equal(t, uint64(0), start)
}

func TestChunkerSmallFile(t *testing.T) {
	b := []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}
	r := bytes.NewReader(b)
	c, err := NewChunker(r, ChunkSizeMinDefault, ChunkSizeAvgDefault, ChunkSizeMaxDefault)
	require.NoError(t, err)

	start, buf, err := c.Next()
	require.NoError(t, err)
	require.Len(t, buf, len(b))
	require.Equal(t, uint64(0), start)
}

// There are no chunk boundaries when all data is nil, make sure we get the
// max chunk size
func TestChunkerNoBoundary(t *testing.T) {
	b := make([]byte, 1024*1024)
	r := bytes.NewReader(b)
	c, err := NewChunker(r, ChunkSizeMinDefault, ChunkSizeAvgDefault, ChunkSizeMaxDefault)
	require.NoError(t, err)
	for {
		start, buf, err := c.Next()
		require.NoError(t, err)
		if len(buf) == 0 {
			break
		}
		require.Equal(t, ChunkSizeMaxDefault, uint64(len(buf)))
		require.Zero(t, start%ChunkSizeMaxDefault, "unexpected start position %d", start)
	}
}

// Test with exactly min, avg, max chunk size of data
func TestChunkerBounds(t *testing.T) {
	for _, c := range []struct {
		name string
		size uint64
	}{
		{"chunker with exactly min chunk size data", ChunkSizeMinDefault},
		{"chunker with exactly avg chunk size data", ChunkSizeAvgDefault},
		{"chunker with exactly max chunk size data", ChunkSizeMaxDefault},
	} {
		t.Run(c.name, func(t *testing.T) {
			b := make([]byte, c.size)
			r := bytes.NewReader(b)
			c, err := NewChunker(r, ChunkSizeMinDefault, ChunkSizeAvgDefault, ChunkSizeMaxDefault)
			require.NoError(t, err)

			start, buf, err := c.Next()
			require.NoError(t, err)
			require.Len(t, buf, len(b))
			require.Equal(t, uint64(0), start)
		})
	}
}

// Test to confirm advancing through the input without producing chunks works.
func TestChunkerAdvance(t *testing.T) {
	// Build an input slice that is NullChunk + <dataA> + Nullchunk + <dataB>.
	// Then skip over the data slices and we should be left with only Null chunks.
	dataA := make([]byte, 128) // Short slice
	for i := range dataA {
		dataA[i] = 'a'
	}

	dataB := make([]byte, 12*ChunkSizeMaxDefault) // Long slice to ensure we read past the chunker-internal buffer
	for i := range dataB {
		dataB[i] = 'b'
	}

	nullChunk := NewNullChunk(ChunkSizeMaxDefault)

	// Build the input slice consisting of Null+dataA+Null+dataB
	input := join(nullChunk.Data, dataA, nullChunk.Data, dataB)

	c, err := NewChunker(bytes.NewReader(input), ChunkSizeMinDefault, ChunkSizeAvgDefault, ChunkSizeMaxDefault)
	require.NoError(t, err)

	// Chunk the first part, this should be a null chunk
	_, buf, err := c.Next()
	require.NoError(t, err)
	require.Equal(t, nullChunk.Data, buf, "expected null chunk")

	// Now skip the dataA slice
	require.NoError(t, c.Advance(len(dataA)))

	// Read the 2nd null chunk
	_, buf, err = c.Next()
	require.NoError(t, err)
	require.Equal(t, nullChunk.Data, buf, "expected null chunk")

	// Skip over dataB
	require.NoError(t, c.Advance(len(dataB)))

	// Should be at the end, nothing more to chunk
	_, buf, err = c.Next()
	require.NoError(t, err)
	require.Empty(t, buf, "expected end of input")
}

// Global vars used for results during the benchmark to prevent optimizer
// from optimizing away some operations
var (
	chunkStart uint64
	chunkBuf   []byte
)

// TestChunkerBoundaryTest verifies that the division-free boundary test built
// from the precomputed constants in NewChunker behaves exactly like the
// plain "hValue % d == d-1" for all edge cases, in particular the values
// where hValue+1 or hValue-(d-1) would wrap around 2^32. An earlier version
// of the fast test declared a false boundary for one hash value per
// discriminator (e.g. 14413 for the default 64KB average chunk size).
func TestChunkerBoundaryTest(t *testing.T) {
	for _, avg := range []uint64{16 * 1024, 64 * 1024, 256 * 1024, 1024 * 1024} {
		c, err := NewChunker(bytes.NewReader(nil), avg/4, avg, avg*4)
		require.NoError(t, err)
		d := c.hDiscriminator

		check := func(h uint32) {
			want := h%d == d-1
			got := bits.RotateLeft32((h+1)*c.hInverseOdd, c.hRot)-c.hQBias <= c.hQMax
			if want != got {
				assert.Equal(t, want, got, "hash value %d, discriminator %d", h, d)
			}
		}

		// The low range covers the wrap-around of hValue-(d-1), the high
		// range covers the wrap-around of hValue+1.
		for h := uint32(0); h < 3*d; h++ {
			check(h)
		}
		for h := ^uint32(0) - 3*d; h != 0; h++ {
			check(h)
		}
	}
}

func BenchmarkChunker(b *testing.B) {
	data, err := os.ReadFile("testdata/chunker.input")
	if err != nil {
		b.Fatal(err)
	}
	b.SetBytes(int64(len(data)))
	b.ResetTimer()
	for b.Loop() {
		c, err := NewChunker(bytes.NewReader(data), ChunkSizeMinDefault, ChunkSizeAvgDefault, ChunkSizeMaxDefault)
		if err != nil {
			b.Fatal(err)
		}
		for {
			start, buf, err := c.Next()
			if err != nil {
				b.Fatal(err)
			}
			if len(buf) == 0 {
				break
			}
			chunkStart = start
			chunkBuf = buf
		}
	}
}

func benchmarkChunkNull(b *testing.B, size int) {
	in := make([]byte, size)
	b.SetBytes(int64(size))
	for b.Loop() {
		c, err := NewChunker(bytes.NewReader(in), ChunkSizeMinDefault, ChunkSizeAvgDefault, ChunkSizeMaxDefault)
		if err != nil {
			b.Fatal(err)
		}
		for {
			start, buf, err := c.Next()
			if err != nil {
				b.Fatal(err)
			}
			if len(buf) == 0 {
				break
			}
			chunkStart = start
			chunkBuf = buf
		}
	}
}

func BenchmarkChunkNull1M(b *testing.B)   { benchmarkChunkNull(b, 1024*1024) }
func BenchmarkChunkNull10M(b *testing.B)  { benchmarkChunkNull(b, 10*1024*1024) }
func BenchmarkChunkNull50M(b *testing.B)  { benchmarkChunkNull(b, 50*1024*1024) }
func BenchmarkChunkNull100M(b *testing.B) { benchmarkChunkNull(b, 100*1024*1024) }