fix: Improve load test external worker timing collection

.github/workflows/build.yml

@@ -257,11 +257,21 @@ jobs:
 
   loadtest:
     runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
     steps:
       - name: Clone repository
         uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3
+      - name: Login to GHCR
+        run: echo "${{ secrets.GITHUB_TOKEN }}" | docker login ghcr.io -u ${{ github.actor }} --password-stdin
       - name: Build loadtest
-        run: docker build -f ./build/package/loadtest.dockerfile .
+        run: docker build -f ./build/package/loadtest.dockerfile -t ghcr.io/hatchet-dev/hatchet/hatchet-loadtest:testing-no-batch .
+      # Mutable "testing" tag for pulling into a staging k8s cluster without a
+      # release - every PR that touches this job overwrites it, so it always
+      # reflects whichever PR pushed most recently, not any one branch.
+      - name: Push loadtest (testing tag)
+        run: docker push ghcr.io/hatchet-dev/hatchet/hatchet-loadtest:testing-no-batch
 
   loadtest-arm:
     runs-on: hatchet-arm64-2

cmd/hatchet-loadtest/do.go

@@ -256,15 +256,22 @@ func do(config LoadTestConfig) error {
 	var phaseSamples chan PhaseSample
 	var phaseResultCh <-chan phaseAccumulator
 	var cancelTiming context.CancelFunc
+	var collector *TimingCollector
 
 	if config.ExternalWorker {
 		workflowIDs := <-resolvedWorkflowIDs
 
+		// Deliberately not derived from `ctx`: `ctx`'s deadline is sized for
+		// registration + emission + config.Wait, with no slack for however
+		// long the collector needs to drain its backlog below. Tying the
+		// collector to it would silently truncate that drain once the
+		// overall test deadline hit, which is exactly the kind of premature
+		// cutoff we're trying to avoid.
 		var timingCtx context.Context
-		timingCtx, cancelTiming = context.WithCancel(ctx)
+		timingCtx, cancelTiming = context.WithCancel(context.Background())
 		defer cancelTiming() // safe to call more than once; guards every return path below
 
-		collector := NewTimingCollector(timingClient, workflowIDs, 2*time.Second)
+		collector = NewTimingCollector(timingClient, workflowIDs, timingPollInterval, config.TimingSampleRate)
 
 		phaseSamples = make(chan PhaseSample, 256)
 		phaseResultCh = accumulatePhases(phaseSamples)
@@ -286,34 +293,81 @@ func do(config LoadTestConfig) error {
 
 	var phases phaseAccumulator
 	if config.ExternalWorker {
-		// Give the engine config.Wait to finish (and the collector to observe) runs that
-		// were still in flight when emission stopped, before tearing down collection -
-		// cancelling immediately here would abort the collector mid-sweep for any run that
-		// completes right around this instant, silently dropping its timing sample instead
-		// of just not counting it (see the "context canceled" warnings this produces).
+		// Give the engine config.Wait to finish runs that were still in flight
+		// when emission stopped.
 		if config.Wait > 0 {
 			l.Info().Msgf("externalWorker: waiting %s for in-flight runs to complete before stopping timing collection...", config.Wait)
 			time.Sleep(config.Wait)
 		}
 
+		// Then keep the collector running until it has actually fetched
+		// every run it has discovered - cancelling on a fixed timer instead
+		// aborts whatever's still queued (however much that is) and
+		// permanently drops those samples. A run only ever leaves the
+		// collector's backlog by being successfully fetched or by aging out
+		// after timingSeenTTL (logged when that happens), so this loop is
+		// guaranteed to terminate on its own.
+		//
+		// Pending() hitting 0 isn't enough by itself: it can be a brief gap
+		// between sweeps, right before the next list pass turns up more
+		// newly-completed runs. Require it to stay at 0 across a full poll
+		// cycle (long enough for at least one more list pass to confirm
+		// there's really nothing left) before calling it drained.
+		l.Info().Msg("externalWorker: waiting for timing collector to fetch all discovered runs...")
+		logTicker := time.NewTicker(5 * time.Second)
+		var zeroSince time.Time
+		for {
+			pending := collector.Pending()
+
+			if pending == 0 {
+				if zeroSince.IsZero() {
+					zeroSince = time.Now()
+				} else if time.Since(zeroSince) >= 2*timingPollInterval {
+					break
+				}
+			} else {
+				zeroSince = time.Time{}
+			}
+
+			select {
+			case <-logTicker.C:
+				l.Info().Msgf("externalWorker: still waiting on %d run(s)...", pending)
+			case <-time.After(200 * time.Millisecond):
+			}
+		}
+		logTicker.Stop()
+
 		cancelTiming()
 		phases = <-phaseResultCh
 	}
 
 	expected := int64(config.EventFanout) * emitted * int64(config.DagSteps)
 
 	if config.ExternalWorker {
+		sampleRate := int64(config.TimingSampleRate)
+		if sampleRate < 1 {
+			sampleRate = 1
+		}
+		expectedSampled := expected / sampleRate
+
 		log.Printf(
-			"ℹ️ pushed %d, using %d events/s (externalWorker: engine-observed samples — queued n=%d, scheduling n=%d, execution n=%d)",
-			emitted, config.Events, phases.queued.count, phases.scheduling.count, phases.execution.count,
+			"ℹ️ pushed %d, using %d events/s (externalWorker: engine-observed samples, 1-in-%d sampled — queued n=%d, scheduling n=%d, execution n=%d)",
+			emitted, config.Events, sampleRate, phases.queued.count, phases.scheduling.count, phases.execution.count,
 		)
 
 		if phases.execution.count == 0 {
 			return fmt.Errorf("❌ no timing samples observed - check that the external SDK worker actually executed tasks for workflow(s) %v", expectedWorkflowNames(timingClient.V0().Namespace(), config.EventFanout))
 		}
 
-		if expected != phases.execution.count {
-			log.Printf("⚠️ warning: pushed and executed-timing-sample counts do not match: expected=%d got=%d", expected, phases.execution.count)
+		// Sampling only targets 1-in-sampleRate, not an exact count, and the
+		// collector already waits for a full drain (see the Pending() loop
+		// above) - so this is a sanity check that the sample landed in the
+		// right ballpark, not a completeness check.
+		if expectedSampled > 0 {
+			lower, upper := expectedSampled/2, expectedSampled+expectedSampled/2
+			if phases.execution.count < lower || phases.execution.count > upper {
+				log.Printf("⚠️ warning: engine-observed sample count is well outside the expected range: expected≈%d (1-in-%d sample of %d pushed) got=%d", expectedSampled, sampleRate, expected, phases.execution.count)
+			}
 		}
 	} else {
 		// NOTE: `emit()` returns successfully pushed events (not merely generated IDs),

cmd/hatchet-loadtest/go/main.go

@@ -65,7 +65,7 @@ func run() error {
 	delayMs := envInt("HATCHET_LOADTEST_DELAY_MS", 0)
 	failureRate := envFloat("HATCHET_LOADTEST_FAILURE_RATE", 0)
 	workerName := envOr("HATCHET_LOADTEST_WORKER_NAME", "load-test-worker")
-	batchTaskName := envOr("HATCHET_LOADTEST_BATCH_WORKFLOW_NAME", "load-test-batch")
+	// batchTaskName := envOr("HATCHET_LOADTEST_BATCH_WORKFLOW_NAME", "load-test-batch")
 
 	task := client.NewStandaloneTask(taskName, func(ctx hatchet.Context, input LoadTestInput) (LoadTestOutput, error) {
 		took := time.Since(input.CreatedAt)
@@ -94,25 +94,25 @@ func run() error {
 	// deliberately doesn't match the "load-test-%d" pattern that cmd/hatchet-loadtest's
 	// expectedWorkflowNames() (do.go) resolves, so the benchmark's TimingCollector never
 	// discovers or polls it and its timings never affect the pass/fail thresholds.
-	batchTask := client.NewStandaloneBatchTask(batchTaskName, func(ctx hatchet.Context, tasks map[string]LoadTestInput) (map[string]LoadTestOutput, error) {
-		out := make(map[string]LoadTestOutput, len(tasks))
-		for id := range tasks {
-			out[id] = LoadTestOutput{
-				Message: "This ran at: " + time.Now().Format(time.RFC3339Nano),
-			}
-		}
-		return out, nil
-	},
-		hatchet.BatchConfig{
-			MaxSize:     10,
-			MaxInterval: durationPtr(500 * time.Millisecond),
-		},
-		hatchet.WithWorkflowEvents(eventKey),
-	)
+	// batchTask := client.NewStandaloneBatchTask(batchTaskName, func(ctx hatchet.Context, tasks map[string]LoadTestInput) (map[string]LoadTestOutput, error) {
+	//	out := make(map[string]LoadTestOutput, len(tasks))
+	//	for id := range tasks {
+	//		out[id] = LoadTestOutput{
+	//			Message: "This ran at: " + time.Now().Format(time.RFC3339Nano),
+	//		}
+	//	}
+	//	return out, nil
+	// },
+	//	hatchet.BatchConfig{
+	//		MaxSize:     10,
+	//		MaxInterval: durationPtr(500 * time.Millisecond),
+	//	},
+	//	hatchet.WithWorkflowEvents(eventKey),
+	//)
 
 	worker, err := client.NewWorker(
 		workerName,
-		hatchet.WithWorkflows(task, batchTask),
+		hatchet.WithWorkflows(task),
 	)
 	if err != nil {
 		return fmt.Errorf("failed to create worker: %w", err)

cmd/hatchet-loadtest/main.go

@@ -45,6 +45,13 @@ type LoadTestConfig struct {
 	// compatible workflow named "load-test-0" (or "load-test-{i}" for
 	// i<EventFanout).
 	ExternalWorker bool
+
+	// TimingSampleRate, in --externalWorker mode, only fetches full task
+	// timings for 1 out of every TimingSampleRate completed runs, rather
+	// than every one - the reported average converges from a uniform sample
+	// just as well as from the full set, at a fraction of the REST load on
+	// the engine. 1 (or less) fetches every run.
+	TimingSampleRate int
 }
 
 func main() {
@@ -98,6 +105,7 @@ func main() {
 	loadtest.Flags().DurationVar(&config.AverageDurationThreshold, "averageDurationThreshold", 100*time.Millisecond, "averageDurationThreshold specifies the threshold for the average duration per executed event to be considered a success")
 	loadtest.Flags().StringVar(&config.PlotDir, "plotDirectory", "", "plotDirectory specifies where to put the generated plots for latency and task duration")
 	loadtest.Flags().BoolVar(&config.ExternalWorker, "externalWorker", false, "externalWorker skips registering a workflow and starting an in-process worker, assuming a separately-running SDK worker (e.g. cmd/hatchet-loadtest/go) has already registered a compatible workflow; worker/workflow flags (slots, dagSteps, eventFanout, rlKeys, workerDelay, failureRate, delay) are ignored in this mode")
+	loadtest.Flags().IntVar(&config.TimingSampleRate, "timingSampleRate", 10, "in --externalWorker mode, fetch full task timings for 1 out of every N completed runs instead of every one - the average still converges, at a fraction of the REST load on the engine; 1 fetches every run")
 	cmd := &cobra.Command{Use: "app"}
 	cmd.AddCommand(loadtest)
 	if err := cmd.Execute(); err != nil {