From 02123df8cc15896f97c1580e49d8e2a991e37ddd Mon Sep 17 00:00:00 2001 From: sangtran Date: Thu, 3 Sep 2026 23:10:33 +0700 Subject: Moving all legacy quacker source code from github to cgit --- .gitignore | 9 + LICENSE | 21 ++ README | 125 ++++++++ fanin/example_test.go | 35 +++ fanin/fanin.go | 142 +++++++++ fanin/fanin_bench_test.go | 36 +++ fanin/fanin_test.go | 281 +++++++++++++++++ fanout/example_test.go | 45 +++ fanout/fanout.go | 203 ++++++++++++ fanout/fanout_bench_test.go | 65 ++++ fanout/fanout_test.go | 412 ++++++++++++++++++++++++ go.mod | 5 + go.sum | 2 + pipeline/batch.go | 64 ++++ pipeline/batch_test.go | 106 +++++++ pipeline/example_test.go | 53 ++++ pipeline/options.go | 34 ++ pipeline/pipeline.go | 21 ++ pipeline/pipeline_bench_test.go | 55 ++++ pipeline/pipeline_test.go | 255 +++++++++++++++ pipeline/stage.go | 310 ++++++++++++++++++ pipeline/stage_test.go | 308 ++++++++++++++++++ workerpool/workerpool.go | 262 ++++++++++++++++ workerpool/workerpool_bench_test.go | 50 +++ workerpool/workerpool_test.go | 608 ++++++++++++++++++++++++++++++++++++ 25 files changed, 3507 insertions(+) create mode 100644 .gitignore create mode 100644 LICENSE create mode 100644 README create mode 100644 fanin/example_test.go create mode 100644 fanin/fanin.go create mode 100644 fanin/fanin_bench_test.go create mode 100644 fanin/fanin_test.go create mode 100644 fanout/example_test.go create mode 100644 fanout/fanout.go create mode 100644 fanout/fanout_bench_test.go create mode 100644 fanout/fanout_test.go create mode 100644 go.mod create mode 100644 go.sum create mode 100644 pipeline/batch.go create mode 100644 pipeline/batch_test.go create mode 100644 pipeline/example_test.go create mode 100644 pipeline/options.go create mode 100644 pipeline/pipeline.go create mode 100644 pipeline/pipeline_bench_test.go create mode 100644 pipeline/pipeline_test.go create mode 100644 pipeline/stage.go create mode 100644 pipeline/stage_test.go create mode 100644 workerpool/workerpool.go create mode 100644 workerpool/workerpool_bench_test.go create mode 100644 workerpool/workerpool_test.go diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..d18c3de --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +.DS_Store + +# Test coverage +*.out +coverage.txt +coverage.html + +# Test binaries +*.test diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..df79de6 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 Tran Quang Sang + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README b/README new file mode 100644 index 0000000..c758b17 --- /dev/null +++ b/README @@ -0,0 +1,125 @@ +QUACKER +======= + +Type-safe, composable concurrency patterns with backpressure for Go. + +Quacker provides composable, channel-native concurrency primitives. +Backpressure propagates naturally through Go's channel semantics — +no custom protocols, no framework lock-in. + + +INSTALLATION +------------ + + go get git.sangtran.io.vn/quacker.git + +Requirements: Go 1.25+ + + +PACKAGES +-------- + +pipeline + Stream processing with backpressure. Compose stages into pipelines. + Each stage transforms a channel into a channel. + + Example: + ctx, cancel := context.WithCancelCause(ctx) + + out := pipeline.Run(ctx, input, + pipeline.Map(3, func(ctx context.Context, e Event) (Event, bool) { + enriched, err := enrich(ctx, e) + if err != nil { + cancel(err) + return e, false + } + return enriched, true + }), + pipeline.Filter(func(e Event) bool { return e.Priority > 0 }), + pipeline.Sink(5, func(ctx context.Context, e Event) { + store(ctx, e) + }), + ) + + for range out {} + if err := context.Cause(ctx); err != nil { + log.Fatal(err) + } + + Stages: Map, Filter, ForEach, Buffer, Tee, Take, Sink, Batch. + +fanout + 1-to-N distribution. Distribute a single stream to multiple + consumers. RoundRobin for load balancing, Broadcast for pub/sub. + + Example: + fo, _ := fanout.NewFanOut[int]( + fanout.WithWorkerCount(3), + fanout.WithStrategy(fanout.Broadcast), + ) + fo.Run(ctx, input) + + for _, ch := range fo.Outputs() { + go consume(ch) + } + +fanin + N-to-1 merge. Merge multiple channels into a single stream. + + Example: + fi, _ := fanin.NewFanIn[int]() + fi.Add(stream1) + fi.Add(stream2) + fi.Add(stream3) + merged := fi.Run(ctx) + +workerpool + Fixed-size worker pool with push-based task queuing, metrics, + and graceful shutdown. + + Example: + pool, _ := workerpool.NewWorkerPool[MyTask](ctx, + workerpool.WithNumWorkers(4), + workerpool.WithTaskQueueSize(100), + ) + pool.Start() + pool.Push(task) + pool.StopAndWait() + + +HOW THEY COMPOSE +---------------- + +The primitives connect through channels — Go's universal connector: + + +-- pipeline --+ + source --> fanout --+-- pipeline --+-- fanin --> pipeline --> sink + +-- pipeline --+ + + Example: + fo.Run(ctx, source) + + fast := pipeline.Run(ctx, fo.Outputs()[0], pipeline.Map(10, enrichFast)) + slow := pipeline.Run(ctx, fo.Outputs()[1], pipeline.Map(2, enrichSlow)) + + fi.Add(fast) + fi.Add(slow) + merged := fi.Run(ctx) + + out := pipeline.Run(ctx, merged, pipeline.Sink(5, store)) + + +DESIGN +------ + +- Backpressure is free. Channels block when full. No custom protocol. +- Errors use stdlib. context.WithCancelCause + context.Cause. +- Shutdown is clear. Close input for drain; cancel context for stop. +- Small interfaces. Observer interfaces are 2-3 methods, optional. +- Zero dependencies. Core logic is stdlib only (golang.org/x/sync). + + +LICENSE +------- + +MIT (see LICENSE) diff --git a/fanin/example_test.go b/fanin/example_test.go new file mode 100644 index 0000000..0edae08 --- /dev/null +++ b/fanin/example_test.go @@ -0,0 +1,35 @@ +package fanin_test + +import ( + "context" + "fmt" + "sort" + + "git.sangtran.io.vn/quacker.git/fanin" +) + +func ExampleNewFanIn() { + fi, _ := fanin.NewFanIn[int]() + + ch1 := make(chan int, 2) + ch1 <- 1 + ch1 <- 2 + close(ch1) + + ch2 := make(chan int, 2) + ch2 <- 3 + ch2 <- 4 + close(ch2) + + fi.Add(ch1) + fi.Add(ch2) + + var results []int + for v := range fi.Run(context.Background()) { + results = append(results, v) + } + + sort.Ints(results) + fmt.Println(results) + // Output: [1 2 3 4] +} diff --git a/fanin/fanin.go b/fanin/fanin.go new file mode 100644 index 0000000..913927f --- /dev/null +++ b/fanin/fanin.go @@ -0,0 +1,142 @@ +// Package fanin merges multiple input channels into a single output channel. +package fanin + +import ( + "context" + "fmt" + "sync" + "sync/atomic" +) + +// FanInObserver defines callbacks for monitoring fan-in lifecycle events. +type FanInObserver interface { + OnInputAdded() + OnInputClosed() +} + +// FanInConfig holds configuration options for creating a FanIn instance. +type FanInConfig struct { + BufferSize int + Observer FanInObserver +} + +// Option is a functional option for configuring a FanIn instance. +type Option func(*FanInConfig) + +// FanIn merges multiple input channels of type T into a single output channel. +type FanIn[T any] struct { + cfg *FanInConfig + mu sync.Mutex + inputs []<-chan T + running atomic.Bool + done chan struct{} +} + +// NewFanIn creates a new FanIn instance with the provided options. +func NewFanIn[T any](opts ...Option) (*FanIn[T], error) { + cfg := &FanInConfig{} + + for _, opt := range opts { + opt(cfg) + } + + if cfg.BufferSize < 0 { + return nil, fmt.Errorf("fanin: buffer size must not be negative, got %d", cfg.BufferSize) + } + + return &FanIn[T]{ + cfg: cfg, + }, nil +} + +// WithBufferSize sets the buffer capacity for the merged output channel. +func WithBufferSize(size int) Option { + return func(f *FanInConfig) { + f.BufferSize = size + } +} + +// WithObserver attaches a FanInObserver to receive lifecycle event callbacks. +func WithObserver(observer FanInObserver) Option { + return func(f *FanInConfig) { + f.Observer = observer + } +} + +// Add registers an input channel to be merged into the output. +// It must be called before Run, and is safe for concurrent use. +func (f *FanIn[T]) Add(ch <-chan T) { + if ch == nil { + panic("fanin: cannot add nil channel") + } + + if f.running.Load() { + panic("fanin: Add() called after Run()") + } + + f.mu.Lock() + defer f.mu.Unlock() + + f.inputs = append(f.inputs, ch) + + if f.cfg.Observer != nil { + safeObserve(f.cfg.Observer.OnInputAdded) + } +} + +// Run starts merging the registered input channels into a single output channel. +func (f *FanIn[T]) Run(ctx context.Context) <-chan T { + if !f.running.CompareAndSwap(false, true) { + panic("fanin: Run() called multiple times") + } + + f.mu.Lock() + inputs := append([]<-chan T(nil), f.inputs...) + out := make(chan T, f.cfg.BufferSize) + f.done = make(chan struct{}) + f.mu.Unlock() + + var wg sync.WaitGroup + for _, ch := range inputs { + wg.Go(func() { + for { + select { + case v, ok := <-ch: + if !ok { + if f.cfg.Observer != nil { + safeObserve(f.cfg.Observer.OnInputClosed) + } + return + } + select { + case out <- v: + case <-ctx.Done(): + return + } + case <-ctx.Done(): + return + } + } + }) + } + + go func() { + wg.Wait() + close(out) + close(f.done) + }() + + return out +} + +func safeObserve(fn func()) { + defer func() { recover() }() + fn() +} + +// Done returns a channel that is closed when all inputs are processed or the context is cancelled. +func (f *FanIn[T]) Done() <-chan struct{} { + f.mu.Lock() + defer f.mu.Unlock() + return f.done +} diff --git a/fanin/fanin_bench_test.go b/fanin/fanin_bench_test.go new file mode 100644 index 0000000..1141224 --- /dev/null +++ b/fanin/fanin_bench_test.go @@ -0,0 +1,36 @@ +package fanin + +import ( + "testing" +) + +func BenchmarkFanIn(b *testing.B) { + b.ReportAllocs() + + fi, _ := NewFanIn[int](WithBufferSize(64)) + + producers := 4 + chs := make([]chan int, producers) + for i := range chs { + chs[i] = make(chan int, 64) + fi.Add(chs[i]) + } + + out := fi.Run(b.Context()) + + // Drain output + go func() { + for range out { + } + }() + + for b.Loop() { + chs[b.N%producers] <- 1 + } + + for _, ch := range chs { + close(ch) + } + + <-fi.Done() +} diff --git a/fanin/fanin_test.go b/fanin/fanin_test.go new file mode 100644 index 0000000..bf2a413 --- /dev/null +++ b/fanin/fanin_test.go @@ -0,0 +1,281 @@ +package fanin + +import ( + "context" + "sync" + "testing" + "time" +) + +type mockObserver struct { + name string + t *testing.T +} + +func (o *mockObserver) OnInputAdded() { + o.t.Logf("%s: OnInputAdded", o.name) +} + +func (o *mockObserver) OnInputClosed() { + o.t.Logf("%s: OnInputClosed", o.name) +} + +func TestNewFanIn(t *testing.T) { + t.Parallel() + + obs := new(mockObserver) + + tests := []struct { + name string + cfg *FanInConfig + shouldErr bool + }{ + { + name: "valid config", + cfg: &FanInConfig{ + BufferSize: 1, + Observer: obs, + }, + shouldErr: false, + }, + { + name: "invalid buffer size", + cfg: &FanInConfig{ + BufferSize: -1, + Observer: obs, + }, + shouldErr: true, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + _, err := NewFanIn[int](WithBufferSize(test.cfg.BufferSize), WithObserver(obs)) + + if (err != nil) != test.shouldErr { + t.Errorf("validation error = %v, want err %v", err, test.shouldErr) + } + }) + } +} + +func TestFanIn_Add(t *testing.T) { + t.Parallel() + fi, _ := NewFanIn[int]() + + inputs := make([]chan int, 10) + for i := range inputs { + inputs[i] = make(chan int, 1) + inputs[i] <- i + close(inputs[i]) + } + + for _, ch := range inputs { + fi.Add(ch) + } + + count := 0 + for range fi.Run(t.Context()) { + count++ + } + + if count != 10 { + t.Fatalf("expected 10 items, got %d", count) + } +} + +func TestFanIn_Run(t *testing.T) { + t.Parallel() + fi, _ := NewFanIn[int]() + + ch := make(chan int) + go func() { + ch <- 1 + ch <- 2 + ch <- 3 + close(ch) + }() + + fi.Add(ch) + + for v := range fi.Run(t.Context()) { + t.Logf("value received %d", v) + } + + <-fi.Done() +} + +func TestFanIn_ContextCancellation(t *testing.T) { + t.Parallel() + fi, _ := NewFanIn[int]() + ctx, cancel := context.WithTimeout(t.Context(), time.Millisecond) + defer cancel() + + inputs := make([]chan int, 10) + for i := range inputs { + inputs[i] = make(chan int, 1) + inputs[i] <- i + close(inputs[i]) + } + + for _, ch := range inputs { + fi.Add(ch) + } + + select { + case <-fi.Run(ctx): + cancel() + case <-fi.Done(): + cancel() + } +} + +func TestFanIn_ShouldPanicRunMultiple(t *testing.T) { + t.Parallel() + + defer func() { + if r := recover(); r == nil { + t.Errorf("expected panic for multiple Run calls") + } + }() + + ctx := t.Context() + fi, _ := NewFanIn[int]() + + inputs := make([]chan int, 10) + for i := range inputs { + inputs[i] = make(chan int, 1) + inputs[i] <- i + close(inputs[i]) + } + + for _, ch := range inputs { + fi.Add(ch) + } + + for v := range fi.Run(ctx) { + t.Logf("value received %d", v) + } + + fi.Run(ctx) +} + +func TestFanIn_ShouldPanicAddAfterRun(t *testing.T) { + t.Parallel() + + defer func() { + if r := recover(); r == nil { + t.Errorf("expected panic for Add after Run") + } + }() + + fi, _ := NewFanIn[int]() + + inputs := make([]chan int, 10) + for i := range inputs { + inputs[i] = make(chan int, 1) + inputs[i] <- i + close(inputs[i]) + } + + for _, ch := range inputs { + fi.Add(ch) + } + + for v := range fi.Run(t.Context()) { + t.Logf("value received %d", v) + } + + fi.Add(make(chan int)) +} + +func TestFanIn_Observer(t *testing.T) { + t.Parallel() + obs := &mockObserver{t: t} + + fi, _ := NewFanIn[int](WithObserver(obs)) + + inputs := make([]chan int, 10) + for i := range inputs { + inputs[i] = make(chan int, 1) + inputs[i] <- i + close(inputs[i]) + } + + for _, ch := range inputs { + fi.Add(ch) + } + + for v := range fi.Run(t.Context()) { + t.Logf("value received %d", v) + } + + <-fi.Done() +} + +func TestFanIn_AddConcurrency(t *testing.T) { + t.Parallel() + + fi, _ := NewFanIn[int]() + var wg sync.WaitGroup + numGoroutines := 10 + channelPerGoroutine := 10 + + for range numGoroutines { + wg.Go(func() { + inputs := make([]chan int, channelPerGoroutine) + for i := range inputs { + inputs[i] = make(chan int, 1) + inputs[i] <- i + close(inputs[i]) + fi.Add(inputs[i]) + } + }) + } + wg.Wait() + + count := 0 + for range fi.Run(t.Context()) { + count++ + } + + <-fi.Done() + + expectedCount := numGoroutines * channelPerGoroutine + if count != expectedCount { + t.Errorf("expected %d items, got %d", expectedCount, count) + } +} + +func TestFanIn_NilChannel(t *testing.T) { + t.Parallel() + + fi, _ := NewFanIn[int]() + + defer func() { + if r := recover(); r == nil { + t.Error("expected panic when adding nil channel") + } + }() + + fi.Add(nil) +} + +func TestFanIn_RunWithNoInputs(t *testing.T) { + t.Parallel() + + fi, _ := NewFanIn[int]() + out := fi.Run(t.Context()) + + count := 0 + for range out { + count++ + } + + if count != 0 { + t.Errorf("expected 0 values from empty fan-in, got %d", count) + } + + <-fi.Done() +} diff --git a/fanout/example_test.go b/fanout/example_test.go new file mode 100644 index 0000000..50e7e96 --- /dev/null +++ b/fanout/example_test.go @@ -0,0 +1,45 @@ +package fanout_test + +import ( + "context" + "fmt" + "sort" + "sync" + + "git.sangtran.io.vn/quacker.git/fanout" +) + +func ExampleNewFanOut() { + fo, _ := fanout.NewFanOut[int]( + fanout.WithWorkerCount(2), + fanout.WithBufferSize(10), + fanout.WithStrategy(fanout.RoundRobin), + ) + + input := make(chan int, 4) + for i := 1; i <= 4; i++ { + input <- i + } + close(input) + + fo.Run(context.Background(), input) + + var mu sync.Mutex + var all []int + var wg sync.WaitGroup + + for _, ch := range fo.Outputs() { + wg.Go(func() { + for v := range ch { + mu.Lock() + all = append(all, v) + mu.Unlock() + } + }) + } + + wg.Wait() + sort.Ints(all) + fmt.Println(all) + // Output: [1 2 3 4] +} diff --git a/fanout/fanout.go b/fanout/fanout.go new file mode 100644 index 0000000..1cd2043 --- /dev/null +++ b/fanout/fanout.go @@ -0,0 +1,203 @@ +// Package fanout distributes values from a single input channel to multiple output channels. +package fanout + +import ( + "context" + "errors" + "fmt" + "sync" + "sync/atomic" +) + +// FanOutObserver defines callbacks for monitoring fan-out lifecycle events. +type FanOutObserver interface { + OnDistribute() + OnOutputClosed() +} + +// FanOutStrategy defines how values are distributed across output channels. +type FanOutStrategy int + +const ( + // RoundRobin distributes values evenly across output channels in rotation. + RoundRobin FanOutStrategy = iota + // Broadcast sends each value to all output channels. + Broadcast +) + +// ErrInvalidStrategy is returned when an unknown FanOutStrategy is provided. +var ErrInvalidStrategy = errors.New("fanout: unknown distribution strategy") + +// FanOutConfig holds configuration options for creating a FanOut instance. +type FanOutConfig struct { + WorkerCount int + BufferSize int + Strategy FanOutStrategy + Observer FanOutObserver +} + +// Option is a functional option for configuring a FanOut instance. +type Option func(*FanOutConfig) + +// FanOut distributes values from a single input channel to multiple output channels. +type FanOut[T any] struct { + cfg *FanOutConfig + outputs []chan T + running atomic.Bool + cachedOutputs []<-chan T + once sync.Once +} + +// NewFanOut creates a new FanOut instance with the provided options. +func NewFanOut[T any](opts ...Option) (*FanOut[T], error) { + cfg := &FanOutConfig{ + WorkerCount: 1, + BufferSize: 0, + Observer: nil, + Strategy: RoundRobin, + } + + for _, opt := range opts { + opt(cfg) + } + + if cfg.WorkerCount < 1 { + return nil, fmt.Errorf("fanout: worker count must be at least 1, got %d", cfg.WorkerCount) + } + + if cfg.BufferSize < 0 { + return nil, fmt.Errorf("fanout: buffer size must not be negative, got %d", cfg.BufferSize) + } + + if cfg.Strategy != RoundRobin && cfg.Strategy != Broadcast { + return nil, ErrInvalidStrategy + } + + outputs := make([]chan T, cfg.WorkerCount) + for i := range outputs { + outputs[i] = make(chan T, cfg.BufferSize) + } + + return &FanOut[T]{ + cfg: cfg, + outputs: outputs, + }, nil +} + +// Run starts distributing values from the input channel to all output channels. +func (f *FanOut[T]) Run(ctx context.Context, input <-chan T) { + if input == nil { + panic("fanout: input channel must not be nil") + } + + if !f.running.CompareAndSwap(false, true) { + panic("fanout: Run() called multiple times") + } + + go func() { + idx := 0 + defer f.running.Store(false) + defer f.closeAllOutput() + for { + select { + case value, ok := <-input: + if !ok { + return + } + switch f.cfg.Strategy { + case RoundRobin: + f.roundRobin(ctx, value, &idx) + case Broadcast: + f.broadcast(ctx, value) + } + case <-ctx.Done(): + return + } + } + }() +} + +// Outputs returns receive-only channel references for the distributed streams. +func (f *FanOut[T]) Outputs() []<-chan T { + f.once.Do(func() { + f.cachedOutputs = make([]<-chan T, len(f.outputs)) + for i, ch := range f.outputs { + f.cachedOutputs[i] = ch + } + }) + + return f.cachedOutputs +} + +func safeObserve(fn func()) { + defer func() { recover() }() + fn() +} + +func (f *FanOut[T]) roundRobin(ctx context.Context, value T, index *int) { + target := *index % f.cfg.WorkerCount + select { + case f.outputs[target] <- value: + *index = (*index + 1) % f.cfg.WorkerCount + if f.cfg.Observer != nil { + safeObserve(f.cfg.Observer.OnDistribute) + } + case <-ctx.Done(): + return + } +} + +func (f *FanOut[T]) broadcast(ctx context.Context, value T) { + var wg sync.WaitGroup + for _, ch := range f.outputs { + wg.Go(func() { + select { + case ch <- value: + if f.cfg.Observer != nil { + safeObserve(f.cfg.Observer.OnDistribute) + } + case <-ctx.Done(): + return + } + }) + } + wg.Wait() +} + +func (f *FanOut[T]) closeAllOutput() { + for _, ch := range f.outputs { + close(ch) + } + + if f.cfg.Observer != nil { + safeObserve(f.cfg.Observer.OnOutputClosed) + } +} + +// WithWorkerCount sets the number of output channels for the fan-out. +func WithWorkerCount(n int) Option { + return func(cfg *FanOutConfig) { + cfg.WorkerCount = n + } +} + +// WithBufferSize sets the buffer capacity for each output channel. +func WithBufferSize(size int) Option { + return func(cfg *FanOutConfig) { + cfg.BufferSize = size + } +} + +// WithStrategy sets the distribution strategy for the fan-out. +func WithStrategy(strategy FanOutStrategy) Option { + return func(cfg *FanOutConfig) { + cfg.Strategy = strategy + } +} + +// WithObserver attaches a FanOutObserver to receive lifecycle event callbacks. +func WithObserver(observer FanOutObserver) Option { + return func(cfg *FanOutConfig) { + cfg.Observer = observer + } +} diff --git a/fanout/fanout_bench_test.go b/fanout/fanout_bench_test.go new file mode 100644 index 0000000..0b88d87 --- /dev/null +++ b/fanout/fanout_bench_test.go @@ -0,0 +1,65 @@ +package fanout + +import ( + "context" + "testing" +) + +func BenchmarkFanOut_RoundRobin(b *testing.B) { + b.ReportAllocs() + + fo, _ := NewFanOut[int]( + WithWorkerCount(4), + WithBufferSize(64), + WithStrategy(RoundRobin), + ) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + input := make(chan int, 64) + fo.Run(ctx, input) + + // Drain outputs + for _, ch := range fo.Outputs() { + go func() { + for range ch { + } + }() + } + + for b.Loop() { + input <- 1 + } + + close(input) +} + +func BenchmarkFanOut_Broadcast(b *testing.B) { + b.ReportAllocs() + + fo, _ := NewFanOut[int]( + WithWorkerCount(4), + WithBufferSize(64), + WithStrategy(Broadcast), + ) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + input := make(chan int, 64) + fo.Run(ctx, input) + + for _, ch := range fo.Outputs() { + go func() { + for range ch { + } + }() + } + + for b.Loop() { + input <- 1 + } + + close(input) +} diff --git a/fanout/fanout_test.go b/fanout/fanout_test.go new file mode 100644 index 0000000..3f15217 --- /dev/null +++ b/fanout/fanout_test.go @@ -0,0 +1,412 @@ +package fanout + +import ( + "context" + "sync" + "testing" + "time" +) + +type mockObserver struct { + name string + t *testing.T +} + +func (m *mockObserver) OnDistribute() { + m.t.Logf("fanout: %s OnDistribute", m.name) +} + +func (m *mockObserver) OnOutputClosed() { + m.t.Logf("fanout: %s OnOutputClosed", m.name) +} + +func TestFanOut_NewFanOut(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + opts []Option + shouldErr bool + }{ + { + name: "worker size is invalid", + opts: []Option{WithWorkerCount(-10)}, + shouldErr: true, + }, + { + name: "buffer size is invalid", + opts: []Option{WithBufferSize(-10), WithWorkerCount(1)}, + shouldErr: true, + }, + { + name: "valid options", + opts: []Option{WithBufferSize(1), WithWorkerCount(1)}, + shouldErr: false, + }, + { + name: "worker count is zero", + opts: []Option{WithWorkerCount(0)}, + shouldErr: true, + }, + { + name: "valid broadcast strategy", + opts: []Option{WithWorkerCount(5), WithBufferSize(10), WithStrategy(Broadcast)}, + shouldErr: false, + }, + { + name: "invalid strategy", + opts: []Option{WithWorkerCount(1), WithStrategy(FanOutStrategy(99))}, + shouldErr: true, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + _, err := NewFanOut[int](test.opts...) + if (err != nil) != test.shouldErr { + t.Errorf("validation error = %v, want err %v", err, test.shouldErr) + } + }) + } +} + +func TestFanOut_ConcurrentStress(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + + workerCount := 5 + fo, err := NewFanOut[int]( + WithWorkerCount(workerCount), + WithBufferSize(100), + ) + if err != nil { + t.Fatalf("new fanout error: %v", err) + } + + input := make(chan int, 1000) + fo.Run(ctx, input) + + var wg sync.WaitGroup + + for i := 0; i < 10; i++ { + wg.Go(func() { + outs := fo.Outputs() + if len(outs) != workerCount { + t.Errorf("expected %d outputs, got %d", workerCount, len(outs)) + } + }) + } + + for _, ch := range fo.Outputs() { + wg.Go(func() { + for range ch { + } + }) + } + + producers := 10 + messagesPerProducer := 100 + + var producerWg sync.WaitGroup + producerWg.Add(producers) + + for i := 0; i < producers; i++ { + go func() { + defer producerWg.Done() + for j := 0; j < messagesPerProducer; j++ { + select { + case input <- j: + case <-ctx.Done(): + return + } + } + }() + } + + go func() { + producerWg.Wait() + close(input) + }() + + wg.Wait() +} + +func TestFanOut_RunBroadcast(t *testing.T) { + t.Parallel() + var wg sync.WaitGroup + workerCount := 2 + size := 10 + obs := &mockObserver{ + name: "fo broadcast", + t: t, + } + fo, _ := NewFanOut[int]( + WithWorkerCount(workerCount), + WithStrategy(Broadcast), + WithObserver(obs), + ) + + ch := make(chan int, size) + wg.Go(func() { + for i := range size { + ch <- i + } + close(ch) + }) + + fo.Run(t.Context(), ch) + + res := make([]int, workerCount) + for i, cha := range fo.Outputs() { + wg.Go(func() { + count := 0 + for v := range cha { + t.Logf("%v", v) + count++ + res[i] = count + } + }) + } + + wg.Wait() + + expected := size + for i, count := range res { + if count != expected { + t.Errorf("worker %d received %d items, expected %d", i, count, expected) + } + } +} + +func TestFanOut_RunRoundRobin(t *testing.T) { + t.Parallel() + var wg sync.WaitGroup + size := 10 + workerCount := 2 + obs := &mockObserver{ + name: "fo round robin", + t: t, + } + fo, _ := NewFanOut[int]( + WithWorkerCount(workerCount), + WithStrategy(RoundRobin), + WithObserver(obs), + ) + + ch := make(chan int, size) + wg.Go(func() { + for i := range size { + ch <- i + } + close(ch) + }) + + fo.Run(t.Context(), ch) + res := make([]int, workerCount) + for i, cha := range fo.Outputs() { + wg.Go(func() { + count := 0 + for v := range cha { + t.Logf("%v", v) + count++ + res[i] = count + } + }) + } + + wg.Wait() + expected := size / workerCount + + for i, count := range res { + if count != expected { + t.Errorf("worker %d received %d items, expected %d", i, count, expected) + } + } +} + +func TestFanOut_ContextTimeOut(t *testing.T) { + t.Parallel() + fo, _ := NewFanOut[int](WithWorkerCount(2)) + ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond) + defer cancel() + + input := make(chan int, 10) + fo.Run(ctx, input) + + var wg sync.WaitGroup + done := make(chan struct{}) + for _, ch := range fo.Outputs() { + wg.Go(func() { + for v := range ch { + t.Logf("received %d", v) + } + }) + } + + go func() { + wg.Wait() + close(done) + }() + + go func() { + i := 0 + for { + select { + case input <- i: + i++ + time.Sleep(time.Millisecond) + case <-ctx.Done(): + return + } + } + }() + + select { + case <-done: + case <-time.After(200 * time.Millisecond): + t.Fatal("timeout waiting for cleanup") + } +} + +func TestFanOut_ContextCancel(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(t.Context()) + + fo, err := NewFanOut[int](WithWorkerCount(3)) + if err != nil { + t.Fatalf("failed to create fanout: %v", err) + } + + input := make(chan int, 10) + fo.Run(ctx, input) + + var wg sync.WaitGroup + for _, ch := range fo.Outputs() { + wg.Go(func() { + for range ch { + } + }) + } + + cancel() + close(input) + + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("timeout waiting for shutdown") + } +} + +func TestFanOut_ContextDoneRoundRobin(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(t.Context()) + + fo, err := NewFanOut[int]( + WithWorkerCount(3), + WithBufferSize(0), + WithStrategy(RoundRobin), + ) + if err != nil { + t.Fatalf("new fanout error: %v", err) + } + + input := make(chan int, 1) + input <- 1 + fo.Run(ctx, input) + + cancel() + + for i, ch := range fo.Outputs() { + select { + case _, ok := <-ch: + if ok { + for range ch { + } + } + case <-time.After(time.Second): + t.Fatalf("worker %d: timeout waiting for close", i) + } + } +} + +func TestFanOut_ContextDoneBroadcast(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(t.Context()) + + fo, err := NewFanOut[int]( + WithWorkerCount(3), + WithStrategy(Broadcast), + WithBufferSize(0), + ) + if err != nil { + t.Fatalf("new fanout error: %v", err) + } + + input := make(chan int, 1) + input <- 1 + fo.Run(ctx, input) + + cancel() + + for i, ch := range fo.Outputs() { + select { + case _, ok := <-ch: + if ok { + for range ch { + } + } + case <-time.After(time.Second): + t.Fatalf("worker %d: timeout waiting for close", i) + } + } +} + +func TestFanOut_RunMultipleTime(t *testing.T) { + t.Parallel() + + ctx := t.Context() + fo, err := NewFanOut[int]( + WithWorkerCount(3), + WithBufferSize(0), + ) + if err != nil { + t.Fatalf("new fanout error: %v", err) + } + + input := make(chan int, 10) + fo.Run(ctx, input) + + defer func() { + if r := recover(); r == nil { + t.Error("expected panic when calling Run() multiple times") + } + }() + + fo.Run(ctx, input) +} + +func TestFanOut_NilInput(t *testing.T) { + t.Parallel() + + fo, err := NewFanOut[int](WithWorkerCount(1)) + if err != nil { + t.Fatalf("failed to create fanout: %v", err) + } + + defer func() { + if r := recover(); r == nil { + t.Error("expected panic when passing nil input channel") + } + }() + + fo.Run(t.Context(), nil) +} diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..df07984 --- /dev/null +++ b/go.mod @@ -0,0 +1,5 @@ +module git.sangtran.io.vn/quacker.git + +go 1.25.1 + +require golang.org/x/sync v0.19.0 diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..159532a --- /dev/null +++ b/go.sum @@ -0,0 +1,2 @@ +golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4= +golang.org/x/sync v0.19.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= diff --git a/pipeline/batch.go b/pipeline/batch.go new file mode 100644 index 0000000..74b08f0 --- /dev/null +++ b/pipeline/batch.go @@ -0,0 +1,64 @@ +package pipeline + +import ( + "context" + "time" +) + +// Batch collects items into slices of up to size, flushing when full or when timeout elapses. +func Batch[T any](ctx context.Context, in <-chan T, size int, timeout time.Duration) <-chan []T { + if in == nil { + panic("pipeline: Batch input channel must not be nil") + } + if size < 1 { + panic("pipeline: Batch size must be at least 1") + } + + out := make(chan []T) + + go func() { + defer close(out) + + buf := make([]T, 0, size) + ticker := time.NewTicker(timeout) + defer ticker.Stop() + + flush := func() bool { + if len(buf) == 0 { + return true + } + batch := buf + buf = make([]T, 0, size) + select { + case out <- batch: + return true + case <-ctx.Done(): + return false + } + } + + for { + select { + case v, ok := <-in: + if !ok { + flush() + return + } + buf = append(buf, v) + if len(buf) >= size { + if !flush() { + return + } + } + case <-ticker.C: + if !flush() { + return + } + case <-ctx.Done(): + return + } + } + }() + + return out +} diff --git a/pipeline/batch_test.go b/pipeline/batch_test.go new file mode 100644 index 0000000..4d79ddc --- /dev/null +++ b/pipeline/batch_test.go @@ -0,0 +1,106 @@ +package pipeline + +import ( + "context" + "testing" + "time" +) + +func TestBatch_Full(t *testing.T) { + t.Parallel() + + input := make(chan int, 10) + for i := 1; i <= 10; i++ { + input <- i + } + close(input) + + out := Batch(t.Context(), input, 3, time.Second) + + var batches [][]int + for b := range out { + batches = append(batches, b) + } + + if len(batches) != 4 { + t.Fatalf("expected 4 batches, got %d: %v", len(batches), batches) + } + if len(batches[0]) != 3 { + t.Fatalf("expected first batch size 3, got %d", len(batches[0])) + } + if len(batches[3]) != 1 { + t.Fatalf("expected last batch size 1, got %d", len(batches[3])) + } +} + +func TestBatch_TimeoutFlush(t *testing.T) { + t.Parallel() + + input := make(chan int) + + out := Batch(t.Context(), input, 100, 50*time.Millisecond) + + input <- 1 + input <- 2 + + select { + case batch := <-out: + if len(batch) != 2 { + t.Fatalf("expected batch of 2 from timeout flush, got %d", len(batch)) + } + case <-time.After(time.Second): + t.Fatal("timeout waiting for batch flush") + } + + close(input) + for range out { + } +} + +func TestBatch_ContextCancel(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(t.Context()) + + input := make(chan int) + out := Batch(ctx, input, 100, time.Second) + + cancel() + + done := make(chan struct{}) + go func() { + for range out { + } + close(done) + }() + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("timeout waiting for shutdown") + } +} + +func TestBatch_NilInput(t *testing.T) { + t.Parallel() + + defer func() { + if r := recover(); r == nil { + t.Error("expected panic for nil input") + } + }() + + Batch[int](t.Context(), nil, 10, time.Second) +} + +func TestBatch_PanicOnZeroSize(t *testing.T) { + t.Parallel() + + defer func() { + if r := recover(); r == nil { + t.Error("expected panic for batch size < 1") + } + }() + + input := make(chan int) + Batch(t.Context(), input, 0, time.Second) +} diff --git a/pipeline/example_test.go b/pipeline/example_test.go new file mode 100644 index 0000000..e322529 --- /dev/null +++ b/pipeline/example_test.go @@ -0,0 +1,53 @@ +package pipeline_test + +import ( + "context" + "fmt" + "sort" + + "git.sangtran.io.vn/quacker.git/pipeline" +) + +func ExampleRun() { + input := make(chan int, 5) + for i := 1; i <= 5; i++ { + input <- i + } + close(input) + + out := pipeline.Run(context.Background(), input, + pipeline.Filter(func(v int) bool { return v%2 == 0 }), + pipeline.Map(1, func(_ context.Context, v int) (int, bool) { + return v * 10, true + }), + ) + + var results []int + for v := range out { + results = append(results, v) + } + sort.Ints(results) + fmt.Println(results) + // Output: [20 40] +} + +func ExampleSink() { + input := make(chan string, 3) + input <- "hello" + input <- "world" + input <- "!" + close(input) + + var collected []string + out := pipeline.Run(context.Background(), input, + pipeline.Sink(1, func(_ context.Context, v string) { + collected = append(collected, v) + }), + ) + + for range out { + } + + fmt.Println(collected) + // Output: [hello world !] +} diff --git a/pipeline/options.go b/pipeline/options.go new file mode 100644 index 0000000..c8c616e --- /dev/null +++ b/pipeline/options.go @@ -0,0 +1,34 @@ +package pipeline + +// StageObserver receives callbacks for stage lifecycle events. +type StageObserver interface { + OnItem() + OnDrop() + OnDone() +} + +// StageOption configures a stage's behavior. +type StageOption func(*stageConfig) + +type stageConfig struct { + observer StageObserver + bufferSize int +} + +// WithObserver attaches an observer to a stage. +func WithObserver(obs StageObserver) StageOption { + return func(cfg *stageConfig) { cfg.observer = obs } +} + +// WithBufferSize sets the output channel buffer size for a stage. +func WithBufferSize(size int) StageOption { + return func(cfg *stageConfig) { cfg.bufferSize = size } +} + +func applyOpts(opts []StageOption) stageConfig { + var cfg stageConfig + for _, opt := range opts { + opt(&cfg) + } + return cfg +} diff --git a/pipeline/pipeline.go b/pipeline/pipeline.go new file mode 100644 index 0000000..170a04e --- /dev/null +++ b/pipeline/pipeline.go @@ -0,0 +1,21 @@ +// Package pipeline provides composable, type-safe stream processing with backpressure. +package pipeline + +import "context" + +// Stage transforms an input channel into an output channel. +// The stage owns the output channel and closes it when the input +// is exhausted or the context is cancelled. +type Stage[T any] func(ctx context.Context, in <-chan T) <-chan T + +// Run connects stages into a pipeline and returns the final output channel. +func Run[T any](ctx context.Context, input <-chan T, stages ...Stage[T]) <-chan T { + if input == nil { + panic("pipeline: input channel must not be nil") + } + ch := input + for _, stage := range stages { + ch = stage(ctx, ch) + } + return ch +} diff --git a/pipeline/pipeline_bench_test.go b/pipeline/pipeline_bench_test.go new file mode 100644 index 0000000..f4f81a0 --- /dev/null +++ b/pipeline/pipeline_bench_test.go @@ -0,0 +1,55 @@ +package pipeline + +import ( + "context" + "testing" +) + +func BenchmarkMap(b *testing.B) { + b.ReportAllocs() + + input := make(chan int, 256) + + out := Run(b.Context(), input, + Map(4, func(_ context.Context, v int) (int, bool) { + return v * 2, true + }), + ) + + go func() { + for range out { + } + }() + + for b.Loop() { + input <- 1 + } + + close(input) +} + +func BenchmarkPipeline_Composition(b *testing.B) { + b.ReportAllocs() + + input := make(chan int, 256) + + out := Run(b.Context(), input, + Filter(func(v int) bool { return v > 0 }), + Map(4, func(_ context.Context, v int) (int, bool) { + return v * 2, true + }), + Buffer[int](64), + Sink(2, func(_ context.Context, v int) {}), + ) + + go func() { + for range out { + } + }() + + for b.Loop() { + input <- 1 + } + + close(input) +} diff --git a/pipeline/pipeline_test.go b/pipeline/pipeline_test.go new file mode 100644 index 0000000..9593356 --- /dev/null +++ b/pipeline/pipeline_test.go @@ -0,0 +1,255 @@ +package pipeline + +import ( + "context" + "fmt" + "sync/atomic" + "testing" + "time" +) + +func TestRun_Passthrough(t *testing.T) { + t.Parallel() + + input := make(chan int, 3) + input <- 1 + input <- 2 + input <- 3 + close(input) + + out := Run(t.Context(), input) + + var got []int + for v := range out { + got = append(got, v) + } + + if len(got) != 3 { + t.Fatalf("expected 3 items, got %d", len(got)) + } +} + +func TestRun_NilInput(t *testing.T) { + t.Parallel() + + defer func() { + if r := recover(); r == nil { + t.Error("expected panic for nil input") + } + }() + + Run[int](t.Context(), nil) +} + +func TestRun_Composition(t *testing.T) { + t.Parallel() + + input := make(chan int, 10) + for i := 1; i <= 10; i++ { + input <- i + } + close(input) + + out := Run(t.Context(), input, + Filter(func(v int) bool { return v%2 == 0 }), + Map(1, func(_ context.Context, v int) (int, bool) { + return v * 2, true + }), + ) + + var got []int + for v := range out { + got = append(got, v) + } + + if len(got) != 5 { + t.Fatalf("expected 5 items, got %d: %v", len(got), got) + } + + sum := 0 + for _, v := range got { + sum += v + } + if sum != 60 { + t.Fatalf("expected sum 60, got %d", sum) + } +} + +func TestRun_Backpressure(t *testing.T) { + t.Parallel() + + input := make(chan int) + consumed := make(chan int, 10) + + out := Run(t.Context(), input, + Map(1, func(_ context.Context, v int) (int, bool) { + consumed <- v + return v, true + }), + ) + + input <- 42 + v := <-consumed + if v != 42 { + t.Fatalf("expected 42, got %d", v) + } + + select { + case input <- 99: + t.Fatal("producer was not blocked") + case <-time.After(50 * time.Millisecond): + } + + go func() { + for range out { + } + }() + close(input) +} + +func TestRun_GracefulDrain(t *testing.T) { + t.Parallel() + + input := make(chan int, 5) + for i := 1; i <= 5; i++ { + input <- i + } + close(input) + + var count atomic.Int32 + out := Run(t.Context(), input, + Sink(2, func(_ context.Context, v int) { + count.Add(1) + }), + ) + + for range out { + } + + if count.Load() != 5 { + t.Fatalf("expected all 5 items drained, got %d", count.Load()) + } +} + +func TestRun_Cancellation(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(t.Context()) + + input := make(chan int) + out := Run(ctx, input, + Map(1, func(_ context.Context, v int) (int, bool) { + return v, true + }), + ) + + go func() { + select { + case input <- 1: + case <-ctx.Done(): + } + }() + + cancel() + + done := make(chan struct{}) + go func() { + for range out { + } + close(done) + }() + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("pipeline did not shut down after cancellation") + } +} + +func TestRun_ErrorPropagation(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancelCause(t.Context()) + defer cancel(nil) + + input := make(chan int, 5) + for i := 1; i <= 5; i++ { + input <- i + } + close(input) + + sentinel := fmt.Errorf("bad item") + + out := Run(ctx, input, + Map(1, func(_ context.Context, v int) (int, bool) { + if v == 3 { + cancel(sentinel) + return v, false + } + return v * 2, true + }), + ) + + for range out { + } + + if cause := context.Cause(ctx); cause != sentinel { + t.Fatalf("expected sentinel error, got: %v", cause) + } +} + +func TestObserver(t *testing.T) { + t.Parallel() + + var items, drops atomic.Int32 + doneCh := make(chan struct{}) + + obs := &testObserver{ + onItem: func() { items.Add(1) }, + onDrop: func() { drops.Add(1) }, + onDone: func() { close(doneCh) }, + } + + input := make(chan int, 6) + for i := 1; i <= 6; i++ { + input <- i + } + close(input) + + out := Run(t.Context(), input, + Filter(func(v int) bool { return v%2 == 0 }, WithObserver(obs)), + ) + + for range out { + } + <-doneCh + + if items.Load() != 3 { + t.Errorf("expected 3 items, got %d", items.Load()) + } + if drops.Load() != 3 { + t.Errorf("expected 3 drops, got %d", drops.Load()) + } +} + +type testObserver struct { + onItem func() + onDrop func() + onDone func() +} + +func (o *testObserver) OnItem() { + if o.onItem != nil { + o.onItem() + } +} + +func (o *testObserver) OnDrop() { + if o.onDrop != nil { + o.onDrop() + } +} + +func (o *testObserver) OnDone() { + if o.onDone != nil { + o.onDone() + } +} diff --git a/pipeline/stage.go b/pipeline/stage.go new file mode 100644 index 0000000..e0716cf --- /dev/null +++ b/pipeline/stage.go @@ -0,0 +1,310 @@ +package pipeline + +import ( + "context" + "sync" +) + +// Map applies fn to each item with bounded concurrency. +// If keep is false, the item is dropped. +func Map[T any](concurrency int, fn func(context.Context, T) (T, bool), opts ...StageOption) Stage[T] { + if concurrency < 1 { + panic("pipeline: Map concurrency must be at least 1") + } + cfg := applyOpts(opts) + + return func(ctx context.Context, in <-chan T) <-chan T { + out := make(chan T, cfg.bufferSize) + var wg sync.WaitGroup + + for range concurrency { + wg.Go(func() { + for { + select { + case v, ok := <-in: + if !ok { + return + } + result, keep := fn(ctx, v) + if !keep { + if cfg.observer != nil { + cfg.observer.OnDrop() + } + continue + } + if cfg.observer != nil { + cfg.observer.OnItem() + } + select { + case out <- result: + case <-ctx.Done(): + return + } + case <-ctx.Done(): + return + } + } + }) + } + + go func() { + wg.Wait() + close(out) + if cfg.observer != nil { + cfg.observer.OnDone() + } + }() + + return out + } +} + +// Filter forwards items for which fn returns true. +func Filter[T any](fn func(T) bool, opts ...StageOption) Stage[T] { + cfg := applyOpts(opts) + + return func(ctx context.Context, in <-chan T) <-chan T { + out := make(chan T, cfg.bufferSize) + + go func() { + defer close(out) + defer func() { + if cfg.observer != nil { + cfg.observer.OnDone() + } + }() + + for { + select { + case v, ok := <-in: + if !ok { + return + } + if !fn(v) { + if cfg.observer != nil { + cfg.observer.OnDrop() + } + continue + } + if cfg.observer != nil { + cfg.observer.OnItem() + } + select { + case out <- v: + case <-ctx.Done(): + return + } + case <-ctx.Done(): + return + } + } + }() + + return out + } +} + +// ForEach calls fn on each item and forwards it unchanged. +func ForEach[T any](fn func(T), opts ...StageOption) Stage[T] { + cfg := applyOpts(opts) + + return func(ctx context.Context, in <-chan T) <-chan T { + out := make(chan T, cfg.bufferSize) + + go func() { + defer close(out) + defer func() { + if cfg.observer != nil { + cfg.observer.OnDone() + } + }() + + for { + select { + case v, ok := <-in: + if !ok { + return + } + fn(v) + if cfg.observer != nil { + cfg.observer.OnItem() + } + select { + case out <- v: + case <-ctx.Done(): + return + } + case <-ctx.Done(): + return + } + } + }() + + return out + } +} + +// Buffer inserts a buffered channel of the specified capacity between stages. +func Buffer[T any](size int) Stage[T] { + if size < 1 { + panic("pipeline: Buffer size must be at least 1") + } + + return func(ctx context.Context, in <-chan T) <-chan T { + out := make(chan T, size) + + go func() { + defer close(out) + for { + select { + case v, ok := <-in: + if !ok { + return + } + select { + case out <- v: + case <-ctx.Done(): + return + } + case <-ctx.Done(): + return + } + } + }() + + return out + } +} + +// Tee sends each item to a side channel while forwarding it to the next stage. +func Tee[T any](side chan<- T, opts ...StageOption) Stage[T] { + if side == nil { + panic("pipeline: Tee side channel must not be nil") + } + cfg := applyOpts(opts) + + return func(ctx context.Context, in <-chan T) <-chan T { + out := make(chan T, cfg.bufferSize) + + go func() { + defer close(out) + defer func() { + if cfg.observer != nil { + cfg.observer.OnDone() + } + }() + + for { + select { + case v, ok := <-in: + if !ok { + return + } + if cfg.observer != nil { + cfg.observer.OnItem() + } + select { + case side <- v: + case <-ctx.Done(): + return + } + select { + case out <- v: + case <-ctx.Done(): + return + } + case <-ctx.Done(): + return + } + } + }() + + return out + } +} + +// Take forwards up to n items, then closes the output channel. +func Take[T any](n int, opts ...StageOption) Stage[T] { + if n < 0 { + panic("pipeline: Take count must not be negative") + } + cfg := applyOpts(opts) + + return func(ctx context.Context, in <-chan T) <-chan T { + out := make(chan T, cfg.bufferSize) + + go func() { + defer close(out) + defer func() { + if cfg.observer != nil { + cfg.observer.OnDone() + } + }() + + count := 0 + for count < n { + select { + case v, ok := <-in: + if !ok { + return + } + if cfg.observer != nil { + cfg.observer.OnItem() + } + select { + case out <- v: + count++ + case <-ctx.Done(): + return + } + case <-ctx.Done(): + return + } + } + }() + + return out + } +} + +// Sink consumes items with bounded concurrency without forwarding them. +func Sink[T any](concurrency int, fn func(context.Context, T), opts ...StageOption) Stage[T] { + if concurrency < 1 { + panic("pipeline: Sink concurrency must be at least 1") + } + cfg := applyOpts(opts) + + return func(ctx context.Context, in <-chan T) <-chan T { + done := make(chan T) + var wg sync.WaitGroup + + for range concurrency { + wg.Go(func() { + for { + select { + case v, ok := <-in: + if !ok { + return + } + fn(ctx, v) + if cfg.observer != nil { + cfg.observer.OnItem() + } + case <-ctx.Done(): + return + } + } + }) + } + + go func() { + wg.Wait() + close(done) + if cfg.observer != nil { + cfg.observer.OnDone() + } + }() + + return done + } +} diff --git a/pipeline/stage_test.go b/pipeline/stage_test.go new file mode 100644 index 0000000..4579158 --- /dev/null +++ b/pipeline/stage_test.go @@ -0,0 +1,308 @@ +package pipeline + +import ( + "context" + "sort" + "sync" + "sync/atomic" + "testing" +) + +func TestMap_Basic(t *testing.T) { + t.Parallel() + + input := make(chan int, 3) + input <- 1 + input <- 2 + input <- 3 + close(input) + + out := Map(1, func(_ context.Context, v int) (int, bool) { + return v * 10, true + })(t.Context(), input) + + var got []int + for v := range out { + got = append(got, v) + } + + sort.Ints(got) + if len(got) != 3 || got[0] != 10 || got[1] != 20 || got[2] != 30 { + t.Fatalf("expected [10 20 30], got %v", got) + } +} + +func TestMap_Concurrent(t *testing.T) { + t.Parallel() + + input := make(chan int, 100) + for i := range 100 { + input <- i + } + close(input) + + out := Map(5, func(_ context.Context, v int) (int, bool) { + return v, true + })(t.Context(), input) + + var got []int + for v := range out { + got = append(got, v) + } + + if len(got) != 100 { + t.Fatalf("expected 100 items, got %d", len(got)) + } + + sort.Ints(got) + for i, v := range got { + if v != i { + t.Fatalf("missing item %d, got %v at position %d", i, v, i) + } + } +} + +func TestMap_Drop(t *testing.T) { + t.Parallel() + + input := make(chan int, 5) + for i := 1; i <= 5; i++ { + input <- i + } + close(input) + + out := Map(1, func(_ context.Context, v int) (int, bool) { + return v, v > 3 + })(t.Context(), input) + + var got []int + for v := range out { + got = append(got, v) + } + + if len(got) != 2 { + t.Fatalf("expected 2 items, got %d: %v", len(got), got) + } +} + +func TestMap_PanicOnZeroConcurrency(t *testing.T) { + t.Parallel() + + defer func() { + if r := recover(); r == nil { + t.Error("expected panic for concurrency < 1") + } + }() + + Map(0, func(_ context.Context, v int) (int, bool) { return v, true }) +} + +func TestFilter(t *testing.T) { + t.Parallel() + + input := make(chan int, 10) + for i := 1; i <= 10; i++ { + input <- i + } + close(input) + + out := Filter(func(v int) bool { return v%3 == 0 })(t.Context(), input) + + var got []int + for v := range out { + got = append(got, v) + } + + if len(got) != 3 { + t.Fatalf("expected [3 6 9], got %v", got) + } +} + +func TestForEach(t *testing.T) { + t.Parallel() + + input := make(chan int, 3) + input <- 1 + input <- 2 + input <- 3 + close(input) + + var sum atomic.Int32 + out := ForEach(func(v int) { + sum.Add(int32(v)) + })(t.Context(), input) + + var got []int + for v := range out { + got = append(got, v) + } + + if sum.Load() != 6 { + t.Fatalf("expected side effect sum 6, got %d", sum.Load()) + } + if len(got) != 3 { + t.Fatalf("expected 3 forwarded items, got %d", len(got)) + } +} + +func TestBuffer(t *testing.T) { + t.Parallel() + + input := make(chan int, 5) + for i := 1; i <= 5; i++ { + input <- i + } + close(input) + + out := Buffer[int](10)(t.Context(), input) + + var got []int + for v := range out { + got = append(got, v) + } + + if len(got) != 5 { + t.Fatalf("expected 5 items, got %d", len(got)) + } +} + +func TestBuffer_PanicOnZero(t *testing.T) { + t.Parallel() + + defer func() { + if r := recover(); r == nil { + t.Error("expected panic for buffer size < 1") + } + }() + + Buffer[int](0) +} + +func TestTee(t *testing.T) { + t.Parallel() + + input := make(chan int, 3) + input <- 1 + input <- 2 + input <- 3 + close(input) + + side := make(chan int, 10) + out := Tee(side)(t.Context(), input) + + var got []int + for v := range out { + got = append(got, v) + } + close(side) + + var sideGot []int + for v := range side { + sideGot = append(sideGot, v) + } + + if len(got) != 3 { + t.Fatalf("expected 3 forwarded items, got %d", len(got)) + } + if len(sideGot) != 3 { + t.Fatalf("expected 3 side items, got %d", len(sideGot)) + } +} + +func TestTee_NilSide(t *testing.T) { + t.Parallel() + + defer func() { + if r := recover(); r == nil { + t.Error("expected panic for nil side channel") + } + }() + + Tee[int](nil) +} + +func TestTake(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + + input := make(chan int, 10) + for i := 1; i <= 10; i++ { + input <- i + } + close(input) + + out := Take[int](3)(ctx, input) + + var got []int + for v := range out { + got = append(got, v) + } + + if len(got) != 3 { + t.Fatalf("expected 3 items, got %d: %v", len(got), got) + } +} + +func TestTake_Zero(t *testing.T) { + t.Parallel() + + input := make(chan int, 5) + for i := 1; i <= 5; i++ { + input <- i + } + close(input) + + out := Take[int](0)(t.Context(), input) + + count := 0 + for range out { + count++ + } + + if count != 0 { + t.Fatalf("expected 0 items from Take(0), got %d", count) + } +} + +func TestSink(t *testing.T) { + t.Parallel() + + input := make(chan int, 5) + for i := 1; i <= 5; i++ { + input <- i + } + close(input) + + var mu sync.Mutex + var collected []int + + out := Sink(2, func(_ context.Context, v int) { + mu.Lock() + collected = append(collected, v) + mu.Unlock() + })(t.Context(), input) + + for range out { + t.Fatal("Sink should not forward items") + } + + mu.Lock() + defer mu.Unlock() + + if len(collected) != 5 { + t.Fatalf("expected 5 items consumed, got %d", len(collected)) + } +} + +func TestSink_PanicOnZeroConcurrency(t *testing.T) { + t.Parallel() + + defer func() { + if r := recover(); r == nil { + t.Error("expected panic for concurrency < 1") + } + }() + + Sink(0, func(_ context.Context, v int) {}) +} diff --git a/workerpool/workerpool.go b/workerpool/workerpool.go new file mode 100644 index 0000000..44a3edf --- /dev/null +++ b/workerpool/workerpool.go @@ -0,0 +1,262 @@ +// Package workerpool provides a generic worker pool implementation for concurrent task processing. +package workerpool + +import ( + "context" + "errors" + "fmt" + "maps" + "sync" + "sync/atomic" + + "golang.org/x/sync/errgroup" +) + +// ErrQueueFull is returned by Push when the task queue has no remaining capacity. +var ErrQueueFull = errors.New("workerpool: task queue is full") + +// Task represents a unit of work that can be executed by the worker pool. +type Task[T any] interface { + Execute(ctx context.Context) error + GetID() string +} + +// WorkerPoolObserver defines callbacks for monitoring worker pool lifecycle events. +type WorkerPoolObserver interface { + OnWorkerStart(workerID int) + OnWorkerStop(workerID int) + OnTaskDone(workerID int, taskID string, err error) +} + +// WorkerPoolConfig holds configuration options for creating a WorkerPool instance. +type WorkerPoolConfig struct { + Name string + NumWorkers int + TaskQueueSize int + Observer WorkerPoolObserver + ErrorHandler func(workerID int, taskID string, err error) + PanicHandler func(workerID int, taskID string, panicValue any) +} + +// Option is a functional option for configuring a WorkerPool instance. +type Option func(*WorkerPoolConfig) + +// WorkerPool manages a fixed-size pool of workers that process tasks concurrently. +type WorkerPool[T Task[T]] struct { + cfg *WorkerPoolConfig + taskQueue chan T + metrics *metrics + eg *errgroup.Group + once sync.Once + ctx context.Context + started atomic.Bool +} + +type metrics struct { + mu sync.RWMutex + TasksProcessed map[string]int + TasksFailed int + ActiveWorkers int +} + +// MetricsSnapshot provides a point-in-time view of worker pool metrics. +type MetricsSnapshot struct { + TasksProcessed map[string]int + TasksFailed int + ActiveWorkers int +} + +// NewWorkerPool creates a new WorkerPool with the provided options. +func NewWorkerPool[T Task[T]](ctx context.Context, opts ...Option) (*WorkerPool[T], error) { + cfg := &WorkerPoolConfig{ + Name: "workerpool", + NumWorkers: 1, + TaskQueueSize: 1, + } + + for _, opt := range opts { + opt(cfg) + } + + if cfg.NumWorkers < 1 { + return nil, fmt.Errorf("workerpool: number of workers must be at least 1, got %d", cfg.NumWorkers) + } + + if cfg.TaskQueueSize < 1 { + return nil, fmt.Errorf("workerpool: task queue size must be at least 1, got %d", cfg.TaskQueueSize) + } + + eg, ctx := errgroup.WithContext(ctx) + + return &WorkerPool[T]{ + cfg: cfg, + taskQueue: make(chan T, cfg.TaskQueueSize), + ctx: ctx, + eg: eg, + metrics: &metrics{ + TasksProcessed: make(map[string]int), + }, + }, nil +} + +func WithName(name string) Option { + return func(cfg *WorkerPoolConfig) { cfg.Name = name } +} + +func WithNumWorkers(n int) Option { + return func(cfg *WorkerPoolConfig) { cfg.NumWorkers = n } +} + +func WithTaskQueueSize(size int) Option { + return func(cfg *WorkerPoolConfig) { cfg.TaskQueueSize = size } +} + +func WithObserver(observer WorkerPoolObserver) Option { + return func(cfg *WorkerPoolConfig) { cfg.Observer = observer } +} + +func WithErrorHandler(handler func(workerID int, taskID string, err error)) Option { + return func(cfg *WorkerPoolConfig) { cfg.ErrorHandler = handler } +} + +func WithPanicHandler(handler func(workerID int, taskID string, panicValue any)) Option { + return func(cfg *WorkerPoolConfig) { cfg.PanicHandler = handler } +} + +// Start launches worker goroutines. It panics if called multiple times. +func (w *WorkerPool[T]) Start() { + if !w.started.CompareAndSwap(false, true) { + panic("workerpool: Start() called multiple times") + } + + for i := 0; i < w.cfg.NumWorkers; i++ { + workerID := i + workerName := fmt.Sprintf("%s.worker.%d", w.cfg.Name, workerID) + w.eg.Go(func() error { + w.metrics.mu.Lock() + w.metrics.ActiveWorkers++ + w.metrics.mu.Unlock() + + if w.cfg.Observer != nil { + safeObserve(func() { w.cfg.Observer.OnWorkerStart(workerID) }) + } + + defer func() { + w.metrics.mu.Lock() + w.metrics.ActiveWorkers-- + w.metrics.mu.Unlock() + + if w.cfg.Observer != nil { + safeObserve(func() { w.cfg.Observer.OnWorkerStop(workerID) }) + } + }() + + return w.worker(workerID, workerName) + }) + } +} + +// Push adds a task to the queue. Returns ErrQueueFull if full, or context error if cancelled. +func (w *WorkerPool[T]) Push(task T) error { + select { + case w.taskQueue <- task: + return nil + case <-w.ctx.Done(): + return fmt.Errorf("workerpool %q: push: %w", w.cfg.Name, w.ctx.Err()) + default: + return fmt.Errorf("workerpool %q: capacity %d: %w", w.cfg.Name, cap(w.taskQueue), ErrQueueFull) + } +} + +// Stop closes the task queue. +func (w *WorkerPool[T]) Stop() { + w.once.Do(func() { + close(w.taskQueue) + }) +} + +// Wait blocks until all workers finish. +func (w *WorkerPool[T]) Wait() error { + return w.eg.Wait() +} + +// StopAndWait closes the queue and waits for all workers to drain. +func (w *WorkerPool[T]) StopAndWait() error { + w.Stop() + return w.Wait() +} + +// GetMetrics returns a snapshot of current pool metrics. +func (w *WorkerPool[T]) GetMetrics() MetricsSnapshot { + w.metrics.mu.RLock() + defer w.metrics.mu.RUnlock() + + tasksProcessed := make(map[string]int, len(w.metrics.TasksProcessed)) + maps.Copy(tasksProcessed, w.metrics.TasksProcessed) + + return MetricsSnapshot{ + ActiveWorkers: w.metrics.ActiveWorkers, + TasksFailed: w.metrics.TasksFailed, + TasksProcessed: tasksProcessed, + } +} + +func (w *WorkerPool[T]) worker(workerID int, workerName string) error { + for { + select { + case task, ok := <-w.taskQueue: + if !ok { + return nil + } + w.executeSafe(workerID, workerName, task) + case <-w.ctx.Done(): + return w.ctx.Err() + } + } +} + +func (w *WorkerPool[T]) executeSafe(workerID int, workerName string, task T) { + taskID := task.GetID() + + defer func() { + if r := recover(); r != nil { + w.recordTaskFailed() + if w.cfg.PanicHandler != nil { + w.cfg.PanicHandler(workerID, taskID, r) + } + } + }() + + if err := task.Execute(w.ctx); err != nil { + w.recordTaskFailed() + if w.cfg.Observer != nil { + safeObserve(func() { w.cfg.Observer.OnTaskDone(workerID, taskID, err) }) + } + if w.cfg.ErrorHandler != nil { + w.cfg.ErrorHandler(workerID, taskID, err) + } + return + } + + w.recordTaskSuccess(workerName) + if w.cfg.Observer != nil { + safeObserve(func() { w.cfg.Observer.OnTaskDone(workerID, taskID, nil) }) + } +} + +func (w *WorkerPool[T]) recordTaskSuccess(workerName string) { + w.metrics.mu.Lock() + defer w.metrics.mu.Unlock() + w.metrics.TasksProcessed[workerName]++ +} + +func (w *WorkerPool[T]) recordTaskFailed() { + w.metrics.mu.Lock() + defer w.metrics.mu.Unlock() + w.metrics.TasksFailed++ +} + +func safeObserve(fn func()) { + defer func() { recover() }() + fn() +} diff --git a/workerpool/workerpool_bench_test.go b/workerpool/workerpool_bench_test.go new file mode 100644 index 0000000..1fbc672 --- /dev/null +++ b/workerpool/workerpool_bench_test.go @@ -0,0 +1,50 @@ +package workerpool + +import ( + "context" + "fmt" + "testing" +) + +type benchTask struct { + id string +} + +func (t *benchTask) Execute(_ context.Context) error { return nil } +func (t *benchTask) GetID() string { return t.id } + +func BenchmarkWorkerPool(b *testing.B) { + b.ReportAllocs() + + pool, _ := NewWorkerPool[*benchTask](b.Context(), + WithNumWorkers(4), + WithTaskQueueSize(256), + ) + pool.Start() + + task := &benchTask{id: "bench"} + + for b.Loop() { + pool.Push(task) + } + + pool.StopAndWait() +} + +func BenchmarkWorkerPool_Push(b *testing.B) { + b.ReportAllocs() + + pool, _ := NewWorkerPool[*benchTask](b.Context(), + WithNumWorkers(8), + WithTaskQueueSize(1024), + ) + pool.Start() + defer pool.StopAndWait() + + b.RunParallel(func(pb *testing.PB) { + task := &benchTask{id: fmt.Sprintf("bench-%d", b.N)} + for pb.Next() { + pool.Push(task) + } + }) +} diff --git a/workerpool/workerpool_test.go b/workerpool/workerpool_test.go new file mode 100644 index 0000000..851874b --- /dev/null +++ b/workerpool/workerpool_test.go @@ -0,0 +1,608 @@ +package workerpool + +import ( + "context" + "errors" + "fmt" + "sync" + "testing" + "time" +) + +type mockTask struct { + id string + fn func(ctx context.Context) error +} + +func (t *mockTask) Execute(ctx context.Context) error { + if t.fn != nil { + return t.fn(ctx) + } + return nil +} + +func (t *mockTask) GetID() string { + return t.id +} + +type taskResult struct { + taskID string + err error +} + +type mockObserver struct { + mu sync.Mutex + workerStarts []int + workerStops []int + tasksDone []taskResult + t *testing.T +} + +func (o *mockObserver) OnWorkerStart(workerID int) { + o.mu.Lock() + defer o.mu.Unlock() + o.workerStarts = append(o.workerStarts, workerID) + if o.t != nil { + o.t.Logf("observer: worker %d started", workerID) + } +} + +func (o *mockObserver) OnWorkerStop(workerID int) { + o.mu.Lock() + defer o.mu.Unlock() + o.workerStops = append(o.workerStops, workerID) + if o.t != nil { + o.t.Logf("observer: worker %d stopped", workerID) + } +} + +func (o *mockObserver) OnTaskDone(workerID int, taskID string, err error) { + o.mu.Lock() + defer o.mu.Unlock() + o.tasksDone = append(o.tasksDone, taskResult{taskID: taskID, err: err}) + if o.t != nil { + if err != nil { + o.t.Logf("observer: worker %d task %s failed: %v", workerID, taskID, err) + } else { + o.t.Logf("observer: worker %d task %s succeeded", workerID, taskID) + } + } +} + +func TestNewWorkerPool(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + opts []Option + shouldErr bool + }{ + { + name: "valid config with defaults", + opts: []Option{}, + shouldErr: false, + }, + { + name: "valid config with all options", + opts: []Option{ + WithName("test-pool"), + WithNumWorkers(4), + WithTaskQueueSize(100), + }, + shouldErr: false, + }, + { + name: "invalid number of workers (zero)", + opts: []Option{WithNumWorkers(0)}, + shouldErr: true, + }, + { + name: "invalid number of workers (negative)", + opts: []Option{WithNumWorkers(-1)}, + shouldErr: true, + }, + { + name: "invalid task queue size (zero)", + opts: []Option{WithTaskQueueSize(0)}, + shouldErr: true, + }, + { + name: "invalid task queue size (negative)", + opts: []Option{WithTaskQueueSize(-1)}, + shouldErr: true, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + _, err := NewWorkerPool[*mockTask](t.Context(), test.opts...) + if (err != nil) != test.shouldErr { + t.Errorf("NewWorkerPool() error = %v, shouldErr %v", err, test.shouldErr) + } + }) + } +} + +func TestWorkerPool_Observer(t *testing.T) { + t.Parallel() + ctx := t.Context() + + obs := &mockObserver{t: t} + numWorkers := 2 + taskCount := 5 + + pool, err := NewWorkerPool[*mockTask](ctx, + WithName("observer-test"), + WithNumWorkers(numWorkers), + WithTaskQueueSize(10), + WithObserver(obs), + ) + if err != nil { + t.Fatalf("failed to create pool: %v", err) + } + + pool.Start() + + for i := range taskCount { + task := &mockTask{id: fmt.Sprintf("task-%d", i)} + if err := pool.Push(task); err != nil { + t.Errorf("failed to push task: %v", err) + } + } + + pool.StopAndWait() + + obs.mu.Lock() + defer obs.mu.Unlock() + + if len(obs.workerStarts) != numWorkers { + t.Errorf("expected %d worker starts, got %d", numWorkers, len(obs.workerStarts)) + } + if len(obs.workerStops) != numWorkers { + t.Errorf("expected %d worker stops, got %d", numWorkers, len(obs.workerStops)) + } + + if len(obs.tasksDone) != taskCount { + t.Errorf("expected %d tasks done, got %d", taskCount, len(obs.tasksDone)) + } + for _, r := range obs.tasksDone { + if r.err != nil { + t.Errorf("task %s should have succeeded, got error: %v", r.taskID, r.err) + } + } +} + +func TestWorkerPool_ObserverTaskFailed(t *testing.T) { + t.Parallel() + ctx := t.Context() + + obs := &mockObserver{t: t} + + pool, err := NewWorkerPool[*mockTask](ctx, + WithName("observer-fail-test"), + WithNumWorkers(1), + WithTaskQueueSize(10), + WithObserver(obs), + ) + if err != nil { + t.Fatalf("failed to create pool: %v", err) + } + + pool.Start() + + task := &mockTask{ + id: "fail-task", + fn: func(ctx context.Context) error { + return fmt.Errorf("intentional error") + }, + } + pool.Push(task) + + pool.StopAndWait() + + obs.mu.Lock() + defer obs.mu.Unlock() + + failures := 0 + for _, r := range obs.tasksDone { + if r.err != nil { + failures++ + } + } + + if failures != 1 { + t.Errorf("expected 1 failed task, got %d", failures) + } + if len(obs.tasksDone) > 0 && obs.tasksDone[0].taskID != "fail-task" { + t.Errorf("expected failed task id 'fail-task', got %s", obs.tasksDone[0].taskID) + } +} + +func TestWorkerPool_ObserverTaskPanic(t *testing.T) { + t.Parallel() + ctx := t.Context() + + panicCaught := make(chan string, 1) + + pool, err := NewWorkerPool[*mockTask](ctx, + WithName("observer-panic-test"), + WithNumWorkers(1), + WithTaskQueueSize(10), + WithPanicHandler(func(workerID int, taskID string, panicValue any) { + panicCaught <- taskID + }), + ) + if err != nil { + t.Fatalf("failed to create pool: %v", err) + } + + pool.Start() + + task := &mockTask{ + id: "panic-task", + fn: func(ctx context.Context) error { + panic("intentional panic") + }, + } + pool.Push(task) + + pool.StopAndWait() + + select { + case taskID := <-panicCaught: + if taskID != "panic-task" { + t.Errorf("expected panic from 'panic-task', got %s", taskID) + } + default: + t.Error("expected PanicHandler to be called") + } +} + +func TestWorkerPool_ConcurrentPush(t *testing.T) { + t.Parallel() + ctx := t.Context() + var wg sync.WaitGroup + taskCount := 20 + + pool, err := NewWorkerPool[*mockTask](ctx, + WithName("concurrent-push"), + WithNumWorkers(3), + WithTaskQueueSize(taskCount), + ) + if err != nil { + t.Fatalf("failed to create pool: %v", err) + } + + pool.Start() + defer pool.StopAndWait() + + for i := range taskCount { + wg.Add(1) + go func(id int) { + defer wg.Done() + task := &mockTask{id: fmt.Sprintf("task-%d", id)} + if err := pool.Push(task); err != nil { + t.Errorf("push failed: %v", err) + } + }(i) + } + + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("timeout waiting for concurrent push") + } +} + +func TestWorkerPool_TaskQueueFull(t *testing.T) { + t.Parallel() + ctx := t.Context() + + pool, err := NewWorkerPool[*mockTask](ctx, + WithName("queue-full"), + WithNumWorkers(1), + WithTaskQueueSize(2), + ) + if err != nil { + t.Fatalf("failed to create pool: %v", err) + } + + for i := range 2 { + task := &mockTask{id: fmt.Sprintf("task-%d", i)} + if err := pool.Push(task); err != nil { + t.Errorf("unexpected error on push %d: %v", i, err) + } + } + + task := &mockTask{id: "overflow-task"} + err = pool.Push(task) + if err == nil { + t.Fatal("expected error when queue is full, got nil") + } + if !errors.Is(err, ErrQueueFull) { + t.Fatalf("expected ErrQueueFull, got: %v", err) + } +} + +func TestWorkerPool_MultipleStop(t *testing.T) { + t.Parallel() + ctx := t.Context() + + pool, err := NewWorkerPool[*mockTask](ctx, + WithName("multiple-stop"), + WithNumWorkers(2), + WithTaskQueueSize(10), + ) + if err != nil { + t.Fatalf("failed to create pool: %v", err) + } + + pool.Start() + pool.Stop() + + func() { + defer func() { + if r := recover(); r != nil { + t.Errorf("Stop() panicked on second call: %v", r) + } + }() + pool.Stop() + }() + + pool.Wait() +} + +func TestWorkerPool_MetricsThreadSafe(t *testing.T) { + t.Parallel() + ctx := t.Context() + var wg sync.WaitGroup + taskCount := 20 + + pool, err := NewWorkerPool[*mockTask](ctx, + WithName("metrics-safe"), + WithNumWorkers(10), + WithTaskQueueSize(taskCount), + ) + if err != nil { + t.Fatalf("failed to create pool: %v", err) + } + + pool.Start() + defer pool.StopAndWait() + + for i := range taskCount { + wg.Add(1) + go func(id int) { + defer wg.Done() + task := &mockTask{ + id: fmt.Sprintf("task-%d", id), + fn: func(ctx context.Context) error { + time.Sleep(time.Millisecond) + return nil + }, + } + pool.Push(task) + }(i) + } + + concurrentReads := 100 + for range concurrentReads { + wg.Add(1) + go func() { + defer wg.Done() + for range 10 { + pool.GetMetrics() + } + }() + } + + wg.Wait() +} + +func TestWorkerPool_ContextCancellation(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(t.Context()) + + pool, err := NewWorkerPool[*mockTask](ctx, + WithName("context-cancel"), + WithNumWorkers(2), + WithTaskQueueSize(10), + ) + if err != nil { + t.Fatalf("failed to create pool: %v", err) + } + + pool.Start() + + for i := range 5 { + task := &mockTask{id: fmt.Sprintf("task-%d", i)} + pool.Push(task) + } + + cancel() + + err = pool.Wait() + if err != nil && !errors.Is(err, context.Canceled) { + t.Errorf("expected context.Canceled, got: %v", err) + } +} + +func TestWorkerPool_ErrorHandler(t *testing.T) { + t.Parallel() + ctx := t.Context() + + errorHandlerCalled := make(chan struct{}) + + pool, err := NewWorkerPool[*mockTask](ctx, + WithName("error-handler-test"), + WithNumWorkers(1), + WithTaskQueueSize(10), + WithErrorHandler(func(workerID int, taskID string, err error) { + t.Logf("ErrorHandler called: worker=%d, task=%s, err=%v", workerID, taskID, err) + close(errorHandlerCalled) + }), + ) + if err != nil { + t.Fatalf("failed to create pool: %v", err) + } + + pool.Start() + + task := &mockTask{ + id: "error-task", + fn: func(ctx context.Context) error { + return fmt.Errorf("task error") + }, + } + pool.Push(task) + + select { + case <-errorHandlerCalled: + case <-time.After(time.Second): + t.Fatal("ErrorHandler was not called within timeout") + } + + pool.StopAndWait() +} + +func TestWorkerPool_PanicHandler(t *testing.T) { + t.Parallel() + ctx := t.Context() + + panicHandlerCalled := make(chan struct{}) + + pool, err := NewWorkerPool[*mockTask](ctx, + WithName("panic-handler-test"), + WithNumWorkers(1), + WithTaskQueueSize(10), + WithPanicHandler(func(workerID int, taskID string, panicValue any) { + t.Logf("PanicHandler called: worker=%d, task=%s, panic=%v", workerID, taskID, panicValue) + close(panicHandlerCalled) + }), + ) + if err != nil { + t.Fatalf("failed to create pool: %v", err) + } + + pool.Start() + + task := &mockTask{ + id: "panic-task", + fn: func(ctx context.Context) error { + panic("task panic") + }, + } + pool.Push(task) + + select { + case <-panicHandlerCalled: + case <-time.After(time.Second): + t.Fatal("PanicHandler was not called within timeout") + } + + pool.StopAndWait() +} + +func TestWorkerPool_PushContextCancelled(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(t.Context()) + + pool, err := NewWorkerPool[*mockTask](ctx, + WithName("push-ctx-cancelled"), + WithNumWorkers(1), + WithTaskQueueSize(1), + ) + if err != nil { + t.Fatalf("failed to create pool: %v", err) + } + + task1 := &mockTask{id: "task-1"} + if err := pool.Push(task1); err != nil { + t.Fatalf("failed to fill queue: %v", err) + } + + cancel() + + task2 := &mockTask{id: "task-2"} + err = pool.Push(task2) + + if err == nil { + t.Fatal("expected error when pushing to full queue with cancelled context") + } + + if !errors.Is(err, ErrQueueFull) && !errors.Is(err, context.Canceled) { + t.Fatalf("expected ErrQueueFull or context.Canceled, got: %v", err) + } +} + +func TestWorkerPool_NoObserver(t *testing.T) { + t.Parallel() + ctx := t.Context() + + pool, err := NewWorkerPool[*mockTask](ctx, + WithName("no-observer"), + WithNumWorkers(2), + WithTaskQueueSize(10), + ) + if err != nil { + t.Fatalf("failed to create pool: %v", err) + } + + pool.Start() + + task1 := &mockTask{id: "success-task"} + pool.Push(task1) + + task2 := &mockTask{ + id: "fail-task", + fn: func(ctx context.Context) error { + return fmt.Errorf("error") + }, + } + pool.Push(task2) + + task3 := &mockTask{ + id: "panic-task", + fn: func(ctx context.Context) error { + panic("panic") + }, + } + pool.Push(task3) + + pool.StopAndWait() + + metrics := pool.GetMetrics() + if metrics.TasksFailed != 2 { + t.Errorf("expected 2 failed tasks, got %d", metrics.TasksFailed) + } +} + +func TestWorkerPool_StartMultipleTimes(t *testing.T) { + t.Parallel() + + pool, err := NewWorkerPool[*mockTask](t.Context(), + WithNumWorkers(1), + WithTaskQueueSize(1), + ) + if err != nil { + t.Fatalf("failed to create pool: %v", err) + } + + pool.Start() + defer pool.StopAndWait() + + defer func() { + if r := recover(); r == nil { + t.Error("expected panic when calling Start() multiple times") + } + }() + + pool.Start() +} -- cgit v1.3.1