Merge pull request #27 from Nitro/relistan/fix-watch-endpoint

Convert /watch to Listener & don't leak goroutines

Author: relistan

catalog/services_state.go

@@ -230,6 +230,23 @@ func (state *ServicesState) AddListener(listener Listener) {
 	log.Debugf("AddListener(): added %s, new count %d", listener.Name(), len(state.listeners))
 }
 
+// Remove an event listener channel by name. This will find the first
+// listener in the list with the specified name and will remove it.
+func (state *ServicesState) RemoveListener(name string) error {
+	state.Lock()
+	defer state.Unlock()
+
+	for i := 0; i < len(state.listeners); i++ {
+		if state.listeners[i].Name() == name {
+			state.listeners = append(state.listeners[:i], state.listeners[i+1:]...)
+			log.Debugf("RemoveListener(): removed %s, new count %d", name, len(state.listeners))
+			return nil
+		}
+	}
+
+	return fmt.Errorf("No listener found with the name: %s", name)
+}
+
 // Take a service and merge it into our state. Correctly handle
 // timestamps so we only add things newer than what we already
 // know about. Retransmits updates to cluster peers.

catalog/services_state_test.go

@@ -461,6 +461,20 @@ func Test_Listeners(t *testing.T) {
 			So(len(state.listeners), ShouldEqual, 1)
 		})
 
+		Convey("Removing listeners results in them being removed from the list", func() {
+			state.AddListener(listener)
+			So(len(state.listeners), ShouldEqual, 1)
+
+			err := state.RemoveListener("listener1")
+			So(len(state.listeners), ShouldEqual, 0)
+			So(err, ShouldBeNil)
+		})
+
+		Convey("Removing a listener that doesn't exist returns an error", func() {
+			err := state.RemoveListener("foo")
+			So(err, ShouldNotBeNil)
+		})
+
 		Convey("A major state change event notifies all listeners", func() {
 			var result ChangeEvent
 			var result2 ChangeEvent

http.go

@@ -30,6 +30,31 @@ type ApiServices struct {
 	ClusterName    string
 }
 
+// A ServicesState.Listener that we use for the /watch endpoint
+type HttpListener struct {
+	eventChan chan catalog.ChangeEvent
+	name      string
+}
+
+func NewHttpListener() *HttpListener {
+	return &HttpListener{
+		// This should be fine enough granularity for practical purposes
+		name: fmt.Sprintf("httpListener-%d", time.Now().UTC().UnixNano()),
+		// Listeners must have buffered channels. We'll use a
+		// somewhat larger buffer here because of the slow link
+		// problem with http
+		eventChan: make(chan catalog.ChangeEvent, 50),
+	}
+}
+
+func (h *HttpListener) Chan() chan catalog.ChangeEvent {
+	return h.eventChan
+}
+
+func (h *HttpListener) Name() string {
+	return h.name
+}
+
 func makeHandler(fn func(http.ResponseWriter, *http.Request,
 	*memberlist.Memberlist, *catalog.ServicesState),
 	list *memberlist.Memberlist, state *catalog.ServicesState) http.HandlerFunc {
@@ -44,29 +69,29 @@ func watchHandler(response http.ResponseWriter, req *http.Request, list *memberl
 
 	response.Header().Set("Content-Type", "application/json")
 
-	lastChange := time.Unix(0, 0)
 	var jsonBytes []byte
 	var err error
 
-	for {
-		var changed bool
+	listener := NewHttpListener()
 
-		func() { // Wrap critical section
-			state.RLock()
-			defer state.RUnlock()
+	// Find out when the http connection closed so we can stop
+	notify := response.(http.CloseNotifier).CloseNotify()
 
-			if state.LastChanged.After(lastChange) {
-				lastChange = state.LastChanged
-				jsonBytes, err = json.Marshal(state.ByService())
-				if err != nil {
-					log.Errorf("Error marshaling state in watchHandler: %s", err.Error())
-					return
-				}
-				changed = true // Trigger sending new encoding
-			}
-		}()
+	// Let's subscribe to state change events
+	state.AddListener(listener)
+	defer state.RemoveListener(listener.Name())
 
-		if changed {
+	for {
+		select {
+		case <-notify:
+			break
+
+		case <-listener.Chan():
+			jsonBytes, err = json.Marshal(state.ByService())
+			if err != nil {
+				log.Errorf("Error marshaling state in watchHandler: %s", err.Error())
+				return
+			}
 			// In order to flush immediately, we have to cast to a Flusher.
 			// The normal HTTP library supports this but not all do, so we
 			// check just in case.
@@ -75,7 +100,6 @@ func watchHandler(response http.ResponseWriter, req *http.Request, list *memberl
 				f.Flush()
 			}
 		}
-		time.Sleep(250 * time.Millisecond)
 	}
 }