diff --git a/bench/internal/roles/config.go b/bench/internal/roles/config.go index 657c334..dd4a3f5 100644 --- a/bench/internal/roles/config.go +++ b/bench/internal/roles/config.go @@ -27,12 +27,13 @@ type KQConfig struct { } type TopologyConfig struct { - ReadyPartitions int `json:"ready_partitions,omitempty"` - RetryPartitions int `json:"retry_partitions,omitempty"` - Producers int `json:"producers"` - Workers int `json:"workers"` - WorkerConcurrency int `json:"worker_concurrency"` - Movers int `json:"movers"` + ReadyPartitions int `json:"ready_partitions,omitempty"` + RetryPartitions int `json:"retry_partitions,omitempty"` + Producers int `json:"producers"` + ProducerConcurrency int `json:"producer_concurrency"` + Workers int `json:"workers"` + WorkerConcurrency int `json:"worker_concurrency"` + Movers int `json:"movers"` } type WorkloadConfig struct { @@ -84,6 +85,9 @@ func (c Config) Validate() error { if c.Topology.Producers <= 0 { return errors.New("roles: producer count must be positive") } + if c.Topology.ProducerConcurrency <= 0 { + return errors.New("roles: producer concurrency must be positive") + } if c.Topology.Workers <= 0 { return errors.New("roles: worker count must be positive") } diff --git a/bench/internal/roles/producer.go b/bench/internal/roles/producer.go index 3bb0f1d..c7a12c5 100644 --- a/bench/internal/roles/producer.go +++ b/bench/internal/roles/producer.go @@ -3,6 +3,8 @@ package roles import ( "context" "encoding/json" + "errors" + "sync" "time" "github.com/raydatray/kq" @@ -36,23 +38,82 @@ func RunProducer(ctx context.Context, config Config, output *events.Writer, warm return runWarmup(ctx, client, config, workloadConfig, output) } + err = runMeasuredProducer(ctx, client, config, workloadConfig, output, started) + _ = output.Write(events.ProcessStopped(events.RoleProducer, config.ProcessIndex, err)) + return err +} + +func runMeasuredProducer( + ctx context.Context, + client *kq.Client, + config Config, + workloadConfig workload.Config, + output *events.Writer, + started time.Time, +) error { + pool := newProducerPool(ctx, config.Topology.ProducerConcurrency, func(ctx context.Context, job workload.Job) error { + return enqueueOne(ctx, client, config.ProcessIndex, job, output) + }) + + var dispatchErr error total := config.TotalTasks() for id := uint64(config.ProcessIndex); id < total; id += uint64(config.Topology.Producers) { if err := waitUntil(ctx, config.DueAt(started, id)); err != nil { - _ = output.Write(events.ProcessStopped(events.RoleProducer, config.ProcessIndex, err)) - return err + dispatchErr = err + break } job := workload.Generate(workloadConfig, id) job.EnqueuedAt = time.Now() - if err := enqueueOne(ctx, client, config.ProcessIndex, job, output); err != nil { - _ = output.Write(events.ProcessStopped(events.RoleProducer, config.ProcessIndex, err)) - return err + if err := pool.Submit(ctx, job); err != nil { + dispatchErr = err + break } } - _ = output.Write(events.ProcessStopped(events.RoleProducer, config.ProcessIndex, nil)) - return nil + return errors.Join(dispatchErr, pool.Close()) +} + +type producerPool struct { + jobs chan workload.Job + workers sync.WaitGroup + errOnce sync.Once + err error +} + +func newProducerPool( + ctx context.Context, + concurrency int, + enqueue func(context.Context, workload.Job) error, +) *producerPool { + pool := &producerPool{jobs: make(chan workload.Job, concurrency)} + for range concurrency { + pool.workers.Add(1) + go func() { + defer pool.workers.Done() + for job := range pool.jobs { + if err := enqueue(ctx, job); err != nil { + pool.errOnce.Do(func() { pool.err = err }) + } + } + }() + } + return pool +} + +func (p *producerPool) Submit(ctx context.Context, job workload.Job) error { + select { + case p.jobs <- job: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (p *producerPool) Close() error { + close(p.jobs) + p.workers.Wait() + return p.err } func runWarmup(ctx context.Context, client *kq.Client, config Config, workloadConfig workload.Config, output *events.Writer) error { diff --git a/bench/internal/roles/producer_test.go b/bench/internal/roles/producer_test.go index c2823d3..3ee464d 100644 --- a/bench/internal/roles/producer_test.go +++ b/bench/internal/roles/producer_test.go @@ -1,8 +1,12 @@ package roles import ( + "context" + "sync/atomic" "testing" "time" + + "github.com/raydatray/kq/bench/internal/workload" ) func testProducerConfig() Config { @@ -16,10 +20,11 @@ func testProducerConfig() Config { RetryGridPartitions: 4, }, Topology: TopologyConfig{ - Producers: 2, - Workers: 1, - WorkerConcurrency: 1, - Movers: 1, + Producers: 2, + ProducerConcurrency: 1, + Workers: 1, + WorkerConcurrency: 1, + Movers: 1, }, Workload: WorkloadConfig{ Arrival: ArrivalConfig{ @@ -125,6 +130,11 @@ func TestConfigValidation(t *testing.T) { t.Fatal("expected producer count error") } bad = good + bad.Topology.ProducerConcurrency = 0 + if err := bad.Validate(); err == nil { + t.Fatal("expected producer concurrency error") + } + bad = good bad.Topology.WorkerConcurrency = 0 if err := bad.Validate(); err == nil { t.Fatal("expected worker concurrency error") @@ -141,6 +151,48 @@ func TestConfigValidation(t *testing.T) { } } +func TestProducerPoolBoundsConcurrency(t *testing.T) { + const concurrency = 4 + started := make(chan struct{}, concurrency) + release := make(chan struct{}) + var active atomic.Int32 + var maximum atomic.Int32 + var completed atomic.Int32 + pool := newProducerPool(context.Background(), concurrency, func(context.Context, workload.Job) error { + current := active.Add(1) + defer active.Add(-1) + for { + prior := maximum.Load() + if current <= prior || maximum.CompareAndSwap(prior, current) { + break + } + } + started <- struct{}{} + <-release + completed.Add(1) + return nil + }) + + for id := range 2 * concurrency { + if err := pool.Submit(context.Background(), workload.Job{WorkloadID: uint64(id)}); err != nil { + t.Fatal(err) + } + } + for range concurrency { + <-started + } + if got := maximum.Load(); got != concurrency { + t.Fatalf("maximum concurrency = %d, want %d", got, concurrency) + } + close(release) + if err := pool.Close(); err != nil { + t.Fatal(err) + } + if got := completed.Load(); got != 2*concurrency { + t.Fatalf("completed = %d, want %d", got, 2*concurrency) + } +} + func TestToKQConfig(t *testing.T) { config := testProducerConfig() kqConfig, err := config.KQ.ToKQConfig() diff --git a/bench/kqbench/config.py b/bench/kqbench/config.py index 218b52d..9060e31 100644 --- a/bench/kqbench/config.py +++ b/bench/kqbench/config.py @@ -77,6 +77,7 @@ def resolve_run(args: argparse.Namespace) -> dict[str, Any]: "ready_partitions": topology_profile.get("ready_partitions", 1), "retry_partitions": topology_profile.get("retry_partitions", 4), "producers": topology_profile.get("producers", 1), + "producer_concurrency": topology_profile.get("producer_concurrency", 1), "workers": topology_profile.get("workers", 1), "worker_concurrency": topology_profile.get("worker_concurrency", 1), "movers": topology_profile.get("movers", 1), @@ -211,7 +212,7 @@ def validate_resolved_config(resolved: dict[str, Any]) -> None: raise ValueError(f"unknown failure mode {failures['mode']!r}") if not 0 <= float(failures["rate"]) <= 1: raise ValueError("failures.rate must be between 0 and 1") - for key in ("producers", "workers", "worker_concurrency", "movers"): + for key in ("producers", "producer_concurrency", "workers", "worker_concurrency", "movers"): if int(topology[key]) <= 0: raise ValueError(f"topology.{key} must be positive") if not kq.get("brokers"): diff --git a/bench/kqbench/test_config.py b/bench/kqbench/test_config.py index c63f721..bb1e5da 100644 --- a/bench/kqbench/test_config.py +++ b/bench/kqbench/test_config.py @@ -17,6 +17,8 @@ def test_profiles_resolve_worker_concurrency(self) -> None: self.assertEqual(baseline["topology"]["worker_concurrency"], 1) self.assertEqual(smoke["topology"]["worker_concurrency"], 2) + self.assertEqual(baseline["topology"]["producer_concurrency"], 1) + self.assertEqual(smoke["topology"]["producer_concurrency"], 1) def test_override_worker_concurrency(self) -> None: resolved = resolve_run(parse_args([ @@ -34,3 +36,20 @@ def test_rejects_non_positive_worker_concurrency(self) -> None: with self.assertRaisesRegex(ValueError, "worker_concurrency must be positive"): resolve_run(args) + + def test_override_producer_concurrency(self) -> None: + resolved = resolve_run(parse_args([ + "run", "--scenario", "ready-success-v1", "--profile", "local-baseline", + "--set", "topology.producer_concurrency=32", + ])) + + self.assertEqual(resolved["topology"]["producer_concurrency"], 32) + + def test_rejects_non_positive_producer_concurrency(self) -> None: + args = parse_args([ + "run", "--scenario", "ready-success-v1", "--profile", "local-baseline", + "--set", "topology.producer_concurrency=0", + ]) + + with self.assertRaisesRegex(ValueError, "producer_concurrency must be positive"): + resolve_run(args) diff --git a/bench/profiles/ci-smoke.json b/bench/profiles/ci-smoke.json index 2e93097..44b16a6 100644 --- a/bench/profiles/ci-smoke.json +++ b/bench/profiles/ci-smoke.json @@ -6,6 +6,7 @@ }, "topology": { "producers": 1, + "producer_concurrency": 1, "workers": 1, "worker_concurrency": 2, "movers": 1 diff --git a/bench/profiles/local-baseline.json b/bench/profiles/local-baseline.json index 0ae1f1b..8f28c3e 100644 --- a/bench/profiles/local-baseline.json +++ b/bench/profiles/local-baseline.json @@ -6,6 +6,7 @@ }, "topology": { "producers": 2, + "producer_concurrency": 1, "workers": 2, "worker_concurrency": 1, "movers": 1 diff --git a/docs/performance/runs/_templates/run.json b/docs/performance/runs/_templates/run.json index 54b804e..a1cab9d 100644 --- a/docs/performance/runs/_templates/run.json +++ b/docs/performance/runs/_templates/run.json @@ -20,6 +20,7 @@ "ready_partitions": 0, "retry_partitions": 0, "producers": 0, + "producer_concurrency": 0, "workers": 0, "worker_concurrency": 0, "movers": 0