optimize thumper

.dockerignore

@@ -0,0 +1,24 @@
+# Keep the build context small and reproducible. The build only needs the Go
+# sources, go.mod/go.sum, and the scripts/ directory.
+.git
+.github
+.idea
+dist
+Dockerfile
+.dockerignore
+
+# Docs and repo metadata — not needed to build the binary.
+*.md
+LICENSE
+NOTICE
+DCO
+CONTRIBUTING.md
+CODE-OF-CONDUCT.md
+
+# Linters and tooling config.
+.golangci.yaml
+.goreleaser.yaml
+.goreleaser.nightly.yaml
+.markdownlint.yaml
+.pre-commit-config.yaml
+.yamllint

Dockerfile

@@ -0,0 +1,30 @@
+# syntax=docker/dockerfile:1
+
+# Build stage runs natively on the builder's architecture ($BUILDPLATFORM) and
+# cross-compiles to the requested target. thumper is a pure-Go, CGO-free static
+# binary, so cross-compilation needs no emulation or C toolchain.
+FROM --platform=$BUILDPLATFORM cgr.dev/chainguard/go:latest AS build
+ARG TARGETOS
+ARG TARGETARCH
+WORKDIR /work
+
+# Download modules in their own layer so they're cached across source changes.
+COPY go.mod go.sum ./
+RUN go mod download
+
+COPY . .
+RUN CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH \
+    go build -ldflags="-s -w" -o /out/thumper ./cmd/thumper/
+
+# Minimal distroless runtime: no shell or package manager, runs as nonroot, and
+# ships CA certificates and tzdata.
+FROM cgr.dev/chainguard/static:latest
+COPY --from=build /out/thumper /usr/bin/thumper
+
+# Bundle the bundled scripts at the ko data location so that the released-image
+# UX works here too: `thumper run scripts/example.yaml` resolves the file via
+# KO_DATA_PATH (see internal/config/load.go findFile).
+COPY scripts/ /var/run/ko/scripts/
+ENV KO_DATA_PATH=/var/run/ko
+
+ENTRYPOINT ["/usr/bin/thumper"]

internal/cmd/run.go

@@ -2,6 +2,8 @@ package cmd
 
 import (
 	"fmt"
+	"net/http"
+	"net/http/pprof"
 	"sync"
 	"time"
 
@@ -30,17 +32,36 @@ var (
 		cobrahttp.WithFlagPrefix("metrics"),
 		cobrahttp.WithDefaultEnabled(true),
 		cobrahttp.WithHandler(promhttp.Handler()))
+	// PProfServerBuilder serves the net/http/pprof profiling endpoints. It is
+	// disabled by default and enabled with --pprof-enabled.
+	PProfServerBuilder = cobrahttp.New("pprof",
+		cobrahttp.WithDefaultAddress(":6060"),
+		cobrahttp.WithFlagPrefix("pprof"),
+		cobrahttp.WithDefaultEnabled(false),
+		cobrahttp.WithHandler(pprofHandler()))
 )
 
