aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorsangtran <tranquangsang12.7@gmail.com>2026-09-03 23:10:33 +0700
committersangtran <tranquangsang12.7@gmail.com>2026-09-03 23:10:33 +0700
commit02123df8cc15896f97c1580e49d8e2a991e37ddd (patch)
treed0c00bc3f65e5f5493a9dff8034c80640cb12960
signatureVerified sangtran <tranquangsang12.7@gmail.com> Key: 8DD79823D235E1BB ultimate trust
downloadquacker-02123df8cc15896f97c1580e49d8e2a991e37ddd.tar.gz
quacker-02123df8cc15896f97c1580e49d8e2a991e37ddd.zip
Verified Moving all legacy quacker source code from github to cgitHEADmain
Diffstat
-rw-r--r--.gitignore9
-rw-r--r--LICENSE21
-rw-r--r--README125
-rw-r--r--fanin/example_test.go35
-rw-r--r--fanin/fanin.go142
-rw-r--r--fanin/fanin_bench_test.go36
-rw-r--r--fanin/fanin_test.go281
-rw-r--r--fanout/example_test.go45
-rw-r--r--fanout/fanout.go203
-rw-r--r--fanout/fanout_bench_test.go65
-rw-r--r--fanout/fanout_test.go412
-rw-r--r--go.mod5
-rw-r--r--go.sum2
-rw-r--r--pipeline/batch.go64
-rw-r--r--pipeline/batch_test.go106
-rw-r--r--pipeline/example_test.go53
-rw-r--r--pipeline/options.go34
-rw-r--r--pipeline/pipeline.go21
-rw-r--r--pipeline/pipeline_bench_test.go55
-rw-r--r--pipeline/pipeline_test.go255
-rw-r--r--pipeline/stage.go310
-rw-r--r--pipeline/stage_test.go308
-rw-r--r--workerpool/workerpool.go262
-rw-r--r--workerpool/workerpool_bench_test.go50
-rw-r--r--workerpool/workerpool_test.go608
25 files changed, 3507 insertions, 0 deletions
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..d18c3de
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,9 @@
1.DS_Store
2
3# Test coverage
4*.out
5coverage.txt
6coverage.html
7
8# Test binaries
9*.test
diff --git a/LICENSE b/LICENSE
new file mode 100644
index 0000000..df79de6
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,21 @@
1MIT License
2
3Copyright (c) 2025 Tran Quang Sang
4
5Permission is hereby granted, free of charge, to any person obtaining a copy
6of this software and associated documentation files (the "Software"), to deal
7in the Software without restriction, including without limitation the rights
8to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9copies of the Software, and to permit persons to whom the Software is
10furnished to do so, subject to the following conditions:
11
12The above copyright notice and this permission notice shall be included in all
13copies or substantial portions of the Software.
14
15THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21SOFTWARE.
diff --git a/README b/README
new file mode 100644
index 0000000..c758b17
--- /dev/null
+++ b/README
@@ -0,0 +1,125 @@
1QUACKER
2=======
3
4Type-safe, composable concurrency patterns with backpressure for Go.
5
6Quacker provides composable, channel-native concurrency primitives.
7Backpressure propagates naturally through Go's channel semantics —
8no custom protocols, no framework lock-in.
9
10
11INSTALLATION
12------------
13
14 go get git.sangtran.io.vn/quacker.git
15
16Requirements: Go 1.25+
17
18
19PACKAGES
20--------
21
22pipeline
23 Stream processing with backpressure. Compose stages into pipelines.
24 Each stage transforms a channel into a channel.
25
26 Example:
27 ctx, cancel := context.WithCancelCause(ctx)
28
29 out := pipeline.Run(ctx, input,
30 pipeline.Map(3, func(ctx context.Context, e Event) (Event, bool) {
31 enriched, err := enrich(ctx, e)
32 if err != nil {
33 cancel(err)
34 return e, false
35 }
36 return enriched, true
37 }),
38 pipeline.Filter(func(e Event) bool { return e.Priority > 0 }),
39 pipeline.Sink(5, func(ctx context.Context, e Event) {
40 store(ctx, e)
41 }),
42 )
43
44 for range out {}
45 if err := context.Cause(ctx); err != nil {
46 log.Fatal(err)
47 }
48
49 Stages: Map, Filter, ForEach, Buffer, Tee, Take, Sink, Batch.
50
51fanout
52 1-to-N distribution. Distribute a single stream to multiple
53 consumers. RoundRobin for load balancing, Broadcast for pub/sub.
54
55 Example:
56 fo, _ := fanout.NewFanOut[int](
57 fanout.WithWorkerCount(3),
58 fanout.WithStrategy(fanout.Broadcast),
59 )
60 fo.Run(ctx, input)
61
62 for _, ch := range fo.Outputs() {
63 go consume(ch)
64 }
65
66fanin
67 N-to-1 merge. Merge multiple channels into a single stream.
68
69 Example:
70 fi, _ := fanin.NewFanIn[int]()
71 fi.Add(stream1)
72 fi.Add(stream2)
73 fi.Add(stream3)
74 merged := fi.Run(ctx)
75
76workerpool
77 Fixed-size worker pool with push-based task queuing, metrics,
78 and graceful shutdown.
79
80 Example:
81 pool, _ := workerpool.NewWorkerPool[MyTask](ctx,
82 workerpool.WithNumWorkers(4),
83 workerpool.WithTaskQueueSize(100),
84 )
85 pool.Start()
86 pool.Push(task)
87 pool.StopAndWait()
88
89
90HOW THEY COMPOSE
91----------------
92
93The primitives connect through channels — Go's universal connector:
94
95 +-- pipeline --+
96 source --> fanout --+-- pipeline --+-- fanin --> pipeline --> sink
97 +-- pipeline --+
98
99 Example:
100 fo.Run(ctx, source)
101
102 fast := pipeline.Run(ctx, fo.Outputs()[0], pipeline.Map(10, enrichFast))
103 slow := pipeline.Run(ctx, fo.Outputs()[1], pipeline.Map(2, enrichSlow))
104
105 fi.Add(fast)
106 fi.Add(slow)
107 merged := fi.Run(ctx)
108
109 out := pipeline.Run(ctx, merged, pipeline.Sink(5, store))
110
111
112DESIGN
113------
114
115- Backpressure is free. Channels block when full. No custom protocol.
116- Errors use stdlib. context.WithCancelCause + context.Cause.
117- Shutdown is clear. Close input for drain; cancel context for stop.
118- Small interfaces. Observer interfaces are 2-3 methods, optional.
119- Zero dependencies. Core logic is stdlib only (golang.org/x/sync).
120
121
122LICENSE
123-------
124
125MIT (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 @@
1package fanin_test
2
3import (
4 "context"
5 "fmt"
6 "sort"
7
8 "git.sangtran.io.vn/quacker.git/fanin"
9)
10
11func ExampleNewFanIn() {
12 fi, _ := fanin.NewFanIn[int]()
13
14 ch1 := make(chan int, 2)
15 ch1 <- 1
16 ch1 <- 2
17 close(ch1)
18
19 ch2 := make(chan int, 2)
20 ch2 <- 3
21 ch2 <- 4
22 close(ch2)
23
24 fi.Add(ch1)
25 fi.Add(ch2)
26
27 var results []int
28 for v := range fi.Run(context.Background()) {
29 results = append(results, v)
30 }
31
32 sort.Ints(results)
33 fmt.Println(results)
34 // Output: [1 2 3 4]
35}
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 @@
1// Package fanin merges multiple input channels into a single output channel.
2package fanin
3
4import (
5 "context"
6 "fmt"
7 "sync"
8 "sync/atomic"
9)
10
11// FanInObserver defines callbacks for monitoring fan-in lifecycle events.
12type FanInObserver interface {
13 OnInputAdded()
14 OnInputClosed()
15}
16
17// FanInConfig holds configuration options for creating a FanIn instance.
18type FanInConfig struct {
19 BufferSize int
20 Observer FanInObserver
21}
22
23// Option is a functional option for configuring a FanIn instance.
24type Option func(*FanInConfig)
25
26// FanIn merges multiple input channels of type T into a single output channel.
27type FanIn[T any] struct {
28 cfg *FanInConfig
29 mu sync.Mutex
30 inputs []<-chan T
31 running atomic.Bool
32 done chan struct{}
33}
34
35// NewFanIn creates a new FanIn instance with the provided options.
36func NewFanIn[T any](opts ...Option) (*FanIn[T], error) {
37 cfg := &FanInConfig{}
38
39 for _, opt := range opts {
40 opt(cfg)
41 }
42
43 if cfg.BufferSize < 0 {
44 return nil, fmt.Errorf("fanin: buffer size must not be negative, got %d", cfg.BufferSize)
45 }
46
47 return &FanIn[T]{
48 cfg: cfg,
49 }, nil
50}
51
52// WithBufferSize sets the buffer capacity for the merged output channel.
53func WithBufferSize(size int) Option {
54 return func(f *FanInConfig) {
55 f.BufferSize = size
56 }
57}
58
59// WithObserver attaches a FanInObserver to receive lifecycle event callbacks.
60func WithObserver(observer FanInObserver) Option {
61 return func(f *FanInConfig) {
62 f.Observer = observer
63 }
64}
65
66// Add registers an input channel to be merged into the output.
67// It must be called before Run, and is safe for concurrent use.
68func (f *FanIn[T]) Add(ch <-chan T) {
69 if ch == nil {
70 panic("fanin: cannot add nil channel")
71 }
72
73 if f.running.Load() {
74 panic("fanin: Add() called after Run()")
75 }
76
77 f.mu.Lock()
78 defer f.mu.Unlock()
79
80 f.inputs = append(f.inputs, ch)
81
82 if f.cfg.Observer != nil {
83 safeObserve(f.cfg.Observer.OnInputAdded)
84 }
85}
86
87// Run starts merging the registered input channels into a single output channel.
88func (f *FanIn[T]) Run(ctx context.Context) <-chan T {
89 if !f.running.CompareAndSwap(false, true) {
90 panic("fanin: Run() called multiple times")
91 }
92
93 f.mu.Lock()
94 inputs := append([]<-chan T(nil), f.inputs...)
95 out := make(chan T, f.cfg.BufferSize)
96 f.done = make(chan struct{})
97 f.mu.Unlock()
98
99 var wg sync.WaitGroup
100 for _, ch := range inputs {
101 wg.Go(func() {
102 for {
103 select {
104 case v, ok := <-ch:
105 if !ok {
106 if f.cfg.Observer != nil {
107 safeObserve(f.cfg.Observer.OnInputClosed)
108 }
109 return
110 }
111 select {
112 case out <- v:
113 case <-ctx.Done():
114 return
115 }
116 case <-ctx.Done():
117 return
118 }
119 }
120 })
121 }
122
123 go func() {
124 wg.Wait()
125 close(out)
126 close(f.done)
127 }()
128
129 return out
130}
131
132func safeObserve(fn func()) {
133 defer func() { recover() }()
134 fn()
135}
136
137// Done returns a channel that is closed when all inputs are processed or the context is cancelled.
138func (f *FanIn[T]) Done() <-chan struct{} {
139 f.mu.Lock()
140 defer f.mu.Unlock()
141 return f.done
142}
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 @@
1package fanin
2
3import (
4 "testing"
5)
6
7func BenchmarkFanIn(b *testing.B) {
8 b.ReportAllocs()
9
10 fi, _ := NewFanIn[int](WithBufferSize(64))
11
12 producers := 4
13 chs := make([]chan int, producers)
14 for i := range chs {
15 chs[i] = make(chan int, 64)
16 fi.Add(chs[i])
17 }
18
19 out := fi.Run(b.Context())
20
21 // Drain output
22 go func() {
23 for range out {
24 }
25 }()
26
27 for b.Loop() {
28 chs[b.N%producers] <- 1
29 }
30
31 for _, ch := range chs {
32 close(ch)
33 }
34
35 <-fi.Done()
36}
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 @@
1package fanin
2
3import (
4 "context"
5 "sync"
6 "testing"
7 "time"
8)
9
10type mockObserver struct {
11 name string
12 t *testing.T
13}
14
15func (o *mockObserver) OnInputAdded() {
16 o.t.Logf("%s: OnInputAdded", o.name)
17}
18
19func (o *mockObserver) OnInputClosed() {
20 o.t.Logf("%s: OnInputClosed", o.name)
21}
22
23func TestNewFanIn(t *testing.T) {
24 t.Parallel()
25
26 obs := new(mockObserver)
27
28 tests := []struct {
29 name string
30 cfg *FanInConfig
31 shouldErr bool
32 }{
33 {
34 name: "valid config",
35 cfg: &FanInConfig{
36 BufferSize: 1,
37 Observer: obs,
38 },
39 shouldErr: false,
40 },
41 {
42 name: "invalid buffer size",
43 cfg: &FanInConfig{
44 BufferSize: -1,
45 Observer: obs,
46 },
47 shouldErr: true,
48 },
49 }
50
51 for _, test := range tests {
52 t.Run(test.name, func(t *testing.T) {
53 t.Parallel()
54 _, err := NewFanIn[int](WithBufferSize(test.cfg.BufferSize), WithObserver(obs))
55
56 if (err != nil) != test.shouldErr {
57 t.Errorf("validation error = %v, want err %v", err, test.shouldErr)
58 }
59 })
60 }
61}
62
63func TestFanIn_Add(t *testing.T) {
64 t.Parallel()
65 fi, _ := NewFanIn[int]()
66
67 inputs := make([]chan int, 10)
68 for i := range inputs {
69 inputs[i] = make(chan int, 1)
70 inputs[i] <- i
71 close(inputs[i])
72 }
73
74 for _, ch := range inputs {
75 fi.Add(ch)
76 }
77
78 count := 0
79 for range fi.Run(t.Context()) {
80 count++
81 }
82
83 if count != 10 {
84 t.Fatalf("expected 10 items, got %d", count)
85 }
86}
87
88func TestFanIn_Run(t *testing.T) {
89 t.Parallel()
90 fi, _ := NewFanIn[int]()
91
92 ch := make(chan int)
93 go func() {
94 ch <- 1
95 ch <- 2
96 ch <- 3
97 close(ch)
98 }()
99
100 fi.Add(ch)
101
102 for v := range fi.Run(t.Context()) {
103 t.Logf("value received %d", v)
104 }
105
106 <-fi.Done()
107}
108
109func TestFanIn_ContextCancellation(t *testing.T) {
110 t.Parallel()
111 fi, _ := NewFanIn[int]()
112 ctx, cancel := context.WithTimeout(t.Context(), time.Millisecond)
113 defer cancel()
114
115 inputs := make([]chan int, 10)
116 for i := range inputs {
117 inputs[i] = make(chan int, 1)
118 inputs[i] <- i
119 close(inputs[i])
120 }
121
122 for _, ch := range inputs {
123 fi.Add(ch)
124 }
125
126 select {
127 case <-fi.Run(ctx):
128 cancel()
129 case <-fi.Done():
130 cancel()
131 }
132}
133
134func TestFanIn_ShouldPanicRunMultiple(t *testing.T) {
135 t.Parallel()
136
137 defer func() {
138 if r := recover(); r == nil {
139 t.Errorf("expected panic for multiple Run calls")
140 }
141 }()
142
143 ctx := t.Context()
144 fi, _ := NewFanIn[int]()
145
146 inputs := make([]chan int, 10)
147 for i := range inputs {
148 inputs[i] = make(chan int, 1)
149 inputs[i] <- i
150 close(inputs[i])
151 }
152
153 for _, ch := range inputs {
154 fi.Add(ch)
155 }
156
157 for v := range fi.Run(ctx) {
158 t.Logf("value received %d", v)
159 }
160
161 fi.Run(ctx)
162}
163
164func TestFanIn_ShouldPanicAddAfterRun(t *testing.T) {
165 t.Parallel()
166
167 defer func() {
168 if r := recover(); r == nil {
169 t.Errorf("expected panic for Add after Run")
170 }
171 }()
172
173 fi, _ := NewFanIn[int]()
174
175 inputs := make([]chan int, 10)
176 for i := range inputs {
177 inputs[i] = make(chan int, 1)
178 inputs[i] <- i
179 close(inputs[i])
180 }
181
182 for _, ch := range inputs {
183 fi.Add(ch)
184 }
185
186 for v := range fi.Run(t.Context()) {
187 t.Logf("value received %d", v)
188 }
189
190 fi.Add(make(chan int))
191}
192
193func TestFanIn_Observer(t *testing.T) {
194 t.Parallel()
195 obs := &mockObserver{t: t}
196
197 fi, _ := NewFanIn[int](WithObserver(obs))
198
199 inputs := make([]chan int, 10)
200 for i := range inputs {
201 inputs[i] = make(chan int, 1)
202 inputs[i] <- i
203 close(inputs[i])
204 }
205
206 for _, ch := range inputs {
207 fi.Add(ch)
208 }
209
210 for v := range fi.Run(t.Context()) {
211 t.Logf("value received %d", v)
212 }
213
214 <-fi.Done()
215}
216
217func TestFanIn_AddConcurrency(t *testing.T) {
218 t.Parallel()
219
220 fi, _ := NewFanIn[int]()
221 var wg sync.WaitGroup
222 numGoroutines := 10
223 channelPerGoroutine := 10
224
225 for range numGoroutines {
226 wg.Go(func() {
227 inputs := make([]chan int, channelPerGoroutine)
228 for i := range inputs {
229 inputs[i] = make(chan int, 1)
230 inputs[i] <- i
231 close(inputs[i])
232 fi.Add(inputs[i])
233 }
234 })
235 }
236 wg.Wait()
237
238 count := 0
239 for range fi.Run(t.Context()) {
240 count++
241 }
242
243 <-fi.Done()
244
245 expectedCount := numGoroutines * channelPerGoroutine
246 if count != expectedCount {
247 t.Errorf("expected %d items, got %d", expectedCount, count)
248 }
249}
250
251func TestFanIn_NilChannel(t *testing.T) {
252 t.Parallel()
253
254 fi, _ := NewFanIn[int]()
255
256 defer func() {
257 if r := recover(); r == nil {
258 t.Error("expected panic when adding nil channel")
259 }
260 }()
261
262 fi.Add(nil)
263}
264
265func TestFanIn_RunWithNoInputs(t *testing.T) {
266 t.Parallel()
267
268 fi, _ := NewFanIn[int]()
269 out := fi.Run(t.Context())
270
271 count := 0
272 for range out {
273 count++
274 }
275
276 if count != 0 {
277 t.Errorf("expected 0 values from empty fan-in, got %d", count)
278 }
279
280 <-fi.Done()
281}
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 @@
1package fanout_test
2
3import (
4 "context"
5 "fmt"
6 "sort"
7 "sync"
8
9 "git.sangtran.io.vn/quacker.git/fanout"
10)
11
12func ExampleNewFanOut() {
13 fo, _ := fanout.NewFanOut[int](
14 fanout.WithWorkerCount(2),
15 fanout.WithBufferSize(10),
16 fanout.WithStrategy(fanout.RoundRobin),
17 )
18
19 input := make(chan int, 4)
20 for i := 1; i <= 4; i++ {
21 input <- i
22 }
23 close(input)
24
25 fo.Run(context.Background(), input)
26
27 var mu sync.Mutex
28 var all []int
29 var wg sync.WaitGroup
30
31 for _, ch := range fo.Outputs() {
32 wg.Go(func() {
33 for v := range ch {
34 mu.Lock()
35 all = append(all, v)
36 mu.Unlock()
37 }
38 })
39 }
40
41 wg.Wait()
42 sort.Ints(all)
43 fmt.Println(all)
44 // Output: [1 2 3 4]
45}
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 @@
1// Package fanout distributes values from a single input channel to multiple output channels.
2package fanout
3
4import (
5 "context"
6 "errors"
7 "fmt"
8 "sync"
9 "sync/atomic"
10)
11
12// FanOutObserver defines callbacks for monitoring fan-out lifecycle events.
13type FanOutObserver interface {
14 OnDistribute()
15 OnOutputClosed()
16}
17
18// FanOutStrategy defines how values are distributed across output channels.
19type FanOutStrategy int
20
21const (
22 // RoundRobin distributes values evenly across output channels in rotation.
23 RoundRobin FanOutStrategy = iota
24 // Broadcast sends each value to all output channels.
25 Broadcast
26)
27
28// ErrInvalidStrategy is returned when an unknown FanOutStrategy is provided.
29var ErrInvalidStrategy = errors.New("fanout: unknown distribution strategy")
30
31// FanOutConfig holds configuration options for creating a FanOut instance.
32type FanOutConfig struct {
33 WorkerCount int
34 BufferSize int
35 Strategy FanOutStrategy
36 Observer FanOutObserver
37}
38
39// Option is a functional option for configuring a FanOut instance.
40type Option func(*FanOutConfig)
41
42// FanOut distributes values from a single input channel to multiple output channels.
43type FanOut[T any] struct {
44 cfg *FanOutConfig
45 outputs []chan T
46 running atomic.Bool
47 cachedOutputs []<-chan T
48 once sync.Once
49}
50
51// NewFanOut creates a new FanOut instance with the provided options.
52func NewFanOut[T any](opts ...Option) (*FanOut[T], error) {
53 cfg := &FanOutConfig{
54 WorkerCount: 1,
55 BufferSize: 0,
56 Observer: nil,
57 Strategy: RoundRobin,
58 }
59
60 for _, opt := range opts {
61 opt(cfg)
62 }
63
64 if cfg.WorkerCount < 1 {
65 return nil, fmt.Errorf("fanout: worker count must be at least 1, got %d", cfg.WorkerCount)
66 }
67
68 if cfg.BufferSize < 0 {
69 return nil, fmt.Errorf("fanout: buffer size must not be negative, got %d", cfg.BufferSize)
70 }
71
72 if cfg.Strategy != RoundRobin && cfg.Strategy != Broadcast {
73 return nil, ErrInvalidStrategy
74 }
75
76 outputs := make([]chan T, cfg.WorkerCount)
77 for i := range outputs {
78 outputs[i] = make(chan T, cfg.BufferSize)
79 }
80
81 return &FanOut[T]{
82 cfg: cfg,
83 outputs: outputs,
84 }, nil
85}
86
87// Run starts distributing values from the input channel to all output channels.
88func (f *FanOut[T]) Run(ctx context.Context, input <-chan T) {
89 if input == nil {
90 panic("fanout: input channel must not be nil")
91 }
92
93 if !f.running.CompareAndSwap(false, true) {
94 panic("fanout: Run() called multiple times")
95 }
96
97 go func() {
98 idx := 0
99 defer f.running.Store(false)
100 defer f.closeAllOutput()
101 for {
102 select {
103 case value, ok := <-input:
104 if !ok {
105 return
106 }
107 switch f.cfg.Strategy {
108 case RoundRobin:
109 f.roundRobin(ctx, value, &idx)
110 case Broadcast:
111 f.broadcast(ctx, value)
112 }
113 case <-ctx.Done():
114 return
115 }
116 }
117 }()
118}
119
120// Outputs returns receive-only channel references for the distributed streams.
121func (f *FanOut[T]) Outputs() []<-chan T {
122 f.once.Do(func() {
123 f.cachedOutputs = make([]<-chan T, len(f.outputs))
124 for i, ch := range f.outputs {
125 f.cachedOutputs[i] = ch
126 }
127 })
128
129 return f.cachedOutputs
130}
131
132func safeObserve(fn func()) {
133 defer func() { recover() }()
134 fn()
135}
136
137func (f *FanOut[T]) roundRobin(ctx context.Context, value T, index *int) {
138 target := *index % f.cfg.WorkerCount
139 select {
140 case f.outputs[target] <- value:
141 *index = (*index + 1) % f.cfg.WorkerCount
142 if f.cfg.Observer != nil {
143 safeObserve(f.cfg.Observer.OnDistribute)
144 }
145 case <-ctx.Done():
146 return
147 }
148}
149
150func (f *FanOut[T]) broadcast(ctx context.Context, value T) {
151 var wg sync.WaitGroup
152 for _, ch := range f.outputs {
153 wg.Go(func() {
154 select {
155 case ch <- value:
156 if f.cfg.Observer != nil {
157 safeObserve(f.cfg.Observer.OnDistribute)
158 }
159 case <-ctx.Done():
160 return
161 }
162 })
163 }
164 wg.Wait()
165}
166
167func (f *FanOut[T]) closeAllOutput() {
168 for _, ch := range f.outputs {
169 close(ch)
170 }
171
172 if f.cfg.Observer != nil {
173 safeObserve(f.cfg.Observer.OnOutputClosed)
174 }
175}
176
177// WithWorkerCount sets the number of output channels for the fan-out.
178func WithWorkerCount(n int) Option {
179 return func(cfg *FanOutConfig) {
180 cfg.WorkerCount = n
181 }
182}
183
184// WithBufferSize sets the buffer capacity for each output channel.
185func WithBufferSize(size int) Option {
186 return func(cfg *FanOutConfig) {
187 cfg.BufferSize = size
188 }
189}
190
191// WithStrategy sets the distribution strategy for the fan-out.
192func WithStrategy(strategy FanOutStrategy) Option {
193 return func(cfg *FanOutConfig) {
194 cfg.Strategy = strategy
195 }
196}
197
198// WithObserver attaches a FanOutObserver to receive lifecycle event callbacks.
199func WithObserver(observer FanOutObserver) Option {
200 return func(cfg *FanOutConfig) {
201 cfg.Observer = observer
202 }
203}
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 @@
1package fanout
2
3import (
4 "context"
5 "testing"
6)
7
8func BenchmarkFanOut_RoundRobin(b *testing.B) {
9 b.ReportAllocs()
10
11 fo, _ := NewFanOut[int](
12 WithWorkerCount(4),
13 WithBufferSize(64),
14 WithStrategy(RoundRobin),
15 )
16
17 ctx, cancel := context.WithCancel(context.Background())
18 defer cancel()
19
20 input := make(chan int, 64)
21 fo.Run(ctx, input)
22
23 // Drain outputs
24 for _, ch := range fo.Outputs() {
25 go func() {
26 for range ch {
27 }
28 }()
29 }
30
31 for b.Loop() {
32 input <- 1
33 }
34
35 close(input)
36}
37
38func BenchmarkFanOut_Broadcast(b *testing.B) {
39 b.ReportAllocs()
40
41 fo, _ := NewFanOut[int](
42 WithWorkerCount(4),
43 WithBufferSize(64),
44 WithStrategy(Broadcast),
45 )
46
47 ctx, cancel := context.WithCancel(context.Background())
48 defer cancel()
49
50 input := make(chan int, 64)
51 fo.Run(ctx, input)
52
53 for _, ch := range fo.Outputs() {
54 go func() {
55 for range ch {
56 }
57 }()
58 }
59
60 for b.Loop() {
61 input <- 1
62 }
63
64 close(input)
65}
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 @@
1package fanout
2
3import (
4 "context"
5 "sync"
6 "testing"
7 "time"
8)
9
10type mockObserver struct {
11 name string
12 t *testing.T
13}
14
15func (m *mockObserver) OnDistribute() {
16 m.t.Logf("fanout: %s OnDistribute", m.name)
17}
18
19func (m *mockObserver) OnOutputClosed() {
20 m.t.Logf("fanout: %s OnOutputClosed", m.name)
21}
22
23func TestFanOut_NewFanOut(t *testing.T) {
24 t.Parallel()
25
26 tests := []struct {
27 name string
28 opts []Option
29 shouldErr bool
30 }{
31 {
32 name: "worker size is invalid",
33 opts: []Option{WithWorkerCount(-10)},
34 shouldErr: true,
35 },
36 {
37 name: "buffer size is invalid",
38 opts: []Option{WithBufferSize(-10), WithWorkerCount(1)},
39 shouldErr: true,
40 },
41 {
42 name: "valid options",
43 opts: []Option{WithBufferSize(1), WithWorkerCount(1)},
44 shouldErr: false,
45 },
46 {
47 name: "worker count is zero",
48 opts: []Option{WithWorkerCount(0)},
49 shouldErr: true,
50 },
51 {
52 name: "valid broadcast strategy",
53 opts: []Option{WithWorkerCount(5), WithBufferSize(10), WithStrategy(Broadcast)},
54 shouldErr: false,
55 },
56 {
57 name: "invalid strategy",
58 opts: []Option{WithWorkerCount(1), WithStrategy(FanOutStrategy(99))},
59 shouldErr: true,
60 },
61 }
62
63 for _, test := range tests {
64 t.Run(test.name, func(t *testing.T) {
65 t.Parallel()
66 _, err := NewFanOut[int](test.opts...)
67 if (err != nil) != test.shouldErr {
68 t.Errorf("validation error = %v, want err %v", err, test.shouldErr)
69 }
70 })
71 }
72}
73
74func TestFanOut_ConcurrentStress(t *testing.T) {
75 t.Parallel()
76 ctx, cancel := context.WithCancel(t.Context())
77 defer cancel()
78
79 workerCount := 5
80 fo, err := NewFanOut[int](
81 WithWorkerCount(workerCount),
82 WithBufferSize(100),
83 )
84 if err != nil {
85 t.Fatalf("new fanout error: %v", err)
86 }
87
88 input := make(chan int, 1000)
89 fo.Run(ctx, input)
90
91 var wg sync.WaitGroup
92
93 for i := 0; i < 10; i++ {
94 wg.Go(func() {
95 outs := fo.Outputs()
96 if len(outs) != workerCount {
97 t.Errorf("expected %d outputs, got %d", workerCount, len(outs))
98 }
99 })
100 }
101
102 for _, ch := range fo.Outputs() {
103 wg.Go(func() {
104 for range ch {
105 }
106 })
107 }
108
109 producers := 10
110 messagesPerProducer := 100
111
112 var producerWg sync.WaitGroup
113 producerWg.Add(producers)
114
115 for i := 0; i < producers; i++ {
116 go func() {
117 defer producerWg.Done()
118 for j := 0; j < messagesPerProducer; j++ {
119 select {
120 case input <- j:
121 case <-ctx.Done():
122 return
123 }
124 }
125 }()
126 }
127
128 go func() {
129 producerWg.Wait()
130 close(input)
131 }()
132
133 wg.Wait()
134}
135
136func TestFanOut_RunBroadcast(t *testing.T) {
137 t.Parallel()
138 var wg sync.WaitGroup
139 workerCount := 2
140 size := 10
141 obs := &mockObserver{
142 name: "fo broadcast",
143 t: t,
144 }
145 fo, _ := NewFanOut[int](
146 WithWorkerCount(workerCount),
147 WithStrategy(Broadcast),
148 WithObserver(obs),
149 )
150
151 ch := make(chan int, size)
152 wg.Go(func() {
153 for i := range size {
154 ch <- i
155 }
156 close(ch)
157 })
158
159 fo.Run(t.Context(), ch)
160
161 res := make([]int, workerCount)
162 for i, cha := range fo.Outputs() {
163 wg.Go(func() {
164 count := 0
165 for v := range cha {
166 t.Logf("%v", v)
167 count++
168 res[i] = count
169 }
170 })
171 }
172
173 wg.Wait()
174
175 expected := size
176 for i, count := range res {
177 if count != expected {
178 t.Errorf("worker %d received %d items, expected %d", i, count, expected)
179 }
180 }
181}
182
183func TestFanOut_RunRoundRobin(t *testing.T) {
184 t.Parallel()
185 var wg sync.WaitGroup
186 size := 10
187 workerCount := 2
188 obs := &mockObserver{
189 name: "fo round robin",
190 t: t,
191 }
192 fo, _ := NewFanOut[int](
193 WithWorkerCount(workerCount),
194 WithStrategy(RoundRobin),
195 WithObserver(obs),
196 )
197
198 ch := make(chan int, size)
199 wg.Go(func() {
200 for i := range size {
201 ch <- i
202 }
203 close(ch)
204 })
205
206 fo.Run(t.Context(), ch)
207 res := make([]int, workerCount)
208 for i, cha := range fo.Outputs() {
209 wg.Go(func() {
210 count := 0
211 for v := range cha {
212 t.Logf("%v", v)
213 count++
214 res[i] = count
215 }
216 })
217 }
218
219 wg.Wait()
220 expected := size / workerCount
221
222 for i, count := range res {
223 if count != expected {
224 t.Errorf("worker %d received %d items, expected %d", i, count, expected)
225 }
226 }
227}
228
229func TestFanOut_ContextTimeOut(t *testing.T) {
230 t.Parallel()
231 fo, _ := NewFanOut[int](WithWorkerCount(2))
232 ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond)
233 defer cancel()
234
235 input := make(chan int, 10)
236 fo.Run(ctx, input)
237
238 var wg sync.WaitGroup
239 done := make(chan struct{})
240 for _, ch := range fo.Outputs() {
241 wg.Go(func() {
242 for v := range ch {
243 t.Logf("received %d", v)
244 }
245 })
246 }
247
248 go func() {
249 wg.Wait()
250 close(done)
251 }()
252
253 go func() {
254 i := 0
255 for {
256 select {
257 case input <- i:
258 i++
259 time.Sleep(time.Millisecond)
260 case <-ctx.Done():
261 return
262 }
263 }
264 }()
265
266 select {
267 case <-done:
268 case <-time.After(200 * time.Millisecond):
269 t.Fatal("timeout waiting for cleanup")
270 }
271}
272
273func TestFanOut_ContextCancel(t *testing.T) {
274 t.Parallel()
275 ctx, cancel := context.WithCancel(t.Context())
276
277 fo, err := NewFanOut[int](WithWorkerCount(3))
278 if err != nil {
279 t.Fatalf("failed to create fanout: %v", err)
280 }
281
282 input := make(chan int, 10)
283 fo.Run(ctx, input)
284
285 var wg sync.WaitGroup
286 for _, ch := range fo.Outputs() {
287 wg.Go(func() {
288 for range ch {
289 }
290 })
291 }
292
293 cancel()
294 close(input)
295
296 done := make(chan struct{})
297 go func() {
298 wg.Wait()
299 close(done)
300 }()
301
302 select {
303 case <-done:
304 case <-time.After(time.Second):
305 t.Fatal("timeout waiting for shutdown")
306 }
307}
308
309func TestFanOut_ContextDoneRoundRobin(t *testing.T) {
310 t.Parallel()
311 ctx, cancel := context.WithCancel(t.Context())
312
313 fo, err := NewFanOut[int](
314 WithWorkerCount(3),
315 WithBufferSize(0),
316 WithStrategy(RoundRobin),
317 )
318 if err != nil {
319 t.Fatalf("new fanout error: %v", err)
320 }
321
322 input := make(chan int, 1)
323 input <- 1
324 fo.Run(ctx, input)
325
326 cancel()
327
328 for i, ch := range fo.Outputs() {
329 select {
330 case _, ok := <-ch:
331 if ok {
332 for range ch {
333 }
334 }
335 case <-time.After(time.Second):
336 t.Fatalf("worker %d: timeout waiting for close", i)
337 }
338 }
339}
340
341func TestFanOut_ContextDoneBroadcast(t *testing.T) {
342 t.Parallel()
343 ctx, cancel := context.WithCancel(t.Context())
344
345 fo, err := NewFanOut[int](
346 WithWorkerCount(3),
347 WithStrategy(Broadcast),
348 WithBufferSize(0),
349 )
350 if err != nil {
351 t.Fatalf("new fanout error: %v", err)
352 }
353
354 input := make(chan int, 1)
355 input <- 1
356 fo.Run(ctx, input)
357
358 cancel()
359
360 for i, ch := range fo.Outputs() {
361 select {
362 case _, ok := <-ch:
363 if ok {
364 for range ch {
365 }
366 }
367 case <-time.After(time.Second):
368 t.Fatalf("worker %d: timeout waiting for close", i)
369 }
370 }
371}
372
373func TestFanOut_RunMultipleTime(t *testing.T) {
374 t.Parallel()
375
376 ctx := t.Context()
377 fo, err := NewFanOut[int](
378 WithWorkerCount(3),
379 WithBufferSize(0),
380 )
381 if err != nil {
382 t.Fatalf("new fanout error: %v", err)
383 }
384
385 input := make(chan int, 10)
386 fo.Run(ctx, input)
387
388 defer func() {
389 if r := recover(); r == nil {
390 t.Error("expected panic when calling Run() multiple times")
391 }
392 }()
393
394 fo.Run(ctx, input)
395}
396
397func TestFanOut_NilInput(t *testing.T) {
398 t.Parallel()
399
400 fo, err := NewFanOut[int](WithWorkerCount(1))
401 if err != nil {
402 t.Fatalf("failed to create fanout: %v", err)
403 }
404
405 defer func() {
406 if r := recover(); r == nil {
407 t.Error("expected panic when passing nil input channel")
408 }
409 }()
410
411 fo.Run(t.Context(), nil)
412}
diff --git a/go.mod b/go.mod
new file mode 100644
index 0000000..df07984
--- /dev/null
+++ b/go.mod
@@ -0,0 +1,5 @@
1module git.sangtran.io.vn/quacker.git
2
3go 1.25.1
4
5require 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 @@
1golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4=
2golang.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 @@
1package pipeline
2
3import (
4 "context"
5 "time"
6)
7
8// Batch collects items into slices of up to size, flushing when full or when timeout elapses.
9func Batch[T any](ctx context.Context, in <-chan T, size int, timeout time.Duration) <-chan []T {
10 if in == nil {
11 panic("pipeline: Batch input channel must not be nil")
12 }
13 if size < 1 {
14 panic("pipeline: Batch size must be at least 1")
15 }
16
17 out := make(chan []T)
18
19 go func() {
20 defer close(out)
21
22 buf := make([]T, 0, size)
23 ticker := time.NewTicker(timeout)
24 defer ticker.Stop()
25
26 flush := func() bool {
27 if len(buf) == 0 {
28 return true
29 }
30 batch := buf
31 buf = make([]T, 0, size)
32 select {
33 case out <- batch:
34 return true
35 case <-ctx.Done():
36 return false
37 }
38 }
39
40 for {
41 select {
42 case v, ok := <-in:
43 if !ok {
44 flush()
45 return
46 }
47 buf = append(buf, v)
48 if len(buf) >= size {
49 if !flush() {
50 return
51 }
52 }
53 case <-ticker.C:
54 if !flush() {
55 return
56 }
57 case <-ctx.Done():
58 return
59 }
60 }
61 }()
62
63 return out
64}
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 @@
1package pipeline
2
3import (
4 "context"
5 "testing"
6 "time"
7)
8
9func TestBatch_Full(t *testing.T) {
10 t.Parallel()
11
12 input := make(chan int, 10)
13 for i := 1; i <= 10; i++ {
14 input <- i
15 }
16 close(input)
17
18 out := Batch(t.Context(), input, 3, time.Second)
19
20 var batches [][]int
21 for b := range out {
22 batches = append(batches, b)
23 }
24
25 if len(batches) != 4 {
26 t.Fatalf("expected 4 batches, got %d: %v", len(batches), batches)
27 }
28 if len(batches[0]) != 3 {
29 t.Fatalf("expected first batch size 3, got %d", len(batches[0]))
30 }
31 if len(batches[3]) != 1 {
32 t.Fatalf("expected last batch size 1, got %d", len(batches[3]))
33 }
34}
35
36func TestBatch_TimeoutFlush(t *testing.T) {
37 t.Parallel()
38
39 input := make(chan int)
40
41 out := Batch(t.Context(), input, 100, 50*time.Millisecond)
42
43 input <- 1
44 input <- 2
45
46 select {
47 case batch := <-out:
48 if len(batch) != 2 {
49 t.Fatalf("expected batch of 2 from timeout flush, got %d", len(batch))
50 }
51 case <-time.After(time.Second):
52 t.Fatal("timeout waiting for batch flush")
53 }
54
55 close(input)
56 for range out {
57 }
58}
59
60func TestBatch_ContextCancel(t *testing.T) {
61 t.Parallel()
62 ctx, cancel := context.WithCancel(t.Context())
63
64 input := make(chan int)
65 out := Batch(ctx, input, 100, time.Second)
66
67 cancel()
68
69 done := make(chan struct{})
70 go func() {
71 for range out {
72 }
73 close(done)
74 }()
75
76 select {
77 case <-done:
78 case <-time.After(time.Second):
79 t.Fatal("timeout waiting for shutdown")
80 }
81}
82
83func TestBatch_NilInput(t *testing.T) {
84 t.Parallel()
85
86 defer func() {
87 if r := recover(); r == nil {
88 t.Error("expected panic for nil input")
89 }
90 }()
91
92 Batch[int](t.Context(), nil, 10, time.Second)
93}
94
95func TestBatch_PanicOnZeroSize(t *testing.T) {
96 t.Parallel()
97
98 defer func() {
99 if r := recover(); r == nil {
100 t.Error("expected panic for batch size < 1")
101 }
102 }()
103
104 input := make(chan int)
105 Batch(t.Context(), input, 0, time.Second)
106}
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 @@
1package pipeline_test
2
3import (
4 "context"
5 "fmt"
6 "sort"
7
8 "git.sangtran.io.vn/quacker.git/pipeline"
9)
10
11func ExampleRun() {
12 input := make(chan int, 5)
13 for i := 1; i <= 5; i++ {
14 input <- i
15 }
16 close(input)
17
18 out := pipeline.Run(context.Background(), input,
19 pipeline.Filter(func(v int) bool { return v%2 == 0 }),
20 pipeline.Map(1, func(_ context.Context, v int) (int, bool) {
21 return v * 10, true
22 }),
23 )
24
25 var results []int
26 for v := range out {
27 results = append(results, v)
28 }
29 sort.Ints(results)
30 fmt.Println(results)
31 // Output: [20 40]
32}
33
34func ExampleSink() {
35 input := make(chan string, 3)
36 input <- "hello"
37 input <- "world"
38 input <- "!"
39 close(input)
40
41 var collected []string
42 out := pipeline.Run(context.Background(), input,
43 pipeline.Sink(1, func(_ context.Context, v string) {
44 collected = append(collected, v)
45 }),
46 )
47
48 for range out {
49 }
50
51 fmt.Println(collected)
52 // Output: [hello world !]
53}
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 @@
1package pipeline
2
3// StageObserver receives callbacks for stage lifecycle events.
4type StageObserver interface {
5 OnItem()
6 OnDrop()
7 OnDone()
8}
9
10// StageOption configures a stage's behavior.
11type StageOption func(*stageConfig)
12
13type stageConfig struct {
14 observer StageObserver
15 bufferSize int
16}
17
18// WithObserver attaches an observer to a stage.
19func WithObserver(obs StageObserver) StageOption {
20 return func(cfg *stageConfig) { cfg.observer = obs }
21}
22
23// WithBufferSize sets the output channel buffer size for a stage.
24func WithBufferSize(size int) StageOption {
25 return func(cfg *stageConfig) { cfg.bufferSize = size }
26}
27
28func applyOpts(opts []StageOption) stageConfig {
29 var cfg stageConfig
30 for _, opt := range opts {
31 opt(&cfg)
32 }
33 return cfg
34}
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 @@
1// Package pipeline provides composable, type-safe stream processing with backpressure.
2package pipeline
3
4import "context"
5
6// Stage transforms an input channel into an output channel.
7// The stage owns the output channel and closes it when the input
8// is exhausted or the context is cancelled.
9type Stage[T any] func(ctx context.Context, in <-chan T) <-chan T
10
11// Run connects stages into a pipeline and returns the final output channel.
12func Run[T any](ctx context.Context, input <-chan T, stages ...Stage[T]) <-chan T {
13 if input == nil {
14 panic("pipeline: input channel must not be nil")
15 }
16 ch := input
17 for _, stage := range stages {
18 ch = stage(ctx, ch)
19 }
20 return ch
21}
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 @@
1package pipeline
2
3import (
4 "context"
5 "testing"
6)
7
8func BenchmarkMap(b *testing.B) {
9 b.ReportAllocs()
10
11 input := make(chan int, 256)
12
13 out := Run(b.Context(), input,
14 Map(4, func(_ context.Context, v int) (int, bool) {
15 return v * 2, true
16 }),
17 )
18
19 go func() {
20 for range out {
21 }
22 }()
23
24 for b.Loop() {
25 input <- 1
26 }
27
28 close(input)
29}
30
31func BenchmarkPipeline_Composition(b *testing.B) {
32 b.ReportAllocs()
33
34 input := make(chan int, 256)
35
36 out := Run(b.Context(), input,
37 Filter(func(v int) bool { return v > 0 }),
38 Map(4, func(_ context.Context, v int) (int, bool) {
39 return v * 2, true
40 }),
41 Buffer[int](64),
42 Sink(2, func(_ context.Context, v int) {}),
43 )
44
45 go func() {
46 for range out {
47 }
48 }()
49
50 for b.Loop() {
51 input <- 1
52 }
53
54 close(input)
55}
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 @@
1package pipeline
2
3import (
4 "context"
5 "fmt"
6 "sync/atomic"
7 "testing"
8 "time"
9)
10
11func TestRun_Passthrough(t *testing.T) {
12 t.Parallel()
13
14 input := make(chan int, 3)
15 input <- 1
16 input <- 2
17 input <- 3
18 close(input)
19
20 out := Run(t.Context(), input)
21
22 var got []int
23 for v := range out {
24 got = append(got, v)
25 }
26
27 if len(got) != 3 {
28 t.Fatalf("expected 3 items, got %d", len(got))
29 }
30}
31
32func TestRun_NilInput(t *testing.T) {
33 t.Parallel()
34
35 defer func() {
36 if r := recover(); r == nil {
37 t.Error("expected panic for nil input")
38 }
39 }()
40
41 Run[int](t.Context(), nil)
42}
43
44func TestRun_Composition(t *testing.T) {
45 t.Parallel()
46
47 input := make(chan int, 10)
48 for i := 1; i <= 10; i++ {
49 input <- i
50 }
51 close(input)
52
53 out := Run(t.Context(), input,
54 Filter(func(v int) bool { return v%2 == 0 }),
55 Map(1, func(_ context.Context, v int) (int, bool) {
56 return v * 2, true
57 }),
58 )
59
60 var got []int
61 for v := range out {
62 got = append(got, v)
63 }
64
65 if len(got) != 5 {
66 t.Fatalf("expected 5 items, got %d: %v", len(got), got)
67 }
68
69 sum := 0
70 for _, v := range got {
71 sum += v
72 }
73 if sum != 60 {
74 t.Fatalf("expected sum 60, got %d", sum)
75 }
76}
77
78func TestRun_Backpressure(t *testing.T) {
79 t.Parallel()
80
81 input := make(chan int)
82 consumed := make(chan int, 10)
83
84 out := Run(t.Context(), input,
85 Map(1, func(_ context.Context, v int) (int, bool) {
86 consumed <- v
87 return v, true
88 }),
89 )
90
91 input <- 42
92 v := <-consumed
93 if v != 42 {
94 t.Fatalf("expected 42, got %d", v)
95 }
96
97 select {
98 case input <- 99:
99 t.Fatal("producer was not blocked")
100 case <-time.After(50 * time.Millisecond):
101 }
102
103 go func() {
104 for range out {
105 }
106 }()
107 close(input)
108}
109
110func TestRun_GracefulDrain(t *testing.T) {
111 t.Parallel()
112
113 input := make(chan int, 5)
114 for i := 1; i <= 5; i++ {
115 input <- i
116 }
117 close(input)
118
119 var count atomic.Int32
120 out := Run(t.Context(), input,
121 Sink(2, func(_ context.Context, v int) {
122 count.Add(1)
123 }),
124 )
125
126 for range out {
127 }
128
129 if count.Load() != 5 {
130 t.Fatalf("expected all 5 items drained, got %d", count.Load())
131 }
132}
133
134func TestRun_Cancellation(t *testing.T) {
135 t.Parallel()
136 ctx, cancel := context.WithCancel(t.Context())
137
138 input := make(chan int)
139 out := Run(ctx, input,
140 Map(1, func(_ context.Context, v int) (int, bool) {
141 return v, true
142 }),
143 )
144
145 go func() {
146 select {
147 case input <- 1:
148 case <-ctx.Done():
149 }
150 }()
151
152 cancel()
153
154 done := make(chan struct{})
155 go func() {
156 for range out {
157 }
158 close(done)
159 }()
160
161 select {
162 case <-done:
163 case <-time.After(time.Second):
164 t.Fatal("pipeline did not shut down after cancellation")
165 }
166}
167
168func TestRun_ErrorPropagation(t *testing.T) {
169 t.Parallel()
170 ctx, cancel := context.WithCancelCause(t.Context())
171 defer cancel(nil)
172
173 input := make(chan int, 5)
174 for i := 1; i <= 5; i++ {
175 input <- i
176 }
177 close(input)
178
179 sentinel := fmt.Errorf("bad item")
180
181 out := Run(ctx, input,
182 Map(1, func(_ context.Context, v int) (int, bool) {
183 if v == 3 {
184 cancel(sentinel)
185 return v, false
186 }
187 return v * 2, true
188 }),
189 )
190
191 for range out {
192 }
193
194 if cause := context.Cause(ctx); cause != sentinel {
195 t.Fatalf("expected sentinel error, got: %v", cause)
196 }
197}
198
199func TestObserver(t *testing.T) {
200 t.Parallel()
201
202 var items, drops atomic.Int32
203 doneCh := make(chan struct{})
204
205 obs := &testObserver{
206 onItem: func() { items.Add(1) },
207 onDrop: func() { drops.Add(1) },
208 onDone: func() { close(doneCh) },
209 }
210
211 input := make(chan int, 6)
212 for i := 1; i <= 6; i++ {
213 input <- i
214 }
215 close(input)
216
217 out := Run(t.Context(), input,
218 Filter(func(v int) bool { return v%2 == 0 }, WithObserver(obs)),
219 )
220
221 for range out {
222 }
223 <-doneCh
224
225 if items.Load() != 3 {
226 t.Errorf("expected 3 items, got %d", items.Load())
227 }
228 if drops.Load() != 3 {
229 t.Errorf("expected 3 drops, got %d", drops.Load())
230 }
231}
232
233type testObserver struct {
234 onItem func()
235 onDrop func()
236 onDone func()
237}
238
239func (o *testObserver) OnItem() {
240 if o.onItem != nil {
241 o.onItem()
242 }
243}
244
245func (o *testObserver) OnDrop() {
246 if o.onDrop != nil {
247 o.onDrop()
248 }
249}
250
251func (o *testObserver) OnDone() {
252 if o.onDone != nil {
253 o.onDone()
254 }
255}
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 @@
1package pipeline
2
3import (
4 "context"
5 "sync"
6)
7
8// Map applies fn to each item with bounded concurrency.
9// If keep is false, the item is dropped.
10func Map[T any](concurrency int, fn func(context.Context, T) (T, bool), opts ...StageOption) Stage[T] {
11 if concurrency < 1 {
12 panic("pipeline: Map concurrency must be at least 1")
13 }
14 cfg := applyOpts(opts)
15
16 return func(ctx context.Context, in <-chan T) <-chan T {
17 out := make(chan T, cfg.bufferSize)
18 var wg sync.WaitGroup
19
20 for range concurrency {
21 wg.Go(func() {
22 for {
23 select {
24 case v, ok := <-in:
25 if !ok {
26 return
27 }
28 result, keep := fn(ctx, v)
29 if !keep {
30 if cfg.observer != nil {
31 cfg.observer.OnDrop()
32 }
33 continue
34 }
35 if cfg.observer != nil {
36 cfg.observer.OnItem()
37 }
38 select {
39 case out <- result:
40 case <-ctx.Done():
41 return
42 }
43 case <-ctx.Done():
44 return
45 }
46 }
47 })
48 }
49
50 go func() {
51 wg.Wait()
52 close(out)
53 if cfg.observer != nil {
54 cfg.observer.OnDone()
55 }
56 }()
57
58 return out
59 }
60}
61
62// Filter forwards items for which fn returns true.
63func Filter[T any](fn func(T) bool, opts ...StageOption) Stage[T] {
64 cfg := applyOpts(opts)
65
66 return func(ctx context.Context, in <-chan T) <-chan T {
67 out := make(chan T, cfg.bufferSize)
68
69 go func() {
70 defer close(out)
71 defer func() {
72 if cfg.observer != nil {
73 cfg.observer.OnDone()
74 }
75 }()
76
77 for {
78 select {
79 case v, ok := <-in:
80 if !ok {
81 return
82 }
83 if !fn(v) {
84 if cfg.observer != nil {
85 cfg.observer.OnDrop()
86 }
87 continue
88 }
89 if cfg.observer != nil {
90 cfg.observer.OnItem()
91 }
92 select {
93 case out <- v:
94 case <-ctx.Done():
95 return
96 }
97 case <-ctx.Done():
98 return
99 }
100 }
101 }()
102
103 return out
104 }
105}
106
107// ForEach calls fn on each item and forwards it unchanged.
108func ForEach[T any](fn func(T), opts ...StageOption) Stage[T] {
109 cfg := applyOpts(opts)
110
111 return func(ctx context.Context, in <-chan T) <-chan T {
112 out := make(chan T, cfg.bufferSize)
113
114 go func() {
115 defer close(out)
116 defer func() {
117 if cfg.observer != nil {
118 cfg.observer.OnDone()
119 }
120 }()
121
122 for {
123 select {
124 case v, ok := <-in:
125 if !ok {
126 return
127 }
128 fn(v)
129 if cfg.observer != nil {
130 cfg.observer.OnItem()
131 }
132 select {
133 case out <- v:
134 case <-ctx.Done():
135 return
136 }
137 case <-ctx.Done():
138 return
139 }
140 }
141 }()
142
143 return out
144 }
145}
146
147// Buffer inserts a buffered channel of the specified capacity between stages.
148func Buffer[T any](size int) Stage[T] {
149 if size < 1 {
150 panic("pipeline: Buffer size must be at least 1")
151 }
152
153 return func(ctx context.Context, in <-chan T) <-chan T {
154 out := make(chan T, size)
155
156 go func() {
157 defer close(out)
158 for {
159 select {
160 case v, ok := <-in:
161 if !ok {
162 return
163 }
164 select {
165 case out <- v:
166 case <-ctx.Done():
167 return
168 }
169 case <-ctx.Done():
170 return
171 }
172 }
173 }()
174
175 return out
176 }
177}
178
179// Tee sends each item to a side channel while forwarding it to the next stage.
180func Tee[T any](side chan<- T, opts ...StageOption) Stage[T] {
181 if side == nil {
182 panic("pipeline: Tee side channel must not be nil")
183 }
184 cfg := applyOpts(opts)
185
186 return func(ctx context.Context, in <-chan T) <-chan T {
187 out := make(chan T, cfg.bufferSize)
188
189 go func() {
190 defer close(out)
191 defer func() {
192 if cfg.observer != nil {
193 cfg.observer.OnDone()
194 }
195 }()
196
197 for {
198 select {
199 case v, ok := <-in:
200 if !ok {
201 return
202 }
203 if cfg.observer != nil {
204 cfg.observer.OnItem()
205 }
206 select {
207 case side <- v:
208 case <-ctx.Done():
209 return
210 }
211 select {
212 case out <- v:
213 case <-ctx.Done():
214 return
215 }
216 case <-ctx.Done():
217 return
218 }
219 }
220 }()
221
222 return out
223 }
224}
225
226// Take forwards up to n items, then closes the output channel.
227func Take[T any](n int, opts ...StageOption) Stage[T] {
228 if n < 0 {
229 panic("pipeline: Take count must not be negative")
230 }
231 cfg := applyOpts(opts)
232
233 return func(ctx context.Context, in <-chan T) <-chan T {
234 out := make(chan T, cfg.bufferSize)
235
236 go func() {
237 defer close(out)
238 defer func() {
239 if cfg.observer != nil {
240 cfg.observer.OnDone()
241 }
242 }()
243
244 count := 0
245 for count < n {
246 select {
247 case v, ok := <-in:
248 if !ok {
249 return
250 }
251 if cfg.observer != nil {
252 cfg.observer.OnItem()
253 }
254 select {
255 case out <- v:
256 count++
257 case <-ctx.Done():
258 return
259 }
260 case <-ctx.Done():
261 return
262 }
263 }
264 }()
265
266 return out
267 }
268}
269
270// Sink consumes items with bounded concurrency without forwarding them.
271func Sink[T any](concurrency int, fn func(context.Context, T), opts ...StageOption) Stage[T] {
272 if concurrency < 1 {
273 panic("pipeline: Sink concurrency must be at least 1")
274 }
275 cfg := applyOpts(opts)
276
277 return func(ctx context.Context, in <-chan T) <-chan T {
278 done := make(chan T)
279 var wg sync.WaitGroup
280
281 for range concurrency {
282 wg.Go(func() {
283 for {
284 select {
285 case v, ok := <-in:
286 if !ok {
287 return
288 }
289 fn(ctx, v)
290 if cfg.observer != nil {
291 cfg.observer.OnItem()
292 }
293 case <-ctx.Done():
294 return
295 }
296 }
297 })
298 }
299
300 go func() {
301 wg.Wait()
302 close(done)
303 if cfg.observer != nil {
304 cfg.observer.OnDone()
305 }
306 }()
307
308 return done
309 }
310}
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 @@
1package pipeline
2
3import (
4 "context"
5 "sort"
6 "sync"
7 "sync/atomic"
8 "testing"
9)
10
11func TestMap_Basic(t *testing.T) {
12 t.Parallel()
13
14 input := make(chan int, 3)
15 input <- 1
16 input <- 2
17 input <- 3
18 close(input)
19
20 out := Map(1, func(_ context.Context, v int) (int, bool) {
21 return v * 10, true
22 })(t.Context(), input)
23
24 var got []int
25 for v := range out {
26 got = append(got, v)
27 }
28
29 sort.Ints(got)
30 if len(got) != 3 || got[0] != 10 || got[1] != 20 || got[2] != 30 {
31 t.Fatalf("expected [10 20 30], got %v", got)
32 }
33}
34
35func TestMap_Concurrent(t *testing.T) {
36 t.Parallel()
37
38 input := make(chan int, 100)
39 for i := range 100 {
40 input <- i
41 }
42 close(input)
43
44 out := Map(5, func(_ context.Context, v int) (int, bool) {
45 return v, true
46 })(t.Context(), input)
47
48 var got []int
49 for v := range out {
50 got = append(got, v)
51 }
52
53 if len(got) != 100 {
54 t.Fatalf("expected 100 items, got %d", len(got))
55 }
56
57 sort.Ints(got)
58 for i, v := range got {
59 if v != i {
60 t.Fatalf("missing item %d, got %v at position %d", i, v, i)
61 }
62 }
63}
64
65func TestMap_Drop(t *testing.T) {
66 t.Parallel()
67
68 input := make(chan int, 5)
69 for i := 1; i <= 5; i++ {
70 input <- i
71 }
72 close(input)
73
74 out := Map(1, func(_ context.Context, v int) (int, bool) {
75 return v, v > 3
76 })(t.Context(), input)
77
78 var got []int
79 for v := range out {
80 got = append(got, v)
81 }
82
83 if len(got) != 2 {
84 t.Fatalf("expected 2 items, got %d: %v", len(got), got)
85 }
86}
87
88func TestMap_PanicOnZeroConcurrency(t *testing.T) {
89 t.Parallel()
90
91 defer func() {
92 if r := recover(); r == nil {
93 t.Error("expected panic for concurrency < 1")
94 }
95 }()
96
97 Map(0, func(_ context.Context, v int) (int, bool) { return v, true })
98}
99
100func TestFilter(t *testing.T) {
101 t.Parallel()
102
103 input := make(chan int, 10)
104 for i := 1; i <= 10; i++ {
105 input <- i
106 }
107 close(input)
108
109 out := Filter(func(v int) bool { return v%3 == 0 })(t.Context(), input)
110
111 var got []int
112 for v := range out {
113 got = append(got, v)
114 }
115
116 if len(got) != 3 {
117 t.Fatalf("expected [3 6 9], got %v", got)
118 }
119}
120
121func TestForEach(t *testing.T) {
122 t.Parallel()
123
124 input := make(chan int, 3)
125 input <- 1
126 input <- 2
127 input <- 3
128 close(input)
129
130 var sum atomic.Int32
131 out := ForEach(func(v int) {
132 sum.Add(int32(v))
133 })(t.Context(), input)
134
135 var got []int
136 for v := range out {
137 got = append(got, v)
138 }
139
140 if sum.Load() != 6 {
141 t.Fatalf("expected side effect sum 6, got %d", sum.Load())
142 }
143 if len(got) != 3 {
144 t.Fatalf("expected 3 forwarded items, got %d", len(got))
145 }
146}
147
148func TestBuffer(t *testing.T) {
149 t.Parallel()
150
151 input := make(chan int, 5)
152 for i := 1; i <= 5; i++ {
153 input <- i
154 }
155 close(input)
156
157 out := Buffer[int](10)(t.Context(), input)
158
159 var got []int
160 for v := range out {
161 got = append(got, v)
162 }
163
164 if len(got) != 5 {
165 t.Fatalf("expected 5 items, got %d", len(got))
166 }
167}
168
169func TestBuffer_PanicOnZero(t *testing.T) {
170 t.Parallel()
171
172 defer func() {
173 if r := recover(); r == nil {
174 t.Error("expected panic for buffer size < 1")
175 }
176 }()
177
178 Buffer[int](0)
179}
180
181func TestTee(t *testing.T) {
182 t.Parallel()
183
184 input := make(chan int, 3)
185 input <- 1
186 input <- 2
187 input <- 3
188 close(input)
189
190 side := make(chan int, 10)
191 out := Tee(side)(t.Context(), input)
192
193 var got []int
194 for v := range out {
195 got = append(got, v)
196 }
197 close(side)
198
199 var sideGot []int
200 for v := range side {
201 sideGot = append(sideGot, v)
202 }
203
204 if len(got) != 3 {
205 t.Fatalf("expected 3 forwarded items, got %d", len(got))
206 }
207 if len(sideGot) != 3 {
208 t.Fatalf("expected 3 side items, got %d", len(sideGot))
209 }
210}
211
212func TestTee_NilSide(t *testing.T) {
213 t.Parallel()
214
215 defer func() {
216 if r := recover(); r == nil {
217 t.Error("expected panic for nil side channel")
218 }
219 }()
220
221 Tee[int](nil)
222}
223
224func TestTake(t *testing.T) {
225 t.Parallel()
226 ctx, cancel := context.WithCancel(t.Context())
227 defer cancel()
228
229 input := make(chan int, 10)
230 for i := 1; i <= 10; i++ {
231 input <- i
232 }
233 close(input)
234
235 out := Take[int](3)(ctx, input)
236
237 var got []int
238 for v := range out {
239 got = append(got, v)
240 }
241
242 if len(got) != 3 {
243 t.Fatalf("expected 3 items, got %d: %v", len(got), got)
244 }
245}
246
247func TestTake_Zero(t *testing.T) {
248 t.Parallel()
249
250 input := make(chan int, 5)
251 for i := 1; i <= 5; i++ {
252 input <- i
253 }
254 close(input)
255
256 out := Take[int](0)(t.Context(), input)
257
258 count := 0
259 for range out {
260 count++
261 }
262
263 if count != 0 {
264 t.Fatalf("expected 0 items from Take(0), got %d", count)
265 }
266}
267
268func TestSink(t *testing.T) {
269 t.Parallel()
270
271 input := make(chan int, 5)
272 for i := 1; i <= 5; i++ {
273 input <- i
274 }
275 close(input)
276
277 var mu sync.Mutex
278 var collected []int
279
280 out := Sink(2, func(_ context.Context, v int) {
281 mu.Lock()
282 collected = append(collected, v)
283 mu.Unlock()
284 })(t.Context(), input)
285
286 for range out {
287 t.Fatal("Sink should not forward items")
288 }
289
290 mu.Lock()
291 defer mu.Unlock()
292
293 if len(collected) != 5 {
294 t.Fatalf("expected 5 items consumed, got %d", len(collected))
295 }
296}
297
298func TestSink_PanicOnZeroConcurrency(t *testing.T) {
299 t.Parallel()
300
301 defer func() {
302 if r := recover(); r == nil {
303 t.Error("expected panic for concurrency < 1")
304 }
305 }()
306
307 Sink(0, func(_ context.Context, v int) {})
308}
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 @@
1// Package workerpool provides a generic worker pool implementation for concurrent task processing.
2package workerpool
3
4import (
5 "context"
6 "errors"
7 "fmt"
8 "maps"
9 "sync"
10 "sync/atomic"
11
12 "golang.org/x/sync/errgroup"
13)
14
15// ErrQueueFull is returned by Push when the task queue has no remaining capacity.
16var ErrQueueFull = errors.New("workerpool: task queue is full")
17
18// Task represents a unit of work that can be executed by the worker pool.
19type Task[T any] interface {
20 Execute(ctx context.Context) error
21 GetID() string
22}
23
24// WorkerPoolObserver defines callbacks for monitoring worker pool lifecycle events.
25type WorkerPoolObserver interface {
26 OnWorkerStart(workerID int)
27 OnWorkerStop(workerID int)
28 OnTaskDone(workerID int, taskID string, err error)
29}
30
31// WorkerPoolConfig holds configuration options for creating a WorkerPool instance.
32type WorkerPoolConfig struct {
33 Name string
34 NumWorkers int
35 TaskQueueSize int
36 Observer WorkerPoolObserver
37 ErrorHandler func(workerID int, taskID string, err error)
38 PanicHandler func(workerID int, taskID string, panicValue any)
39}
40
41// Option is a functional option for configuring a WorkerPool instance.
42type Option func(*WorkerPoolConfig)
43
44// WorkerPool manages a fixed-size pool of workers that process tasks concurrently.
45type WorkerPool[T Task[T]] struct {
46 cfg *WorkerPoolConfig
47 taskQueue chan T
48 metrics *metrics
49 eg *errgroup.Group
50 once sync.Once
51 ctx context.Context
52 started atomic.Bool
53}
54
55type metrics struct {
56 mu sync.RWMutex
57 TasksProcessed map[string]int
58 TasksFailed int
59 ActiveWorkers int
60}
61
62// MetricsSnapshot provides a point-in-time view of worker pool metrics.
63type MetricsSnapshot struct {
64 TasksProcessed map[string]int
65 TasksFailed int
66 ActiveWorkers int
67}
68
69// NewWorkerPool creates a new WorkerPool with the provided options.
70func NewWorkerPool[T Task[T]](ctx context.Context, opts ...Option) (*WorkerPool[T], error) {
71 cfg := &WorkerPoolConfig{
72 Name: "workerpool",
73 NumWorkers: 1,
74 TaskQueueSize: 1,
75 }
76
77 for _, opt := range opts {
78 opt(cfg)
79 }
80
81 if cfg.NumWorkers < 1 {
82 return nil, fmt.Errorf("workerpool: number of workers must be at least 1, got %d", cfg.NumWorkers)
83 }
84
85 if cfg.TaskQueueSize < 1 {
86 return nil, fmt.Errorf("workerpool: task queue size must be at least 1, got %d", cfg.TaskQueueSize)
87 }
88
89 eg, ctx := errgroup.WithContext(ctx)
90
91 return &WorkerPool[T]{
92 cfg: cfg,
93 taskQueue: make(chan T, cfg.TaskQueueSize),
94 ctx: ctx,
95 eg: eg,
96 metrics: &metrics{
97 TasksProcessed: make(map[string]int),
98 },
99 }, nil
100}
101
102func WithName(name string) Option {
103 return func(cfg *WorkerPoolConfig) { cfg.Name = name }
104}
105
106func WithNumWorkers(n int) Option {
107 return func(cfg *WorkerPoolConfig) { cfg.NumWorkers = n }
108}
109
110func WithTaskQueueSize(size int) Option {
111 return func(cfg *WorkerPoolConfig) { cfg.TaskQueueSize = size }
112}
113
114func WithObserver(observer WorkerPoolObserver) Option {
115 return func(cfg *WorkerPoolConfig) { cfg.Observer = observer }
116}
117
118func WithErrorHandler(handler func(workerID int, taskID string, err error)) Option {
119 return func(cfg *WorkerPoolConfig) { cfg.ErrorHandler = handler }
120}
121
122func WithPanicHandler(handler func(workerID int, taskID string, panicValue any)) Option {
123 return func(cfg *WorkerPoolConfig) { cfg.PanicHandler = handler }
124}
125
126// Start launches worker goroutines. It panics if called multiple times.
127func (w *WorkerPool[T]) Start() {
128 if !w.started.CompareAndSwap(false, true) {
129 panic("workerpool: Start() called multiple times")
130 }
131
132 for i := 0; i < w.cfg.NumWorkers; i++ {
133 workerID := i
134 workerName := fmt.Sprintf("%s.worker.%d", w.cfg.Name, workerID)
135 w.eg.Go(func() error {
136 w.metrics.mu.Lock()
137 w.metrics.ActiveWorkers++
138 w.metrics.mu.Unlock()
139
140 if w.cfg.Observer != nil {
141 safeObserve(func() { w.cfg.Observer.OnWorkerStart(workerID) })
142 }
143
144 defer func() {
145 w.metrics.mu.Lock()
146 w.metrics.ActiveWorkers--
147 w.metrics.mu.Unlock()
148
149 if w.cfg.Observer != nil {
150 safeObserve(func() { w.cfg.Observer.OnWorkerStop(workerID) })
151 }
152 }()
153
154 return w.worker(workerID, workerName)
155 })
156 }
157}
158
159// Push adds a task to the queue. Returns ErrQueueFull if full, or context error if cancelled.
160func (w *WorkerPool[T]) Push(task T) error {
161 select {
162 case w.taskQueue <- task:
163 return nil
164 case <-w.ctx.Done():
165 return fmt.Errorf("workerpool %q: push: %w", w.cfg.Name, w.ctx.Err())
166 default:
167 return fmt.Errorf("workerpool %q: capacity %d: %w", w.cfg.Name, cap(w.taskQueue), ErrQueueFull)
168 }
169}
170
171// Stop closes the task queue.
172func (w *WorkerPool[T]) Stop() {
173 w.once.Do(func() {
174 close(w.taskQueue)
175 })
176}
177
178// Wait blocks until all workers finish.
179func (w *WorkerPool[T]) Wait() error {
180 return w.eg.Wait()
181}
182
183// StopAndWait closes the queue and waits for all workers to drain.
184func (w *WorkerPool[T]) StopAndWait() error {
185 w.Stop()
186 return w.Wait()
187}
188
189// GetMetrics returns a snapshot of current pool metrics.
190func (w *WorkerPool[T]) GetMetrics() MetricsSnapshot {
191 w.metrics.mu.RLock()
192 defer w.metrics.mu.RUnlock()
193
194 tasksProcessed := make(map[string]int, len(w.metrics.TasksProcessed))
195 maps.Copy(tasksProcessed, w.metrics.TasksProcessed)
196
197 return MetricsSnapshot{
198 ActiveWorkers: w.metrics.ActiveWorkers,
199 TasksFailed: w.metrics.TasksFailed,
200 TasksProcessed: tasksProcessed,
201 }
202}
203
204func (w *WorkerPool[T]) worker(workerID int, workerName string) error {
205 for {
206 select {
207 case task, ok := <-w.taskQueue:
208 if !ok {
209 return nil
210 }
211 w.executeSafe(workerID, workerName, task)
212 case <-w.ctx.Done():
213 return w.ctx.Err()
214 }
215 }
216}
217
218func (w *WorkerPool[T]) executeSafe(workerID int, workerName string, task T) {
219 taskID := task.GetID()
220
221 defer func() {
222 if r := recover(); r != nil {
223 w.recordTaskFailed()
224 if w.cfg.PanicHandler != nil {
225 w.cfg.PanicHandler(workerID, taskID, r)
226 }
227 }
228 }()
229
230 if err := task.Execute(w.ctx); err != nil {
231 w.recordTaskFailed()
232 if w.cfg.Observer != nil {
233 safeObserve(func() { w.cfg.Observer.OnTaskDone(workerID, taskID, err) })
234 }
235 if w.cfg.ErrorHandler != nil {
236 w.cfg.ErrorHandler(workerID, taskID, err)
237 }
238 return
239 }
240
241 w.recordTaskSuccess(workerName)
242 if w.cfg.Observer != nil {
243 safeObserve(func() { w.cfg.Observer.OnTaskDone(workerID, taskID, nil) })
244 }
245}
246
247func (w *WorkerPool[T]) recordTaskSuccess(workerName string) {
248 w.metrics.mu.Lock()
249 defer w.metrics.mu.Unlock()
250 w.metrics.TasksProcessed[workerName]++
251}
252
253func (w *WorkerPool[T]) recordTaskFailed() {
254 w.metrics.mu.Lock()
255 defer w.metrics.mu.Unlock()
256 w.metrics.TasksFailed++
257}
258
259func safeObserve(fn func()) {
260 defer func() { recover() }()
261 fn()
262}
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 @@
1package workerpool
2
3import (
4 "context"
5 "fmt"
6 "testing"
7)
8
9type benchTask struct {
10 id string
11}
12
13func (t *benchTask) Execute(_ context.Context) error { return nil }
14func (t *benchTask) GetID() string { return t.id }
15
16func BenchmarkWorkerPool(b *testing.B) {
17 b.ReportAllocs()
18
19 pool, _ := NewWorkerPool[*benchTask](b.Context(),
20 WithNumWorkers(4),
21 WithTaskQueueSize(256),
22 )
23 pool.Start()
24
25 task := &benchTask{id: "bench"}
26
27 for b.Loop() {
28 pool.Push(task)
29 }
30
31 pool.StopAndWait()
32}
33
34func BenchmarkWorkerPool_Push(b *testing.B) {
35 b.ReportAllocs()
36
37 pool, _ := NewWorkerPool[*benchTask](b.Context(),
38 WithNumWorkers(8),
39 WithTaskQueueSize(1024),
40 )
41 pool.Start()
42 defer pool.StopAndWait()
43
44 b.RunParallel(func(pb *testing.PB) {
45 task := &benchTask{id: fmt.Sprintf("bench-%d", b.N)}
46 for pb.Next() {
47 pool.Push(task)
48 }
49 })
50}
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 @@
1package workerpool
2
3import (
4 "context"
5 "errors"
6 "fmt"
7 "sync"
8 "testing"
9 "time"
10)
11
12type mockTask struct {
13 id string
14 fn func(ctx context.Context) error
15}
16
17func (t *mockTask) Execute(ctx context.Context) error {
18 if t.fn != nil {
19 return t.fn(ctx)
20 }
21 return nil
22}
23
24func (t *mockTask) GetID() string {
25 return t.id
26}
27
28type taskResult struct {
29 taskID string
30 err error
31}
32
33type mockObserver struct {
34 mu sync.Mutex
35 workerStarts []int
36 workerStops []int
37 tasksDone []taskResult
38 t *testing.T
39}
40
41func (o *mockObserver) OnWorkerStart(workerID int) {
42 o.mu.Lock()
43 defer o.mu.Unlock()
44 o.workerStarts = append(o.workerStarts, workerID)
45 if o.t != nil {
46 o.t.Logf("observer: worker %d started", workerID)
47 }
48}
49
50func (o *mockObserver) OnWorkerStop(workerID int) {
51 o.mu.Lock()
52 defer o.mu.Unlock()
53 o.workerStops = append(o.workerStops, workerID)
54 if o.t != nil {
55 o.t.Logf("observer: worker %d stopped", workerID)
56 }
57}
58
59func (o *mockObserver) OnTaskDone(workerID int, taskID string, err error) {
60 o.mu.Lock()
61 defer o.mu.Unlock()
62 o.tasksDone = append(o.tasksDone, taskResult{taskID: taskID, err: err})
63 if o.t != nil {
64 if err != nil {
65 o.t.Logf("observer: worker %d task %s failed: %v", workerID, taskID, err)
66 } else {
67 o.t.Logf("observer: worker %d task %s succeeded", workerID, taskID)
68 }
69 }
70}
71
72func TestNewWorkerPool(t *testing.T) {
73 t.Parallel()
74
75 tests := []struct {
76 name string
77 opts []Option
78 shouldErr bool
79 }{
80 {
81 name: "valid config with defaults",
82 opts: []Option{},
83 shouldErr: false,
84 },
85 {
86 name: "valid config with all options",
87 opts: []Option{
88 WithName("test-pool"),
89 WithNumWorkers(4),
90 WithTaskQueueSize(100),
91 },
92 shouldErr: false,
93 },
94 {
95 name: "invalid number of workers (zero)",
96 opts: []Option{WithNumWorkers(0)},
97 shouldErr: true,
98 },
99 {
100 name: "invalid number of workers (negative)",
101 opts: []Option{WithNumWorkers(-1)},
102 shouldErr: true,
103 },
104 {
105 name: "invalid task queue size (zero)",
106 opts: []Option{WithTaskQueueSize(0)},
107 shouldErr: true,
108 },
109 {
110 name: "invalid task queue size (negative)",
111 opts: []Option{WithTaskQueueSize(-1)},
112 shouldErr: true,
113 },
114 }
115
116 for _, test := range tests {
117 t.Run(test.name, func(t *testing.T) {
118 t.Parallel()
119 _, err := NewWorkerPool[*mockTask](t.Context(), test.opts...)
120 if (err != nil) != test.shouldErr {
121 t.Errorf("NewWorkerPool() error = %v, shouldErr %v", err, test.shouldErr)
122 }
123 })
124 }
125}
126
127func TestWorkerPool_Observer(t *testing.T) {
128 t.Parallel()
129 ctx := t.Context()
130
131 obs := &mockObserver{t: t}
132 numWorkers := 2
133 taskCount := 5
134
135 pool, err := NewWorkerPool[*mockTask](ctx,
136 WithName("observer-test"),
137 WithNumWorkers(numWorkers),
138 WithTaskQueueSize(10),
139 WithObserver(obs),
140 )
141 if err != nil {
142 t.Fatalf("failed to create pool: %v", err)
143 }
144
145 pool.Start()
146
147 for i := range taskCount {
148 task := &mockTask{id: fmt.Sprintf("task-%d", i)}
149 if err := pool.Push(task); err != nil {
150 t.Errorf("failed to push task: %v", err)
151 }
152 }
153
154 pool.StopAndWait()
155
156 obs.mu.Lock()
157 defer obs.mu.Unlock()
158
159 if len(obs.workerStarts) != numWorkers {
160 t.Errorf("expected %d worker starts, got %d", numWorkers, len(obs.workerStarts))
161 }
162 if len(obs.workerStops) != numWorkers {
163 t.Errorf("expected %d worker stops, got %d", numWorkers, len(obs.workerStops))
164 }
165
166 if len(obs.tasksDone) != taskCount {
167 t.Errorf("expected %d tasks done, got %d", taskCount, len(obs.tasksDone))
168 }
169 for _, r := range obs.tasksDone {
170 if r.err != nil {
171 t.Errorf("task %s should have succeeded, got error: %v", r.taskID, r.err)
172 }
173 }
174}
175
176func TestWorkerPool_ObserverTaskFailed(t *testing.T) {
177 t.Parallel()
178 ctx := t.Context()
179
180 obs := &mockObserver{t: t}
181
182 pool, err := NewWorkerPool[*mockTask](ctx,
183 WithName("observer-fail-test"),
184 WithNumWorkers(1),
185 WithTaskQueueSize(10),
186 WithObserver(obs),
187 )
188 if err != nil {
189 t.Fatalf("failed to create pool: %v", err)
190 }
191
192 pool.Start()
193
194 task := &mockTask{
195 id: "fail-task",
196 fn: func(ctx context.Context) error {
197 return fmt.Errorf("intentional error")
198 },
199 }
200 pool.Push(task)
201
202 pool.StopAndWait()
203
204 obs.mu.Lock()
205 defer obs.mu.Unlock()
206
207 failures := 0
208 for _, r := range obs.tasksDone {
209 if r.err != nil {
210 failures++
211 }
212 }
213
214 if failures != 1 {
215 t.Errorf("expected 1 failed task, got %d", failures)
216 }
217 if len(obs.tasksDone) > 0 && obs.tasksDone[0].taskID != "fail-task" {
218 t.Errorf("expected failed task id 'fail-task', got %s", obs.tasksDone[0].taskID)
219 }
220}
221
222func TestWorkerPool_ObserverTaskPanic(t *testing.T) {
223 t.Parallel()
224 ctx := t.Context()
225
226 panicCaught := make(chan string, 1)
227
228 pool, err := NewWorkerPool[*mockTask](ctx,
229 WithName("observer-panic-test"),
230 WithNumWorkers(1),
231 WithTaskQueueSize(10),
232 WithPanicHandler(func(workerID int, taskID string, panicValue any) {
233 panicCaught <- taskID
234 }),
235 )
236 if err != nil {
237 t.Fatalf("failed to create pool: %v", err)
238 }
239
240 pool.Start()
241
242 task := &mockTask{
243 id: "panic-task",
244 fn: func(ctx context.Context) error {
245 panic("intentional panic")
246 },
247 }
248 pool.Push(task)
249
250 pool.StopAndWait()
251
252 select {
253 case taskID := <-panicCaught:
254 if taskID != "panic-task" {
255 t.Errorf("expected panic from 'panic-task', got %s", taskID)
256 }
257 default:
258 t.Error("expected PanicHandler to be called")
259 }
260}
261
262func TestWorkerPool_ConcurrentPush(t *testing.T) {
263 t.Parallel()
264 ctx := t.Context()
265 var wg sync.WaitGroup
266 taskCount := 20
267
268 pool, err := NewWorkerPool[*mockTask](ctx,
269 WithName("concurrent-push"),
270 WithNumWorkers(3),
271 WithTaskQueueSize(taskCount),
272 )
273 if err != nil {
274 t.Fatalf("failed to create pool: %v", err)
275 }
276
277 pool.Start()
278 defer pool.StopAndWait()
279
280 for i := range taskCount {
281 wg.Add(1)
282 go func(id int) {
283 defer wg.Done()
284 task := &mockTask{id: fmt.Sprintf("task-%d", id)}
285 if err := pool.Push(task); err != nil {
286 t.Errorf("push failed: %v", err)
287 }
288 }(i)
289 }
290
291 done := make(chan struct{})
292 go func() {
293 wg.Wait()
294 close(done)
295 }()
296
297 select {
298 case <-done:
299 case <-time.After(5 * time.Second):
300 t.Fatal("timeout waiting for concurrent push")
301 }
302}
303
304func TestWorkerPool_TaskQueueFull(t *testing.T) {
305 t.Parallel()
306 ctx := t.Context()
307
308 pool, err := NewWorkerPool[*mockTask](ctx,
309 WithName("queue-full"),
310 WithNumWorkers(1),
311 WithTaskQueueSize(2),
312 )
313 if err != nil {
314 t.Fatalf("failed to create pool: %v", err)
315 }
316
317 for i := range 2 {
318 task := &mockTask{id: fmt.Sprintf("task-%d", i)}
319 if err := pool.Push(task); err != nil {
320 t.Errorf("unexpected error on push %d: %v", i, err)
321 }
322 }
323
324 task := &mockTask{id: "overflow-task"}
325 err = pool.Push(task)
326 if err == nil {
327 t.Fatal("expected error when queue is full, got nil")
328 }
329 if !errors.Is(err, ErrQueueFull) {
330 t.Fatalf("expected ErrQueueFull, got: %v", err)
331 }
332}
333
334func TestWorkerPool_MultipleStop(t *testing.T) {
335 t.Parallel()
336 ctx := t.Context()
337
338 pool, err := NewWorkerPool[*mockTask](ctx,
339 WithName("multiple-stop"),
340 WithNumWorkers(2),
341 WithTaskQueueSize(10),
342 )
343 if err != nil {
344 t.Fatalf("failed to create pool: %v", err)
345 }
346
347 pool.Start()
348 pool.Stop()
349
350 func() {
351 defer func() {
352 if r := recover(); r != nil {
353 t.Errorf("Stop() panicked on second call: %v", r)
354 }
355 }()
356 pool.Stop()
357 }()
358
359 pool.Wait()
360}
361
362func TestWorkerPool_MetricsThreadSafe(t *testing.T) {
363 t.Parallel()
364 ctx := t.Context()
365 var wg sync.WaitGroup
366 taskCount := 20
367
368 pool, err := NewWorkerPool[*mockTask](ctx,
369 WithName("metrics-safe"),
370 WithNumWorkers(10),
371 WithTaskQueueSize(taskCount),
372 )
373 if err != nil {
374 t.Fatalf("failed to create pool: %v", err)
375 }
376
377 pool.Start()
378 defer pool.StopAndWait()
379
380 for i := range taskCount {
381 wg.Add(1)
382 go func(id int) {
383 defer wg.Done()
384 task := &mockTask{
385 id: fmt.Sprintf("task-%d", id),
386 fn: func(ctx context.Context) error {
387 time.Sleep(time.Millisecond)
388 return nil
389 },
390 }
391 pool.Push(task)
392 }(i)
393 }
394
395 concurrentReads := 100
396 for range concurrentReads {
397 wg.Add(1)
398 go func() {
399 defer wg.Done()
400 for range 10 {
401 pool.GetMetrics()
402 }
403 }()
404 }
405
406 wg.Wait()
407}
408
409func TestWorkerPool_ContextCancellation(t *testing.T) {
410 t.Parallel()
411 ctx, cancel := context.WithCancel(t.Context())
412
413 pool, err := NewWorkerPool[*mockTask](ctx,
414 WithName("context-cancel"),
415 WithNumWorkers(2),
416 WithTaskQueueSize(10),
417 )
418 if err != nil {
419 t.Fatalf("failed to create pool: %v", err)
420 }
421
422 pool.Start()
423
424 for i := range 5 {
425 task := &mockTask{id: fmt.Sprintf("task-%d", i)}
426 pool.Push(task)
427 }
428
429 cancel()
430
431 err = pool.Wait()
432 if err != nil && !errors.Is(err, context.Canceled) {
433 t.Errorf("expected context.Canceled, got: %v", err)
434 }
435}
436
437func TestWorkerPool_ErrorHandler(t *testing.T) {
438 t.Parallel()
439 ctx := t.Context()
440
441 errorHandlerCalled := make(chan struct{})
442
443 pool, err := NewWorkerPool[*mockTask](ctx,
444 WithName("error-handler-test"),
445 WithNumWorkers(1),
446 WithTaskQueueSize(10),
447 WithErrorHandler(func(workerID int, taskID string, err error) {
448 t.Logf("ErrorHandler called: worker=%d, task=%s, err=%v", workerID, taskID, err)
449 close(errorHandlerCalled)
450 }),
451 )
452 if err != nil {
453 t.Fatalf("failed to create pool: %v", err)
454 }
455
456 pool.Start()
457
458 task := &mockTask{
459 id: "error-task",
460 fn: func(ctx context.Context) error {
461 return fmt.Errorf("task error")
462 },
463 }
464 pool.Push(task)
465
466 select {
467 case <-errorHandlerCalled:
468 case <-time.After(time.Second):
469 t.Fatal("ErrorHandler was not called within timeout")
470 }
471
472 pool.StopAndWait()
473}
474
475func TestWorkerPool_PanicHandler(t *testing.T) {
476 t.Parallel()
477 ctx := t.Context()
478
479 panicHandlerCalled := make(chan struct{})
480
481 pool, err := NewWorkerPool[*mockTask](ctx,
482 WithName("panic-handler-test"),
483 WithNumWorkers(1),
484 WithTaskQueueSize(10),
485 WithPanicHandler(func(workerID int, taskID string, panicValue any) {
486 t.Logf("PanicHandler called: worker=%d, task=%s, panic=%v", workerID, taskID, panicValue)
487 close(panicHandlerCalled)
488 }),
489 )
490 if err != nil {
491 t.Fatalf("failed to create pool: %v", err)
492 }
493
494 pool.Start()
495
496 task := &mockTask{
497 id: "panic-task",
498 fn: func(ctx context.Context) error {
499 panic("task panic")
500 },
501 }
502 pool.Push(task)
503
504 select {
505 case <-panicHandlerCalled:
506 case <-time.After(time.Second):
507 t.Fatal("PanicHandler was not called within timeout")
508 }
509
510 pool.StopAndWait()
511}
512
513func TestWorkerPool_PushContextCancelled(t *testing.T) {
514 t.Parallel()
515 ctx, cancel := context.WithCancel(t.Context())
516
517 pool, err := NewWorkerPool[*mockTask](ctx,
518 WithName("push-ctx-cancelled"),
519 WithNumWorkers(1),
520 WithTaskQueueSize(1),
521 )
522 if err != nil {
523 t.Fatalf("failed to create pool: %v", err)
524 }
525
526 task1 := &mockTask{id: "task-1"}
527 if err := pool.Push(task1); err != nil {
528 t.Fatalf("failed to fill queue: %v", err)
529 }
530
531 cancel()
532
533 task2 := &mockTask{id: "task-2"}
534 err = pool.Push(task2)
535
536 if err == nil {
537 t.Fatal("expected error when pushing to full queue with cancelled context")
538 }
539
540 if !errors.Is(err, ErrQueueFull) && !errors.Is(err, context.Canceled) {
541 t.Fatalf("expected ErrQueueFull or context.Canceled, got: %v", err)
542 }
543}
544
545func TestWorkerPool_NoObserver(t *testing.T) {
546 t.Parallel()
547 ctx := t.Context()
548
549 pool, err := NewWorkerPool[*mockTask](ctx,
550 WithName("no-observer"),
551 WithNumWorkers(2),
552 WithTaskQueueSize(10),
553 )
554 if err != nil {
555 t.Fatalf("failed to create pool: %v", err)
556 }
557
558 pool.Start()
559
560 task1 := &mockTask{id: "success-task"}
561 pool.Push(task1)
562
563 task2 := &mockTask{
564 id: "fail-task",
565 fn: func(ctx context.Context) error {
566 return fmt.Errorf("error")
567 },
568 }
569 pool.Push(task2)
570
571 task3 := &mockTask{
572 id: "panic-task",
573 fn: func(ctx context.Context) error {
574 panic("panic")
575 },
576 }
577 pool.Push(task3)
578
579 pool.StopAndWait()
580
581 metrics := pool.GetMetrics()
582 if metrics.TasksFailed != 2 {
583 t.Errorf("expected 2 failed tasks, got %d", metrics.TasksFailed)
584 }
585}
586
587func TestWorkerPool_StartMultipleTimes(t *testing.T) {
588 t.Parallel()
589
590 pool, err := NewWorkerPool[*mockTask](t.Context(),
591 WithNumWorkers(1),
592 WithTaskQueueSize(1),
593 )
594 if err != nil {
595 t.Fatalf("failed to create pool: %v", err)
596 }
597
598 pool.Start()
599 defer pool.StopAndWait()
600
601 defer func() {
602 if r := recover(); r == nil {
603 t.Error("expected panic when calling Start() multiple times")
604 }
605 }()
606
607 pool.Start()
608}
generated by cgit v1.3.1 (git 2.54.0) at 2026-10-17 18:43:26 +0000