Fix concurrent panic

melody.go

@@ -174,13 +174,14 @@ func (m *Melody) HandleRequestWithKeys(w http.ResponseWriter, r *http.Request, k
 	}
 
 	session := &Session{
-		Request: r,
+		Request:    r,
-		Keys:    keys,
+		Keys:       keys,
-		conn:    conn,
+		conn:       conn,
-		output:  make(chan *envelope, m.Config.MessageBufferSize),
+		output:     make(chan *envelope, m.Config.MessageBufferSize),
-		melody:  m,
+		outputDone: make(chan struct{}),
-		open:    true,
+		melody:     m,
-		rwmutex: &sync.RWMutex{},
+		open:       true,
+		rwmutex:    &sync.RWMutex{},
 	}
 
 	m.hub.register <- session

session.go

@@ -11,13 +11,14 @@ import (
 
 // Session wrapper around websocket connections.
 type Session struct {
-	Request *http.Request
+	Request    *http.Request
-	Keys    map[string]interface{}
+	Keys       map[string]interface{}
-	conn    *websocket.Conn
+	conn       *websocket.Conn
-	output  chan *envelope
+	output     chan *envelope
-	melody  *Melody
+	outputDone chan struct{}
-	open    bool
+	melody     *Melody
-	rwmutex *sync.RWMutex
+	open       bool
+	rwmutex    *sync.RWMutex
 }
 
 func (s *Session) writeMessage(message *envelope) {
@@ -56,12 +57,13 @@ func (s *Session) closed() bool {
 }
 
 func (s *Session) close() {
-	if !s.closed() {
+	s.rwmutex.Lock()
-		s.rwmutex.Lock()
+	open := s.open
-		s.open = false
+	s.open = false
+	s.rwmutex.Unlock()
+	if open {
 		s.conn.Close()
-		close(s.output)
+		close(s.outputDone)
-		s.rwmutex.Unlock()
 	}
 }
 
@@ -76,11 +78,7 @@ func (s *Session) writePump() {
 loop:
 	for {
 		select {
-		case msg, ok := <-s.output:
+		case msg := <-s.output:
-			if !ok {
-				break loop
-			}
-
 			err := s.writeRaw(msg)
 
 			if err != nil {
@@ -101,6 +99,10 @@ loop:
 			}
 		case <-ticker.C:
 			s.ping()
+		case _, ok := <-s.outputDone:
+			if !ok {
+				break loop
+			}
 		}
 	}
 }