packer/rpc: no more muxconn

packer/rpc/build.go

@@ -9,14 +9,14 @@ import (
 // over an RPC connection.
 type build struct {
 	client *rpc.Client
-	mux    *MuxConn
+	mux    *muxBroker
 }
 
 // BuildServer wraps a packer.Build implementation and makes it exportable
 // as part of a Golang RPC server.
 type BuildServer struct {
 	build packer.Build
-	mux   *MuxConn
+	mux   *muxBroker
 }
 
 type BuildPrepareResponse struct {

packer/rpc/builder.go

@@ -10,14 +10,14 @@ import (
 // over an RPC connection.
 type builder struct {
 	client *rpc.Client
-	mux    *MuxConn
+	mux    *muxBroker
 }
 
 // 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
+	mux     *muxBroker
 }
 
 type BuilderPrepareArgs struct {

packer/rpc/client.go

@@ -12,22 +12,29 @@ import (
 // Establishing a connection is up to the user, the Client can just
 // communicate over any ReadWriteCloser.
 type Client struct {
-	mux      *MuxConn
+	mux      *muxBroker
 	client   *rpc.Client
 	closeMux bool
 }
 
 func NewClient(rwc io.ReadWriteCloser) (*Client, error) {
-	result, err := newClientWithMux(NewMuxConn(rwc), 0)
+	mux, err := newMuxBrokerClient(rwc)
 	if err != nil {
 		return nil, err
 	}
+	go mux.Run()
+
+	result, err := newClientWithMux(mux, 0)
+	if err != nil {
+		mux.Close()
+		return nil, err
+	}
 
 	result.closeMux = true
 	return result, err
 }
 
-func newClientWithMux(mux *MuxConn, streamId uint32) (*Client, error) {
+func newClientWithMux(mux *muxBroker, streamId uint32) (*Client, error) {
 	clientConn, err := mux.Dial(streamId)
 	if err != nil {
 		return nil, err

packer/rpc/command.go

@@ -9,14 +9,14 @@ import (
 // command is actually executed over an RPC connection.
 type command struct {
 	client *rpc.Client
-	mux    *MuxConn
+	mux    *muxBroker
 }
 
 // A CommandServer wraps a packer.Command and makes it exportable as part
 // of a Golang RPC server.
 type CommandServer struct {
 	command packer.Command
-	mux     *MuxConn
+	mux     *muxBroker
 }
 
 type CommandRunArgs struct {

packer/rpc/communicator.go

@@ -12,14 +12,14 @@ import (
 // executed over an RPC connection.
 type communicator struct {
 	client *rpc.Client
-	mux    *MuxConn
+	mux    *muxBroker
 }
 
 // CommunicatorServer wraps a packer.Communicator implementation and makes
 // it exportable as part of a Golang RPC server.
 type CommunicatorServer struct {
 	c   packer.Communicator
-	mux *MuxConn
+	mux *muxBroker
 }
 
 type CommandFinished struct {
@@ -252,7 +252,7 @@ func (c *CommunicatorServer) Download(args *CommunicatorDownloadArgs, reply *int
 	return
 }
 
-func serveSingleCopy(name string, mux *MuxConn, id uint32, dst io.Writer, src io.Reader) {
+func serveSingleCopy(name string, mux *muxBroker, id uint32, dst io.Writer, src io.Reader) {
 	conn, err := mux.Accept(id)
 	if err != nil {
 		log.Printf("[ERR] '%s' accept error: %s", name, err)

packer/rpc/environment.go

@@ -10,14 +10,14 @@ import (
 // where the actual environment is executed over an RPC connection.
 type Environment struct {
 	client *rpc.Client
-	mux    *MuxConn
+	mux    *muxBroker
 }
 
 // A EnvironmentServer wraps a packer.Environment and makes it exportable
 // as part of a Golang RPC server.
 type EnvironmentServer struct {
 	env packer.Environment
-	mux *MuxConn
+	mux *muxBroker
 }
 
 type EnvironmentCliArgs struct {

packer/rpc/hook.go

@@ -10,14 +10,14 @@ import (
 // over an RPC connection.
 type hook struct {
 	client *rpc.Client
-	mux    *MuxConn
+	mux    *muxBroker
 }
 
 // HookServer wraps a packer.Hook implementation and makes it exportable
 // as part of a Golang RPC server.
 type HookServer struct {
 	hook packer.Hook
-	mux  *MuxConn
+	mux  *muxBroker
 }
 
 type HookRunArgs struct {

packer/rpc/mux_broker.go

@@ -3,8 +3,10 @@ package rpc
 import (
 	"encoding/binary"
 	"fmt"
+	"io"
 	"net"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/hashicorp/yamux"
@@ -16,6 +18,7 @@ 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
 
@@ -34,6 +37,24 @@ func newMuxBroker(s *yamux.Session) *muxBroker {
 	}
 }
 
+func newMuxBrokerClient(rwc io.ReadWriteCloser) (*muxBroker, error) {
+	s, err := yamux.Client(rwc, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	return newMuxBroker(s), nil
+}
+
+func newMuxBrokerServer(rwc io.ReadWriteCloser) (*muxBroker, error) {
+	s, err := yamux.Server(rwc, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	return newMuxBroker(s), nil
+}
+
 // Accept accepts a connection by ID.
 //
 // This should not be called multiple times with the same ID at one time.
@@ -60,6 +81,11 @@ func (m *muxBroker) Accept(id uint32) (net.Conn, error) {
 	return c, nil
 }
 
+// Close closes the connection and all sub-connections.
+func (m *muxBroker) Close() error {
+	return m.session.Close()
+}
+
 // Dial opens a connection by ID.
 func (m *muxBroker) Dial(id uint32) (net.Conn, error) {
 	// Open the stream
@@ -88,6 +114,11 @@ func (m *muxBroker) Dial(id uint32) (net.Conn, error) {
 	return stream, nil
 }
 
+// NextId returns a unique ID to use next.
+func (m *muxBroker) NextId() uint32 {
+	return atomic.AddUint32(&m.nextId, 1)
+}
+
 // Run starts the brokering and should be executed in a goroutine, since it
 // blocks forever, or until the session closes.
 func (m *muxBroker) Run() {

packer/rpc/muxconn_test.go

-package rpc
-
-import (
-	"io"
-	"net"
-	"sync"
-	"testing"
-)
-
-func readStream(t *testing.T, s io.Reader) string {
-	var data [1024]byte
-	n, err := s.Read(data[:])
-	if err != nil {
-		t.Fatalf("err: %s", err)
-	}
-
-	return string(data[0:n])
-}
-
-func testMux(t *testing.T) (client *MuxConn, server *MuxConn) {
-	l, err := net.Listen("tcp", "127.0.0.1:0")
-	if err != nil {
-		t.Fatalf("err: %s", err)
-	}
-
-	// Server side
-	doneCh := make(chan struct{})
-	go func() {
-		defer close(doneCh)
-		conn, err := l.Accept()
-		l.Close()
-		if err != nil {
-			t.Fatalf("err: %s", err)
-		}
-
-		server = NewMuxConn(conn)
-	}()
-
-	// Client side
-	conn, err := net.Dial("tcp", l.Addr().String())
-	if err != nil {
-		t.Fatalf("err: %s", err)
-	}
-	client = NewMuxConn(conn)
-
-	// Wait for the server
-	<-doneCh
-
-	return
-}
-
-func TestMuxConn(t *testing.T) {
-	client, server := testMux(t)
-	defer client.Close()
-	defer server.Close()
-
-	// When the server is done
-	doneCh := make(chan struct{})
-
-	// The server side
-	go func() {
-		defer close(doneCh)
-
-		s0, err := server.Accept(0)
-		if err != nil {
-			t.Fatalf("err: %s", err)
-		}
-
-		s1, err := server.Dial(1)
-		if err != nil {
-			t.Fatalf("err: %s", err)
-		}
-
-		var wg sync.WaitGroup
-		wg.Add(2)
-
-		go func() {
-			defer wg.Done()
-			defer s1.Close()
-			data := readStream(t, s1)
-			if data != "another" {
-				t.Fatalf("bad: %#v", data)
-			}
-		}()
-
-		go func() {
-			defer wg.Done()
-			defer s0.Close()
-			data := readStream(t, s0)
-			if data != "hello" {
-				t.Fatalf("bad: %#v", data)
-			}
-		}()
-
-		wg.Wait()
-	}()
-
-	s0, err := client.Dial(0)
-	if err != nil {
-		t.Fatalf("err: %s", err)
-	}
-
-	s1, err := client.Accept(1)
-	if err != nil {
-		t.Fatalf("err: %s", err)
-	}
-
-	if _, err := s0.Write([]byte("hello")); err != nil {
-		t.Fatalf("err: %s", err)
-	}
-	if _, err := s1.Write([]byte("another")); err != nil {
-		t.Fatalf("err: %s", err)
-	}
-
-	s0.Close()
-	s1.Close()
-
-	// Wait for the server to be done
-	<-doneCh
-}
-
-func TestMuxConn_lotsOfData(t *testing.T) {
-	client, server := testMux(t)
-	defer client.Close()
-	defer server.Close()
-
-	// When the server is done
-	doneCh := make(chan struct{})
-
-	// The server side
-	go func() {
-		defer close(doneCh)
-
-		s0, err := server.Accept(0)
-		if err != nil {
-			t.Fatalf("err: %s", err)
-		}
-
-		var data [1024]byte
-		for {
-			n, err := s0.Read(data[:])
-			if err == io.EOF {
-				break
-			}
-
-			dataString := string(data[0:n])
-			if dataString != "hello" {
-				t.Fatalf("bad: %#v", dataString)
-			}
-		}
-
-		s0.Close()
-	}()
-
-	s0, err := client.Dial(0)
-	if err != nil {
-		t.Fatalf("err: %s", err)
-	}
-
-	for i := 0; i < 4096*4; i++ {
-		if _, err := s0.Write([]byte("hello")); err != nil {
-			t.Fatalf("err: %s", err)
-		}
-	}
-
-	if err := s0.Close(); err != nil {
-		t.Fatalf("err: %s", err)
-	}
-
-	// Wait for the server to be done
-	<-doneCh
-}
-
-// This tests that even when the client end is closed, data can be
-// read from the server.
-func TestMuxConn_clientCloseRead(t *testing.T) {
-	client, server := testMux(t)
-	defer client.Close()
-	defer server.Close()
-
-	// This channel will be closed when we close
-	waitCh := make(chan struct{})
-
-	go func() {
-		conn, err := server.Accept(0)
-		if err != nil {
-			t.Fatalf("err: %s", err)
-		}
-
-		<-waitCh
-
-		_, err = conn.Write([]byte("foo"))
-		if err != nil {
-			t.Fatalf("err: %s", err)
-		}
-
-		conn.Close()
-	}()
-
-	s0, err := client.Dial(0)
-	if err != nil {
-		t.Fatalf("err: %s", err)
-	}
-
-	if err := s0.Close(); err != nil {
-		t.Fatalf("bad: %s", err)
-	}
-
-	// Close this to continue on on the server-side
-	close(waitCh)
-
-	var data [1024]byte
-	n, err := s0.Read(data[:])
-	if string(data[:n]) != "foo" {
-		t.Fatalf("bad: %#v", string(data[:n]))
-	}
-}
-
-func TestMuxConn_socketClose(t *testing.T) {
-	client, server := testMux(t)
-	defer client.Close()
-	defer server.Close()
-
-	go func() {
-		_, err := server.Accept(0)
-		if err != nil {
-			t.Fatalf("err: %s", err)
-		}
-
-		server.rwc.Close()
-	}()
-
-	s0, err := client.Dial(0)
-	if err != nil {
-		t.Fatalf("err: %s", err)
-	}
-
-	var data [1024]byte
-	_, err = s0.Read(data[:])
-	if err != io.EOF {
-		t.Fatalf("err: %s", err)
-	}
-}
-
-func TestMuxConn_clientClosesStreams(t *testing.T) {
-	client, server := testMux(t)
-	defer client.Close()
-	defer server.Close()
-
-	go func() {
-		conn, err := server.Accept(0)
-		if err != nil {
-			t.Fatalf("err: %s", err)
-		}
-		conn.Close()
-	}()
-
-	s0, err := client.Dial(0)
-	if err != nil {
-		t.Fatalf("err: %s", err)
-	}
-
-	var data [1024]byte
-	_, err = s0.Read(data[:])
-	if err != io.EOF {
-		t.Fatalf("err: %s", err)
-	}
-}
-
-func TestMuxConn_serverClosesStreams(t *testing.T) {
-	client, server := testMux(t)
-	defer client.Close()
-	defer server.Close()
-	go server.Accept(0)
-
-	s0, err := client.Dial(0)
-	if err != nil {
-		t.Fatalf("err: %s", err)
-	}
-
-	if err := server.Close(); err != nil {
-		t.Fatalf("err: %s", err)
-	}
-
-	// This should block forever since we never write onto this stream.
-	var data [1024]byte
-	_, err = s0.Read(data[:])
-	if err != io.EOF {
-		t.Fatalf("err: %s", err)
-	}
-}
-
-func TestMuxConnNextId(t *testing.T) {
-	client, server := testMux(t)
-	defer client.Close()
-	defer server.Close()
-
-	a := client.NextId()
-	b := client.NextId()
-
-	if a != 1 || b != 2 {
-		t.Fatalf("IDs should increment")
-	}
-
-	a = server.NextId()
-	b = server.NextId()
-
-	if a != 1 || b != 2 {
-		t.Fatalf("IDs should increment: %d %d", a, b)
-	}
-}

packer/rpc/post_processor.go

@@ -9,14 +9,14 @@ import (
 // executed over an RPC connection.
 type postProcessor struct {
 	client *rpc.Client
-	mux    *MuxConn
+	mux    *muxBroker
 }
 
 // PostProcessorServer wraps a packer.PostProcessor implementation and makes it
 // exportable as part of a Golang RPC server.
 type PostProcessorServer struct {
 	client *rpc.Client
-	mux    *MuxConn
+	mux    *muxBroker
 	p      packer.PostProcessor
 }
 

packer/rpc/provisioner.go

@@ -10,14 +10,14 @@ import (
 // executed over an RPC connection.
 type provisioner struct {
 	client *rpc.Client
-	mux    *MuxConn
+	mux    *muxBroker
 }
 
 // ProvisionerServer wraps a packer.Provisioner implementation and makes it
 // exportable as part of a Golang RPC server.
 type ProvisionerServer struct {
 	p   packer.Provisioner
-	mux *MuxConn
+	mux *muxBroker
 }
 
 type ProvisionerPrepareArgs struct {

packer/rpc/server.go

@@ -29,7 +29,7 @@ const (
 // Server represents an RPC server for Packer. This must be paired on
 // the other side with a Client.
 type Server struct {
-	mux      *MuxConn
+	mux      *muxBroker
 	streamId uint32
 	server   *rpc.Server
 	closeMux bool
@@ -37,12 +37,14 @@ type Server struct {
 
 // NewServer returns a new Packer RPC server.
 func NewServer(conn io.ReadWriteCloser) *Server {
-	result := newServerWithMux(NewMuxConn(conn), 0)
+	mux, _ := newMuxBrokerServer(conn)
+	result := newServerWithMux(mux, 0)
 	result.closeMux = true
+	go mux.Run()
 	return result
 }
 
-func newServerWithMux(mux *MuxConn, streamId uint32) *Server {
+func newServerWithMux(mux *muxBroker, streamId uint32) *Server {
 	return &Server{
 		mux:      mux,
 		streamId: streamId,
@@ -140,11 +142,11 @@ func (s *Server) Serve() {
 	// Accept a connection on stream ID 0, which is always used for
 	// normal client to server connections.
 	stream, err := s.mux.Accept(s.streamId)
-	defer stream.Close()
 	if err != nil {
 		log.Printf("[ERR] Error retrieving stream for serving: %s", err)
 		return
 	}
+	defer stream.Close()
 
 	var h codec.MsgpackHandle
 	rpcCodec := codec.GoRpc.ServerCodec(stream, &h)