Fix request throttling per host (#192)

* Fix host limiting

* Refactor

* Remove comment

* Rename a field

* Test host throttler

* Fix data race

* Test pool

* Test more

* Calculate rate after gettign semaphore

command.go

@@ -59,6 +59,7 @@ func (c *command) runWithError(ss []string) (bool, error) {
 		),
 		args.RateLimit,
 		args.MaxConnections,
+		args.MaxConnectionsPerHost,
 	)
 
 	pp := newPageParser(newLinkFinder(args.ExcludedPatterns))

host_throttler.go

@@ -0,0 +1,27 @@
+package main
+
+import "go.uber.org/ratelimit"
+
+type hostThrottler struct {
+	limiter     ratelimit.Limiter
+	connections semaphore
+}
+
+func newHostThrottler(requestPerSecond, maxConnectionsPerHost int) *hostThrottler {
+	l := ratelimit.NewUnlimited()
+
+	if requestPerSecond > 0 {
+		l = ratelimit.New(requestPerSecond)
+	}
+
+	return &hostThrottler{l, newSemaphore(maxConnectionsPerHost)}
+}
+
+func (t *hostThrottler) Request() {
+	t.connections.Request()
+	t.limiter.Take()
+}
+
+func (t *hostThrottler) Release() {
+	t.connections.Release()
+}

host_throttler_pool.go

@@ -0,0 +1,23 @@
+package main
+
+import "sync"
+
+type hostThrottlerPool struct {
+	requestPerSecond, maxConnectionsPerHost int
+	hostMap                                 sync.Map
+}
+
+func newHostThrottlerPool(requestPerSecond, maxConnectionsPerHost int) *hostThrottlerPool {
+	return &hostThrottlerPool{requestPerSecond, maxConnectionsPerHost, sync.Map{}}
+}
+
+func (p *hostThrottlerPool) Get(name string) *hostThrottler {
+	t := newHostThrottler(p.requestPerSecond, p.maxConnectionsPerHost)
+	x, ok := p.hostMap.LoadOrStore(name, t)
+
+	if ok {
+		t = x.(*hostThrottler)
+	}
+
+	return t
+}

host_throttler_pool_test.go

@@ -0,0 +1,55 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewHostThrottlerPool(t *testing.T) {
+	newHostThrottlerPool(1, 1)
+}
+
+func TestHostThrottlerPoolGetHost(t *testing.T) {
+	c := make(chan struct{}, 100)
+	s := newHostThrottlerPool(1000000, 1)
+
+	for i := 0; i < 2; i++ {
+		go func() {
+			s.Get("foo").Request()
+			c <- struct{}{}
+		}()
+	}
+
+	<-c
+
+	assert.Equal(t, 0, len(c))
+
+	s.Get("foo").Release()
+	<-c
+}
+
+func TestHostThrottlerPoolGetHosts(t *testing.T) {
+	hosts := []string{"foo", "bar"}
+	c := make(chan struct{}, 100)
+	s := newHostThrottlerPool(1000000, 1)
+
+	for _, host := range hosts {
+		for i := 0; i < 2; i++ {
+			go func(host string) {
+				s.Get(host).Request()
+				c <- struct{}{}
+			}(host)
+		}
+	}
+
+	<-c
+	<-c
+
+	assert.Equal(t, 0, len(c))
+
+	for _, host := range hosts {
+		s.Get(host).Release()
+		<-c
+	}
+}

host_throttler_test.go

@@ -0,0 +1,30 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewHostThrottler(t *testing.T) {
+	newHostThrottler(1, 1)
+}
+
+func TestHostThrottlerRequest(t *testing.T) {
+	c := make(chan struct{}, 100)
+	s := newHostThrottler(1000000, 1)
+
+	for i := 0; i < 2; i++ {
+		go func() {
+			s.Request()
+			c <- struct{}{}
+		}()
+	}
+
+	<-c
+
+	assert.Equal(t, 0, len(c))
+
+	s.Release()
+	<-c
+}

throttled_http_client.go

@@ -2,32 +2,29 @@ package main
 
 import (
 	"net/url"
-
-	"go.uber.org/ratelimit"
 )
 
-// TODO Throttle requests for each host.
 type throttledHttpClient struct {
-	client    httpClient
-	limiter   ratelimit.Limiter
-	semaphore semaphore
+	client            httpClient
+	connections       semaphore
+	hostThrottlerPool *hostThrottlerPool
 }
 
-func newThrottledHttpClient(c httpClient, rps int, maxConnections int) httpClient {
-	l := ratelimit.NewUnlimited()
-
-	if rps > 0 {
-		l = ratelimit.New(rps)
+func newThrottledHttpClient(c httpClient, requestPerSecond int, maxConnections, maxConnectionsPerHost int) httpClient {
+	return &throttledHttpClient{
+		c,
+		newSemaphore(maxConnections),
+		newHostThrottlerPool(requestPerSecond, maxConnectionsPerHost),
 	}
-
-	return &throttledHttpClient{c, l, newSemaphore(maxConnections)}
 }
 
 func (c *throttledHttpClient) Get(u *url.URL) (httpResponse, error) {
-	c.semaphore.Request()
-	defer c.semaphore.Release()
+	c.connections.Request()
+	defer c.connections.Release()
 
-	c.limiter.Take()
+	t := c.hostThrottlerPool.Get(u.Hostname())
+	t.Request()
+	defer t.Release()
 
 	return c.client.Get(u)
 }