aboutsummaryrefslogtreecommitdiff
path: root/pipeline/stage.go
diff options
context:
space:
mode:
Diffstat (limited to 'pipeline/stage.go')
-rw-r--r--pipeline/stage.go310
1 files changed, 310 insertions, 0 deletions
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}
generated by cgit v1.3.1 (git 2.54.0) at 2026-10-17 19:36:58 +0000