refactor(flow): rework parallel processing of stream elements (#178)

Author: reugn

flow/filter.go

@@ -2,6 +2,7 @@ package flow
 
 import (
 	"fmt"
+	"sync"
 
 	"github.com/reugn/go-streams"
 )
@@ -31,20 +32,27 @@ var _ streams.Flow = (*Filter[any])(nil)
 // NewFilter returns a new Filter operator.
 // T specifies the incoming and the outgoing element type.
 //
-// filterPredicate is the boolean-valued filter function.
-// parallelism is the flow parallelism factor. In case the events order matters, use parallelism = 1.
-// If the parallelism argument is not positive, NewFilter will panic.
+// filterPredicate is a function that accepts an element of type T and returns true
+// if the element should be included in the output stream, and false if it should be
+// filtered out.
+// parallelism specifies the number of goroutines to use for parallel processing. If
+// the order of elements in the output stream must be preserved, set parallelism to 1.
+//
+// NewFilter will panic if parallelism is less than 1.
 func NewFilter[T any](filterPredicate FilterPredicate[T], parallelism int) *Filter[T] {
 	if parallelism < 1 {
 		panic(fmt.Sprintf("nonpositive Filter parallelism: %d", parallelism))
 	}
+
 	filter := &Filter[T]{
 		filterPredicate: filterPredicate,
 		in:              make(chan any),
 		out:             make(chan any),
 		parallelism:     parallelism,
 	}
-	go filter.doStream()
+
+	// start processing stream elements
+	go filter.stream()
 
 	return filter
 }
@@ -79,20 +87,26 @@ func (f *Filter[T]) transmit(inlet streams.Inlet) {
 	close(inlet.In())
 }
 
-// doStream discards items that don't match the filter predicate.
-func (f *Filter[T]) doStream() {
-	sem := make(chan struct{}, f.parallelism)
-	for elem := range f.in {
-		sem <- struct{}{}
-		go func(element T) {
-			defer func() { <-sem }()
-			if f.filterPredicate(element) {
-				f.out <- element
-			}
-		}(elem.(T))
-	}
+// stream reads elements from the input channel, filters them using the
+// filterPredicate, and sends the filtered elements to the output channel.
+// It uses a pool of goroutines to process elements in parallel.
+func (f *Filter[T]) stream() {
+	var wg sync.WaitGroup
+	// create a pool of worker goroutines
 	for i := 0; i < f.parallelism; i++ {
-		sem <- struct{}{}
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for element := range f.in {
+				if f.filterPredicate(element.(T)) {
+					f.out <- element
+				}
+			}
+		}()
 	}
+
+	// wait for worker goroutines to finish processing inbound elements
+	wg.Wait()
+	// close the output channel
 	close(f.out)
 }

flow/flat_map.go

@@ -2,6 +2,7 @@ package flow
 
 import (
 	"fmt"
+	"sync"
 
 	"github.com/reugn/go-streams"
 )
@@ -30,19 +31,24 @@ var _ streams.Flow = (*FlatMap[any, any])(nil)
 // T specifies the incoming element type, and the outgoing element type is []R.
 //
 // flatMapFunction is the FlatMap transformation function.
-// parallelism is the flow parallelism factor. In case the events order matters, use parallelism = 1.
-// If the parallelism argument is not positive, NewFlatMap will panic.
+// parallelism specifies the number of goroutines to use for parallel processing. If
+// the order of elements in the output stream must be preserved, set parallelism to 1.
+//
+// NewFlatMap will panic if parallelism is less than 1.
 func NewFlatMap[T, R any](flatMapFunction FlatMapFunction[T, R], parallelism int) *FlatMap[T, R] {
 	if parallelism < 1 {
 		panic(fmt.Sprintf("nonpositive FlatMap parallelism: %d", parallelism))
 	}
+
 	flatMap := &FlatMap[T, R]{
 		flatMapFunction: flatMapFunction,
 		in:              make(chan any),
 		out:             make(chan any),
 		parallelism:     parallelism,
 	}
-	go flatMap.doStream()
+
+	// start processing stream elements
+	go flatMap.stream()
 
 	return flatMap
 }
@@ -77,20 +83,27 @@ func (fm *FlatMap[T, R]) transmit(inlet streams.Inlet) {
 	close(inlet.In())
 }
 
-func (fm *FlatMap[T, R]) doStream() {
-	sem := make(chan struct{}, fm.parallelism)
-	for elem := range fm.in {
-		sem <- struct{}{}
-		go func(element T) {
-			defer func() { <-sem }()
-			result := fm.flatMapFunction(element)
-			for _, item := range result {
-				fm.out <- item
-			}
-		}(elem.(T))
-	}
+// stream reads elements from the input channel, applies the flatMapFunction
+// to each element, and sends the resulting elements to the output channel.
+// It uses a pool of goroutines to process elements in parallel.
+func (fm *FlatMap[T, R]) stream() {
+	var wg sync.WaitGroup
+	// create a pool of worker goroutines
 	for i := 0; i < fm.parallelism; i++ {
-		sem <- struct{}{}
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for element := range fm.in {
+				result := fm.flatMapFunction(element.(T))
+				for _, item := range result {
+					fm.out <- item
+				}
+			}
+		}()
 	}
+
+	// wait for worker goroutines to finish processing inbound elements
+	wg.Wait()
+	// close the output channel
 	close(fm.out)
 }

flow/fold.go

@@ -38,7 +38,9 @@ func NewFold[T, R any](init R, foldFunction FoldFunction[T, R]) *Fold[T, R] {
 		in:           make(chan any),
 		out:          make(chan any),
 	}
-	go foldFlow.doStream()
+
+	// start processing stream elements
+	go foldFlow.stream()
 
 	return foldFlow
 }
