From 972437cd4a4eca9dcbc02269b93cf98eee504870 Mon Sep 17 00:00:00 2001
From: Kane York <rikingcoding@gmail.com>
Date: Fri, 8 Jul 2016 19:16:10 -0700
Subject: [PATCH] Add /get_sub_count

---
 socketserver/server/handlecore.go    |  1 +
 socketserver/server/publisher.go     | 17 +++++++++++++++++
 socketserver/server/subscriptions.go | 17 +++++++++++++++++
 3 files changed, 35 insertions(+)

diff --git a/socketserver/server/handlecore.go b/socketserver/server/handlecore.go
index 5b426c63..723f4fb2 100644
--- a/socketserver/server/handlecore.go
+++ b/socketserver/server/handlecore.go
@@ -102,6 +102,7 @@ func SetupServerAndHandle(config *ConfigFile, serveMux *http.ServeMux) {
 	serveMux.HandleFunc("/drop_backlog", HTTPBackendDropBacklog)
 	serveMux.HandleFunc("/uncached_pub", HTTPBackendUncachedPublish)
 	serveMux.HandleFunc("/cached_pub", HTTPBackendCachedPublish)
+	serveMux.HandleFunc("/get_sub_count", HTTPGetSubscriberCount)
 
 	announceForm, err := Backend.SealRequest(url.Values{
 		"startup": []string{"1"},
diff --git a/socketserver/server/publisher.go b/socketserver/server/publisher.go
index c5120156..85912dd6 100644
--- a/socketserver/server/publisher.go
+++ b/socketserver/server/publisher.go
@@ -220,3 +220,20 @@ func HTTPBackendUncachedPublish(w http.ResponseWriter, r *http.Request) {
 	}
 	fmt.Fprint(w, count)
 }
+
+// HTTPGetSubscriberCount handles the /get_sub_count route.
+// It replies with the number of clients subscribed to a pub/sub topic.
+// A "global" option is not available, use fetch(/stats).CurrentClientCount instead.
+func HTTPGetSubscriberCount(w http.ResponseWriter, r *http.Request) {
+	r.ParseForm()
+	formData, err := Backend.UnsealRequest(r.Form)
+	if err != nil {
+		w.WriteHeader(403)
+		fmt.Fprintf(w, "Error: %v", err)
+		return
+	}
+
+	channel := formData.Get("channel")
+
+	fmt.Fprint(w, CountSubscriptions(strings.Split(channel, ",")))
+}
\ No newline at end of file
diff --git a/socketserver/server/subscriptions.go b/socketserver/server/subscriptions.go
index e55db08d..30bc4112 100644
--- a/socketserver/server/subscriptions.go
+++ b/socketserver/server/subscriptions.go
@@ -19,6 +19,23 @@ var ChatSubscriptionLock sync.RWMutex
 var GlobalSubscriptionInfo []*ClientInfo
 var GlobalSubscriptionLock sync.RWMutex
 
+func CountSubscriptions(channels []string) int {
+	ChatSubscriptionLock.RLock()
+	defer ChatSubscriptionLock.RUnlock()
+
+	count := 0
+	for _, channelName := range channels {
+		list := ChatSubscriptionInfo[channelName]
+		if list != nil {
+			list.RLock()
+			count += len(list.Members)
+			list.RUnlock()
+		}
+	}
+
+	return count
+}
+
 func SubscribeChannel(client *ClientInfo, channelName string) {
 	ChatSubscriptionLock.RLock()
 	_subscribeWhileRlocked(channelName, client.MessageChannel)