Fix concurrent panic

melody.go

@@ -178,6 +178,7 @@ func (m *Melody) HandleRequestWithKeys(w http.ResponseWriter, r *http.Request, k
 		Keys:       keys,
 		conn:       conn,
 		output:     make(chan *envelope, m.Config.MessageBufferSize),
+		outputDone: make(chan struct{}),
 		melody:     m,
 		open:       true,
 		rwmutex:    &sync.RWMutex{},

session.go

@@ -15,6 +15,7 @@ type Session struct {
 	Keys       map[string]interface{}
 	conn       *websocket.Conn
 	output     chan *envelope
+	outputDone chan struct{}
 	melody     *Melody
 	open       bool
 	rwmutex    *sync.RWMutex
@@ -56,12 +57,13 @@ func (s *Session) closed() bool {
 }
 
 func (s *Session) close() {
-	if !s.closed() {
 	s.rwmutex.Lock()
+	open := s.open
 	s.open = false
-		s.conn.Close()
-		close(s.output)
 	s.rwmutex.Unlock()
+	if open {
+		s.conn.Close()
+		close(s.outputDone)
 	}
 }
 
@@ -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
+			}
 		}
 	}
 }