@@ -73,7 +75,12 @@ func (m *Fold[T, R]) transmit(inlet streams.Inlet) {
 	close(inlet.In())
 }
 
-func (m *Fold[T, R]) doStream() {
+// stream consumes elements from the input channel, applies the foldFunction to
+// each element along with the previously accumulated value, and emits the updated
+// value into the output channel. All input elements are assumed to be of type T.
+// The processing is done sequentially, ensuring that the order of accumulation is
+// maintained.
+func (m *Fold[T, R]) stream() {
 	lastFolded := m.init
 	for element := range m.in {
 		lastFolded = m.foldFunction(element.(T), lastFolded)

flow/map.go

@@ -2,6 +2,7 @@ package flow
 
 import (
 	"fmt"
+	"sync"
 
 	"github.com/reugn/go-streams"
 )
@@ -30,19 +31,24 @@ var _ streams.Flow = (*Map[any, any])(nil)
 // T specifies the incoming element type, and the outgoing element type is R.
 //
 // mapFunction is the Map transformation function.
-// parallelism is the flow parallelism factor. In case the events order matters, use parallelism = 1.
-// If the parallelism argument is not positive, NewMap will panic.
+// parallelism specifies the number of goroutines to use for parallel processing. If
+// the order of elements in the output stream must be preserved, set parallelism to 1.
+//
+// NewMap will panic if parallelism is less than 1.
 func NewMap[T, R any](mapFunction MapFunction[T, R], parallelism int) *Map[T, R] {
 	if parallelism < 1 {
 		panic(fmt.Sprintf("nonpositive Map parallelism: %d", parallelism))
 	}
+
 	mapFlow := &Map[T, R]{
 		mapFunction: mapFunction,
 		in:          make(chan any),
 		out:         make(chan any),
 		parallelism: parallelism,
 	}
-	go mapFlow.doStream()
+
+	// start processing stream elements
+	go mapFlow.stream()
 
 	return mapFlow
 }
@@ -77,18 +83,24 @@ func (m *Map[T, R]) transmit(inlet streams.Inlet) {
 	close(inlet.In())
 }
 
-func (m *Map[T, R]) doStream() {
-	sem := make(chan struct{}, m.parallelism)
-	for elem := range m.in {
-		sem <- struct{}{}
-		go func(element T) {
-			defer func() { <-sem }()
-			result := m.mapFunction(element)
-			m.out <- result
-		}(elem.(T))
-	}
+// stream reads elements from the input channel, applies the mapFunction
+// to each element, and sends the resulting element to the output channel.
+// It uses a pool of goroutines to process elements in parallel.
+func (m *Map[T, R]) stream() {
+	var wg sync.WaitGroup
+	// create a pool of worker goroutines
 	for i := 0; i < m.parallelism; i++ {
-		sem <- struct{}{}
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for element := range m.in {
+				m.out <- m.mapFunction(element.(T))
+			}
+		}()
 	}
+
+	// wait for worker goroutines to finish processing inbound elements
+	wg.Wait()
+	// close the output channel
 	close(m.out)
 }

flow/pass_through.go

@@ -23,7 +23,9 @@ func NewPassThrough() *PassThrough {
 		in:  make(chan any),
 		out: make(chan any),
 	}
-	go passThrough.doStream()
+
+	// start processing stream elements
+	go passThrough.stream()
 
 	return passThrough
 }
@@ -58,7 +60,7 @@ func (pt *PassThrough) transmit(inlet streams.Inlet) {
 	close(inlet.In())
 }
 
-func (pt *PassThrough) doStream() {
+func (pt *PassThrough) stream() {
 	for element := range pt.in {
 		pt.out <- element
 	}

flow/reduce.go

@@ -34,7 +34,9 @@ func NewReduce[T any](reduceFunction ReduceFunction[T]) *Reduce[T] {
 		in:             make(chan any),
 		out:            make(chan any),
 	}
-	go reduce.doStream()
+
+	// start processing stream elements
+	go reduce.stream()
 
 	return reduce
 }
@@ -69,7 +71,13 @@ func (r *Reduce[T]) transmit(inlet streams.Inlet) {
 	close(inlet.In())
 }
 
-func (r *Reduce[T]) doStream() {
+// stream consumes elements from the input channel, applies the reduceFunction to
+// each element along with the previously reduced value, and emits the updated
+// value into the output channel. The first element received becomes the initial
+// reduced value. Subsequent elements are combined with the accumulated result.
+// All input elements are assumed to be of type T. The processing is done sequentially,
+// ensuring that the order of accumulation is maintained.
+func (r *Reduce[T]) stream() {
 	var lastReduced any
 	for element := range r.in {
 		if lastReduced == nil {