Repository: raviqqe/muffet
Pull request #192 (2021-11-10 05:27:35)

Fix request throttling per host
Related to #172.


codecov[bot]: # [Codecov](https://codecov.io/gh/raviqqe/muffet/pull/192?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yota+Toyama) Report
> Merging [#192](https://codecov.io/gh/raviqqe/muffet/pull/192?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yota+Toyama) (384b08d) into [master](https://codecov.io/gh/raviqqe/muffet/commit/0a58e1e116c673428fdd3e467869c6880864e32c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yota+Toyama) (0a58e1e) will **increase** coverage by `0.79%`.
> The diff coverage is `90.90%`.

[![Impacted file tree graph](https://codecov.io/gh/raviqqe/muffet/pull/192/graphs/tree.svg?width=650&height=150&src=pr&token=qXJG42bGNw&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yota+Toyama)](https://codecov.io/gh/raviqqe/muffet/pull/192?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yota+Toyama)

```diff
@@            Coverage Diff             @@
##           master     #192      +/-   ##
==========================================
+ Coverage   77.08%   77.88%   +0.79%     
==========================================
  Files          23       23              
  Lines         419      434      +15     
==========================================
+ Hits          323      338      +15     
  Misses         80       80              
  Partials       16       16              
```


| [Impacted Files](https://codecov.io/gh/raviqqe/muffet/pull/192?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yota+Toyama) | Coverage Δ | |
|---|---|---|
| [throttled\_http\_client.go](https://codecov.io/gh/raviqqe/muffet/pull/192/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yota+Toyama#diff-dGhyb3R0bGVkX2h0dHBfY2xpZW50Lmdv) | `90.90% <90.47%> (+15.90%)` | :arrow_up: |
| [command.go](https://codecov.io/gh/raviqqe/muffet/pull/192/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yota+Toyama#diff-Y29tbWFuZC5nbw==) | `84.44% <100.00%> (+0.17%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/raviqqe/muffet/pull/192?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yota+Toyama).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yota+Toyama)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/raviqqe/muffet/pull/192?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yota+Toyama). Last update [0a58e1e...384b08d](https://codecov.io/gh/raviqqe/muffet/pull/192?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yota+Toyama). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Yota+Toyama).


Patch:
diff --git a/command.go b/command.go
index 58a79e3..7dc9da8 100644
--- a/command.go
+++ b/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))
diff --git a/configuration.go b/configuration.go
index a32ed04..57531b5 100644
--- a/configuration.go
+++ b/configuration.go
@@ -3,7 +3,7 @@ package main
 import "time"
 
 const (
-	version     = "2.4.4"
+	version     = "2.4.5"
 	agentName   = "muffet"
 	concurrency = 1024
 	tcpTimeout  = 5 * time.Second
diff --git a/host_throttler.go b/host_throttler.go
new file mode 100644
index 0000000..eac913d
--- /dev/null
+++ b/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()
+}
diff --git a/host_throttler_pool.go b/host_throttler_pool.go
new file mode 100644
index 0000000..2709eb6
--- /dev/null
+++ b/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
+}
diff --git a/throttled_http_client.go b/throttled_http_client.go
index 2d30f8a..fecc323 100644
--- a/throttled_http_client.go
+++ b/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)
 }


Test patch:
diff --git a/host_throttler_pool_test.go b/host_throttler_pool_test.go
new file mode 100644
index 0000000..c82738d
--- /dev/null
+++ b/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
+	}
+}
diff --git a/host_throttler_test.go b/host_throttler_test.go
new file mode 100644
index 0000000..b52f66d
--- /dev/null
+++ b/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
+}