packer/rpc: Builder

packer/rpc/builder.go

@@ -1,8 +1,6 @@
 package rpc
 
 import (
-	"encoding/gob"
-	"fmt"
 	"github.com/mitchellh/packer/packer"
 	"log"
 	"net/rpc"
@@ -12,35 +10,27 @@ import (
 // over an RPC connection.
 type builder struct {
 	client *rpc.Client
+	mux *MuxConn
 }
 
 // BuilderServer wraps a packer.Builder implementation and makes it exportable
 // as part of a Golang RPC server.
 type BuilderServer struct {
 	builder packer.Builder
+	mux *MuxConn
 }
 
 type BuilderPrepareArgs struct {
 	Configs []interface{}
 }
 
-type BuilderRunArgs struct {
-	RPCAddress      string
-	ResponseAddress string
-}
-
 type BuilderPrepareResponse struct {
 	Warnings []string
 	Error    error
 }
 
-type BuilderRunResponse struct {
-	Err        error
-	RPCAddress string
-}
-
 func Builder(client *rpc.Client) *builder {
-	return &builder{client}
+	return &builder{client: client}
 }
 
 func (b *builder) Prepare(config ...interface{}) ([]string, error) {
@@ -54,58 +44,28 @@ func (b *builder) Prepare(config ...interface{}) ([]string, error) {
 }
 
 func (b *builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packer.Artifact, error) {
-	// Create and start the server for the Build and UI
+	nextId := b.mux.NextId()
-	server := rpc.NewServer()
+	server := NewServerWithMux(b.mux, nextId)
-	RegisterCache(server, cache)
+	server.RegisterCache(cache)
-	RegisterHook(server, hook)
+	server.RegisterHook(hook)
-	RegisterUi(server, ui)
+	server.RegisterUi(ui)
+	go server.Serve()
 
-	// Create a server for the response
+	var responseId uint32
-	responseL := netListenerInRange(portRangeMin, portRangeMax)
+	if err := b.client.Call("Builder.Run", nextId, &responseId); err != nil {
-	runResponseCh := make(chan *BuilderRunResponse)
-	go func() {
-		defer responseL.Close()
-
-		var response BuilderRunResponse
-		defer func() { runResponseCh <- &response }()
-
-		conn, err := responseL.Accept()
-		if err != nil {
-			response.Err = err
-			return
-		}
-		defer conn.Close()
-
-		decoder := gob.NewDecoder(conn)
-		if err := decoder.Decode(&response); err != nil {
-			response.Err = fmt.Errorf("Error waiting for Run: %s", err)
-		}
-	}()
-
-	args := &BuilderRunArgs{
-		serveSingleConn(server),
-		responseL.Addr().String(),
-	}
-
-	if err := b.client.Call("Builder.Run", args, new(interface{})); err != nil {
 		return nil, err
 	}
 
-	response := <-runResponseCh
+	if responseId == 0 {
-	if response.Err != nil {
-		return nil, response.Err
-	}
-
-	if response.RPCAddress == "" {
 		return nil, nil
 	}
 
-	client, err := rpcDial(response.RPCAddress)
+	client, err := NewClientWithMux(b.mux, responseId)
 	if err != nil {
 		return nil, err
 	}
 
-	return Artifact(client), nil
+	return client.Artifact(), nil
 }
 
 func (b *builder) Cancel() {
@@ -127,46 +87,27 @@ func (b *BuilderServer) Prepare(args *BuilderPrepareArgs, reply *BuilderPrepareR
 	return nil
 }
 
-func (b *BuilderServer) Run(args *BuilderRunArgs, reply *interface{}) error {
+func (b *BuilderServer) Run(streamId uint32, reply *uint32) error {
-	client, err := rpcDial(args.RPCAddress)
+	client, err := NewClientWithMux(b.mux, streamId)
 	if err != nil {
-		return err
+		return NewBasicError(err)
+	}
+	defer client.Close()
+
+	artifact, err := b.builder.Run(client.Ui(), client.Hook(), client.Cache())
+	if err != nil {
+		return NewBasicError(err)
 	}
 
-	responseC, err := tcpDial(args.ResponseAddress)
+	*reply = 0
-	if err != nil {
+	if artifact != nil {
-		return err
+		streamId = b.mux.NextId()
+		server := NewServerWithMux(b.mux, streamId)
+		server.RegisterArtifact(artifact)
+		go server.Serve()
+		*reply = streamId
 	}
 
-	responseWriter := gob.NewEncoder(responseC)
-
-	// Run the build in a goroutine so we don't block the RPC connection
-	go func() {
-		defer responseC.Close()
-
-		cache := Cache(client)
-		hook := Hook(client)
-		ui := &Ui{client: client}
-		artifact, responseErr := b.builder.Run(ui, hook, cache)
-		responseAddress := ""
-
-		if responseErr == nil && artifact != nil {
-			// Wrap the artifact
-			server := rpc.NewServer()
-			RegisterArtifact(server, artifact)
-			responseAddress = serveSingleConn(server)
-		}
-
-		if responseErr != nil {
-			responseErr = NewBasicError(responseErr)
-		}
-
-		err := responseWriter.Encode(&BuilderRunResponse{responseErr, responseAddress})
-		if err != nil {
-			log.Printf("BuildServer.Run error: %s", err)
-		}
-	}()
-
 	return nil
 }
 

packer/rpc/builder_test.go

@@ -2,31 +2,19 @@ package rpc
 
 import (
 	"github.com/mitchellh/packer/packer"
-	"net/rpc"
 	"reflect"
 	"testing"
 )
 
 var testBuilderArtifact = &packer.MockArtifact{}
 
-func builderRPCClient(t *testing.T) (*packer.MockBuilder, packer.Builder) {
-	b := new(packer.MockBuilder)
-
-	// Start the server
-	server := rpc.NewServer()
-	RegisterBuilder(server, b)
-	address := serveSingleConn(server)
-
-	// Create the client over RPC and run some methods to verify it works
-	client, err := rpc.Dial("tcp", address)
-	if err != nil {
-		t.Fatalf("err: %s", err)
-	}
-	return b, Builder(client)
-}
-
 func TestBuilderPrepare(t *testing.T) {
-	b, bClient := builderRPCClient(t)
+	b := new(packer.MockBuilder)
+	client, server := testClientServer(t)
+	defer client.Close()
+	defer server.Close()
+	server.RegisterBuilder(b)
+	bClient := client.Builder()
 
 	// Test Prepare
 	config := 42
@@ -48,7 +36,12 @@ func TestBuilderPrepare(t *testing.T) {
 }
 
 func TestBuilderPrepare_Warnings(t *testing.T) {
-	b, bClient := builderRPCClient(t)
+	b := new(packer.MockBuilder)
+	client, server := testClientServer(t)
+	defer client.Close()
+	defer server.Close()
+	server.RegisterBuilder(b)
+	bClient := client.Builder()
 
 	expected := []string{"foo"}
 	b.PrepareWarnings = expected
@@ -64,7 +57,12 @@ func TestBuilderPrepare_Warnings(t *testing.T) {
 }
 
 func TestBuilderRun(t *testing.T) {
-	b, bClient := builderRPCClient(t)
+	b := new(packer.MockBuilder)
+	client, server := testClientServer(t)
+	defer client.Close()
+	defer server.Close()
+	server.RegisterBuilder(b)
+	bClient := client.Builder()
 
 	// Test Run
 	cache := new(testCache)
@@ -79,34 +77,21 @@ func TestBuilderRun(t *testing.T) {
 		t.Fatal("run should be called")
 	}
 
-	b.RunCache.Lock("foo")
-	if !cache.lockCalled {
-		t.Fatal("should be called")
-	}
-
-	b.RunHook.Run("foo", nil, nil, nil)
-	if !hook.RunCalled {
-		t.Fatal("should be called")
-	}
-
-	b.RunUi.Say("format")
-	if !ui.sayCalled {
-		t.Fatal("say should be called")
-	}
-
-	if ui.sayMessage != "format" {
-		t.Fatalf("bad: %s", ui.sayMessage)
-	}
-
 	if artifact.Id() != testBuilderArtifact.Id() {
 		t.Fatalf("bad: %s", artifact.Id())
 	}
 }
 
 func TestBuilderRun_nilResult(t *testing.T) {
-	b, bClient := builderRPCClient(t)
+	b := new(packer.MockBuilder)
 	b.RunNilResult = true
 
+	client, server := testClientServer(t)
+	defer client.Close()
+	defer server.Close()
+	server.RegisterBuilder(b)
+	bClient := client.Builder()
+
 	cache := new(testCache)
 	hook := &packer.MockHook{}
 	ui := &testUi{}
@@ -120,7 +105,13 @@ func TestBuilderRun_nilResult(t *testing.T) {
 }
 
 func TestBuilderRun_ErrResult(t *testing.T) {
-	b, bClient := builderRPCClient(t)
+	b := new(packer.MockBuilder)
+	client, server := testClientServer(t)
+	defer client.Close()
+	defer server.Close()
+	server.RegisterBuilder(b)
+	bClient := client.Builder()
+
 	b.RunErrResult = true
 
 	cache := new(testCache)
@@ -136,7 +127,12 @@ func TestBuilderRun_ErrResult(t *testing.T) {
 }
 
 func TestBuilderCancel(t *testing.T) {
-	b, bClient := builderRPCClient(t)
+	b := new(packer.MockBuilder)
+	client, server := testClientServer(t)
+	defer client.Close()
+	defer server.Close()
+	server.RegisterBuilder(b)
+	bClient := client.Builder()
 
 	bClient.Cancel()
 	if !b.CancelCalled {

packer/rpc/client.go

@@ -45,6 +45,13 @@ func (c *Client) Artifact() packer.Artifact {
 	}
 }
 
+func (c *Client) Builder() packer.Builder {
+	return &builder{
+		client: c.client,
+		mux:    c.mux,
+	}
+}
+
 func (c *Client) Cache() packer.Cache {
 	return &cache{
 		client: c.client,

packer/rpc/server.go

@@ -20,7 +20,7 @@ func RegisterBuild(s *rpc.Server, b packer.Build) {
 // Registers the appropriate endpoint on an RPC server to serve a
 // Packer Builder.
 func RegisterBuilder(s *rpc.Server, b packer.Builder) {
-	registerComponent(s, "Builder", &BuilderServer{b}, false)
+	registerComponent(s, "Builder", &BuilderServer{builder: b}, false)
 }
 
 // Registers the appropriate endpoint on an RPC server to serve a

packer/rpc/server_new.go

@@ -13,6 +13,7 @@ var endpointId uint64
 
 const (
 	DefaultArtifactEndpoint      string = "Artifact"
+	DefaultBuilderEndpoint              = "Builder"
 	DefaultCacheEndpoint                = "Cache"
 	DefaultCommandEndpoint              = "Command"
 	DefaultCommunicatorEndpoint         = "Communicator"
@@ -53,6 +54,13 @@ func (s *Server) RegisterArtifact(a packer.Artifact) {
 	})
 }
 
+func (s *Server) RegisterBuilder(b packer.Builder) {
+	s.server.RegisterName(DefaultBuilderEndpoint, &BuilderServer{
+		builder: b,
+		mux:     s.mux,
+	})
+}
+
 func (s *Server) RegisterCache(c packer.Cache) {
 	s.server.RegisterName(DefaultCacheEndpoint, &CacheServer{
 		cache: c,