Moving listener count to APIEventHandler

Signed-off-by: Nishant Totla <nishanttotla@gmail.com>

Author: nishanttotla

api/handlers.go

@@ -13,7 +13,6 @@ import (
 	"sort"
 	"strconv"
 	"strings"
-	"sync/atomic"
 	"time"
 
 	log "github.com/Sirupsen/logrus"
@@ -42,7 +41,7 @@ var (
 func getInfo(c *context, w http.ResponseWriter, r *http.Request) {
 	info := apitypes.Info{
 		Images:            len(c.cluster.Images().Filter(cluster.ImageFilterOptions{})),
-		NEventsListener:   int(atomic.LoadUint64(c.listenerCount)),
+		NEventsListener:   c.eventsHandler.Size(),
 		Debug:             c.debug,
 		MemoryLimit:       true,
 		SwapLimit:         true,
@@ -893,9 +892,6 @@ func getEvents(c *context, w http.ResponseWriter, r *http.Request) {
 	eventsChan, cancelFunc := c.eventsHandler.Watch()
 	defer cancelFunc()
 
-	atomic.AddUint64(c.listenerCount, 1)
-	defer atomic.AddUint64(c.listenerCount, ^uint64(0))
-
 	// create timer for --until
 	var (
 		timer   *time.Timer

api/primary.go

@@ -15,7 +15,6 @@ import (
 type context struct {
 	cluster       cluster.Cluster
 	eventsHandler *cluster.APIEventHandler
-	listenerCount *uint64
 	statusHandler StatusHandler
 	debug         bool
 	tlsConfig     *tls.Config
@@ -121,14 +120,11 @@ func NewPrimary(cluster cluster.Cluster, tlsConfig *tls.Config, status StatusHan
 
 	// eventsHandler is the handler for API events
 	eventsHandler := cluster.NewAPIEventHandler()
-	// listenerCount keeps track of the number of API events connections
-	listenerCount := uint64(0)
 	cluster.RegisterEventHandler(eventsHandler)
 
 	context := &context{
 		cluster:       cluster,
 		eventsHandler: eventsHandler,
-		listenerCount: &listenerCount,
 		statusHandler: status,
 		tlsConfig:     tlsConfig,
 	}

cluster/api_events.go

@@ -1,6 +1,7 @@
 package cluster
 
 import (
+	"sync/atomic"
 	"time"
 
 	"github.com/docker/go-events"
@@ -14,20 +15,32 @@ const (
 
 // APIEventsHandler broadcasts events to multiple client listeners.
 type APIEventHandler struct {
-	watchQueue *watch.Queue
+	listenerCount *uint64
+	watchQueue    *watch.Queue
 }
 
 // NewAPIEventHandler creates a new APIEventsHandler for a cluster.
 // The new eventsHandler is initialized with no writers or channels.
 func NewAPIEventHandler() *APIEventHandler {
+	count := uint64(0)
 	return &APIEventHandler{
-		watchQueue: watch.NewQueue(watch.WithTimeout(defaultEventQueueTimeout), watch.WithLimit(defaultEventQueueLimit), watch.WithCloseOutChan()),
+		listenerCount: &count,
+		watchQueue:    watch.NewQueue(watch.WithTimeout(defaultEventQueueTimeout), watch.WithLimit(defaultEventQueueLimit), watch.WithCloseOutChan()),
 	}
 }
 
 // Add adds the writer and a new channel for the remote address.
-func (eh *APIEventHandler) Watch() (eventq chan events.Event, cancel func()) {
-	eventq, cancel = eh.watchQueue.Watch()
+func (eh *APIEventHandler) Watch() (chan events.Event, func()) {
+	// create a new queue and subscribe to it
+	eventq, cancelFunc := eh.watchQueue.Watch()
+	// increment counter
+	atomic.AddUint64(eh.listenerCount, 1)
+
+	cancel := func() {
+		// decrement counter
+		atomic.AddUint64(eh.listenerCount, ^uint64(0))
+		cancelFunc()
+	}
 	return eventq, cancel
 }
 
@@ -41,3 +54,8 @@ func (eh *APIEventHandler) Handle(e *Event) error {
 	eh.watchQueue.Publish(e)
 	return nil
 }
+
+// Size returns the number of event queues currently listening for events
+func (eh *APIEventHandler) Size() int {
+	return int(atomic.LoadUint64(eh.listenerCount))
+}