channel: add serial yamux channel close timeout

bergwolf · bergwolf · commit 1d559a7cf00a · 2018-09-05T22:53:29.000+08:00
So that if serial yamux session fails to close somehow, we can catch it. Fixes: #359 Signed-off-by: Peng Tao <bergwolf@gmail.com>
diff --git a/channel.go b/channel.go
@@ -25,6 +25,7 @@ import (
 var (
 	channelExistMaxTries   = 200
 	channelExistWaitTime   = 50 * time.Millisecond
+	channelCloseTimeout    = 5 * time.Second
 	isAFVSockSupportedFunc = isAFVSockSupported
 )
 
@@ -209,7 +210,13 @@ func (c *serialChannel) listen() (net.Listener, error) {
 func (c *serialChannel) teardown() error {
 	// wait for the session to be fully shutdown first
 	if c.waitCh != nil {
-		<-c.waitCh
+		t := time.NewTimer(channelCloseTimeout)
+		select {
+		case <-c.waitCh:
+			t.Stop()
+		case <-t.C:
+			return fmt.Errorf("timeout waiting for yamux channel to close")
+		}
 	}
 	return c.serialConn.Close()
 }
diff --git a/channel_test.go b/channel_test.go
@@ -12,6 +12,7 @@ import (
 	"os"
 	"path/filepath"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -75,6 +76,20 @@ func TestTeardownSerialChannel(t *testing.T) {
 	assert.Nil(t, err, "%v", err)
 }
 
+func TestTeardownSerialChannelTimeout(t *testing.T) {
+	_, f, err := os.Pipe()
+	assert.Nil(t, err, "%v", err)
+	channelCloseTimeout = 1 * time.Microsecond
+
+	c := &serialChannel{
+		serialConn: f,
+		waitCh:     make(chan struct{}),
+	}
+
+	err = c.teardown()
+	assert.NotNil(t, err, "channel close should timeout")
+}
+
 func TestNewChannel(t *testing.T) {
 	assert := assert.New(t)
 
