make Bridge not block in the first stream that not closed (#288)

* not block in the first channel

* make Bridge not block in the first stream that not closed

* Bridge with test
This commit is contained in:
Tuuuuuuuuu
2025-01-14 16:19:45 +08:00
committed by GitHub
parent a0d97cf38e
commit e27df00fa8
3 changed files with 31 additions and 19 deletions

View File

@@ -157,10 +157,10 @@ func (c *Channel[T]) Tee(ctx context.Context, in <-chan T) (<-chan T, <-chan T)
// Play: https://go.dev/play/p/qmWSy1NVF-Y
func (c *Channel[T]) Bridge(ctx context.Context, chanStream <-chan <-chan T) <-chan T {
valStream := make(chan T)
go func() {
defer close(valStream)
wg := sync.WaitGroup{}
defer wg.Wait()
for {
var stream <-chan T
select {
@@ -169,19 +169,22 @@ func (c *Channel[T]) Bridge(ctx context.Context, chanStream <-chan <-chan T) <-c
return
}
stream = maybeStream
wg.Add(1)
case <-ctx.Done():
return
}
for val := range c.OrDone(ctx, stream) {
select {
case valStream <- val:
case <-ctx.Done():
go func() {
defer wg.Done()
for val := range c.OrDone(ctx, stream) {
select {
case valStream <- val:
case <-ctx.Done():
}
}
}
}()
}
}()
return valStream
}

View File

@@ -168,7 +168,8 @@ func ExampleChannel_Tee() {
func ExampleChannel_Bridge() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
m1 := make(map[int]int)
m2 := make(map[int]int)
c := NewChannel[int]()
genVals := func() <-chan <-chan int {
out := make(chan (<-chan int))
@@ -177,6 +178,7 @@ func ExampleChannel_Bridge() {
for i := 1; i <= 5; i++ {
stream := make(chan int, 1)
stream <- i
m1[i]++
close(stream)
out <- stream
}
@@ -185,12 +187,15 @@ func ExampleChannel_Bridge() {
}
for v := range c.Bridge(ctx, genVals()) {
fmt.Println(v)
m2[v]++
}
for k, v := range m1 {
fmt.Println(m2[k] == v)
}
// Output:
// 1
// 2
// 3
// 4
// 5
// true
// true
// true
// true
// true
}

View File

@@ -169,7 +169,8 @@ func TestTee(t *testing.T) {
func TestBridge(t *testing.T) {
t.Parallel()
assert := internal.NewAssert(t, "TestBridge")
m1 := make(map[int]int)
m2 := make(map[int]int)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
@@ -181,6 +182,7 @@ func TestBridge(t *testing.T) {
for i := 0; i < 10; i++ {
stream := make(chan int, 1)
stream <- i
m1[i]++
close(stream)
chanStream <- stream
}
@@ -188,9 +190,11 @@ func TestBridge(t *testing.T) {
return chanStream
}
index := 0
for val := range c.Bridge(ctx, genVals()) {
assert.Equal(index, val)
index++
m2[val]++
}
for k, v := range m1 {
assert.Equal(m2[k], v)
}
}