Merge pull request Imgur#67 from Imgur/dev-redbtn-longpoll

Add button to disable longpoll (via redis)

config.go

@@ -27,6 +27,7 @@ func NewConfig(configFilePath string) {
 		ConfigOption("datadog_host", "127.0.0.1")
 	}
 
+	ConfigOption("longpoll_killswitch", "longpoll_killswitch")
 	ConfigOption("redis_enabled", false)
 
 	if viper.GetBool("redis_enabled") {

config.yml

@@ -16,6 +16,9 @@ datadog_enabled: false
 # Datadog host if datadog enabled
 datadog_host: "127.0.0.1"
 
+# Redis key to monitor. If it exists, longpolling is disabled, for performance reasons.
+longpoll_killswitch: "longpoll_killswitch"
+
 # ----- Redis Support -----
 
 # Bool; Redis must be enabled if running Incus in a cluster.

incus/main.go

@@ -76,6 +76,8 @@ func main() {
 	go server.ListenFromRedis()
 	go server.ListenFromSockets()
 	go server.ListenFromLongpoll()
+	go server.MonitorLongpollKillswitch()
+
 	go server.ListenForHTTPPings()
 	go server.SendHeartbeatsPeriodically(20 * time.Second)
 

redis_store.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/garyburd/redigo/redis"
+	"github.com/spf13/viper"
 )
 
 const ClientsKey = "SocketClients"
@@ -229,6 +230,40 @@ func (this *RedisStore) QueryIsUserActive(user string, nowTimestamp int64) (bool
 	}
 }
 
+func (this *RedisStore) GetIsLongpollKillswitchActive() (bool, error) {
+	killswitchKey := viper.Get("longpoll_killswitch")
+
+	result := this.redisPendingQueue.RunAsyncTimeout(5*time.Second, func(conn redis.Conn) (result interface{}, err error) {
+		return conn.Do("TTL", killswitchKey)
+	})
+
+	if result.Error == nil {
+		if result.Value.(int64) >= -1 {
+			return true, nil
+		} else {
+			return false, nil
+		}
+	}
+
+	return false, timedOut
+}
+
+func (this *RedisStore) ActivateLongpollKillswitch(seconds int64) error {
+	killswitchKey := viper.Get("longpoll_killswitch")
+
+	return this.redisPendingQueue.RunAsyncTimeout(5*time.Second, func(conn redis.Conn) (result interface{}, err error) {
+		return conn.Do("SETEX", killswitchKey, seconds, "1")
+	}).Error
+}
+
+func (this *RedisStore) DeactivateLongpollKillswitch() error {
+	killswitchKey := viper.Get("longpoll_killswitch")
+
+	return this.redisPendingQueue.RunAsyncTimeout(5*time.Second, func(conn redis.Conn) (result interface{}, err error) {
+		return conn.Do("DEL", killswitchKey)
+	}).Error
+}
+
 func (this *RedisStore) Publish(channel string, message string) {
 	publisher, err := this.GetConn()
 	if err != nil {

redis_store_test.go

@@ -254,3 +254,40 @@ func TestUserPresenceIsImmediatelyRemovedUponMarkingInactive(t *testing.T) {
 		t.Fatalf("Expected 'bazbar' to be inactive after affirmatively marking as inactive")
 	}
 }
+
+func TestKillswitch(t *testing.T) {
+	store := newTestRedisStore()
+
+	store.DeactivateLongpollKillswitch()
+
+	active, err := store.GetIsLongpollKillswitchActive()
+
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err.Error())
+	}
+	if active {
+		t.Fatalf("Expected precondition that killswitch is inactive")
+	}
+
+	store.ActivateLongpollKillswitch(3)
+
+	active, err = store.GetIsLongpollKillswitchActive()
+
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err.Error())
+	}
+	if !active {
+		t.Fatalf("Expected killswitch to be active")
+	}
+
+	time.Sleep(4 * time.Second)
+
+	active, err = store.GetIsLongpollKillswitchActive()
+
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err.Error())
+	}
+	if active {
+		t.Fatalf("Expected killswitch to be inactive")
+	}
+}

server.go

@@ -9,6 +9,7 @@ import (
 	"net/http"
 	"os"
 	"os/signal"
+	"sync/atomic"
 	"syscall"
 	"time"
 
@@ -31,6 +32,10 @@ const (
 	closeCodeUnexpectedError = 1011
 )
 
+var (
+	disableLongpoll atomic.Value
+)
+
 type GCMClient interface {
 	Send(*gcm.Message, int) (*gcm.Response, error)
 }
@@ -162,13 +167,21 @@ func (this *Server) ListenFromLongpoll() {
 			}
 		}()
 
-		sock := newSocket(nil, w, this, "")
 		w.Header().Set("Access-Control-Allow-Origin", "*")
 		w.Header().Set("Content-Type", "application/json")
 		w.Header().Set("Cache-Control", "private, no-store, no-cache, must-revalidate, post-check=0, pre-check=0")
 		w.Header().Set("Connection", "keep-alive")
 		//w.Header().Set("Content-Encoding", "gzip")
 
+		longpollIsDisabled := disableLongpoll.Load()
+		if longpollIsDisabled != nil && longpollIsDisabled.(bool) == true {
+			w.Header().Set("Connection", "close")
+			w.WriteHeader(503)
+			return
+		}
+
+		sock := newSocket(nil, w, this, "")
+
 		if DEBUG {
 			log.Printf("Long poll connected via \n")
 		}
@@ -316,3 +329,19 @@ func (this *Server) GetAPNSClient(build string) apns.APNSClient {
 func (this *Server) GetGCMClient() GCMClient {
 	return this.gcmProvider()
 }
+
+func (this *Server) MonitorLongpollKillswitch() {
+	if !viper.GetBool("redis_enabled") {
+		return
+	}
+
+	for {
+		longpollSwitchedOff, err := this.Store.redis.GetIsLongpollKillswitchActive()
+
+		if err == nil {
+			disableLongpoll.Store(longpollSwitchedOff)
+		}
+
+		time.Sleep(5 * time.Second)
+	}
+}