diff options
Diffstat (limited to 'pipeline')
| -rw-r--r-- | pipeline/batch.go | 64 | ||||
| -rw-r--r-- | pipeline/batch_test.go | 106 | ||||
| -rw-r--r-- | pipeline/example_test.go | 53 | ||||
| -rw-r--r-- | pipeline/options.go | 34 | ||||
| -rw-r--r-- | pipeline/pipeline.go | 21 | ||||
| -rw-r--r-- | pipeline/pipeline_bench_test.go | 55 | ||||
| -rw-r--r-- | pipeline/pipeline_test.go | 255 | ||||
| -rw-r--r-- | pipeline/stage.go | 310 | ||||
| -rw-r--r-- | pipeline/stage_test.go | 308 |
9 files changed, 1206 insertions, 0 deletions
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 @@ | |||
| 1 | package pipeline | ||
| 2 | |||
| 3 | import ( | ||
| 4 | "context" | ||
| 5 | "time" | ||
| 6 | ) | ||
| 7 | |||
| 8 | // Batch collects items into slices of up to size, flushing when full or when timeout elapses. | ||
| 9 | func 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 @@ | |||
| 1 | package pipeline | ||
| 2 | |||
| 3 | import ( | ||
| 4 | "context" | ||
| 5 | "testing" | ||
| 6 | "time" | ||
| 7 | ) | ||
| 8 | |||
| 9 | func 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 | |||
| 36 | func 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 | |||
| 60 | func 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 | |||
| 83 | func 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 | |||
| 95 | func 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 @@ | |||
| 1 | package pipeline_test | ||
| 2 | |||
| 3 | import ( | ||
| 4 | "context" | ||
| 5 | "fmt" | ||
| 6 | "sort" | ||
| 7 | |||
| 8 | "git.sangtran.io.vn/quacker.git/pipeline" | ||
| 9 | ) | ||
| 10 | |||
| 11 | func 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 | |||
| 34 | func 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 @@ | |||
| 1 | package pipeline | ||
| 2 | |||
| 3 | // StageObserver receives callbacks for stage lifecycle events. | ||
| 4 | type StageObserver interface { | ||
| 5 | OnItem() | ||
| 6 | OnDrop() | ||
| 7 | OnDone() | ||
| 8 | } | ||
| 9 | |||
| 10 | // StageOption configures a stage's behavior. | ||
| 11 | type StageOption func(*stageConfig) | ||
| 12 | |||
| 13 | type stageConfig struct { | ||
| 14 | observer StageObserver | ||
| 15 | bufferSize int | ||
| 16 | } | ||
| 17 | |||
| 18 | // WithObserver attaches an observer to a stage. | ||
| 19 | func 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. | ||
| 24 | func WithBufferSize(size int) StageOption { | ||
| 25 | return func(cfg *stageConfig) { cfg.bufferSize = size } | ||
| 26 | } | ||
| 27 | |||
| 28 | func 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. | ||
| 2 | package pipeline | ||
| 3 | |||
| 4 | import "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. | ||
| 9 | type 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. | ||
| 12 | func 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 @@ | |||
| 1 | package pipeline | ||
| 2 | |||
| 3 | import ( | ||
| 4 | "context" | ||
| 5 | "testing" | ||
| 6 | ) | ||
| 7 | |||
| 8 | func 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 | |||
| 31 | func 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 @@ | |||
| 1 | package pipeline | ||
| 2 | |||
| 3 | import ( | ||
| 4 | "context" | ||
| 5 | "fmt" | ||
| 6 | "sync/atomic" | ||
| 7 | "testing" | ||
| 8 | "time" | ||
| 9 | ) | ||
| 10 | |||
| 11 | func 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 | |||
| 32 | func 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 | |||
| 44 | func 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 | |||
| 78 | func 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 | |||
| 110 | func 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 | |||
| 134 | func 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 | |||
| 168 | func 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 | |||
| 199 | func 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 | |||
| 233 | type testObserver struct { | ||
| 234 | onItem func() | ||
| 235 | onDrop func() | ||
| 236 | onDone func() | ||
| 237 | } | ||
| 238 | |||
| 239 | func (o *testObserver) OnItem() { | ||
| 240 | if o.onItem != nil { | ||
| 241 | o.onItem() | ||
| 242 | } | ||
| 243 | } | ||
| 244 | |||
| 245 | func (o *testObserver) OnDrop() { | ||
| 246 | if o.onDrop != nil { | ||
| 247 | o.onDrop() | ||
| 248 | } | ||
| 249 | } | ||
| 250 | |||
| 251 | func (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 @@ | |||
| 1 | package pipeline | ||
| 2 | |||
| 3 | import ( | ||
| 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. | ||
| 10 | func 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. | ||
| 63 | func 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. | ||
| 108 | func 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. | ||
| 148 | func 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. | ||
| 180 | func 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. | ||
| 227 | func 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. | ||
| 271 | func 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 @@ | |||
| 1 | package pipeline | ||
| 2 | |||
| 3 | import ( | ||
| 4 | "context" | ||
| 5 | "sort" | ||
| 6 | "sync" | ||
| 7 | "sync/atomic" | ||
| 8 | "testing" | ||
| 9 | ) | ||
| 10 | |||
| 11 | func 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 | |||
| 35 | func 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 | |||
| 65 | func 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 | |||
| 88 | func 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 | |||
| 100 | func 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 | |||
| 121 | func 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 | |||
| 148 | func 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 | |||
| 169 | func 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 | |||
| 181 | func 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 | |||
| 212 | func 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 | |||
| 224 | func 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 | |||
| 247 | func 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 | |||
| 268 | func 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 | |||
| 298 | func 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 | } | ||
