From 425dbc4797f153666f3e3cb5d21a3573f7335d71 Mon Sep 17 00:00:00 2001
From: Qian Wang <cbeuw.andy@gmail.com>
Date: Sat, 23 Mar 2019 23:45:12 +1100
Subject: [PATCH] Session times out 30s after all streams have been closed

---
 internal/multiplex/session.go | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/internal/multiplex/session.go b/internal/multiplex/session.go
index 0a37283..c56fd04 100644
--- a/internal/multiplex/session.go
+++ b/internal/multiplex/session.go
@@ -6,6 +6,7 @@ import (
 	"net"
 	"sync"
 	"sync/atomic"
+	"time"
 )
 
 const (
@@ -28,7 +29,7 @@ type Session struct {
 	// atomic
 	nextStreamID uint32
 
-	streamsM sync.RWMutex
+	streamsM sync.Mutex
 	streams  map[uint32]*Stream
 
 	// Switchboard manages all connections to remote
@@ -90,6 +91,9 @@ func (sesh *Session) AcceptStream() (*Stream, error) {
 func (sesh *Session) delStream(id uint32) {
 	sesh.streamsM.Lock()
 	delete(sesh.streams, id)
+	if len(sesh.streams) == 0 {
+		go sesh.timeoutAfter(30 * time.Second)
+	}
 	sesh.streamsM.Unlock()
 }
 
@@ -142,3 +146,14 @@ func (sesh *Session) Close() error {
 func (sesh *Session) IsBroken() bool {
 	return atomic.LoadUint32(&sesh.broken) == 1
 }
+
+func (sesh *Session) timeoutAfter(to time.Duration) {
+	time.Sleep(to)
+	sesh.streamsM.Lock()
+	if len(sesh.streams) == 0 && !sesh.IsBroken() {
+		sesh.streamsM.Unlock()
+		sesh.Close()
+	} else {
+		sesh.streamsM.Unlock()
+	}
+}