From eb139b7a84fdaa193437ab155996e18b54e381d5 Mon Sep 17 00:00:00 2001
From: Klaus Post <klauspost@gmail.com>
Date: Tue, 31 Dec 2024 15:48:48 +0100
Subject: [PATCH] Add CLI benchmarks

Add in-memory benchmark of stream encoding.

Specify `-bench=n` to run n benchmarks on every input with provided settings.
---
 cmd/lz4c/compress.go   | 42 ++++++++++++++++++++++++++++++++++++++++++
 cmd/lz4c/uncompress.go | 32 +++++++++++++++++++++++++++++++-
 2 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/cmd/lz4c/compress.go b/cmd/lz4c/compress.go
index 9ce26377..5cfff1e1 100644
--- a/cmd/lz4c/compress.go
+++ b/cmd/lz4c/compress.go
@@ -1,11 +1,14 @@
 package main
 
 import (
+	"bytes"
 	"flag"
 	"fmt"
 	"io"
 	"os"
+	"runtime"
 	"sync/atomic"
+	"time"
 
 	"code.cloudfoundry.org/bytefmt"
 	"github.com/schollz/progressbar/v3"
@@ -26,6 +29,7 @@ func Compress(fs *flag.FlagSet) cmdflag.Handler {
 	fs.UintVar(&level, "l", 0, "compression level (0=fastest)")
 	var concurrency int
 	fs.IntVar(&concurrency, "c", -1, "concurrency (default=all CPUs")
+	bench := fs.Int("bench", 0, "Run benchmark n times. No output will be written")
 
 	var lvl lz4.CompressionLevel
 	switch level {
@@ -87,6 +91,33 @@ func Compress(fs *flag.FlagSet) cmdflag.Handler {
 			if err != nil {
 				return fidx, err
 			}
+			if *bench > 0 {
+				fmt.Print("Reading ", filename, "...")
+				input, err := io.ReadAll(file)
+				if err != nil {
+					return fidx, err
+				}
+				file.Close()
+				for i := 0; i < *bench; i++ {
+					fmt.Print("\nCompressing...")
+					runtime.GC()
+					start := time.Now()
+					counter := wCounter{out: io.Discard}
+					zw.Reset(&counter)
+					_, err := io.Copy(zw, bytes.NewReader(input))
+					if err != nil {
+						return fidx, err
+					}
+					output := counter.n
+					elapsed := time.Since(start)
+					ms := elapsed.Round(time.Millisecond)
+					mbPerSec := (float64(len(input)) / 1e6) / (float64(elapsed) / (float64(time.Second)))
+					pct := float64(output) * 100 / float64(len(input))
+					fmt.Printf(" %d -> %d [%.02f%%]; %v, %.01fMB/s", len(input), output, pct, ms, mbPerSec)
+				}
+				fmt.Println("")
+				continue
+			}
 			finfo, err := file.Stat()
 			if err != nil {
 				return fidx, err
@@ -147,3 +178,14 @@ func Compress(fs *flag.FlagSet) cmdflag.Handler {
 		return len(args), nil
 	}
 }
+
+type wCounter struct {
+	n   int
+	out io.Writer
+}
+
+func (w *wCounter) Write(p []byte) (n int, err error) {
+	n, err = w.out.Write(p)
+	w.n += n
+	return n, err
+}
diff --git a/cmd/lz4c/uncompress.go b/cmd/lz4c/uncompress.go
index 04624dbe..5bc2e4f6 100644
--- a/cmd/lz4c/uncompress.go
+++ b/cmd/lz4c/uncompress.go
@@ -1,11 +1,14 @@
 package main
 
 import (
+	"bytes"
 	"flag"
 	"fmt"
 	"io"
 	"os"
+	"runtime"
 	"strings"
+	"time"
 
 	"github.com/schollz/progressbar/v3"
 
@@ -14,7 +17,8 @@ import (
 )
 
 // Uncompress uncompresses a set of files or from stdin to stdout.
-func Uncompress(_ *flag.FlagSet) cmdflag.Handler {
+func Uncompress(fs *flag.FlagSet) cmdflag.Handler {
+	bench := fs.Int("bench", 0, "Run benchmark n times. No output will be written")
 	return func(args ...string) (int, error) {
 		zr := lz4.NewReader(nil)
 
@@ -31,6 +35,32 @@ func Uncompress(_ *flag.FlagSet) cmdflag.Handler {
 			if err != nil {
 				return fidx, err
 			}
+
+			if *bench > 0 {
+				fmt.Print("Reading ", zfilename, "...")
+				compressed, err := io.ReadAll(zfile)
+				if err != nil {
+					return fidx, err
+				}
+				zfile.Close()
+				for i := 0; i < *bench; i++ {
+					fmt.Print("\nDecompressing...")
+					runtime.GC()
+					start := time.Now()
+					zr.Reset(bytes.NewReader(compressed))
+					output, err := io.Copy(io.Discard, zr)
+					if err != nil {
+						return fidx, err
+					}
+					elapsed := time.Since(start)
+					ms := elapsed.Round(time.Millisecond)
+					mbPerSec := (float64(output) / 1e6) / (float64(elapsed) / (float64(time.Second)))
+					pct := float64(output) * 100 / float64(len(compressed))
+					fmt.Printf(" %d -> %d [%.02f%%]; %v, %.01fMB/s", len(compressed), output, pct, ms, mbPerSec)
+				}
+				fmt.Println("")
+				continue
+			}
 			zinfo, err := zfile.Stat()
 			if err != nil {
 				return fidx, err