From ad2488a989366d7b0e6424ff7fdda955e1b38031 Mon Sep 17 00:00:00 2001
From: Klaus Post <klauspost@gmail.com>
Date: Wed, 16 Dec 2020 13:26:22 -0800
Subject: [PATCH] tests: Add random data stream compression (#302)

Add benchmark with random stream compression.
---
 s2/encode_test.go | 44 ++++++++++++++++++++++++++++++++++++++++----
 s2/s2_test.go     |  4 ++--
 2 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/s2/encode_test.go b/s2/encode_test.go
index 2d83095d0d..b4490b0411 100644
--- a/s2/encode_test.go
+++ b/s2/encode_test.go
@@ -16,7 +16,7 @@ import (
 	"github.com/klauspost/compress/zip"
 )
 
-func testOptions(t *testing.T) map[string][]WriterOption {
+func testOptions(t testing.TB) map[string][]WriterOption {
 	var testOptions = map[string][]WriterOption{
 		"default": {},
 		"better":  {WriterBetterCompression()},
@@ -47,16 +47,25 @@ func testOptions(t *testing.T) map[string][]WriterOption {
 	x = make(map[string][]WriterOption)
 	for name, opt := range testOptions {
 		x[name] = opt
-		x[name+"-pad-min"] = cloneAdd(opt, WriterPadding(2), WriterPaddingSrc(rand.New(rand.NewSource(0))))
+		x[name+"-pad-min"] = cloneAdd(opt, WriterPadding(2), WriterPaddingSrc(zeroReader{}))
 		if !testing.Short() {
-			x[name+"-pad-8000"] = cloneAdd(opt, WriterPadding(8000), WriterPaddingSrc(rand.New(rand.NewSource(0))))
-			x[name+"-pad-max"] = cloneAdd(opt, WriterPadding(4<<20), WriterPaddingSrc(rand.New(rand.NewSource(0))))
+			x[name+"-pad-8000"] = cloneAdd(opt, WriterPadding(8000), WriterPaddingSrc(zeroReader{}))
+			x[name+"-pad-max"] = cloneAdd(opt, WriterPadding(4<<20), WriterPaddingSrc(zeroReader{}))
 		}
 	}
 	testOptions = x
 	return testOptions
 }
 
+type zeroReader struct{}
+
+func (zeroReader) Read(p []byte) (int, error) {
+	for i := range p {
+		p[i] = 0
+	}
+	return len(p), nil
+}
+
 func TestEncoderRegression(t *testing.T) {
 	data, err := ioutil.ReadFile("testdata/enc_regressions.zip")
 	if err != nil {
@@ -265,3 +274,30 @@ func TestWriterPadding(t *testing.T) {
 		}
 	}
 }
+
+func BenchmarkWriterRandom(b *testing.B) {
+	rng := rand.New(rand.NewSource(1))
+	// Make max window so we never get matches.
+	data := make([]byte, 4<<20)
+	for i := range data {
+		data[i] = uint8(rng.Intn(256))
+	}
+
+	for name, opts := range testOptions(b) {
+		w := NewWriter(ioutil.Discard, opts...)
+		b.Run(name, func(b *testing.B) {
+			b.ResetTimer()
+			b.ReportAllocs()
+			b.SetBytes(int64(len(data)))
+			for i := 0; i < b.N; i++ {
+				err := w.EncodeBuffer(data)
+				if err != nil {
+					b.Fatal(err)
+				}
+			}
+			// Flush output
+			w.Flush()
+		})
+		w.Close()
+	}
+}
diff --git a/s2/s2_test.go b/s2/s2_test.go
index 8d7ff3fb8a..41b502c9e5 100644
--- a/s2/s2_test.go
+++ b/s2/s2_test.go
@@ -1391,7 +1391,7 @@ func BenchmarkWordsEncode1e4(b *testing.B) { benchWords(b, 1e4, false) }
 func BenchmarkWordsEncode1e5(b *testing.B) { benchWords(b, 1e5, false) }
 func BenchmarkWordsEncode1e6(b *testing.B) { benchWords(b, 1e6, false) }
 
-func BenchmarkRandomEncode(b *testing.B) {
+func BenchmarkRandomEncodeBlock1MB(b *testing.B) {
 	rng := rand.New(rand.NewSource(1))
 	data := make([]byte, 1<<20)
 	for i := range data {
@@ -1400,7 +1400,7 @@ func BenchmarkRandomEncode(b *testing.B) {
 	benchEncode(b, data)
 }
 
-func BenchmarkRandomEncodeBetter(b *testing.B) {
+func BenchmarkRandomEncodeBetterBlock1MB(b *testing.B) {
 	rng := rand.New(rand.NewSource(1))
 	data := make([]byte, 1<<20)
 	for i := range data {