Fix lingering transports after dmsg.Server disconnects.

skycoin · Jun 26, 2019 · f724286 · f724286
1 parent e68e923
commit f724286
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 13 deletions.
diff --git a/go.sum b/go.sum
@@ -118,6 +118,8 @@ golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223/go.mod h1:STP8DvDyc/dI5b8T5h
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20190422165155-953cdadca894 h1:Cz4ceDQGXuKRnVBDTS23GTn/pU5OE2C0WrNTOYK1Uuc=
 golang.org/x/sys v0.0.0-20190422165155-953cdadca894/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20190624142023-c5567b49c5d0 h1:HyfiK1WMnHj5FXFXatD+Qs1A/xC2Run6RzeW1SyHxpc=
+golang.org/x/sys v0.0.0-20190624142023-c5567b49c5d0/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/text v0.3.0 h1:g61tztE5qeGQ89tm6NTjjM9VPIm088od1l6aSorWRWg=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.2 h1:tW2bmiBqwgJj/UpqtC8EpXEZVYOwU0yG4iWbprSVAcs=

diff --git a/pkg/dmsg/client.go b/pkg/dmsg/client.go
@@ -183,14 +183,11 @@ func (c *ClientConn) Serve(ctx context.Context, accept chan<- *Transport) (err e
 	log := c.log.WithField("remoteServer", c.remoteSrv)
 	log.WithField("connCount", incrementServeCount()).Infoln("ServingConn")
 	defer func() {
+		c.close(true)
 		log.WithError(err).WithField("connCount", decrementServeCount()).Infoln("ConnectionClosed")
 		c.wg.Done()
 	}()
 
-	closeConn := func(log *logrus.Entry) {
-		log.WithError(c.Close()).Warn("ClosingConnection")
-	}
-
 	for {
 		f, err := readFrame(c.Conn)
 		if err != nil {
@@ -224,7 +221,8 @@ func (c *ClientConn) Serve(ctx context.Context, accept chan<- *Transport) (err e
 				if err != nil {
 					log.WithField("remoteClient", initPK).WithError(err).Infoln("Rejected [REQUEST]")
 					if isWriteError(err) || err == ErrClientClosed {
-						closeConn(log)
+						err := c.Close()
+						log.WithError(err).Warn("ClosingConnection")
 					}
 					return
 				}
@@ -258,21 +256,37 @@ func (c *ClientConn) DialTransport(ctx context.Context, clientPK cipher.PubKey)
 	return tp, nil
 }
 
-// Close closes the connection to dms_server.
-func (c *ClientConn) Close() error {
+// If 'isSrvDisconnect' is set, we don't need to send CLOSE frames or close the underlying TCP connection.
+func (c *ClientConn) close(isSrvDisconnect bool) (closed bool) {
 	c.once.Do(func() {
+		closed = true
 		c.log.WithField("remoteServer", c.remoteSrv).Infoln("ClosingConnection")
 		close(c.done)
 		c.mx.Lock()
-		for _, tp := range c.tps {
-			if tp != nil {
-				go tp.Close() //nolint:errcheck
+		if isSrvDisconnect {
+			for _, tp := range c.tps {
+				if tp != nil {
+					go tp.close()
+				}
+			}
+		} else {
+			for _, tp := range c.tps {
+				if tp != nil {
+					go tp.Close() //nolint:errcheck
+				}
 			}
+			_ = c.Conn.Close() //nolint:errcheck
 		}
-		_ = c.Conn.Close() //nolint:errcheck
 		c.mx.Unlock()
-		c.wg.Wait()
 	})
+	return closed
+}
+
+// Close closes the connection to dms_server.
+func (c *ClientConn) Close() error {
+	if c.close(false) {
+		c.wg.Wait()
+	}
 	return nil
 }
 

diff --git a/pkg/dmsg/server_test.go b/pkg/dmsg/server_test.go
@@ -14,11 +14,11 @@ import (
 	"testing"
 	"time"
 
+	"github.com/skycoin/skycoin/src/util/logging"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/net/nettest"
 
-	"github.com/skycoin/skycoin/src/util/logging"
 	"github.com/skycoin/skywire/internal/noise"
 	"github.com/skycoin/skywire/pkg/cipher"
 	"github.com/skycoin/skywire/pkg/messaging-discovery/client"