-
Notifications
You must be signed in to change notification settings - Fork 17
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
simplify fanin code: drop tomb package and don't use reflection
- Loading branch information
1 parent
5c9956d
commit c13ef2e
Showing
8 changed files
with
78 additions
and
159 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,83 +1,52 @@ | ||
| package rabtap | ||
|
|
||
| import ( | ||
| "reflect" | ||
|
|
||
| tomb "gopkg.in/tomb.v2" | ||
| "context" | ||
| "sync" | ||
| ) | ||
|
|
||
| // Fanin allows to do a select ("fan-in") on an set of channels | ||
| type Fanin struct { | ||
| Ch chan interface{} | ||
| channels []reflect.SelectCase | ||
| t tomb.Tomb | ||
| } | ||
|
|
||
| // NewFanin creates a new Fanin object | ||
| func NewFanin(channels []interface{}) *Fanin { | ||
| fanin := Fanin{Ch: make(chan interface{})} | ||
| fanin.add(fanin.t.Dying()) | ||
| for _, c := range channels { | ||
| fanin.add(c) | ||
| } | ||
|
|
||
| fanin.t.Go(fanin.loop) | ||
| return &fanin | ||
| } | ||
|
|
||
| func (s *Fanin) add(c interface{}) { | ||
| s.channels = append(s.channels, | ||
| reflect.SelectCase{ | ||
| Dir: reflect.SelectRecv, | ||
| Chan: reflect.ValueOf(c)}) | ||
| } | ||
|
|
||
| // Stop stops the fanin go-routine | ||
| func (s *Fanin) Stop() error { | ||
| s.t.Kill(nil) | ||
| return s.t.Wait() | ||
| } | ||
|
|
||
| // Alive returns true if the fanin is running | ||
| func (s *Fanin) Alive() bool { | ||
| return s.t.Alive() | ||
| } | ||
|
|
||
| // Select wait for activity on any of the channels | ||
| func (s *Fanin) loop() error { | ||
|
|
||
| for { | ||
| chosen, message, ok := reflect.Select(s.channels) | ||
|
|
||
| // channels[0] is always the tomb Dying() chan. Request to end fanin. | ||
| if chosen == 0 { | ||
| close(s.Ch) // note: sends nil message on s.Ch | ||
| return nil | ||
| // WrapChan takes a channel of any type and returns a new channel of type interface{} | ||
| func WrapChan[T any](c <-chan T) <-chan interface{} { | ||
| wrapped := make(chan interface{}) | ||
| go func() { | ||
| for m := range c { | ||
| wrapped <- m | ||
| } | ||
| close(wrapped) | ||
| }() | ||
| return wrapped | ||
| } | ||
|
|
||
| if !ok { | ||
| // The chosen channel has been closed, so zero | ||
| // out the channel to disable the case (happens on normal shutdown) | ||
| s.channels[chosen].Chan = reflect.ValueOf(nil) | ||
| // Fanin selects sumultanously from an array of channels and sends received | ||
| // messages to a new channel ("fan-in" of channels) | ||
| func Fanin(ctx context.Context, channels []<-chan interface{}) chan interface{} { | ||
| var wg sync.WaitGroup | ||
| out := make(chan interface{}) | ||
|
|
||
| // if no channels are left, end the fanin. | ||
| active := false | ||
| for i := 1; i < len(s.channels); i++ { | ||
| if s.channels[i].Chan != reflect.ValueOf(nil) { | ||
| active = true | ||
| break | ||
| } | ||
| } | ||
| if !active { | ||
| close(s.Ch) | ||
| return nil | ||
| } | ||
| } else { | ||
| // allow a blocking write to s.Ch to be terminated | ||
| receiver := func(ctx context.Context, c <-chan interface{}) { | ||
| defer wg.Done() | ||
| for { | ||
| select { | ||
| case s.Ch <- message.Interface(): | ||
| case <-s.channels[0].Chan.Interface().(<-chan struct{}): | ||
| case <-ctx.Done(): | ||
| return | ||
| case val, ok := <-c: | ||
| if !ok { | ||
| return | ||
| } | ||
| out <- val | ||
| } | ||
| } | ||
| } | ||
|
|
||
| wg.Add(len(channels)) | ||
| for _, c := range channels { | ||
| go receiver(ctx, c) | ||
| } | ||
|
|
||
| go func() { | ||
| wg.Wait() | ||
| close(out) | ||
| }() | ||
|
|
||
| return out | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,75 +1,47 @@ | ||
| package rabtap | ||
|
|
||
| import ( | ||
| "context" | ||
| "testing" | ||
| "time" | ||
|
|
||
| "github.com/stretchr/testify/assert" | ||
| ) | ||
|
|
||
| func expectIntOnChan(t *testing.T, val int, ch <-chan interface{}) { | ||
| func expectOnChan[T any](t *testing.T, val T, ch <-chan interface{}) { | ||
| t.Helper() | ||
| select { | ||
| case message := <-ch: | ||
| assert.Equal(t, val, message.(int)) | ||
| case <-time.After(3 * time.Second): | ||
| assert.Fail(t, "did not receive message in expected time") | ||
| } | ||
| } | ||
|
|
||
| func expectNilOnChan(t *testing.T, ch <-chan interface{}) { | ||
| select { | ||
| case message := <-ch: | ||
| assert.Nil(t, message) | ||
| case <-time.After(3 * time.Second): | ||
| assert.Equal(t, val, message.(T)) | ||
| case <-time.After(1 * time.Second): | ||
| assert.Fail(t, "did not receive message in expected time") | ||
| } | ||
| } | ||
|
|
||
| func TestFaninReceivesFromMultipleChannels(t *testing.T) { | ||
|
|
||
| // create fanin of 3 int channels | ||
| chan1 := make(chan int) | ||
| chan2 := make(chan int) | ||
| chan3 := make(chan int) | ||
| fanin := NewFanin([]interface{}{chan1, chan2, chan3}) | ||
|
|
||
| assert.True(t, fanin.Alive()) | ||
|
|
||
| go func() { | ||
| chan1 <- 99 | ||
| chan2 <- 100 | ||
| chan3 <- 101 | ||
| }() | ||
| chan1 := make(chan int, 1) | ||
| defer close(chan1) | ||
| chan2 := make(chan string, 1) | ||
| defer close(chan2) | ||
| fanin := Fanin(context.TODO(), []<-chan interface{}{WrapChan(chan1), WrapChan(chan2)}) | ||
|
|
||
| expectIntOnChan(t, 99, fanin.Ch) | ||
| expectIntOnChan(t, 100, fanin.Ch) | ||
| expectIntOnChan(t, 101, fanin.Ch) | ||
|
|
||
| // fanin.Stop() closes fanin channel which in turn sends nil message | ||
| assert.Nil(t, fanin.Stop()) | ||
| expectNilOnChan(t, fanin.Ch) | ||
|
|
||
| assert.False(t, fanin.Alive()) | ||
| chan1 <- 99 | ||
| expectOnChan(t, 99, fanin) | ||
| chan2 <- "hello" | ||
| expectOnChan(t, "hello", fanin) | ||
| } | ||
|
|
||
| func TestFaninClosesChanWhenAllInputsAreClosed(t *testing.T) { | ||
|
|
||
| chan1 := make(chan int) | ||
| chan2 := make(chan int) | ||
| fanin := NewFanin([]interface{}{chan1, chan2}) | ||
|
|
||
| go func() { | ||
| chan1 <- 99 | ||
| chan2 <- 100 | ||
| close(chan1) | ||
| close(chan2) | ||
| }() | ||
| fanin := Fanin(context.TODO(), []<-chan interface{}{WrapChan(chan1), WrapChan(chan2)}) | ||
|
|
||
| expectIntOnChan(t, 99, fanin.Ch) | ||
| expectIntOnChan(t, 100, fanin.Ch) | ||
| close(chan1) | ||
| close(chan2) | ||
|
|
||
| // close of last channel closes fanin channel which in turn sends nil message | ||
| expectNilOnChan(t, fanin.Ch) | ||
| _, ok := <-fanin | ||
|
|
||
| assert.False(t, fanin.Alive()) | ||
| assert.False(t, ok) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters