Use atomic add for connection count

This commit is contained in:
simon 2019-06-23 10:55:56 -04:00
parent 7077c21a30
commit 3ae1089048
2 changed files with 14 additions and 9 deletions

View File

@ -20,8 +20,8 @@ and error handling. Built for automated web scraping.
### Usage ### Usage
```bash ```bash
wget https://simon987.net/data/architeuthis/14_architeuthis.tar.gz wget https://simon987.net/data/architeuthis/15_architeuthis.tar.gz
tar -xzf 11_architeuthis.tar.gz tar -xzf 15_architeuthis.tar.gz
vim config.json # Configure settings here vim config.json # Configure settings here
./architeuthis ./architeuthis

19
main.go
View File

@ -13,6 +13,7 @@ import (
"sort" "sort"
"strings" "strings"
"sync" "sync"
"sync/atomic"
"time" "time"
) )
@ -35,7 +36,7 @@ type Proxy struct {
Url *url.URL Url *url.URL
Limiters []*ExpiringLimiter Limiters []*ExpiringLimiter
HttpClient *http.Client HttpClient *http.Client
Connections int Connections *int32
} }
type RequestCtx struct { type RequestCtx struct {
@ -54,7 +55,7 @@ func (a ByConnectionCount) Swap(i, j int) {
} }
func (a ByConnectionCount) Less(i, j int) bool { 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 { func (p *Proxy) getLimiter(host string) *rate.Limiter {
@ -152,7 +153,7 @@ func New() *Balancer {
logrus.WithFields(logrus.Fields{ logrus.WithFields(logrus.Fields{
"proxy": p.Name, "proxy": p.Name,
"conns": p.Connections, "conns": *p.Connections,
"url": r.URL, "url": r.URL,
}).Trace("Routing request") }).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) { func (p *Proxy) processRequest(r *http.Request) (*http.Response, error) {
p.Connections += 1 atomic.AddInt32(p.Connections, 1)
defer func() { defer func() {
p.Connections -= 1 atomic.AddInt32(p.Connections, -1)
}() }()
retries := 0 retries := 0
additionalRetries := 0 additionalRetries := 0
@ -378,11 +379,15 @@ func NewProxy(name, stringUrl string) (*Proxy, error) {
httpClient.Timeout = config.Timeout httpClient.Timeout = config.Timeout
return &Proxy{ p := &Proxy{
Name: name, Name: name,
Url: parsedUrl, Url: parsedUrl,
HttpClient: httpClient, HttpClient: httpClient,
}, nil }
conns := int32(0)
p.Connections = &conns
return p, nil
} }
func main() { func main() {