helper/resource: StopCh as a helper for provider stopCh + timeout

This commit is contained in:
Mitchell Hashimoto 2016-10-23 18:25:47 -07:00
parent 8c11f137f5
commit 89647745b0
No known key found for this signature in database
GPG Key ID: 744E147AA52F5B0A
2 changed files with 102 additions and 0 deletions

40
helper/resource/stop.go Normal file
View File

@ -0,0 +1,40 @@
package resource
import (
"time"
)
// StopCh is used to construct a channel that is closed when the provider
// stopCh is closed, or a timeout is reached.
//
// The first return value is the channel that is closed when the operation
// should stop. The second return value should be closed when the operation
// completes so that the stop channel is never closed and to cleanup the
// goroutine watching the channels.
func StopCh(stopCh <-chan struct{}, max time.Duration) (<-chan struct{}, chan<- struct{}) {
ch := make(chan struct{})
doneCh := make(chan struct{})
// If we have a max of 0 then it is unlimited. A nil channel blocks on
// receive forever so this ensures that behavior.
var timeCh <-chan time.Time
if max > 0 {
timeCh = time.After(max)
}
// Start a goroutine to watch the various cases of cancellation
go func() {
select {
case <-doneCh:
// If we are done, don't trigger the cancel
return
case <-timeCh:
case <-stopCh:
}
// Trigger the cancel
close(ch)
}()
return ch, doneCh
}

View File

@ -0,0 +1,62 @@
package resource
import (
"testing"
"time"
)
func TestStopCh_stop(t *testing.T) {
stopCh := make(chan struct{})
ch, _ := StopCh(stopCh, 0)
// ch should block
select {
case <-ch:
t.Fatal("ch should block")
case <-time.After(10 * time.Millisecond):
}
// Close the stop channel
close(stopCh)
// ch should return
select {
case <-ch:
case <-time.After(10 * time.Millisecond):
t.Fatal("ch should not block")
}
}
func TestStopCh_done(t *testing.T) {
stopCh := make(chan struct{})
ch, doneCh := StopCh(stopCh, 0)
// ch should block
select {
case <-ch:
t.Fatal("ch should block")
case <-time.After(10 * time.Millisecond):
}
// Close the done channel
close(doneCh)
// ch should block
select {
case <-ch:
t.Fatal("ch should block")
case <-time.After(10 * time.Millisecond):
}
}
func TestStopCh_timeout(t *testing.T) {
stopCh := make(chan struct{})
ch, _ := StopCh(stopCh, 10*time.Millisecond)
// ch should return
select {
case <-ch:
case <-time.After(100 * time.Millisecond):
t.Fatal("ch should not block")
}
}