From 4f7b5b7d9804efb8118080d244d2830dfdbabf02 Mon Sep 17 00:00:00 2001
From: simon <fortier.simon@protonmail.com>
Date: Sat, 1 Jun 2019 15:00:50 -0400
Subject: [PATCH] Add hot config reload

---
 README.md           | 13 +++++++++++--
 config.go           | 27 +++++++++++++++++++++++++++
 gc.go               |  3 +++
 jenkins/Jenkinsfile |  2 +-
 main.go             | 36 ++++++++++++++++++++----------------
 reload.sh           |  3 +++
 6 files changed, 65 insertions(+), 19 deletions(-)
 create mode 100755 reload.sh

diff --git a/README.md b/README.md
index b9ca237..a12571f 100644
--- a/README.md
+++ b/README.md
@@ -19,13 +19,22 @@ and error handling. Built for automated web scraping.
 ### Usage
 
 ```bash
-wget https://simon987.net/data/architeuthis/9_architeuthis.tar.gz
-tar -xzf 9_architeuthis.tar.gz
+wget https://simon987.net/data/architeuthis/11_architeuthis.tar.gz
+tar -xzf 11_architeuthis.tar.gz
 
 vim config.json # Configure settings here
 ./architeuthis
 ```
 
+### Hot config reload
+
+```bash
+# Note: this will reset current rate limiters, if there are many active
+# connections, this might cause a small request spike and go over
+# the rate limits.
+./reload.sh
+```
+
 ### Sample configuration
 
 ```json
diff --git a/config.go b/config.go
index ecaa854..d0e518d 100644
--- a/config.go
+++ b/config.go
@@ -3,6 +3,7 @@ package main
 import (
 	"encoding/json"
 	"fmt"
+	"github.com/sirupsen/logrus"
 	"golang.org/x/time/rate"
 	"io/ioutil"
 	"os"
@@ -98,6 +99,32 @@ func applyConfig(proxy *Proxy) {
 	}
 }
 
+func (b *Balancer) reloadConfig() {
+
+	b.proxyMutex.Lock()
+	loadConfig()
+
+	if b.proxies != nil {
+		b.proxies = b.proxies[:0]
+	}
+
+	for _, proxyConf := range config.Proxies {
+		proxy, err := NewProxy(proxyConf.Name, proxyConf.Url)
+		handleErr(err)
+		b.proxies = append(b.proxies, proxy)
+
+		applyConfig(proxy)
+
+		logrus.WithFields(logrus.Fields{
+			"name": proxy.Name,
+			"url":  proxy.Url,
+		}).Info("Proxy")
+	}
+	b.proxyMutex.Unlock()
+
+	logrus.Info("Reloaded config")
+}
+
 func handleErr(err error) {
 	if err != nil {
 		panic(err)
diff --git a/gc.go b/gc.go
index e4f9023..f66d082 100644
--- a/gc.go
+++ b/gc.go
@@ -24,11 +24,14 @@ func (b *Balancer) setupGarbageCollector() {
 func (b *Balancer) cleanAllExpiredLimits() {
 	before := 0
 	after := 0
+
+	b.proxyMutex.RLock()
 	for _, p := range b.proxies {
 		before += len(p.Limiters)
 		cleanExpiredLimits(p)
 		after += len(p.Limiters)
 	}
+	b.proxyMutex.RUnlock()
 
 	logrus.WithFields(logrus.Fields{
 		"removed": before - after,
diff --git a/jenkins/Jenkinsfile b/jenkins/Jenkinsfile
index 44f05df..8f85688 100644
--- a/jenkins/Jenkinsfile
+++ b/jenkins/Jenkinsfile
@@ -26,7 +26,7 @@ pipeline {
                 sh 'cp *.go "/go/src/github.com/simon987/Architeuthis"'
                 sh 'cd /go/src/github.com/simon987/Architeuthis && go get ./...'
                 sh 'cd /go/src/github.com/simon987/Architeuthis && go build -a -installsuffix cgo -o "${WORKSPACE}/architeuthis" .'
-                sh 'tar -czf ${BUILD_NUMBER}_architeuthis.tar.gz config.json architeuthis'
+                sh 'tar -czf ${BUILD_NUMBER}_architeuthis.tar.gz config.json architeuthis reload.sh'
                 sshPut remote: remote, from: env.BUILD_NUMBER + '_architeuthis.tar.gz', into: 'architeuthis/webroot/'
             }
         }
diff --git a/main.go b/main.go
index 447225e..d0d4e8c 100644
--- a/main.go
+++ b/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"fmt"
 	"github.com/elazarl/goproxy"
 	"github.com/pkg/errors"
 	"github.com/ryanuber/go-glob"
@@ -11,12 +12,14 @@ import (
 	"net/url"
 	"sort"
 	"strings"
+	"sync"
 	"time"
 )
 
 type Balancer struct {
-	server  *goproxy.ProxyHttpServer
-	proxies []*Proxy
+	server     *goproxy.ProxyHttpServer
+	proxies    []*Proxy
+	proxyMutex *sync.RWMutex
 }
 
 type ExpiringLimiter struct {
@@ -113,6 +116,7 @@ func New() *Balancer {
 
 	balancer := new(Balancer)
 
+	balancer.proxyMutex = &sync.RWMutex{}
 	balancer.server = goproxy.NewProxyHttpServer()
 
 	balancer.server.OnRequest().HandleConnect(goproxy.AlwaysMitm)
@@ -120,6 +124,7 @@ func New() *Balancer {
 	balancer.server.OnRequest().DoFunc(
 		func(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
 
+			balancer.proxyMutex.RLock()
 			p := balancer.chooseProxy()
 
 			logrus.WithFields(logrus.Fields{
@@ -129,6 +134,7 @@ func New() *Balancer {
 			}).Trace("Routing request")
 
 			resp, err := p.processRequest(r)
+			balancer.proxyMutex.RUnlock()
 
 			if err != nil {
 				logrus.WithError(err).Trace("Could not complete request")
@@ -137,6 +143,17 @@ func New() *Balancer {
 
 			return nil, resp
 		})
+
+	balancer.server.NonproxyHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+
+		if r.URL.Path == "/reload" {
+			balancer.reloadConfig()
+			_, _ = fmt.Fprint(w, "Reloaded\n")
+		} else {
+			w.Header().Set("Content-Type", "application/json")
+			_, _ = fmt.Fprint(w, "{\"name\":\"Architeuthis\",\"version\":1.0}")
+		}
+	})
 	return balancer
 }
 
@@ -275,21 +292,8 @@ func NewProxy(name, stringUrl string) (*Proxy, error) {
 func main() {
 	logrus.SetLevel(logrus.TraceLevel)
 
-	loadConfig()
 	balancer := New()
-
-	for _, proxyConf := range config.Proxies {
-		proxy, err := NewProxy(proxyConf.Name, proxyConf.Url)
-		handleErr(err)
-		balancer.proxies = append(balancer.proxies, proxy)
-
-		applyConfig(proxy)
-
-		logrus.WithFields(logrus.Fields{
-			"name": proxy.Name,
-			"url":  proxy.Url,
-		}).Info("Proxy")
-	}
+	balancer.reloadConfig()
 
 	balancer.setupGarbageCollector()
 	balancer.Run()
diff --git a/reload.sh b/reload.sh
new file mode 100755
index 0000000..406116a
--- /dev/null
+++ b/reload.sh
@@ -0,0 +1,3 @@
+#!/usr/bin/env bash
+
+curl -X GET localhost:5050/reload
\ No newline at end of file