diff --git a/nursery.go b/nursery.go index f9cce93..635134a 100644 --- a/nursery.go +++ b/nursery.go @@ -29,6 +29,7 @@ type nursery struct { limiter limiter goRoutine chan Routine routinesCount atomic.Int32 + bufSize int } func newNursery() *nursery { @@ -39,6 +40,7 @@ func newNursery() *nursery { errors: make(chan error), limiter: nil, goRoutine: make(chan func() error), + bufSize: 0, } return n diff --git a/options.go b/options.go index 0f5cf8b..cdb3b04 100644 --- a/options.go +++ b/options.go @@ -88,3 +88,16 @@ func WithMaxGoroutines(max int) BlockOption { n.limiter = make(chan struct{}, max+1) } } + +// WithOutputBuffer returns a nursery block option that sets the output buffer +// size for AsyncMap operations. This option is only used by AsyncMap and is +// ignored by other functions. If size is zero or negative, an unbuffered +// channel is used (default behavior). +func WithOutputBuffer(size int) BlockOption { + return func(n *nursery) { + if size < 0 { + size = 0 + } + n.bufSize = size + } +} diff --git a/utils.go b/utils.go index 4aabc71..3331b85 100644 --- a/utils.go +++ b/utils.go @@ -153,3 +153,108 @@ func Map2[K comparable, V any](input map[K]V, f func(context.Context, K, V) (K, return nil }, opts...) } + +// AsyncMap concurrently transforms an input sequence into an output sequence. +// The order of the output sequence is non-deterministic. Items are delivered as they finish. +// +// Within the parent nursery, a goroutine is launched which creates a nested nursery. +// The sequence is processed entirely within this nested nursery. By default, the parent +// nursery is used as the context for the nested nursery. +// +// If a transform function returns an error, that error is propagated to the parent nursery +// and the nested nursery's context is canceled, stopping processing of all remaining items. +// +// By default, AsyncMap uses an unbuffered channel for outputs. For high-throughput scenarios +// where producers might temporarily outpace consumers, use WithOutputBuffer to specify +// a buffer size. +// +// If consuming from the output sequence stops before it's fully drained (e.g., +// breaking out of the range loop early), transformation goroutines may block indefinitely. +// To prevent this, cancel the parent nursery's context when done consuming. For multi-stage +// pipelines requiring finer control, provide a separate cancellable context via WithContext. +// +// Example: +// +// _ = Block(func(n Nursery) error { +// inputs := slices.Values([]int{1, 2, 3, 4, 5}) +// outputs := conc.AsyncMap(n, inputs, func(_ context.Context, i int) (string, error) { +// return fmt.Sprintf("item-%d", i), nil +// }) +// +// for output := range outputs { +// fmt.Println(output) +// } +// }) +// +// Note: If supplying a WithContext option, it will override the default parent context. +// Ensure any custom context derives from the parent nursery context to maintain proper +// cancellation propagation. +func AsyncMap[I any, O any]( + parent Nursery, + inputs iter.Seq[I], + transform func(context.Context, I) (O, error), + opts ...BlockOption, +) iter.Seq[O] { + // Create a temporary nursery to get correct output buffer size + tn := newNursery() + for _, opt := range opts { + opt(tn) + } + var outputs chan O + if tn.bufSize > 0 { + outputs = make(chan O, tn.bufSize) + } else { + outputs = make(chan O) + } + allOpts := append([]BlockOption{WithContext(parent)}, opts...) + + // Core processing loop + process := func(nested Nursery) error { + done := nested.Done() + for input := range inputs { + select { + case <-done: + return nil + default: + input := input + nested.Go(func() error { + output, err := transform(nested, input) + if err != nil { + return err + } + select { + case outputs <- output: + return nil + case <-done: + return nested.Err() + } + }) + } + } + return nil + } + + // Launch + parent.Go(func() error { + defer func() { + close(outputs) + }() + if err := Block(process, allOpts...); err != nil { + return err + } + return nil + }) + + return chanSeq(outputs) +} + +// chanSeq returns an iterator that yields the values of ch until it is closed. +func chanSeq[T any](ch <-chan T) iter.Seq[T] { + return func(yield func(T) bool) { + for v := range ch { + if !yield(v) { + return + } + } + } +} diff --git a/utils_test.go b/utils_test.go index 6b168b6..4e2318f 100644 --- a/utils_test.go +++ b/utils_test.go @@ -3,9 +3,12 @@ package conc import ( "context" "errors" + "fmt" "iter" "reflect" + "slices" "sync" + "sync/atomic" "testing" "time" ) @@ -371,3 +374,575 @@ func TestMap2(t *testing.T) { } }) } + +func TestAsyncMap(t *testing.T) { + t.Run("Basic", func(t *testing.T) { + inputs := []int{70, 10, 33, 20, 50, 80} + var collected []string + + err := Block(func(n Nursery) error { + outputs := AsyncMap( + n, + slices.Values(inputs), + func(ctx context.Context, i int) (string, error) { + output := fmt.Sprintf("output %02d", i) + Sleep(ctx, time.Millisecond*time.Duration(i)) + return output, nil + }, + ) + for output := range outputs { + collected = append(collected, output) + } + + return nil + }) + if err != nil { + t.Error(err) + } + if len(collected) != len(inputs) { + t.Errorf("expected to collect an output for each input") + } + sortedOutputs := slices.Clone(collected) + slices.Sort(sortedOutputs) + if !slices.Equal(collected, sortedOutputs) { + t.Errorf("expected output slice to be in sorted order") + } + }) + + t.Run("Multistage", func(t *testing.T) { + err := Block(func(n Nursery) error { + // stage A: ints to strings + inputs := []int{7, 5, 3, 1, 2, 4, 6} + outputsA := AsyncMap( + n, + slices.Values(inputs), + func(ctx context.Context, i int) (string, error) { + output := fmt.Sprintf("output %02d", i) + Sleep(ctx, time.Millisecond*time.Duration(i)) + return output, nil + }, + ) + + // stage B: strings to ints + outputsB := AsyncMap( + n, + outputsA, + func(_ context.Context, s string) (int, error) { + return len(s), nil + }, + ) + + for _ = range outputsB { + } + + return nil + }) + if err != nil { + t.Error(err) + } + }) + + t.Run("ErrorHandling", func(t *testing.T) { + inputs := []int{1, 2, 3, 4, 5} + expectedError := errors.New("test error") + var collected []string + var processedCount atomic.Int32 + + err := Block(func(n Nursery) error { + outputs := AsyncMap( + n, + slices.Values(inputs), + func(_ context.Context, i int) (string, error) { + processedCount.Add(1) + time.Sleep(1 * time.Millisecond) + if i == 3 { + return "", expectedError + } + return fmt.Sprintf("value %d", i), nil + }, + ) + + for output := range outputs { + collected = append(collected, output) + } + return nil + }) + + if err == nil { + t.Error("expected an error, got nil") + } + if !errors.Is(err, expectedError) { + t.Errorf("expected error %v, got %v", expectedError, err) + } + + if processedCount.Load() == 0 { + t.Error("expected at least some inputs to be processed") + } + + for _, output := range collected { + if output == "value 3" { + t.Error("output from error-producing input should not be collected") + } + } + }) + + t.Run("EmptyInput", func(t *testing.T) { + var inputs []int + var collected []string + var transformCalled atomic.Bool + + err := Block(func(n Nursery) error { + outputs := AsyncMap( + n, + slices.Values(inputs), + func(_ context.Context, i int) (string, error) { + transformCalled.Store(true) + return fmt.Sprintf("value %d", i), nil + }, + ) + for output := range outputs { + collected = append(collected, output) + } + return nil + }) + + if err != nil { + t.Error(err) + } + if transformCalled.Load() { + t.Error("transform function should not be called for empty input") + } + if len(collected) != 0 { + t.Errorf("expected empty output for empty input, got %d items", len(collected)) + } + }) + + t.Run("Cancellation", func(t *testing.T) { + inputs := []int{1, 2, 3, 4, 5} + var processedCount atomic.Int32 + var collected []string + var contextCancelled atomic.Bool + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + err := Block(func(n Nursery) error { + outputs := AsyncMap( + n, + slices.Values(inputs), + func(ctx context.Context, i int) (string, error) { + processedCount.Add(1) + // Cancel after processing a few items + if i == 3 { + cancel() + // Give a little time for cancellation to propagate + time.Sleep(1 * time.Millisecond) + } + + // Check if context is cancelled + select { + case <-ctx.Done(): + contextCancelled.Store(true) + return "", ctx.Err() + default: + time.Sleep(1 * time.Millisecond) + return fmt.Sprintf("value %d", i), nil + } + }, + ) + + for output := range outputs { + collected = append(collected, output) + } + return nil + }, WithContext(ctx)) + + // The error might not be propagated if the main goroutine completes + // before the error is processed, so we don't strictly check for an error + if err != nil { + t.Logf("Got expected error: %v", err) + } + + if !contextCancelled.Load() { + t.Error("expected context to be cancelled") + } + + // Note: We can't reliably assert that not all inputs were processed + // after cancellation. Depending on timing, all inputs might be processed + // before cancellation fully propagates. This is expected behavior. + t.Logf("Processed %d/%d inputs after cancellation", processedCount.Load(), len(inputs)) + }) + + t.Run("LargeInput", func(t *testing.T) { + // Create a large input slice + const inputSize = 50 // Reduced from 100 to make test faster + inputs := make([]int, inputSize) + for i := 0; i < inputSize; i++ { + inputs[i] = i + } + + var processedCount atomic.Int32 + var collected []int + + err := Block(func(n Nursery) error { + outputs := AsyncMap( + n, + slices.Values(inputs), + func(_ context.Context, i int) (int, error) { + processedCount.Add(1) + return i * 2, nil + }, + ) + for output := range outputs { + collected = append(collected, output) + } + return nil + }) + + if err != nil { + t.Error(err) + } + if processedCount.Load() != inputSize { + t.Errorf("expected all %d inputs to be processed, got %d", inputSize, processedCount.Load()) + } + if len(collected) != inputSize { + t.Errorf("expected to collect %d outputs, got %d", inputSize, len(collected)) + } + }) + + t.Run("MixedProcessingTimes", func(t *testing.T) { + inputs := []int{10, 1, 5, 2, 8, 3} + var processOrder []int + var outputOrder []int + var mu sync.Mutex + + err := Block(func(n Nursery) error { + outputs := AsyncMap( + n, + slices.Values(inputs), + func(ctx context.Context, i int) (int, error) { + // Process in order of input + mu.Lock() + processOrder = append(processOrder, i) + mu.Unlock() + + // Sleep for different durations to ensure different completion times + Sleep(ctx, time.Millisecond*time.Duration(i)) + return i, nil + }, + ) + for output := range outputs { + outputOrder = append(outputOrder, output) + } + return nil + }) + + if err != nil { + t.Error(err) + } + + if len(outputOrder) != len(inputs) { + t.Errorf("expected %d outputs, got %d", len(inputs), len(outputOrder)) + } + + // Verify outputs are not in the same order as inputs (should be ordered by completion time) + // This is a probabilistic test, but with the sleep times we've chosen, it should almost always pass + if slices.Equal(processOrder, outputOrder) { + t.Errorf("expected output order to differ from processing order due to different processing times") + } + }) + + t.Run("DrainAfterPartialConsumption", func(t *testing.T) { + // This test demonstrates that if you stop consuming from the output channel early, + // you must drain the remaining outputs to allow all goroutines to complete + inputs := []int{1, 2, 3, 4, 5} + var processedCount atomic.Int32 + var consumedCount atomic.Int32 + + err := Block(func(n Nursery) error { + outputs := AsyncMap( + n, + slices.Values(inputs), + func(_ context.Context, i int) (int, error) { + processedCount.Add(1) + // Simulate work + time.Sleep(1 * time.Millisecond) + return i, nil + }, + ) + + // Only consume first 3 outputs + count := 0 + for output := range outputs { + consumedCount.Add(1) + count++ + if count >= 3 { + // We've consumed enough outputs, but we need to drain the channel + // to allow all goroutines to complete + go func() { + // Drain remaining outputs in a separate goroutine + for range outputs { + // Just drain, don't count these + } + }() + break + } + _ = output + } + + return nil + }) + + if err != nil { + t.Error(err) + } + + if processedCount.Load() != int32(len(inputs)) { + t.Errorf("expected all %d inputs to be processed, got %d", len(inputs), processedCount.Load()) + } + + if consumedCount.Load() != 3 { + t.Errorf("expected to consume exactly 3 outputs, got %d", consumedCount.Load()) + } + }) + + t.Run("SafeEarlyTerminationWithDefer", func(t *testing.T) { + // This test demonstrates the recommended pattern for safely handling early termination + // using a defer to ensure the channel is always drained + inputs := []int{1, 2, 3, 4, 5} + var processedCount atomic.Int32 + var consumedCount atomic.Int32 + + err := Block(func(n Nursery) error { + outputs := AsyncMap( + n, + slices.Values(inputs), + func(_ context.Context, i int) (int, error) { + processedCount.Add(1) + // Simulate work + time.Sleep(1 * time.Millisecond) + return i, nil + }, + ) + + // Set up a deferred drain to handle early returns or breaks + drainStarted := false + defer func() { + if !drainStarted { + n.Go(func() error { + // Drain any remaining outputs + for range outputs { + // Discard remaining outputs + } + return nil + }) + } + }() + + // Process outputs until some condition + for output := range outputs { + consumedCount.Add(1) + _ = output + + // Simulate early termination after consuming 3 outputs + if consumedCount.Load() >= 3 { + drainStarted = true + go func() { + // Drain remaining outputs + for range outputs { + // Discard remaining outputs + } + }() + break + } + } + + return nil + }) + + if err != nil { + t.Error(err) + } + + if processedCount.Load() != int32(len(inputs)) { + t.Errorf("expected all %d inputs to be processed, got %d", len(inputs), processedCount.Load()) + } + + if consumedCount.Load() != 3 { + t.Errorf("expected to consume exactly 3 outputs, got %d", consumedCount.Load()) + } + }) + + t.Run("ContextCancellationUnblocksProducers", func(t *testing.T) { + // This test verifies that cancelling the context unblocks goroutines + // that are trying to send to the outputs channel + inputs := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10} + var processedCount atomic.Int32 + var blockedCount atomic.Int32 + var unblockedCount atomic.Int32 + var readyToCancel atomic.Bool + var cancelDone atomic.Bool + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + // Use a longer timeout to ensure the test has enough time to complete + ctx, timeoutCancel := context.WithTimeout(ctx, 2*time.Second) + defer timeoutCancel() // Ensure the timeout context is properly cancelled + + err := Block(func(n Nursery) error { + // Create a buffered channel to control when goroutines are blocked + // Buffer size 3 means the first 3 items will be processed without blocking + controlCh := make(chan struct{}, 3) + + outputs := AsyncMap( + n, + slices.Values(inputs), + func(ctx context.Context, i int) (int, error) { + processedCount.Add(1) + + // Simulate slow processing for later items + if i > 3 { + // Signal that we're about to block + blockedCount.Add(1) + + // Wait until we're ready to proceed or context is cancelled + select { + case controlCh <- struct{}{}: + // This will block once the buffer is full + readyToCancel.Store(true) + + // Wait for cancellation + <-ctx.Done() + unblockedCount.Add(1) + return i, n.Err() + case <-ctx.Done(): + // Already cancelled + unblockedCount.Add(1) + return i, n.Err() + } + } + return i, nil + }, + ) + + // Only consume first 3 outputs, then cancel the context + count := 0 + for output := range outputs { + count++ + if count >= 3 { + // Wait until at least one goroutine is blocked + for !readyToCancel.Load() { + time.Sleep(10 * time.Millisecond) + // If we've waited too long, break to avoid hanging + if count := blockedCount.Load(); count > 0 { + break + } + } + + // Cancel the context to unblock any goroutines trying to send + cancel() + cancelDone.Store(true) + break + } + _ = output + } + + // Wait for cancellation to propagate + for i := 0; i < 10 && !cancelDone.Load(); i++ { + time.Sleep(10 * time.Millisecond) + } + + return nil + }, WithContext(ctx)) + + if err == nil { + t.Error("expected an error due to context cancellation") + } + + if blockedCount.Load() == 0 { + t.Error("expected some goroutines to be blocked") + } + + if unblockedCount.Load() == 0 { + t.Error("expected blocked goroutines to be unblocked by context cancellation") + } + + if processedCount.Load() >= int32(len(inputs)) { + t.Logf("All inputs were processed despite cancellation, which is unexpected but not an error") + } + }) + + t.Run("WithOutputBuffer", func(t *testing.T) { + inputs := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10} + var producerUnblocked atomic.Bool + var processedCount atomic.Int32 + + bufferSize := 5 + err := Block(func(n Nursery) error { + // Create a slow consumer scenario + outputs := AsyncMap( + n, + slices.Values(inputs), + func(ctx context.Context, i int) (int, error) { + // Mark this item as processed + processedCount.Add(1) + + // All items should be processed without blocking, + // even though we haven't started consuming yet + if processedCount.Load() > int32(bufferSize) { + producerUnblocked.Store(true) + } + + return i, nil + }, + WithOutputBuffer(bufferSize), + ) + + // Wait a bit to allow producers to process items + time.Sleep(50 * time.Millisecond) + + // Then consume all outputs + var collected []int + for output := range outputs { + collected = append(collected, output) + } + + if len(collected) != len(inputs) { + t.Errorf("expected %d outputs, got %d", len(inputs), len(collected)) + } + + return nil + }) + + if err != nil { + t.Error(err) + } + + // All inputs should be processed + if processedCount.Load() != int32(len(inputs)) { + t.Errorf("expected all %d inputs to be processed, got %d", len(inputs), processedCount.Load()) + } + + // Verify that producers were able to make progress beyond the buffer size + if !producerUnblocked.Load() { + t.Error("expected producers to be unblocked by the buffered channel") + } + }) + + t.Run("ZeroBufferEqualsUnbuffered", func(t *testing.T) { + tempNursery := newNursery() + WithOutputBuffer(0)(tempNursery) + if tempNursery.bufSize != 0 { + t.Errorf("expected buffer size 0, got %d", tempNursery.bufSize) + } + }) + + t.Run("NegativeBufferEqualsUnbuffered", func(t *testing.T) { + tempNursery := newNursery() + WithOutputBuffer(-5)(tempNursery) + if tempNursery.bufSize != 0 { + t.Errorf("expected buffer size 0, got %d", tempNursery.bufSize) + } + }) +}