From 3ae1089048f1d510d60b1c11c4443e33ebe0aecd Mon Sep 17 00:00:00 2001
From: simon <me@simon987.net>
Date: Sun, 23 Jun 2019 10:55:56 -0400
Subject: [PATCH] Use atomic add for connection count

---
 README.md |  4 ++--
 main.go   | 19 ++++++++++++-------
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index e7eaa77..bda7aa1 100644
--- a/README.md
+++ b/README.md
@@ -20,8 +20,8 @@ and error handling. Built for automated web scraping.
 ### Usage
 
 ```bash
-wget https://simon987.net/data/architeuthis/14_architeuthis.tar.gz
-tar -xzf 11_architeuthis.tar.gz
+wget https://simon987.net/data/architeuthis/15_architeuthis.tar.gz
+tar -xzf 15_architeuthis.tar.gz
 
 vim config.json # Configure settings here
 ./architeuthis
diff --git a/main.go b/main.go
index af0ad14..d6dc3fa 100644
--- a/main.go
+++ b/main.go
@@ -13,6 +13,7 @@ import (
 	"sort"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 )
 
@@ -35,7 +36,7 @@ type Proxy struct {
 	Url         *url.URL
 	Limiters    []*ExpiringLimiter
 	HttpClient  *http.Client
-	Connections int
+	Connections *int32
 }
 
 type RequestCtx struct {
@@ -54,7 +55,7 @@ func (a ByConnectionCount) Swap(i, j int) {
 }
 
 func (a ByConnectionCount) Less(i, j int) bool {
-	return a[i].Connections < a[j].Connections
+	return *a[i].Connections < *a[j].Connections
 }
 
 func (p *Proxy) getLimiter(host string) *rate.Limiter {
@@ -152,7 +153,7 @@ func New() *Balancer {
 
 			logrus.WithFields(logrus.Fields{
 				"proxy": p.Name,
-				"conns": p.Connections,
+				"conns": *p.Connections,
 				"url":   r.URL,
 			}).Trace("Routing request")
 
@@ -237,9 +238,9 @@ func computeRules(ctx *RequestCtx, configs []*HostConfig) (dontRetry, forceRetry
 
 func (p *Proxy) processRequest(r *http.Request) (*http.Response, error) {
 
-	p.Connections += 1
+	atomic.AddInt32(p.Connections, 1)
 	defer func() {
-		p.Connections -= 1
+		atomic.AddInt32(p.Connections, -1)
 	}()
 	retries := 0
 	additionalRetries := 0
@@ -378,11 +379,15 @@ func NewProxy(name, stringUrl string) (*Proxy, error) {
 
 	httpClient.Timeout = config.Timeout
 
-	return &Proxy{
+	p := &Proxy{
 		Name:       name,
 		Url:        parsedUrl,
 		HttpClient: httpClient,
-	}, nil
+	}
+
+	conns := int32(0)
+	p.Connections = &conns
+	return p, nil
 }
 
 func main() {