Close session if all conn are closed

Signed-off-by: Yilun <[email protected]>

client.go

@@ -42,6 +42,7 @@ type TunaSessionClient struct {
 	sessions     map[string]*ncp.Session
 	sessionConns map[string]map[string]net.Conn
 	sharedKeys   map[string]*[sharedKeySize]byte
+	connCount    map[string]int
 	isClosed     bool
 }
 
@@ -62,6 +63,7 @@ func NewTunaSessionClient(clientAccount *nkn.Account, m *nkn.MultiClient, wallet
 		sessions:      make(map[string]*ncp.Session),
 		sessionConns:  make(map[string]map[string]net.Conn),
 		sharedKeys:    make(map[string]*[sharedKeySize]byte),
+		connCount:     make(map[string]int),
 	}
 
 	return c, nil
@@ -302,7 +304,23 @@ func (c *TunaSessionClient) listenNet(i int) {
 				}
 			}
 
+			c.Lock()
+			c.connCount[sessionKey]++
+			c.Unlock()
+
 			c.handleConn(conn, sess, i)
+
+			c.Lock()
+			c.connCount[sessionKey]--
+			shouldClose := c.connCount[sessionKey] == 0
+			if shouldClose {
+				delete(c.connCount, sessionKey)
+			}
+			c.Unlock()
+
+			if shouldClose {
+				sess.Close()
+			}
 		}(conn)
 	}
 }
@@ -460,8 +478,24 @@ func (c *TunaSessionClient) DialWithConfig(remoteAddr string, config *DialConfig
 	for i := 0; i < len(pubAddrs.Addrs); i++ {
 		if conn, ok := conns[connID(i)]; ok {
 			go func(conn net.Conn, i int) {
-				defer conn.Close()
+				c.Lock()
+				c.connCount[sessionKey]++
+				c.Unlock()
+
 				c.handleConn(conn, sess, i)
+				conn.Close()
+
+				c.Lock()
+				c.connCount[sessionKey]--
+				shouldClose := c.connCount[sessionKey] == 0
+				if shouldClose {
+					delete(c.connCount, sessionKey)
+				}
+				c.Unlock()
+
+				if shouldClose {
+					sess.Close()
+				}
 			}(conn, i)
 		}
 	}
@@ -566,7 +600,7 @@ func (c *TunaSessionClient) newSession(remoteAddr string, sessionID []byte, conn
 		if writeTimeout > 0 {
 			err := conn.SetWriteDeadline(time.Now().Add(writeTimeout))
 			if err != nil {
-				return err
+				return ncp.ErrConnClosed
 			}
 		}
 		buf, err := c.encode(buf, remoteAddr)
@@ -580,7 +614,7 @@ func (c *TunaSessionClient) newSession(remoteAddr string, sessionID []byte, conn
 		if writeTimeout > 0 {
 			err = conn.SetWriteDeadline(zeroTime)
 			if err != nil {
-				return err
+				return ncp.ErrConnClosed
 			}
 		}
 		return nil

go.mod

@@ -4,9 +4,9 @@ go 1.13
 
 require (
 	github.com/imdario/mergo v0.3.9
-	github.com/nknorg/ncp-go v1.0.1
+	github.com/nknorg/ncp-go v1.0.2
 	github.com/nknorg/nkn-sdk-go v1.3.0
 	github.com/nknorg/nkn/v2 v2.0.2
-	github.com/nknorg/tuna v0.0.0-20200729232930-42c0c6547187
+	github.com/nknorg/tuna v0.0.0-20200803091628-f86c69d40928
 	golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9
 )

go.sum

@@ -74,8 +74,8 @@ github.com/nknorg/go-nat v1.0.1/go.mod h1:dblX1Ac2j08rTUGs5CKCAfjHGN5eDFhbeqt2rc
 github.com/nknorg/go-portscanner v0.0.0-20181002101859-8493ef01db79 h1:AT5YRim0696HFnGSvFg0uOmSuo1X11roGBQhKB25PHE=
 github.com/nknorg/go-portscanner v0.0.0-20181002101859-8493ef01db79/go.mod h1:lTev9/EZpc4U3w/sp95T8UPeaav/0Hk+JjykVDuPaQs=
 github.com/nknorg/ncp-go v1.0.0/go.mod h1:LfiwrFlKZAAoMwhCBA959jSzVD8E9RGOH9nyfXkQSfY=
-github.com/nknorg/ncp-go v1.0.1 h1:oJKmcaAcVC1oCJiJ4tDvot7kved8uUGBYXldDHxc8vM=
-github.com/nknorg/ncp-go v1.0.1/go.mod h1:LfiwrFlKZAAoMwhCBA959jSzVD8E9RGOH9nyfXkQSfY=
+github.com/nknorg/ncp-go v1.0.2 h1:nZRdfXtbrVEBmhxCfEMu20XZIzk8Wztx4OE47drzIWw=
+github.com/nknorg/ncp-go v1.0.2/go.mod h1:LfiwrFlKZAAoMwhCBA959jSzVD8E9RGOH9nyfXkQSfY=
 github.com/nknorg/nkn-sdk-go v1.3.0 h1:iT/MqUx7VyqVo3ivz3VIBbB+VvaByld0YrQiVIlfdCI=
 github.com/nknorg/nkn-sdk-go v1.3.0/go.mod h1:SeHZ25YcKIXDauhPw/v4XZjprWd+DtLDPNs3/U9/5qE=
 github.com/nknorg/nkn/v2 v2.0.0-20200703083743-f83cde273fe2/go.mod h1:LxKZTpn5VZTAbjsACSFL5nlQAQ4+WQKyIpVZUYopXm4=
@@ -85,8 +85,8 @@ github.com/nknorg/nnet v0.0.0-20200521002812-357d1b11179f h1:JvDKr6D9LdnsdScND+8
 github.com/nknorg/nnet v0.0.0-20200521002812-357d1b11179f/go.mod h1:4DHEQEMhlRGIKGSyhATdjeusdqaHafDatadtpeHBpvI=
 github.com/nknorg/portmapper v0.0.0-20200114081049-1c03cdccc283 h1:uS3/DvxCbi0zZau66ggQAgEjyGmql2mj77UQFgumq1I=
 github.com/nknorg/portmapper v0.0.0-20200114081049-1c03cdccc283/go.mod h1:dL4PQJ4670oTO6LqvkjrBQEkD+iMiOYjlKRBBw55Csg=
-github.com/nknorg/tuna v0.0.0-20200729232930-42c0c6547187 h1:fQyuj5rkyQr8OzG5HdmgAhH+/fq7I3XU7c9vMe89Z6k=
-github.com/nknorg/tuna v0.0.0-20200729232930-42c0c6547187/go.mod h1:GhSzMlgyFy9SvT9TCc/UUEGmMiLxPkyeiWaiM0IWT1U=
+github.com/nknorg/tuna v0.0.0-20200803091628-f86c69d40928 h1:bO40M16Zu4pFxjIyc2PG4oTfu2IUeZcWFVEST1VYizA=
+github.com/nknorg/tuna v0.0.0-20200803091628-f86c69d40928/go.mod h1:GhSzMlgyFy9SvT9TCc/UUEGmMiLxPkyeiWaiM0IWT1U=
 github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
 github.com/onsi/ginkgo v1.7.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
 github.com/onsi/gomega v1.4.3/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY=