Issue #3: Add close method

.travis.yml

@@ -1,6 +1,8 @@
 language: go
+sudo: required
 go:
 - 1.4
+- 1.5
 before_install:
   - go get github.com/axw/gocov/gocov
   - go get github.com/mattn/goveralls

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 2015-09-03
+
+* Add `Close` method to melody instance.
+
 ### 2015-06-10
 
 * Support for binary messages.

hub.go

@@ -5,6 +5,8 @@ type hub struct {
 	broadcast  chan *envelope
 	register   chan *Session
 	unregister chan *Session
+	exit       chan bool
+	open       bool
 }
 
 func newHub() *hub {
@@ -13,10 +15,13 @@ func newHub() *hub {
 		broadcast:  make(chan *envelope),
 		register:   make(chan *Session),
 		unregister: make(chan *Session),
+		exit:       make(chan bool),
+		open:       true,
 	}
 }
 
 func (h *hub) run() {
+loop:
 	for {
 		select {
 		case s := <-h.register:
@@ -24,8 +29,8 @@ func (h *hub) run() {
 		case s := <-h.unregister:
 			if _, ok := h.sessions[s]; ok {
 				delete(h.sessions, s)
-				close(s.output)
 				s.conn.Close()
+				close(s.output)
 			}
 		case m := <-h.broadcast:
 			for s := range h.sessions {
@@ -37,6 +42,14 @@ func (h *hub) run() {
 					go s.writeMessage(m)
 				}
 			}
+		case <-h.exit:
+			for s := range h.sessions {
+				delete(h.sessions, s)
+				s.conn.Close()
+				close(s.output)
+			}
+			h.open = false
+			break loop
 		}
 	}
 }

melody.go

@@ -88,7 +88,9 @@ func (m *Melody) HandleRequest(w http.ResponseWriter, r *http.Request) {
 
 	session.readPump(m.messageHandler, m.messageHandlerBinary, m.errorHandler)
 
-	m.hub.unregister <- session
+	if m.hub.open {
+		m.hub.unregister <- session
+	}
 
 	go m.disconnectHandler(session)
 }
@@ -111,3 +113,8 @@ func (m *Melody) BroadcastOthers(msg []byte, s *Session) {
 		return s != q
 	})
 }
+
+// Closes the melody instance and all connected sessions.
+func (m *Melody) Close() {
+	m.hub.exit <- true
+}

melody_test.go

@@ -325,3 +325,18 @@ func TestBroadcastFilter(t *testing.T) {
 		t.Errorf("should not be false")
 	}
 }
+
+func TestStop(t *testing.T) {
+	noecho := NewTestServer()
+	server := httptest.NewServer(noecho)
+	defer server.Close()
+
+	conn, err := NewDialer(server.URL)
+	defer conn.Close()
+
+	if err != nil {
+		t.Error(err)
+	}
+
+	noecho.m.Close()
+}

session.go

@@ -60,16 +60,17 @@ func (s *Session) writePump(errorHandler handleErrorFunc) {
 	ticker := time.NewTicker(s.config.PingPeriod)
 	defer ticker.Stop()
 
+loop:
 	for {
 		select {
 		case msg, ok := <-s.output:
 			if !ok {
 				s.close()
-				return
+				break loop
 			}
 			if err := s.writeRaw(msg); err != nil {
 				go errorHandler(s, err)
-				return
+				break loop
 			}
 		case <-ticker.C:
 			s.ping()