feature(streamwork): make cancellation more predictable

Author: hectorj-thetreep

streamwork/circuitbreaker.go

@@ -102,6 +102,9 @@ func CircuitBreakerWorker[T any](cb CircuitBreaker) Worker[T, T] {
 		go func() {
 			defer close(chanOut)
 			for {
+				if ctx.Err() != nil {
+					return
+				}
 				select {
 				case <-ctx.Done():
 					return
@@ -110,6 +113,9 @@ func CircuitBreakerWorker[T any](cb CircuitBreaker) Worker[T, T] {
 						return
 					}
 					if cb.IsPaused() {
+						if ctx.Err() != nil {
+							return
+						}
 						select {
 						case <-ctx.Done():
 							return
@@ -118,12 +124,18 @@ func CircuitBreakerWorker[T any](cb CircuitBreaker) Worker[T, T] {
 					}
 				vLoop:
 					for {
+						if ctx.Err() != nil {
+							return
+						}
 						select {
 						case <-ctx.Done():
 							return
 						case chanOut <- v:
 							break vLoop
 						case <-cb.pauseChan():
+							if ctx.Err() != nil {
+								return
+							}
 							select {
 							case <-ctx.Done():
 								return

streamwork/filter.go

@@ -9,6 +9,9 @@ func Filter[TIN any](f func(ctx context.Context, v TIN) bool) Worker[TIN, TIN] {
 		go func() {
 			defer close(chanOut)
 			for {
+				if ctx.Err() != nil {
+					return
+				}
 				select {
 				case <-ctx.Done():
 					return
@@ -19,6 +22,9 @@ func Filter[TIN any](f func(ctx context.Context, v TIN) bool) Worker[TIN, TIN] {
 					if !f(ctx, v) {
 						continue
 					}
+					if ctx.Err() != nil {
+						return
+					}
 					select {
 					case <-ctx.Done():
 						return

streamwork/flatten.go

@@ -11,6 +11,9 @@ func Flatten[T any]() Worker[[]T, T] {
 		go func() {
 			defer close(chanOut)
 			for {
+				if ctx.Err() != nil {
+					return
+				}
 				select {
 				case <-ctx.Done():
 					return
@@ -19,6 +22,9 @@ func Flatten[T any]() Worker[[]T, T] {
 						return
 					}
 					for _, v := range batch {
+						if ctx.Err() != nil {
+							return
+						}
 						select {
 						case <-ctx.Done():
 							return

streamwork/parallel.go

@@ -30,13 +30,19 @@ func Parallelize[TIN any, TOUT any](instancesCount int, worker Worker[TIN, TOUT]
 
 func pipeChannels[T any](ctx context.Context, src <-chan T, dest chan T) {
 	for {
+		if ctx.Err() != nil {
+			return
+		}
 		select {
 		case <-ctx.Done():
 			return
 		case v, ok := <-src:
 			if !ok {
 				return
 			}
+			if ctx.Err() != nil {
+				return
+			}
 			select {
 			case <-ctx.Done():
 				return

streamwork/sources.go

@@ -14,6 +14,9 @@ func ReadSlice[T any](s []T) Source[T] {
 		go func() {
 			defer close(c)
 			for _, v := range s {
+				if ctx.Err() != nil {
+					return
+				}
 				select {
 				case <-ctx.Done():
 					return
@@ -48,6 +51,9 @@ func ReadSeq[T any](s iter.Seq[T], options ...StreamOption) Source[T] {
 		go func() {
 			defer close(c)
 			for v := range s {
+				if ctx.Err() != nil {
+					return
+				}
 				select {
 				case <-ctx.Done():
 					return
@@ -76,6 +82,9 @@ func ReadSeqErr[T any](s iter.Seq2[T, error], options ...StreamOption) Source[T]
 					}
 					continue
 				}
+				if ctx.Err() != nil {
+					return
+				}
 				select {
 				case <-ctx.Done():
 					return

streamwork/worker.go

@@ -15,13 +15,19 @@ func WorkerFunc[TIN any, TOUT any](f func(ctx context.Context, v TIN) TOUT) Work
 		go func() {
 			defer close(chanOut)
 			for {
+				if ctx.Err() != nil {
+					return
+				}
 				select {
 				case <-ctx.Done():
 					return
 				case v, ok := <-source:
 					if !ok {
 						return
 					}
+					if ctx.Err() != nil {
+						return
+					}
 					select {
 					case <-ctx.Done():
 						return
@@ -43,6 +49,9 @@ func WorkerSeq[TIN any, TOUT any](f func(ctx context.Context, v iter.Seq[TIN]) i
 		seqOut := f(
 			ctx, func(yield func(TIN) bool) {
 				for {
+					if ctx.Err() != nil {
+						return
+					}
 					select {
 					case <-ctx.Done():
 						return
@@ -60,6 +69,9 @@ func WorkerSeq[TIN any, TOUT any](f func(ctx context.Context, v iter.Seq[TIN]) i
 		go func() {
 			defer close(chanOut)
 			for vOut := range seqOut {
+				if ctx.Err() != nil {
+					return
+				}
 				select {
 				case <-ctx.Done():
 					return
@@ -80,6 +92,9 @@ func WorkerSeqErr[TIN any, TOUT any](f func(ctx context.Context, v iter.Seq[TIN]
 		seqOut := f(
 			ctx, func(yield func(TIN) bool) {
 				for {
+					if ctx.Err() != nil {
+						return
+					}
 					select {
 					case <-ctx.Done():
 						return
@@ -103,6 +118,9 @@ func WorkerSeqErr[TIN any, TOUT any](f func(ctx context.Context, v iter.Seq[TIN]
 					}
 					continue
 				}
+				if ctx.Err() != nil {
+					return
+				}
 				select {
 				case <-ctx.Done():
 					return
@@ -123,6 +141,9 @@ func WorkerFuncErr[TIN any, TOUT any](f func(ctx context.Context, v TIN) (TOUT,
 		go func() {
 			defer close(chanOut)
 			for {
+				if ctx.Err() != nil {
+					return
+				}
 				select {
 				case <-ctx.Done():
 					return
@@ -137,6 +158,9 @@ func WorkerFuncErr[TIN any, TOUT any](f func(ctx context.Context, v TIN) (TOUT,
 						}
 						continue
 					}
+					if ctx.Err() != nil {
+						return
+					}
 					select {
 					case <-ctx.Done():
 						return