From 8c11f137f5d8cc113a38cec129ac57e226b1c550 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Sun, 23 Oct 2016 18:07:17 -0700
Subject: [PATCH] helper/schema: support Stop()

---
 helper/schema/provider.go      | 44 ++++++++++++++++++++++++++++
 helper/schema/provider_test.go | 53 ++++++++++++++++++++++++++++++++++
 2 files changed, 97 insertions(+)

diff --git a/helper/schema/provider.go b/helper/schema/provider.go
index 546794f406..636ea55182 100644
--- a/helper/schema/provider.go
+++ b/helper/schema/provider.go
@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"sort"
+	"sync"
 
 	"github.com/hashicorp/go-multierror"
 	"github.com/hashicorp/terraform/terraform"
@@ -49,6 +50,10 @@ type Provider struct {
 	ConfigureFunc ConfigureFunc
 
 	meta interface{}
+
+	stopCh   chan struct{} // stopCh is closed when Stop is called
+	stopped  bool          // set to true once stopped to avoid double close of stopCh
+	stopLock sync.Mutex
 }
 
 // ConfigureFunc is the function used to configure a Provider.
@@ -104,6 +109,45 @@ func (p *Provider) SetMeta(v interface{}) {
 	p.meta = v
 }
 
+// Stopped reports whether the provider has been stopped or not.
+func (p *Provider) Stopped() bool {
+	p.stopLock.Lock()
+	defer p.stopLock.Unlock()
+	return p.stopped
+}
+
+// StopCh returns a channel that is closed once the provider is stopped.
+func (p *Provider) StopCh() <-chan struct{} {
+	p.stopLock.Lock()
+	defer p.stopLock.Unlock()
+
+	if p.stopCh == nil {
+		p.stopCh = make(chan struct{})
+	}
+
+	return p.stopCh
+}
+
+// Stop implementation of terraform.ResourceProvider interface.
+func (p *Provider) Stop() error {
+	p.stopLock.Lock()
+	defer p.stopLock.Unlock()
+
+	// Close the stop channel and mark as stopped if we haven't
+	if !p.stopped {
+		// Initialize the stop channel so future calls to StopCh work
+		if p.stopCh == nil {
+			p.stopCh = make(chan struct{})
+		}
+
+		// Close and mark
+		close(p.stopCh)
+		p.stopped = true
+	}
+
+	return nil
+}
+
 // Input implementation of terraform.ResourceProvider interface.
 func (p *Provider) Input(
 	input terraform.UIInput,
diff --git a/helper/schema/provider_test.go b/helper/schema/provider_test.go
index aa8a787bfa..cde72d86bf 100644
--- a/helper/schema/provider_test.go
+++ b/helper/schema/provider_test.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"reflect"
 	"testing"
+	"time"
 
 	"github.com/hashicorp/terraform/config"
 	"github.com/hashicorp/terraform/terraform"
@@ -328,3 +329,55 @@ func TestProviderMeta(t *testing.T) {
 		t.Fatalf("bad: %#v", v)
 	}
 }
+
+func TestProviderStop(t *testing.T) {
+	var p Provider
+
+	if p.Stopped() {
+		t.Fatal("should not be stopped")
+	}
+
+	// Verify stopch blocks
+	ch := p.StopCh()
+	select {
+	case <-ch:
+		t.Fatal("should not be stopped")
+	case <-time.After(10 * time.Millisecond):
+	}
+
+	// Stop it
+	if err := p.Stop(); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	// Verify
+	if !p.Stopped() {
+		t.Fatal("should be stopped")
+	}
+
+	select {
+	case <-ch:
+	case <-time.After(10 * time.Millisecond):
+		t.Fatal("should be stopped")
+	}
+}
+
+func TestProviderStop_stopFirst(t *testing.T) {
+	var p Provider
+
+	// Stop it
+	if err := p.Stop(); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	// Verify
+	if !p.Stopped() {
+		t.Fatal("should be stopped")
+	}
+
+	select {
+	case <-p.StopCh():
+	case <-time.After(10 * time.Millisecond):
+		t.Fatal("should be stopped")
+	}
+}