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 --- pipeline/stage.go | 310 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 310 insertions(+) create mode 100644 pipeline/stage.go (limited to 'pipeline/stage.go') 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 + } +} -- cgit v1.3.1