Merge pull request #338 from bergwolf/yamux

wait on yamux session close

Gopkg.toml

@@ -4,7 +4,7 @@
 
 [[constraint]]
   name = "github.com/hashicorp/yamux"
-  revision = "f5742cb6b85602e7fa834e9d5d91a7d7fa850824"
+  revision = "cc6d2ea263b2471faabce371255777a365bf8306"
 
 [[constraint]]
   name = "github.com/mdlayher/vsock"

channel.go

@@ -103,6 +103,7 @@ func (c *vSockChannel) teardown() error {
 type serialChannel struct {
 	serialPath string
 	serialConn *os.File
+	waitCh     <-chan struct{}
 }
 
 func (c *serialChannel) setup() error {
@@ -200,11 +201,16 @@ func (c *serialChannel) listen() (net.Listener, error) {
 	if err != nil {
 		return nil, err
 	}
+	c.waitCh = session.CloseChan()
 
 	return session, nil
 }
 
 func (c *serialChannel) teardown() error {
+	// wait for the session to be fully shutdown first
+	if c.waitCh != nil {
+		<-c.waitCh
+	}
 	return c.serialConn.Close()
 }
 

channel_test.go

@@ -49,13 +49,20 @@ func TestWaitSerialChannel(t *testing.T) {
 }
 
 func TestListenSerialChannel(t *testing.T) {
-	_, f, err := os.Pipe()
+	f, _, err := os.Pipe()
 	assert.Nil(t, err, "%v", err)
 
 	c := &serialChannel{serialConn: f}
 
-	_, err = c.listen()
+	l, err := c.listen()
+	assert.Nil(t, err, "%v", err)
+	assert.NotNil(t, l, "listen should not return nil listener")
+
+	err = l.Close()
 	assert.Nil(t, err, "%v", err)
+
+	err = c.teardown()
+	assert.Error(t, err, "connection should be already closed")
 }
 
 func TestTeardownSerialChannel(t *testing.T) {

protocols/client/client.go

@@ -8,6 +8,7 @@ package client
 
 import (
 	"context"
+	"fmt"
 	"net"
 	"net/url"
 	"strconv"
@@ -31,6 +32,7 @@ const (
 )
 
 var defaultDialTimeout = 15 * time.Second
+var defaultCloseTimeout = 5 * time.Second
 
 // AgentClient is an agent gRPC client connection wrapper for agentgrpc.AgentServiceClient
 type AgentClient struct {
@@ -45,7 +47,26 @@ type yamuxSessionStream struct {
 }
 
 func (y *yamuxSessionStream) Close() error {
-	return y.session.Close()
+	waitCh := y.session.CloseChan()
+	timeout := time.NewTimer(defaultCloseTimeout)
+
+	if err := y.Conn.Close(); err != nil {
+		return err
+	}
+
+	if err := y.session.Close(); err != nil {
+		return err
+	}
+
+	// block until session is really closed
+	select {
+	case <-waitCh:
+		timeout.Stop()
+	case <-timeout.C:
+		return fmt.Errorf("timeout waiting for session close")
+	}
+
+	return nil
 }
 
 type dialer func(string, time.Duration) (net.Conn, error)