From 2f17841f8579bd5151127aa1839c3e6262c61f18 Mon Sep 17 00:00:00 2001
From: Andy Wang <cbeuw.andy@gmail.com>
Date: Fri, 25 Dec 2020 23:16:57 +0000
Subject: [PATCH] Use Compare-And-Swap for atomic booleans indicating session
 and switchboard closed

---
 internal/multiplex/session.go     | 29 ++++++++++++++---------------
 internal/multiplex/switchboard.go | 17 +++++++----------
 2 files changed, 21 insertions(+), 25 deletions(-)

diff --git a/internal/multiplex/session.go b/internal/multiplex/session.go
index 159172e..0113afa 100644
--- a/internal/multiplex/session.go
+++ b/internal/multiplex/session.go
@@ -66,6 +66,8 @@ type Session struct {
 
 	streamsM sync.Mutex
 	streams  map[uint32]*Stream
+	// For accepting new streams
+	acceptCh chan *Stream
 
 	// a pool of heap allocated frame objects so we don't have to allocate a new one each time we receive a frame
 	recvFramePool sync.Pool
@@ -78,9 +80,6 @@ type Session struct {
 	// Used for LocalAddr() and RemoteAddr() etc.
 	addrs atomic.Value
 
-	// For accepting new streams
-	acceptCh chan *Stream
-
 	closed uint32
 
 	terminalMsg atomic.Value
@@ -181,7 +180,7 @@ func (sesh *Session) Accept() (net.Conn, error) {
 }
 
 func (sesh *Session) closeStream(s *Stream, active bool) error {
-	if atomic.SwapUint32(&s.closed, 1) == 1 {
+	if !atomic.CompareAndSwapUint32(&s.closed, 0, 1) {
 		return fmt.Errorf("closing stream %v: %w", s.id, errRepeatStreamClosing)
 	}
 	_ = s.getRecvBuf().Close() // recvBuf.Close should not return error
@@ -244,6 +243,10 @@ func (sesh *Session) recvDataFromRemote(data []byte) error {
 	}
 
 	sesh.streamsM.Lock()
+	if sesh.IsClosed() {
+		sesh.streamsM.Unlock()
+		return ErrBrokenSession
+	}
 	existingStream, existing := sesh.streams[frame.StreamID]
 	if existing {
 		sesh.streamsM.Unlock()
@@ -255,10 +258,10 @@ func (sesh *Session) recvDataFromRemote(data []byte) error {
 	} else {
 		newStream := makeStream(sesh, frame.StreamID)
 		sesh.streams[frame.StreamID] = newStream
+		sesh.acceptCh <- newStream
 		sesh.streamsM.Unlock()
 		// new stream
 		sesh.streamCountIncr()
-		sesh.acceptCh <- newStream
 		return newStream.recvFrame(frame)
 	}
 }
@@ -276,14 +279,14 @@ func (sesh *Session) TerminalMsg() string {
 	}
 }
 
-func (sesh *Session) closeSession(closeSwitchboard bool) error {
-	if atomic.SwapUint32(&sesh.closed, 1) == 1 {
+func (sesh *Session) closeSession() error {
+	if !atomic.CompareAndSwapUint32(&sesh.closed, 0, 1) {
 		log.Debugf("session %v has already been closed", sesh.id)
 		return errRepeatSessionClosing
 	}
-	sesh.acceptCh <- nil
 
 	sesh.streamsM.Lock()
+	close(sesh.acceptCh)
 	for id, stream := range sesh.streams {
 		if stream == nil {
 			continue
@@ -294,26 +297,23 @@ func (sesh *Session) closeSession(closeSwitchboard bool) error {
 		sesh.streamCountDecr()
 	}
 	sesh.streamsM.Unlock()
-
-	if closeSwitchboard {
-		sesh.sb.closeAll()
-	}
 	return nil
 }
 
 func (sesh *Session) passiveClose() error {
 	log.Debugf("attempting to passively close session %v", sesh.id)
-	err := sesh.closeSession(true)
+	err := sesh.closeSession()
 	if err != nil {
 		return err
 	}
+	sesh.sb.closeAll()
 	log.Debugf("session %v closed gracefully", sesh.id)
 	return nil
 }
 
 func (sesh *Session) Close() error {
 	log.Debugf("attempting to actively close session %v", sesh.id)
-	err := sesh.closeSession(false)
+	err := sesh.closeSession()
 	if err != nil {
 		return err
 	}
@@ -339,7 +339,6 @@ func (sesh *Session) Close() error {
 	if err != nil {
 		return err
 	}
-
 	sesh.sb.closeAll()
 	log.Debugf("session %v closed gracefully", sesh.id)
 	return nil
diff --git a/internal/multiplex/switchboard.go b/internal/multiplex/switchboard.go
index 7d2c93c..ea21308 100644
--- a/internal/multiplex/switchboard.go
+++ b/internal/multiplex/switchboard.go
@@ -71,7 +71,8 @@ func (sb *switchboard) send(data []byte, connId *uint32) (n int, err error) {
 		n, err = conn.Write(d)
 		if err != nil {
 			sb.conns.Delete(*connId)
-			sb.close("failed to write to remote " + err.Error())
+			sb.session.SetTerminalMsg("failed to write to remote " + err.Error())
+			sb.session.passiveClose()
 			return n, err
 		}
 		sb.valve.AddTx(int64(n))
@@ -138,16 +139,11 @@ func (sb *switchboard) pickRandConn() (uint32, net.Conn, error) {
 	return id, conn, nil
 }
 
-func (sb *switchboard) close(terminalMsg string) {
-	atomic.StoreUint32(&sb.broken, 1)
-	if !sb.session.IsClosed() {
-		sb.session.SetTerminalMsg(terminalMsg)
-		sb.session.passiveClose()
-	}
-}
-
 // actively triggered by session.Close()
 func (sb *switchboard) closeAll() {
+	if !atomic.CompareAndSwapUint32(&sb.broken, 0, 1) {
+		return
+	}
 	sb.conns.Range(func(key, connI interface{}) bool {
 		conn := connI.(net.Conn)
 		conn.Close()
@@ -168,7 +164,8 @@ func (sb *switchboard) deplex(connId uint32, conn net.Conn) {
 			log.Debugf("a connection for session %v has closed: %v", sb.session.id, err)
 			sb.conns.Delete(connId)
 			atomic.AddUint32(&sb.numConns, ^uint32(0))
-			sb.close("a connection has dropped unexpectedly")
+			sb.session.SetTerminalMsg("a connection has dropped unexpectedly")
+			sb.session.passiveClose()
 			return
 		}