+// pprofHandler returns a mux serving the standard net/http/pprof endpoints
+// under /debug/pprof/.
+func pprofHandler() http.Handler {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/debug/pprof/", pprof.Index)
+	mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
+	mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
+	mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
+	mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
+	return mux
+}
+
 func RegisterRunFlags(cmd *cobra.Command) {
 	cmd.Flags().Int("qps", 1, "queries per second to generate")
 	cmd.Flags().Duration("step-timeout", 500*time.Millisecond, "maximum time a single step is allowed to run")
 	cmd.Flags().Duration("step-interval", time.Second, "time each worker waits between steps; lower it to drive more steps per worker (effective throughput is qps / step-interval)")
 	cmd.Flags().Bool("randomize-starting-step", false, "randomize the starting script step for each worker")
-	cmd.Flags().Bool("rerender", false, "re-render each script from its source file at the start of every cycle through its steps, regenerating template values such as randomObjectID for each cycle instead of once at startup")
 
 	// Register http flags
 	MetricsServerBuilder.RegisterFlags(cmd.Flags())
+	PProfServerBuilder.RegisterFlags(cmd.Flags())
 }
 
 var RunCmd = &cobra.Command{
@@ -66,7 +87,6 @@ func runCmdFunc(cmd *cobra.Command, args []string) error {
 	stepTimeout := cobrautil.MustGetDuration(cmd, "step-timeout")
 	stepInterval := cobrautil.MustGetDuration(cmd, "step-interval")
 	stepRandomization := cobrautil.MustGetBool(cmd, "randomize-starting-step")
-	rerender := cobrautil.MustGetBool(cmd, "rerender")
 	psName := cobrautil.MustGetString(cmd, "permissions-system")
 	log.Info().Int("qps", qps).Str("permission-system", psName).Msg("starting run command")
 
@@ -78,43 +98,57 @@ func runCmdFunc(cmd *cobra.Command, args []string) error {
 	// Keep track of the total stats for all workers
 	var scriptsForStats []*thumperconf.Script
 
-	scriptCache := make(map[string][]*thumperrunner.ExecutableScript, len(args))
+	// Parse each unique script file exactly once. Scripts that use
+	// randomObjectID keep their placeholders here; each worker resolves them to
+	// fresh values at runtime (see PrepareRenderable), so the YAML is never
+	// re-parsed no matter how long the run lasts. Static scripts are identical
+	// across workers, so they are prepared once here and shared by pointer.
+	type loadedFile struct {
+		scripts    []*thumperconf.Script
+		usedRandom bool
+		prepared   []*thumperrunner.ExecutableScript // set only for static scripts
+	}
+	files := make(map[string]loadedFile, len(args))
+	for _, scriptFilename := range args {
+		if _, ok := files[scriptFilename]; ok {
+			continue
+		}
 
-	// Load the scripts and transform them, one copy per worker
+		scripts, usedRandom, err := thumperconf.Load(scriptFilename, scriptVars)
+		if err != nil {
+			return fmt.Errorf("unable to load script file: %w", err)
+		}
+
+		loaded := loadedFile{scripts: scripts, usedRandom: usedRandom}
+		if !usedRandom {
+			if loaded.prepared, err = thumperrunner.Prepare(scripts); err != nil {
+				return fmt.Errorf("error preparing scripts for execution: %w", err)
+			}
+		}
+		files[scriptFilename] = loaded
+		scriptsForStats = append(scriptsForStats, scripts...)
+	}
+
+	// Build one set of executable scripts per worker. Scripts with random IDs
+	// get their own renderable copy per worker (so a randomized starting step
+	// can't make workers emit identical IDs during their first partial cycle);
+	// static scripts reuse the single shared Prepare result from above.
 	workerScripts := make([][]*thumperrunner.ExecutableScript, 0, qps)
 	for i := 0; i < qps; i++ {
 		var preparedScripts []*thumperrunner.ExecutableScript
 		for _, scriptFilename := range args {
-			if cached, ok := scriptCache[scriptFilename]; ok {
-				preparedScripts = append(preparedScripts, cached...)
+			loaded := files[scriptFilename]
 
-				// Skip actually loading it from disk
+			if !loaded.usedRandom {
+				preparedScripts = append(preparedScripts, loaded.prepared...)
 				continue
 			}
 
-			fileScripts, usedRandom, err := thumperconf.Load(scriptFilename, scriptVars)
-			if err != nil {
-				return fmt.Errorf("unable to load script file: %w", err)
-			}
-
-			if i == 0 {
-				scriptsForStats = append(scriptsForStats, fileScripts...)
-			}
-
-			preparedFileScripts, err := thumperrunner.Prepare(fileScripts)
+			renderable, err := thumperrunner.PrepareRenderable(loaded.scripts)
 			if err != nil {
 				return fmt.Errorf("error preparing scripts for execution: %w", err)
 			}
-
-			if rerender {
-				thumperrunner.EnableRerender(preparedFileScripts, scriptFilename, scriptVars)
-			}
-
-			if !usedRandom {
-				scriptCache[scriptFilename] = preparedFileScripts
-			}
-
-			preparedScripts = append(preparedScripts, preparedFileScripts...)
+			preparedScripts = append(preparedScripts, renderable...)
 		}
 
 		workerScripts = append(workerScripts, preparedScripts)
@@ -157,6 +191,15 @@ func runCmdFunc(cmd *cobra.Command, args []string) error {
 		}
 	}()
 
+	// Start the pprof endpoint. Disabled by default; ListenFromFlags is a no-op
+	// unless --pprof-enabled is set.
+	pprofSrv := PProfServerBuilder.ServerFromFlags(cmd)
+	go func() {
+		if err := PProfServerBuilder.ListenFromFlags(cmd, pprofSrv); err != nil {
+			log.Fatal().Err(err).Msg("failed while serving pprof")
+		}
+	}()
+
 	wg.Wait()
 	log.Info().Msg("terminating")
 

internal/config/load.go

@@ -4,11 +4,11 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
-	"html/template"
 	"io"
 	"math/rand"
 	"os"
 	"path"
+	"text/template"
 
 	"github.com/Masterminds/sprig/v3"
 	"github.com/ccoveille/go-safecast/v2"
@@ -18,9 +18,14 @@ import (
 
 // Load reads a script file, replaces the templated values with the values from
 // the execution environment, and then processes it as a thumper script yaml.
+//
+// randomObjectID does not bake a concrete value at load time; it emits a keyed
+// placeholder that is resolved to a fresh random value at runtime (see
+// ResolveScript). This lets a single parse be reused for every execution cycle
+// while still producing new, uniformly-distributed, correlated IDs each cycle.
+// The returned bool reports whether the script used randomObjectID at all.
 func Load(filename string, vars ScriptVariables) ([]*Script, bool, error) {
 	usedRandom := false
-	randomID := randomObjectID(64)
 
 	// Look for the file in the given path *or* in the kodata dir
 	filepath, err := findFile(filename, os.Getenv("KO_DATA_PATH"))
@@ -41,9 +46,15 @@ func Load(filename string, vars ScriptVariables) ([]*Script, bool, error) {
 			}
 			return indices
 		},
-		"randomObjectID": func() string {
+		// randomObjectID emits a placeholder keyed by its (optional) argument.
+		// Calls with the same key render to the same placeholder, so the same
+		// value can be referenced across steps (e.g. to CREATE and later DELETE
+		// the same relationship); distinct keys render to distinct placeholders.
+		// A keyless call uses a single shared key. The placeholders are resolved
+		// to concrete random values per cycle by ResolveScript.
+		"randomObjectID": func(key ...any) string {
 			usedRandom = true
-			return randomID
+			return randomIDPlaceholder(fmt.Sprint(key...))
 		},
 	}).Funcs(sprig.FuncMap())
 

internal/config/resolve.go

@@ -0,0 +1,138 @@
+package config
+
+import (
+	"strings"
+	"unicode/utf8"
+)
+
+// Random-ID placeholders are delimited by Unicode private-use-area runes: they
+// never appear in real object IDs (so substitution can't hit a real value) and
+// survive YAML decoding untouched, so a placeholder embedded in an ID
+// (e.g. "org_<ph>") arrives intact at ResolveScript.
+const (
+	randomIDStart = '\ue000'
+	randomIDEnd   = '\ue001'
+)
+
+// randomIDPlaceholder returns the placeholder token for a given key. See the
+// randomObjectID template function in Load.
+func randomIDPlaceholder(key string) string {
+	return string(randomIDStart) + key + string(randomIDEnd)
+}
+
+// substituteRandomIDs replaces every random-ID placeholder in s with gen(key),
+// where key is the text the placeholder was created with. gen is expected to
+// return the same value for repeated calls with the same key so that IDs
+// referenced in multiple places stay consistent within one resolution.
+func substituteRandomIDs(s string, gen func(key string) string) string {
+	i := strings.IndexRune(s, randomIDStart)
+	if i < 0 {
+		return s
+	}
+
+	var b strings.Builder
+	b.Grow(len(s)) // resolved values are at least as long as the tokens they replace
+	for {
+		b.WriteString(s[:i])
+
+		rest := s[i+utf8.RuneLen(randomIDStart):]
+		j := strings.IndexRune(rest, randomIDEnd)
+		if j < 0 {
+			// Unterminated placeholder; emit the remainder verbatim rather than
+			// dropping data.
+			b.WriteString(s[i:])
+			return b.String()
+		}
+
+		b.WriteString(gen(rest[:j]))
+		s = rest[j+utf8.RuneLen(randomIDEnd):]
+
+		i = strings.IndexRune(s, randomIDStart)
+		if i < 0 {
+			b.WriteString(s)
+			return b.String()
+		}
+	}
+}
+
+// ResolveScript returns a copy of template with every random-ID placeholder
+// replaced by a freshly generated random object ID. Placeholders that share a
+// key resolve to the same value within the returned script (so, for example, a
+// CREATE and a DELETE step that reference the same key operate on the same
+// relationship), while distinct keys resolve to independent, uniformly
+// distributed values (scattering writes across the keyspace). Each call
+// generates new values, so calling it once per cycle yields fresh data without
+// re-parsing the script.
+//
+// The input template is not modified and may be resolved concurrently.
+func ResolveScript(template *Script) *Script {
+	values := make(map[string]string)
+	gen := func(key string) string {
+		v, ok := values[key]
+		if !ok {
+			v = randomObjectID(64)
+			values[key] = v
+		}
+		return v
+	}
+	return resolveScript(template, gen)
+}
+
+func resolveScript(in *Script, gen func(string) string) *Script {
+	out := &Script{
+		Name:      in.Name,
+		Weight:    in.Weight,
+		RecordTTL: in.RecordTTL,
+	}
+	if len(in.Steps) > 0 {
+		out.Steps = make([]ScriptStep, len(in.Steps))
+		for i, step := range in.Steps {
+			out.Steps[i] = resolveStep(step, gen)
+		}
+	}
+	return out
+}
+
+// resolveStep receives step by value (a shallow copy) and returns it with every
+// placeholder-bearing string field resolved. Slices that are rewritten are
+// reallocated so the input script's backing arrays are never mutated.
+func resolveStep(step ScriptStep, gen func(string) string) ScriptStep {
+	sub := func(p *string) { *p = substituteRandomIDs(*p, gen) }
+
+	// Only the fields that carry object references / IDs are resolved; keyword
+	// fields such as Op and Consistency are left alone. Context (*ProtoStruct)
+	// is intentionally not descended into.
+	sub(&step.Resource)
+	sub(&step.Subject)
+	sub(&step.Permission)
+	sub(&step.Schema)
+
+	if len(step.Updates) > 0 {
+		updates := make([]Update, len(step.Updates))
+		for i, u := range step.Updates {
+			sub(&u.Resource)
+			sub(&u.Subject)
+			sub(&u.Relation)
+			if u.Caveat != nil {
+				caveat := *u.Caveat
+				sub(&caveat.Name)
+				u.Caveat = &caveat
+			}
+			updates[i] = u
+		}
+		step.Updates = updates
+	}
+
+	if len(step.Checks) > 0 {
+		checks := make([]Check, len(step.Checks))
+		for i, c := range step.Checks {
+			sub(&c.Resource)
+			sub(&c.Subject)
+			sub(&c.Permission)
+			checks[i] = c
+		}
+		step.Checks = checks
+	}
+
+	return step
+}