From 2f9d0b51dd27a0e389b6ed4d9bddf482e66b4bc6 Mon Sep 17 00:00:00 2001
From: Martin Hutchinson <mhutchinson@gmail.com>
Date: Thu, 15 Aug 2024 13:34:04 +0100
Subject: [PATCH] Hammer: exit criteria (#141)

The hammer can be configured to stop successfully if a desired number of leaves are written, and to fail after a max runtime is reached. These two can be used in conjunction to run a headless test for write performance.
---
 hammer/README.md | 15 +++++++++++++++
 hammer/hammer.go | 45 ++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/hammer/README.md b/hammer/README.md
index 6065c450..dfd250df 100644
--- a/hammer/README.md
+++ b/hammer/README.md
@@ -29,3 +29,18 @@ go run ./hammer \
   --max_write_ops=42
 ```
 
+For a headless write-only example that could be used for integration tests, this command attempts to write 2500 leaves within 1 minute.
+If the target number of leaves is reached then it exits successfully.
+If the timeout of 1 minute is reached first, then it exits with an exit code of 1.
+
+```shell
+go run ./hammer \
+  --log_public_key=Test-Betty+df84580a+AQQASqPUZoIHcJAF5mBOryctwFdTV1E0GRY4kEAtTzwB \
+  --log_url=http://localhost:2024 \
+  --max_read_ops=0 \
+  --num_writers=512 \
+  --max_write_ops=512 \
+  --max_runtime=1m \
+  --leaf_write_goal=2500 \
+  --show_ui=false
+```
diff --git a/hammer/hammer.go b/hammer/hammer.go
index dfa0fbd3..86a7253c 100644
--- a/hammer/hammer.go
+++ b/hammer/hammer.go
@@ -52,6 +52,9 @@ var (
 	leafMinSize = flag.Int("leaf_min_size", 0, "Minimum size in bytes of individual leaves")
 	dupChance   = flag.Float64("dup_chance", 0.1, "The probability of a generated leaf being a duplicate of a previous value")
 
+	leafWriteGoal = flag.Int64("leaf_write_goal", 0, "Exit after writing this number of leaves, or 0 to keep going indefinitely")
+	maxRunTime    = flag.Duration("max_runtime", 0, "Fail after this amount of time has passed, or 0 to keep going indefinitely")
+
 	showUI = flag.Bool("show_ui", true, "Set to false to disable the text-based UI")
 
 	hc = &http.Client{
@@ -68,7 +71,7 @@ func main() {
 	klog.InitFlags(nil)
 	flag.Parse()
 
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
 
 	logSigV, err := note.NewVerifier(*logPubKey)
 	if err != nil {
@@ -95,6 +98,45 @@ func main() {
 
 	gen := newLeafGenerator(tracker.LatestConsistent.Size, *leafMinSize, *dupChance)
 	hammer := NewHammer(&tracker, f.Fetch, w.Write, gen, ha.seqLeafChan, ha.errChan)
+
+	exitCode := 0
+	if *leafWriteGoal > 0 {
+		go func() {
+			startTime := time.Now()
+			goal := tracker.LatestConsistent.Size + uint64(*leafWriteGoal)
+			klog.Infof("Will exit once tree size is at least %d", goal)
+			tick := time.NewTicker(1 * time.Second)
+			for {
+				select {
+				case <-ctx.Done():
+					return
+				case <-tick.C:
+					if tracker.LatestConsistent.Size >= goal {
+						elapsed := time.Since(startTime)
+						klog.Infof("Reached tree size goal of %d after %s; exiting", goal, elapsed)
+						cancel()
+						return
+					}
+				}
+			}
+		}()
+	}
+	if *maxRunTime > 0 {
+		go func() {
+			klog.Infof("Will fail after %s", *maxRunTime)
+			for {
+				select {
+				case <-ctx.Done():
+					return
+				case <-time.After(*maxRunTime):
+					klog.Infof("Max runtime reached; exiting")
+					exitCode = 1
+					cancel()
+					return
+				}
+			}
+		}()
+	}
 	hammer.Run(ctx)
 
 	if *showUI {
@@ -103,6 +145,7 @@ func main() {
 	} else {
 		<-ctx.Done()
 	}
+	os.Exit(exitCode)
 }
 
 func NewHammer(tracker *client.LogStateTracker, f client.Fetcher, w LeafWriter, gen func() []byte, seqLeafChan chan<- leafTime, errChan chan<- error) *Hammer {