From ff425ac17583fc7395705b19359460abae373f15 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ola=20Holmstr=C3=B6m?= <ola.holmstrom@gmail.com>
Date: Wed, 2 Sep 2015 21:40:26 +0200
Subject: [PATCH] Issue #3: Add close method

---
 .travis.yml    |  2 ++
 CHANGELOG.md   |  4 ++++
 hub.go         | 15 ++++++++++++++-
 melody.go      |  9 ++++++++-
 melody_test.go | 15 +++++++++++++++
 session.go     |  5 +++--
 6 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index e638e9e..c7dde77 100644
--- a/.travis.yml
+++ b/.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
diff --git a/CHANGELOG.md b/CHANGELOG.md
index dbbaaff..45c695a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 2015-09-03
+
+* Add `Close` method to melody instance.
+
 ### 2015-06-10
 
 * Support for binary messages.
diff --git a/hub.go b/hub.go
index 2918fa6..63ed440 100644
--- a/hub.go
+++ b/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
 		}
 	}
 }
diff --git a/melody.go b/melody.go
index 1ec4778..61ef4fd 100644
--- a/melody.go
+++ b/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
+}
diff --git a/melody_test.go b/melody_test.go
index 48c9727..e03f750 100644
--- a/melody_test.go
+++ b/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()
+}
diff --git a/session.go b/session.go
index 72bf9f0..e84b603 100644
--- a/session.go
+++ b/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()