core: close provider/provisioner connections when not used anymore

rpc/mux_broker.go

@@ -3,9 +3,9 @@ package rpc
 import (
 	"encoding/binary"
 	"fmt"
+	"math/rand"
 	"net"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/hashicorp/yamux"
@@ -17,7 +17,6 @@ import (
 // or accept a connection from, and the broker handles the details of
 // holding these channels open while they're being negotiated.
 type muxBroker struct {
-	nextId  uint32
 	session *yamux.Session
 	streams map[uint32]*muxBrokerPending
 
@@ -95,9 +94,12 @@ func (m *muxBroker) Dial(id uint32) (net.Conn, error) {
 	return stream, nil
 }
 
-// NextId returns a unique ID to use next.
+// NextId returns a unique ID to use next. There is no need for seeding the
+// rand package as the returned ID's aren't stored or used anywhere outside
+// the current runtime. So it's perfectly fine to get the same pseudo-random
+// numbers each time terraform is running.
 func (m *muxBroker) NextId() uint32 {
-	return atomic.AddUint32(&m.nextId, 1)
+	return rand.Uint32()
 }
 
 // Run starts the brokering and should be executed in a goroutine, since it

rpc/resource_provider.go

@@ -174,6 +174,10 @@ func (p *ResourceProvider) Resources() []terraform.ResourceType {
 	return result
 }
 
+func (p *ResourceProvider) Close() error {
+	return p.Client.Close()
+}
+
 // ResourceProviderServer is a net/rpc compatible structure for serving
 // a ResourceProvider. This should not be used directly.
 type ResourceProviderServer struct {

rpc/resource_provider_test.go

@@ -488,3 +488,31 @@ func TestResourceProvider_validateResource_warns(t *testing.T) {
 		t.Fatalf("bad: %#v", w)
 	}
 }
+
+func TestResourceProvider_close(t *testing.T) {
+	client, _ := testNewClientServer(t)
+	defer client.Close()
+
+	provider, err := client.ResourceProvider()
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	var p interface{}
+	p = provider
+	pCloser, ok := p.(terraform.ResourceProviderCloser)
+	if !ok {
+		t.Fatal("should be a ResourceProviderCloser")
+	}
+
+	if err := pCloser.Close(); err != nil {
+		t.Fatalf("failed to close provider: %s", err)
+	}
+
+	// The connection should be closed now, so if we to make a
+	// new call we should get an error.
+	err = provider.Configure(&terraform.ResourceConfig{})
+	if err == nil {
+		t.Fatal("should have error")
+	}
+}

rpc/resource_provisioner.go

@@ -63,6 +63,10 @@ func (p *ResourceProvisioner) Apply(
 	return err
 }
 
+func (p *ResourceProvisioner) Close() error {
+	return p.Client.Close()
+}
+
 type ResourceProvisionerValidateArgs struct {
 	Config *terraform.ResourceConfig
 }

rpc/resource_provisioner_test.go

@@ -132,3 +132,34 @@ func TestResourceProvisioner_validate_warns(t *testing.T) {
 		t.Fatalf("bad: %#v", w)
 	}
 }
+
+func TestResourceProvisioner_close(t *testing.T) {
+	client, _ := testNewClientServer(t)
+	defer client.Close()
+
+	provisioner, err := client.ResourceProvisioner()
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	var p interface{}
+	p = provisioner
+	pCloser, ok := p.(terraform.ResourceProvisionerCloser)
+	if !ok {
+		t.Fatal("should be a ResourceProvisionerCloser")
+	}
+
+	if err := pCloser.Close(); err != nil {
+		t.Fatalf("failed to close provisioner: %s", err)
+	}
+
+	// The connection should be closed now, so if we to make a
+	// new call we should get an error.
+	o := &terraform.MockUIOutput{}
+	s := &terraform.InstanceState{}
+	c := &terraform.ResourceConfig{}
+	err = provisioner.Apply(o, s, c)
+	if err == nil {
+		t.Fatal("should have error")
+	}
+}

terraform/eval_context.go

@@ -28,6 +28,9 @@ type EvalContext interface {
 	// initialized) or returns nil if the provider isn't initialized.
 	Provider(string) ResourceProvider
 
+	// CloseProvider closes provider connections that aren't needed anymore.
+	CloseProvider(string) error
+
 	// ConfigureProvider configures the provider with the given
 	// configuration. This is a separate context call because this call
 	// is used to store the provider configuration for inheritance lookups
@@ -51,6 +54,10 @@ type EvalContext interface {
 	// initialized) or returns nil if the provisioner isn't initialized.
 	Provisioner(string) ResourceProvisioner
 
+	// CloseProvisioner closes provisioner connections that aren't needed
+	// anymore.
+	CloseProvisioner(string) error
+
 	// Interpolate takes the given raw configuration and completes
 	// the interpolations, returning the processed ResourceConfig.
 	//

terraform/eval_context_builtin.go

@@ -114,6 +114,28 @@ func (ctx *BuiltinEvalContext) Provider(n string) ResourceProvider {
 	return ctx.ProviderCache[PathCacheKey(providerPath)]
 }
 
+func (ctx *BuiltinEvalContext) CloseProvider(n string) error {
+	ctx.once.Do(ctx.init)
+
+	ctx.ProviderLock.Lock()
+	defer ctx.ProviderLock.Unlock()
+
+	providerPath := make([]string, len(ctx.Path())+1)
+	copy(providerPath, ctx.Path())
+	providerPath[len(providerPath)-1] = n
+
+	var provider interface{}
+	provider = ctx.ProviderCache[PathCacheKey(providerPath)]
+	if provider != nil {
+		if p, ok := provider.(ResourceProviderCloser); ok {
+			delete(ctx.ProviderCache, PathCacheKey(providerPath))
+			return p.Close()
+		}
+	}
+
+	return nil
+}
+
 func (ctx *BuiltinEvalContext) ConfigureProvider(
 	n string, cfg *ResourceConfig) error {
 	p := ctx.Provider(n)
@@ -222,6 +244,28 @@ func (ctx *BuiltinEvalContext) Provisioner(n string) ResourceProvisioner {
 	return ctx.ProvisionerCache[PathCacheKey(provPath)]
 }
 
+func (ctx *BuiltinEvalContext) CloseProvisioner(n string) error {
+	ctx.once.Do(ctx.init)
+
+	ctx.ProvisionerLock.Lock()
+	defer ctx.ProvisionerLock.Unlock()
+
+	provPath := make([]string, len(ctx.Path())+1)
+	copy(provPath, ctx.Path())
+	provPath[len(provPath)-1] = n
+
+	var prov interface{}
+	prov = ctx.ProvisionerCache[PathCacheKey(provPath)]
+	if prov != nil {
+		if p, ok := prov.(ResourceProvisionerCloser); ok {
+			delete(ctx.ProvisionerCache, PathCacheKey(provPath))
+			return p.Close()
+		}
+	}
+
+	return nil
+}
+
 func (ctx *BuiltinEvalContext) Interpolate(
 	cfg *config.RawConfig, r *Resource) (*ResourceConfig, error) {
 	if cfg != nil {

terraform/eval_context_mock.go

@@ -25,6 +25,10 @@ type MockEvalContext struct {
 	ProviderName     string
 	ProviderProvider ResourceProvider
 
+	CloseProviderCalled   bool
+	CloseProviderName     string
+	CloseProviderProvider ResourceProvider
+
 	ProviderInputCalled bool
 	ProviderInputName   string
 	ProviderInputConfig map[string]interface{}
@@ -55,6 +59,10 @@ type MockEvalContext struct {
 	ProvisionerName        string
 	ProvisionerProvisioner ResourceProvisioner
 
+	CloseProvisionerCalled      bool
+	CloseProvisionerName        string
+	CloseProvisionerProvisioner ResourceProvisioner
+
 	InterpolateCalled       bool
 	InterpolateConfig       *config.RawConfig
 	InterpolateResource     *Resource
@@ -105,6 +113,12 @@ func (c *MockEvalContext) Provider(n string) ResourceProvider {
 	return c.ProviderProvider
 }
 
+func (c *MockEvalContext) CloseProvider(n string) error {
+	c.CloseProviderCalled = true
+	c.CloseProviderName = n
+	return nil
+}
+
 func (c *MockEvalContext) ConfigureProvider(n string, cfg *ResourceConfig) error {
 	c.ConfigureProviderCalled = true
 	c.ConfigureProviderName = n
@@ -150,6 +164,12 @@ func (c *MockEvalContext) Provisioner(n string) ResourceProvisioner {
 	return c.ProvisionerProvisioner
 }
 
+func (c *MockEvalContext) CloseProvisioner(n string) error {
+	c.CloseProvisionerCalled = true
+	c.CloseProvisionerName = n
+	return nil
+}
+
 func (c *MockEvalContext) Interpolate(
 	config *config.RawConfig, resource *Resource) (*ResourceConfig, error) {
 	c.InterpolateCalled = true

terraform/eval_provider.go

@@ -71,6 +71,17 @@ func (n *EvalInitProvider) Eval(ctx EvalContext) (interface{}, error) {
 	return ctx.InitProvider(n.Name)
 }
 
+// EvalCloseProvider is an EvalNode implementation that closes provider
+// connections that aren't needed anymore.
+type EvalCloseProvider struct {
+	Name string
+}
+
+func (n *EvalCloseProvider) Eval(ctx EvalContext) (interface{}, error) {
+	ctx.CloseProvider(n.Name)
+	return nil, nil
+}
+
 // EvalGetProvider is an EvalNode implementation that retrieves an already
 // initialized provider instance for the given name.
 type EvalGetProvider struct {

terraform/eval_provider_test.go

@@ -112,6 +112,22 @@ func TestEvalInitProvider(t *testing.T) {
 	}
 }
 
+func TestEvalCloseProvider(t *testing.T) {
+	n := &EvalCloseProvider{Name: "foo"}
+	provider := &MockResourceProvider{}
+	ctx := &MockEvalContext{CloseProviderProvider: provider}
+	if _, err := n.Eval(ctx); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	if !ctx.CloseProviderCalled {
+		t.Fatal("should be called")
+	}
+	if ctx.CloseProviderName != "foo" {
+		t.Fatalf("bad: %#v", ctx.CloseProviderName)
+	}
+}
+
 func TestEvalGetProvider_impl(t *testing.T) {
 	var _ EvalNode = new(EvalGetProvider)
 }

terraform/eval_provisioner.go

@@ -15,6 +15,17 @@ func (n *EvalInitProvisioner) Eval(ctx EvalContext) (interface{}, error) {
 	return ctx.InitProvisioner(n.Name)
 }
 
+// EvalCloseProvisioner is an EvalNode implementation that closes provisioner
+// connections that aren't needed anymore.
+type EvalCloseProvisioner struct {
+	Name string
+}
+
+func (n *EvalCloseProvisioner) Eval(ctx EvalContext) (interface{}, error) {
+	ctx.CloseProvisioner(n.Name)
+	return nil, nil
+}
+
 // EvalGetProvisioner is an EvalNode implementation that retrieves an already
 // initialized provisioner instance for the given name.
 type EvalGetProvisioner struct {

terraform/eval_provisioner_test.go

@@ -24,6 +24,22 @@ func TestEvalInitProvisioner(t *testing.T) {
 	}
 }
 
+func TestEvalCloseProvisioner(t *testing.T) {
+	n := &EvalCloseProvisioner{Name: "foo"}
+	provisioner := &MockResourceProvisioner{}
+	ctx := &MockEvalContext{CloseProvisionerProvisioner: provisioner}
+	if _, err := n.Eval(ctx); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	if !ctx.CloseProvisionerCalled {
+		t.Fatal("should be called")
+	}
+	if ctx.CloseProvisionerName != "foo" {
+		t.Fatalf("bad: %#v", ctx.CloseProvisionerName)
+	}
+}
+
 func TestEvalGetProvisioner_impl(t *testing.T) {
 	var _ EvalNode = new(EvalGetProvisioner)
 }

terraform/evaltree_provider.go

@@ -72,3 +72,9 @@ func ProviderEvalTree(n string, config *config.RawConfig) EvalNode {
 
 	return &EvalSequence{Nodes: seq}
 }
+
+// CloseProviderEvalTree returns the evaluation tree for closing
+// provider connections that aren't needed anymore.
+func CloseProviderEvalTree(n string) EvalNode {
+	return &EvalCloseProvider{Name: n}
+}

terraform/graph_builder.go

@@ -116,12 +116,14 @@ func (b *BuiltinGraphBuilder) Steps(path []string) []GraphTransformer {
 		// Provider-related transformations
 		&MissingProviderTransformer{Providers: b.Providers},
 		&ProviderTransformer{},
+		&CloseProviderTransformer{},
 		&PruneProviderTransformer{},
 		&DisableProviderTransformer{},
 
 		// Provisioner-related transformations
 		&MissingProvisionerTransformer{Provisioners: b.Provisioners},
 		&ProvisionerTransformer{},
+		&CloseProvisionerTransformer{},
 		&PruneProvisionerTransformer{},
 
 		// Run our vertex-level transforms