From 89647745b0985a1dd74d04418f7436c63363f7d7 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Sun, 23 Oct 2016 18:25:47 -0700
Subject: [PATCH] helper/resource: StopCh as a helper for provider stopCh +
 timeout

---
 helper/resource/stop.go      | 40 +++++++++++++++++++++++
 helper/resource/stop_test.go | 62 ++++++++++++++++++++++++++++++++++++
 2 files changed, 102 insertions(+)
 create mode 100644 helper/resource/stop.go
 create mode 100644 helper/resource/stop_test.go

diff --git a/helper/resource/stop.go b/helper/resource/stop.go
new file mode 100644
index 000000000..b3bfa776c
--- /dev/null
+++ b/helper/resource/stop.go
@@ -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
+}
diff --git a/helper/resource/stop_test.go b/helper/resource/stop_test.go
new file mode 100644
index 000000000..839f4af9a
--- /dev/null
+++ b/helper/resource/stop_test.go
@@ -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")
+	}
+}