feat(network): define config for rate limit threshold (#1284)

pactus-project · May 28, 2024 · c76c16e · c76c16e
1 parent 9b857bc
commit c76c16e
Show file tree

Hide file tree

Showing 6 changed files with 16 additions and 10 deletions.
diff --git a/config/example_config.toml b/config/example_config.toml
@@ -75,6 +75,9 @@
   # Default is `false`.
   force_private_network = false
 
+  # `rate_limit_threshold` defines a threshold (message per second) to ignore messages in P2P network.
+  rate_limit_threshold = 3
+
 # `sync` contains configuration of sync module.
 [sync]
 

diff --git a/network/config.go b/network/config.go
@@ -23,6 +23,7 @@ type Config struct {
 	EnableMdns           bool     `toml:"enable_mdns"`
 	EnableMetrics        bool     `toml:"enable_metrics"`
 	ForcePrivateNetwork  bool     `toml:"force_private_network"`
+	RateLimitThreshold   int      `toml:"rate_limit_threshold"`
 
 	// Private configs
 	NetworkName                 string   `toml:"-"`
@@ -48,6 +49,7 @@ func DefaultConfig() *Config {
 		ForcePrivateNetwork:  false,
 		DefaultPort:          0,
 		IsBootstrapper:       false,
+		RateLimitThreshold:   3,
 	}
 }
 

diff --git a/network/gossip.go b/network/gossip.go
@@ -64,7 +64,7 @@ func newGossipService(ctx context.Context, host lp2phost.Host, eventCh chan Even
 		wg:               sync.WaitGroup{},
 		eventCh:          eventCh,
 		logger:           log,
-		rateLimit:        newRateLimit(3, 1*time.Second),
+		rateLimit:        newRateLimit(conf.RateLimitThreshold, 1*time.Second),
 		generalTopicName: generalTopicName,
 	}
 }

diff --git a/network/network_test.go b/network/network_test.go
@@ -50,6 +50,7 @@ func testConfig() *Config {
 		ForcePrivateNetwork:  true,
 		NetworkName:          "test",
 		DefaultPort:          12345,
+		RateLimitThreshold:   124,
 	}
 }
 

diff --git a/network/ratelimit.go b/network/ratelimit.go
@@ -6,12 +6,12 @@ import (
 
 type rateLimit struct {
 	referenceTime time.Time
-	threshold     uint8
-	counter       uint8
+	threshold     int
+	counter       int
 	window        time.Duration
 }
 
-func newRateLimit(threshold uint8, window time.Duration) *rateLimit {
+func newRateLimit(threshold int, window time.Duration) *rateLimit {
 	return &rateLimit{
 		referenceTime: time.Now(),
 		threshold:     threshold,

diff --git a/network/ratelimit_test.go b/network/ratelimit_test.go
@@ -8,16 +8,16 @@ import (
 )
 
 func TestRateLimit(t *testing.T) {
-	threshold := uint8(5)
+	threshold := 5
 	window := 100 * time.Millisecond
 	r := newRateLimit(threshold, window)
 
 	t.Run("InitialState", func(t *testing.T) {
-		assert.Equal(t, uint8(0), r.counter)
+		assert.Equal(t, 0, r.counter)
 	})
 
 	t.Run("IncrementWithinThreshold", func(t *testing.T) {
-		for i := uint8(0); i < threshold; i++ {
+		for i := 0; i < threshold; i++ {
 			assert.True(t, r.increment())
 		}
 		assert.Equal(t, threshold, r.counter)
@@ -30,14 +30,14 @@ func TestRateLimit(t *testing.T) {
 	t.Run("ResetAfterWindow", func(t *testing.T) {
 		time.Sleep(window + 10*time.Millisecond)
 		assert.True(t, r.increment())
-		assert.Equal(t, uint8(1), r.counter)
+		assert.Equal(t, 1, r.counter)
 	})
 
 	t.Run("ResetMethod", func(t *testing.T) {
 		r.reset()
-		assert.Equal(t, uint8(0), r.counter)
+		assert.Equal(t, 0, r.counter)
 		assert.True(t, r.increment())
-		assert.Equal(t, uint8(1), r.counter)
+		assert.Equal(t, 1, r.counter)
 	})
 
 	t.Run("DiffMethod", func(t *testing.T) {