Merge pull request #1456 from mitchellh/f-yamux

packer/rpc: Replace MuxConn with Yamux

packer/plugin/server.go

@@ -33,7 +33,7 @@ const MagicCookieValue = "d602bf8f470bc67ca7faa0386276bbdd4330efaf76d1a219cb4d69
 // The APIVersion is outputted along with the RPC address. The plugin
 // client validates this API version and will show an error if it doesn't
 // know how to speak it.
-const APIVersion = "3"
+const APIVersion = "4"
 
 // Server waits for a connection to this plugin and returns a Packer
 // RPC server that you can use to register components and serve them.

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

+package rpc
+
+import (
+	"encoding/binary"
+	"fmt"
+	"io"
+	"net"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/hashicorp/yamux"
+)
+
+// muxBroker is responsible for brokering multiplexed connections by unique ID.
+//
+// This allows a plugin to request a channel with a specific ID to connect to
+// 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
+
+	sync.Mutex
+}
+
+type muxBrokerPending struct {
+	ch     chan net.Conn
+	doneCh chan struct{}
+}
+
+func newMuxBroker(s *yamux.Session) *muxBroker {
+	return &muxBroker{
+		session: s,
+		streams: make(map[uint32]*muxBrokerPending),
+	}
+}
+
+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.
+func (m *muxBroker) Accept(id uint32) (net.Conn, error) {
+	var c net.Conn
+	p := m.getStream(id)
+	select {
+	case c = <-p.ch:
+		close(p.doneCh)
+	case <-time.After(5 * time.Second):
+		m.Lock()
+		defer m.Unlock()
+		delete(m.streams, id)
+
+		return nil, fmt.Errorf("timeout waiting for accept")
+	}
+
+	// Ack our connection
+	if err := binary.Write(c, binary.LittleEndian, id); err != nil {
+		c.Close()
+		return nil, err
+	}
+
+	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
+	stream, err := m.session.OpenStream()
+	if err != nil {
+		return nil, err
+	}
+
+	// Write the stream ID onto the wire.
+	if err := binary.Write(stream, binary.LittleEndian, id); err != nil {
+		stream.Close()
+		return nil, err
+	}
+
+	// Read the ack that we connected. Then we're off!
+	var ack uint32
+	if err := binary.Read(stream, binary.LittleEndian, &ack); err != nil {
+		stream.Close()
+		return nil, err
+	}
+	if ack != id {
+		stream.Close()
+		return nil, fmt.Errorf("bad ack: %d (expected %d)", ack, id)
+	}
+
+	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() {
+	for {
+		stream, err := m.session.AcceptStream()
+		if err != nil {
+			// Once we receive an error, just exit
+			break
+		}
+
+		// Read the stream ID from the stream
+		var id uint32
+		if err := binary.Read(stream, binary.LittleEndian, &id); err != nil {
+			stream.Close()
+			continue
+		}
+
+		// Initialize the waiter
+		p := m.getStream(id)
+		select {
+		case p.ch <- stream:
+		default:
+		}
+
+		// Wait for a timeout
+		go m.timeoutWait(id, p)
+	}
+}
+
+func (m *muxBroker) getStream(id uint32) *muxBrokerPending {
+	m.Lock()
+	defer m.Unlock()
+
+	p, ok := m.streams[id]
+	if ok {
+		return p
+	}
+
+	m.streams[id] = &muxBrokerPending{
+		ch:     make(chan net.Conn, 1),
+		doneCh: make(chan struct{}),
+	}
+	return m.streams[id]
+}
+
+func (m *muxBroker) timeoutWait(id uint32, p *muxBrokerPending) {
+	// Wait for the stream to either be picked up and connected, or
+	// for a timeout.
+	timeout := false
+	select {
+	case <-p.doneCh:
+	case <-time.After(5 * time.Second):
+		timeout = true
+	}
+
+	m.Lock()
+	defer m.Unlock()
+
+	// Delete the stream so no one else can grab it
+	delete(m.streams, id)
+
+	// If we timed out, then check if we have a channel in the buffer,
+	// and if so, close it.
+	if timeout {
+		select {
+		case s := <-p.ch:
+			s.Close()
+		}
+	}
+}

packer/rpc/mux_broker_test.go

+package rpc
+
+import (
+	"net"
+	"testing"
+
+	"github.com/hashicorp/yamux"
+)
+
+func TestMuxBroker(t *testing.T) {
+	c, s := testYamux(t)
+	defer c.Close()
+	defer s.Close()
+
+	bc := newMuxBroker(c)
+	bs := newMuxBroker(s)
+	go bc.Run()
+	go bs.Run()
+
+	go func() {
+		c, err := bc.Dial(5)
+		if err != nil {
+			t.Fatalf("err: %s", err)
+		}
+
+		if _, err := c.Write([]byte{42}); err != nil {
+			t.Fatalf("err: %s", err)
+		}
+	}()
+
+	client, err := bs.Accept(5)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	var data [1]byte
+	if _, err := client.Read(data[:]); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	if data[0] != 42 {
+		t.Fatalf("bad: %d", data[0])
+	}
+}
+
+func testYamux(t *testing.T) (client *yamux.Session, server *yamux.Session) {
+	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, err = yamux.Server(conn, nil)
+		if err != nil {
+			t.Fatalf("err: %s", err)
+		}
+	}()
+
+	// Client side
+	conn, err := net.Dial("tcp", l.Addr().String())
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	client, err = yamux.Client(conn, nil)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	// Wait for the server
+	<-doneCh
+
+	return
+}

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)