From 9fcd23bd3890e881b666b761b2b35a8d2db05fae Mon Sep 17 00:00:00 2001
From: Alexander Neumann <alexander@bumpern.de>
Date: Mon, 9 Feb 2015 23:39:16 +0100
Subject: [PATCH] Add parallel benchmark for chunk+encrypt

---
 archiver_test.go | 99 +++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 76 insertions(+), 23 deletions(-)

diff --git a/archiver_test.go b/archiver_test.go
index 2e1f66994..64cb16bd0 100644
--- a/archiver_test.go
+++ b/archiver_test.go
@@ -2,7 +2,6 @@ package restic_test
 
 import (
 	"bytes"
-	"crypto/sha256"
 	"io"
 	"math/rand"
 	"testing"
@@ -28,6 +27,38 @@ func get_random(seed, count int) []byte {
 
 const bufSize = chunker.MiB
 
+func benchmarkChunkEncrypt(b testing.TB, rd interface {
+	io.ReadSeeker
+	io.ReaderAt
+}, key *restic.Key) {
+	chunkBuf := make([]byte, restic.CiphertextExtension+chunker.MaxSize)
+
+	ch := restic.GetChunker("BenchmarkChunkEncrypt")
+	rd.Seek(0, 0)
+	ch.Reset(rd)
+
+	for {
+		chunk, err := ch.Next()
+
+		if err == io.EOF {
+			break
+		}
+
+		ok(b, err)
+
+		// reduce length of chunkBuf
+		chunkBuf = chunkBuf[:chunk.Length]
+		n, err := io.ReadFull(chunk.Reader(rd), chunkBuf)
+		ok(b, err)
+		assert(b, uint(n) == chunk.Length, "invalid length: got %d, expected %d", n, chunk.Length)
+
+		_, err = key.Encrypt(chunkBuf, chunkBuf)
+		ok(b, err)
+	}
+
+	restic.FreeChunker("BenchmarkChunkEncrypt", ch)
+}
+
 func BenchmarkChunkEncrypt(b *testing.B) {
 	data := get_random(23, 10<<20) // 10MiB
 	rd := bytes.NewReader(data)
@@ -35,32 +66,54 @@ func BenchmarkChunkEncrypt(b *testing.B) {
 	be := setupBackend(b)
 	defer teardownBackend(b, be)
 	key := setupKey(b, be, "geheim")
-	chunkBuf := make([]byte, restic.CiphertextExtension+chunker.MaxSize)
 
 	b.ResetTimer()
 	b.SetBytes(int64(len(data)))
 
 	for i := 0; i < b.N; i++ {
-		rd.Seek(0, 0)
-		ch := chunker.New(rd, bufSize, sha256.New)
-
-		for {
-			chunk, err := ch.Next()
-
-			if err == io.EOF {
-				break
-			}
-
-			ok(b, err)
-
-			// reduce length of chunkBuf
-			chunkBuf = chunkBuf[:chunk.Length]
-			n, err := io.ReadFull(chunk.Reader(rd), chunkBuf)
-			ok(b, err)
-			assert(b, uint(n) == chunk.Length, "invalid length: got %d, expected %d", n, chunk.Length)
-
-			_, err = key.Encrypt(chunkBuf, chunkBuf)
-			ok(b, err)
-		}
+		benchmarkChunkEncrypt(b, rd, key)
 	}
 }
+
+func benchmarkChunkEncryptP(b *testing.PB, rd interface {
+	io.ReadSeeker
+	io.ReaderAt
+}, key *restic.Key) {
+	chunkBuf := make([]byte, restic.CiphertextExtension+chunker.MaxSize)
+
+	ch := restic.GetChunker("BenchmarkChunkEncryptP")
+	rd.Seek(0, 0)
+	ch.Reset(rd)
+
+	for {
+		chunk, err := ch.Next()
+		if err == io.EOF {
+			break
+		}
+
+		// reduce length of chunkBuf
+		chunkBuf = chunkBuf[:chunk.Length]
+		io.ReadFull(chunk.Reader(rd), chunkBuf)
+		key.Encrypt(chunkBuf, chunkBuf)
+	}
+
+	restic.FreeChunker("BenchmarkChunkEncryptP", ch)
+}
+
+func BenchmarkChunkEncryptParallel(b *testing.B) {
+	be := setupBackend(b)
+	defer teardownBackend(b, be)
+	key := setupKey(b, be, "geheim")
+
+	data := get_random(23, 10<<20) // 10MiB
+
+	b.ResetTimer()
+	b.SetBytes(int64(len(data)))
+
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			rd := bytes.NewReader(data)
+			benchmarkChunkEncryptP(pb, rd, key)
+		}
+	})
+}