Correcting a blocking issue in Fan Out

Author: benjivesterby

stream.go

@@ -13,6 +13,8 @@ package stream
 import (
 	"context"
 	"reflect"
+
+	. "go.structs.dev/gen"
 )
 
 // Pipe accepts an incoming data channel and pipes it to the supplied
@@ -131,25 +133,45 @@ func FanOut[T any](ctx context.Context, in <-chan T, out ...chan<- T) {
 				return
 			}
 
-			for _, o := range out {
-				// Closure to catch panic on closed channel write.
-				// Continue Loop
-				func() {
-					select {
-					case <-ctx.Done():
-						return
-					case o <- v:
-					}
-				}()
+			// Closure to catch panic on closed channel write.
+			selectCases := make([]reflect.SelectCase, 0, len(out)+1)
+
+			// 0 index is context
+			selectCases = append(selectCases, reflect.SelectCase{
+				Dir:  reflect.SelectRecv,
+				Chan: reflect.ValueOf(ctx.Done()),
+			})
+
+			for _, outc := range out {
+				// Skip nil channels until they are non-nil
+				if outc == nil {
+					continue
+				}
+
+				selectCases = append(selectCases, reflect.SelectCase{
+					Dir:  reflect.SelectSend,
+					Chan: reflect.ValueOf(outc),
+					Send: reflect.ValueOf(v),
+				})
+			}
+
+			for len(selectCases) > 1 {
+				chosen, _, _ := reflect.Select(selectCases)
+
+				// The context was cancelled.
+				if chosen == 0 {
+					return
+				}
+
+				selectCases = Exclude(selectCases, selectCases[chosen])
 			}
 		}
 
 	}
 }
 
 // Distribute accepts an incoming data channel and distributes the data among
-// the supplied outgoing data channels. This distribution is done stochastically
-// using the cryptographic random number generator.
+// the supplied outgoing data channels using a dynamic select statement.
 //
 // NOTE: Execute the Distribute function in a goroutine if parallel execution is
 // desired. Cancelling the context or closing the incoming channel is important
@@ -170,24 +192,19 @@ func Distribute[T any](ctx context.Context, in <-chan T, out ...chan<- T) {
 				return
 			}
 
-			// Closure to catch panic on closed channel write.
-			func() {
-				defer recover()
-
-				selectCases := make([]reflect.SelectCase, 0, len(out)+1)
-				for _, outc := range out {
-					selectCases = append(selectCases, reflect.SelectCase{
-						Dir:  reflect.SelectSend,
-						Chan: reflect.ValueOf(outc),
-						Send: reflect.ValueOf(v),
-					})
-				}
+			selectCases := make([]reflect.SelectCase, 0, len(out)+1)
+			for _, outc := range out {
 				selectCases = append(selectCases, reflect.SelectCase{
-					Dir:  reflect.SelectRecv,
-					Chan: reflect.ValueOf(ctx.Done()),
+					Dir:  reflect.SelectSend,
+					Chan: reflect.ValueOf(outc),
+					Send: reflect.ValueOf(v),
 				})
-				_, _, _ = reflect.Select(selectCases)
-			}()
+			}
+			selectCases = append(selectCases, reflect.SelectCase{
+				Dir:  reflect.SelectRecv,
+				Chan: reflect.ValueOf(ctx.Done()),
+			})
+			_, _, _ = reflect.Select(selectCases)
 		}
 
 	}

stream_test.go

@@ -443,6 +443,66 @@ func Test_Distribute_ZeroOut(t *testing.T) {
 	Distribute(ctx, in)
 }
 
+func Test_FanOut(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	c1, c2, c3 := make(chan int), make(chan int), make(chan int)
+	var c4 chan int
+	data := Ints[int](1000)
+
+	go FanOut(ctx, Slice[int](data).Chan(ctx), c1, c2, c3, c4)
+
+	seen := make(map[int]int)
+	for i := 0; i < len(data)*3; i++ {
+		select {
+		case <-ctx.Done():
+			t.Fatal("context cancelled")
+			return
+		case _, ok := <-c1:
+			if !ok {
+				return
+			}
+
+			seen[1]++
+		case _, ok := <-c2:
+			if !ok {
+				return
+			}
+
+			seen[2]++
+		case _, ok := <-c3:
+			if !ok {
+				return
+			}
+
+			seen[3]++
+		case _, ok := <-c4:
+			if !ok {
+				return
+			}
+
+			seen[4]++
+		}
+	}
+
+	if len(seen) != 3 {
+		t.Fatalf("expected %v, got %v", len(data)-1, len(seen))
+	}
+
+	for k, v := range seen {
+		if k == 4 {
+			if v > 0 {
+				t.Fatalf("expected %v, got %v", 0, v)
+			}
+		}
+
+		if v != len(data) {
+			t.Fatalf("Chan C%v: expected %v, got %v", k, len(data), v)
+		}
+	}
+}
+
 func Test_FanOut_ZeroOut(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()