From 02d01141ca25c07c9f07098175e8d671fd98c325 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Sat, 4 Oct 2014 09:54:03 -0700
Subject: [PATCH] rpc: implement provisioner update

---
 rpc/client.go                    |  1 +
 rpc/client_test.go               |  3 ++-
 rpc/resource_provisioner.go      | 35 +++++++++++++++++++++++++++-----
 rpc/resource_provisioner_test.go | 13 +++++++-----
 rpc/server.go                    |  1 +
 5 files changed, 42 insertions(+), 11 deletions(-)

diff --git a/rpc/client.go b/rpc/client.go
index f6d20e3a6..0c80385ee 100644
--- a/rpc/client.go
+++ b/rpc/client.go
@@ -101,6 +101,7 @@ func (c *Client) ResourceProvisioner() (terraform.ResourceProvisioner, error) {
 	}
 
 	return &ResourceProvisioner{
+		Broker: c.broker,
 		Client: rpc.NewClient(conn),
 		Name:   "ResourceProvisioner",
 	}, nil
diff --git a/rpc/client_test.go b/rpc/client_test.go
index c4479cfd1..f8c286fe8 100644
--- a/rpc/client_test.go
+++ b/rpc/client_test.go
@@ -60,9 +60,10 @@ func TestClient_ResourceProvisioner(t *testing.T) {
 	}
 
 	// Apply
+	output := &terraform.MockUIOutput{}
 	state := &terraform.InstanceState{}
 	conf := &terraform.ResourceConfig{}
-	err = provisioner.Apply(state, conf)
+	err = provisioner.Apply(output, state, conf)
 	if !p.ApplyCalled {
 		t.Fatal("apply should be called")
 	}
diff --git a/rpc/resource_provisioner.go b/rpc/resource_provisioner.go
index 5fc45a98c..cf8c00812 100644
--- a/rpc/resource_provisioner.go
+++ b/rpc/resource_provisioner.go
@@ -9,6 +9,7 @@ import (
 // ResourceProvisioner is an implementation of terraform.ResourceProvisioner
 // that communicates over RPC.
 type ResourceProvisioner struct {
+	Broker *muxBroker
 	Client *rpc.Client
 	Name   string
 }
@@ -36,12 +37,19 @@ func (p *ResourceProvisioner) Validate(c *terraform.ResourceConfig) ([]string, [
 }
 
 func (p *ResourceProvisioner) Apply(
+	output terraform.UIOutput,
 	s *terraform.InstanceState,
 	c *terraform.ResourceConfig) error {
+	id := p.Broker.NextId()
+	go acceptAndServe(p.Broker, id, "UIOutput", &UIOutputServer{
+		UIOutput: output,
+	})
+
 	var resp ResourceProvisionerApplyResponse
 	args := &ResourceProvisionerApplyArgs{
-		State:  s,
-		Config: c,
+		OutputId: id,
+		State:    s,
+		Config:   c,
 	}
 
 	err := p.Client.Call(p.Name+".Apply", args, &resp)
@@ -65,8 +73,9 @@ type ResourceProvisionerValidateResponse struct {
 }
 
 type ResourceProvisionerApplyArgs struct {
-	State  *terraform.InstanceState
-	Config *terraform.ResourceConfig
+	OutputId uint32
+	State    *terraform.InstanceState
+	Config   *terraform.ResourceConfig
 }
 
 type ResourceProvisionerApplyResponse struct {
@@ -76,13 +85,29 @@ type ResourceProvisionerApplyResponse struct {
 // ResourceProvisionerServer is a net/rpc compatible structure for serving
 // a ResourceProvisioner. This should not be used directly.
 type ResourceProvisionerServer struct {
+	Broker      *muxBroker
 	Provisioner terraform.ResourceProvisioner
 }
 
 func (s *ResourceProvisionerServer) Apply(
 	args *ResourceProvisionerApplyArgs,
 	result *ResourceProvisionerApplyResponse) error {
-	err := s.Provisioner.Apply(args.State, args.Config)
+	conn, err := s.Broker.Dial(args.OutputId)
+	if err != nil {
+		*result = ResourceProvisionerApplyResponse{
+			Error: NewBasicError(err),
+		}
+		return nil
+	}
+	client := rpc.NewClient(conn)
+	defer client.Close()
+
+	output := &UIOutput{
+		Client: client,
+		Name:   "UIOutput",
+	}
+
+	err = s.Provisioner.Apply(output, args.State, args.Config)
 	*result = ResourceProvisionerApplyResponse{
 		Error: NewBasicError(err),
 	}
diff --git a/rpc/resource_provisioner_test.go b/rpc/resource_provisioner_test.go
index b91252648..32b6c06f3 100644
--- a/rpc/resource_provisioner_test.go
+++ b/rpc/resource_provisioner_test.go
@@ -13,18 +13,21 @@ func TestResourceProvisioner_impl(t *testing.T) {
 }
 
 func TestResourceProvisioner_apply(t *testing.T) {
-	p := new(terraform.MockResourceProvisioner)
-	client, server := testClientServer(t)
-	name, err := Register(server, p)
+	client, server := testNewClientServer(t)
+	defer client.Close()
+
+	p := server.ProvisionerFunc().(*terraform.MockResourceProvisioner)
+
+	provisioner, err := client.ResourceProvisioner()
 	if err != nil {
 		t.Fatalf("err: %s", err)
 	}
-	provisioner := &ResourceProvisioner{Client: client, Name: name}
 
 	// Apply
+	output := &terraform.MockUIOutput{}
 	state := &terraform.InstanceState{}
 	conf := &terraform.ResourceConfig{}
-	err = provisioner.Apply(state, conf)
+	err = provisioner.Apply(output, state, conf)
 	if !p.ApplyCalled {
 		t.Fatal("apply should be called")
 	}
diff --git a/rpc/server.go b/rpc/server.go
index 5ca6c6746..dd1e9b7b0 100644
--- a/rpc/server.go
+++ b/rpc/server.go
@@ -118,6 +118,7 @@ func (d *dispenseServer) ResourceProvisioner(
 		}
 
 		serve(conn, "ResourceProvisioner", &ResourceProvisionerServer{
+			Broker:      d.broker,
 			Provisioner: d.ProvisionerFunc(),
 		})
 	}()