From e696946adf4df4358f8511b70208e96783c6730d Mon Sep 17 00:00:00 2001 From: Pavel Zbitskiy Date: Wed, 20 Nov 2024 08:16:39 -0500 Subject: [PATCH 1/2] ERL: Use IP address instead of IP:port pair --- util/rateLimit.go | 17 ++++++++++------- util/rateLimit_test.go | 28 +++++++++++++++++----------- 2 files changed, 27 insertions(+), 18 deletions(-) diff --git a/util/rateLimit.go b/util/rateLimit.go index 3fea7f8912..a6d7b4dead 100644 --- a/util/rateLimit.go +++ b/util/rateLimit.go @@ -42,7 +42,7 @@ type ElasticRateLimiter struct { MaxCapacity int CapacityPerReservation int sharedCapacity capacityQueue - capacityByClient map[ErlClient]capacityQueue + capacityByClient map[string]capacityQueue clientLock deadlock.RWMutex // CongestionManager and enable flag cm CongestionManager @@ -53,6 +53,7 @@ type ElasticRateLimiter struct { // ErlClient clients must support OnClose for reservation closing type ErlClient interface { OnClose(func()) + RoutingAddr() []byte } // capacity is an empty structure used for loading and draining queues @@ -122,7 +123,7 @@ func NewElasticRateLimiter( ret := ElasticRateLimiter{ MaxCapacity: maxCapacity, CapacityPerReservation: reservedCapacity, - capacityByClient: map[ErlClient]capacityQueue{}, + capacityByClient: map[string]capacityQueue{}, sharedCapacity: capacityQueue(make(chan capacity, maxCapacity)), congestionControlCounter: conmanCount, } @@ -178,7 +179,7 @@ func (erl *ElasticRateLimiter) ConsumeCapacity(c ErlClient) (*ErlCapacityGuard, var isCMEnabled bool // get the client's queue erl.clientLock.RLock() - q, exists = erl.capacityByClient[c] + q, exists = erl.capacityByClient[string(c.RoutingAddr())] isCMEnabled = erl.enableCM erl.clientLock.RUnlock() @@ -234,7 +235,8 @@ func (erl *ElasticRateLimiter) ConsumeCapacity(c ErlClient) (*ErlCapacityGuard, func (erl *ElasticRateLimiter) openReservation(c ErlClient) (capacityQueue, error) { erl.clientLock.Lock() defer erl.clientLock.Unlock() - if _, exists := erl.capacityByClient[c]; exists { + addr := string(c.RoutingAddr()) + if _, exists := erl.capacityByClient[addr]; exists { return capacityQueue(nil), errERLReservationExists } // guard against overprovisioning, if there is less than a reservedCapacity amount left @@ -244,7 +246,7 @@ func (erl *ElasticRateLimiter) openReservation(c ErlClient) (capacityQueue, erro } // make capacity for the provided client q := capacityQueue(make(chan capacity, erl.CapacityPerReservation)) - erl.capacityByClient[c] = q + erl.capacityByClient[addr] = q // create a thread to drain the capacity from sharedCapacity in a blocking way // and move it to the reservation, also in a blocking way go func() { @@ -261,12 +263,13 @@ func (erl *ElasticRateLimiter) openReservation(c ErlClient) (capacityQueue, erro func (erl *ElasticRateLimiter) closeReservation(c ErlClient) { erl.clientLock.Lock() defer erl.clientLock.Unlock() - q, exists := erl.capacityByClient[c] + addr := string(c.RoutingAddr()) + q, exists := erl.capacityByClient[addr] // guard clauses, and preventing the ElasticRateLimiter from draining its own sharedCapacity if !exists || q == erl.sharedCapacity { return } - delete(erl.capacityByClient, c) + delete(erl.capacityByClient, addr) // start a routine to consume capacity from the closed reservation, and return it to the sharedCapacity go func() { for i := 0; i < erl.CapacityPerReservation; i++ { diff --git a/util/rateLimit_test.go b/util/rateLimit_test.go index 8888bfcf4c..938ba6a657 100644 --- a/util/rateLimit_test.go +++ b/util/rateLimit_test.go @@ -38,6 +38,10 @@ func (c mockClient) OnClose(func()) { return } +func (c mockClient) RoutingAddr() []byte { + return []byte(c) +} + func TestNewElasticRateLimiter(t *testing.T) { partitiontest.PartitionTest(t) erl := NewElasticRateLimiter(100, 10, time.Second, nil) @@ -49,6 +53,7 @@ func TestNewElasticRateLimiter(t *testing.T) { func TestElasticRateLimiterCongestionControlled(t *testing.T) { partitiontest.PartitionTest(t) client := mockClient("client") + clientAddr := string(client.RoutingAddr()) erl := NewElasticRateLimiter(3, 2, time.Second, nil) // give the ERL a congestion controler with well defined behavior for testing erl.cm = mockCongestionControl{} @@ -57,24 +62,24 @@ func TestElasticRateLimiterCongestionControlled(t *testing.T) { // because the ERL gives capacity to a reservation, and then asynchronously drains capacity from the share, // wait a moment before testing the size of the sharedCapacity time.Sleep(100 * time.Millisecond) - assert.Equal(t, 1, len(erl.capacityByClient[client])) + assert.Equal(t, 1, len(erl.capacityByClient[clientAddr])) assert.Equal(t, 1, len(erl.sharedCapacity)) assert.NoError(t, err) erl.EnableCongestionControl() _, _, err = erl.ConsumeCapacity(client) - assert.Equal(t, 0, len(erl.capacityByClient[client])) + assert.Equal(t, 0, len(erl.capacityByClient[clientAddr])) assert.Equal(t, 1, len(erl.sharedCapacity)) assert.NoError(t, err) _, _, err = erl.ConsumeCapacity(client) - assert.Equal(t, 0, len(erl.capacityByClient[client])) + assert.Equal(t, 0, len(erl.capacityByClient[clientAddr])) assert.Equal(t, 1, len(erl.sharedCapacity)) assert.Error(t, err) erl.DisableCongestionControl() _, _, err = erl.ConsumeCapacity(client) - assert.Equal(t, 0, len(erl.capacityByClient[client])) + assert.Equal(t, 0, len(erl.capacityByClient[clientAddr])) assert.Equal(t, 0, len(erl.sharedCapacity)) assert.NoError(t, err) } @@ -132,46 +137,47 @@ func TestZeroSizeReservations(t *testing.T) { func TestConsumeReleaseCapacity(t *testing.T) { partitiontest.PartitionTest(t) client := mockClient("client") + clientAddr := string(client.RoutingAddr()) erl := NewElasticRateLimiter(4, 3, time.Second, nil) c1, _, err := erl.ConsumeCapacity(client) // because the ERL gives capacity to a reservation, and then asynchronously drains capacity from the share, // wait a moment before testing the size of the sharedCapacity time.Sleep(100 * time.Millisecond) - assert.Equal(t, 2, len(erl.capacityByClient[client])) + assert.Equal(t, 2, len(erl.capacityByClient[clientAddr])) assert.Equal(t, 1, len(erl.sharedCapacity)) assert.NoError(t, err) _, _, err = erl.ConsumeCapacity(client) - assert.Equal(t, 1, len(erl.capacityByClient[client])) + assert.Equal(t, 1, len(erl.capacityByClient[clientAddr])) assert.Equal(t, 1, len(erl.sharedCapacity)) assert.NoError(t, err) _, _, err = erl.ConsumeCapacity(client) - assert.Equal(t, 0, len(erl.capacityByClient[client])) + assert.Equal(t, 0, len(erl.capacityByClient[clientAddr])) assert.Equal(t, 1, len(erl.sharedCapacity)) assert.NoError(t, err) // remember this capacity, as it is a shared capacity c4, _, err := erl.ConsumeCapacity(client) - assert.Equal(t, 0, len(erl.capacityByClient[client])) + assert.Equal(t, 0, len(erl.capacityByClient[clientAddr])) assert.Equal(t, 0, len(erl.sharedCapacity)) assert.NoError(t, err) _, _, err = erl.ConsumeCapacity(client) - assert.Equal(t, 0, len(erl.capacityByClient[client])) + assert.Equal(t, 0, len(erl.capacityByClient[clientAddr])) assert.Equal(t, 0, len(erl.sharedCapacity)) assert.Error(t, err) // now release the capacity and observe the items return to the correct places err = c1.Release() - assert.Equal(t, 1, len(erl.capacityByClient[client])) + assert.Equal(t, 1, len(erl.capacityByClient[clientAddr])) assert.Equal(t, 0, len(erl.sharedCapacity)) assert.NoError(t, err) // now release the capacity and observe the items return to the correct places err = c4.Release() - assert.Equal(t, 1, len(erl.capacityByClient[client])) + assert.Equal(t, 1, len(erl.capacityByClient[clientAddr])) assert.Equal(t, 1, len(erl.sharedCapacity)) assert.NoError(t, err) From 39911118fa9c3e956b77a6c08d66bce2e476b484 Mon Sep 17 00:00:00 2001 From: Pavel Zbitskiy Date: Wed, 20 Nov 2024 08:27:42 -0500 Subject: [PATCH 2/2] fix linter --- util/rateLimit_test.go | 24 +++++++++++------------- 1 file changed, 11 insertions(+), 13 deletions(-) diff --git a/util/rateLimit_test.go b/util/rateLimit_test.go index 938ba6a657..7c09ee3e9a 100644 --- a/util/rateLimit_test.go +++ b/util/rateLimit_test.go @@ -53,7 +53,6 @@ func TestNewElasticRateLimiter(t *testing.T) { func TestElasticRateLimiterCongestionControlled(t *testing.T) { partitiontest.PartitionTest(t) client := mockClient("client") - clientAddr := string(client.RoutingAddr()) erl := NewElasticRateLimiter(3, 2, time.Second, nil) // give the ERL a congestion controler with well defined behavior for testing erl.cm = mockCongestionControl{} @@ -62,24 +61,24 @@ func TestElasticRateLimiterCongestionControlled(t *testing.T) { // because the ERL gives capacity to a reservation, and then asynchronously drains capacity from the share, // wait a moment before testing the size of the sharedCapacity time.Sleep(100 * time.Millisecond) - assert.Equal(t, 1, len(erl.capacityByClient[clientAddr])) + assert.Equal(t, 1, len(erl.capacityByClient[string(client.RoutingAddr())])) assert.Equal(t, 1, len(erl.sharedCapacity)) assert.NoError(t, err) erl.EnableCongestionControl() _, _, err = erl.ConsumeCapacity(client) - assert.Equal(t, 0, len(erl.capacityByClient[clientAddr])) + assert.Equal(t, 0, len(erl.capacityByClient[string(client.RoutingAddr())])) assert.Equal(t, 1, len(erl.sharedCapacity)) assert.NoError(t, err) _, _, err = erl.ConsumeCapacity(client) - assert.Equal(t, 0, len(erl.capacityByClient[clientAddr])) + assert.Equal(t, 0, len(erl.capacityByClient[string(client.RoutingAddr())])) assert.Equal(t, 1, len(erl.sharedCapacity)) assert.Error(t, err) erl.DisableCongestionControl() _, _, err = erl.ConsumeCapacity(client) - assert.Equal(t, 0, len(erl.capacityByClient[clientAddr])) + assert.Equal(t, 0, len(erl.capacityByClient[string(client.RoutingAddr())])) assert.Equal(t, 0, len(erl.sharedCapacity)) assert.NoError(t, err) } @@ -137,47 +136,46 @@ func TestZeroSizeReservations(t *testing.T) { func TestConsumeReleaseCapacity(t *testing.T) { partitiontest.PartitionTest(t) client := mockClient("client") - clientAddr := string(client.RoutingAddr()) erl := NewElasticRateLimiter(4, 3, time.Second, nil) c1, _, err := erl.ConsumeCapacity(client) // because the ERL gives capacity to a reservation, and then asynchronously drains capacity from the share, // wait a moment before testing the size of the sharedCapacity time.Sleep(100 * time.Millisecond) - assert.Equal(t, 2, len(erl.capacityByClient[clientAddr])) + assert.Equal(t, 2, len(erl.capacityByClient[string(client.RoutingAddr())])) assert.Equal(t, 1, len(erl.sharedCapacity)) assert.NoError(t, err) _, _, err = erl.ConsumeCapacity(client) - assert.Equal(t, 1, len(erl.capacityByClient[clientAddr])) + assert.Equal(t, 1, len(erl.capacityByClient[string(client.RoutingAddr())])) assert.Equal(t, 1, len(erl.sharedCapacity)) assert.NoError(t, err) _, _, err = erl.ConsumeCapacity(client) - assert.Equal(t, 0, len(erl.capacityByClient[clientAddr])) + assert.Equal(t, 0, len(erl.capacityByClient[string(client.RoutingAddr())])) assert.Equal(t, 1, len(erl.sharedCapacity)) assert.NoError(t, err) // remember this capacity, as it is a shared capacity c4, _, err := erl.ConsumeCapacity(client) - assert.Equal(t, 0, len(erl.capacityByClient[clientAddr])) + assert.Equal(t, 0, len(erl.capacityByClient[string(client.RoutingAddr())])) assert.Equal(t, 0, len(erl.sharedCapacity)) assert.NoError(t, err) _, _, err = erl.ConsumeCapacity(client) - assert.Equal(t, 0, len(erl.capacityByClient[clientAddr])) + assert.Equal(t, 0, len(erl.capacityByClient[string(client.RoutingAddr())])) assert.Equal(t, 0, len(erl.sharedCapacity)) assert.Error(t, err) // now release the capacity and observe the items return to the correct places err = c1.Release() - assert.Equal(t, 1, len(erl.capacityByClient[clientAddr])) + assert.Equal(t, 1, len(erl.capacityByClient[string(client.RoutingAddr())])) assert.Equal(t, 0, len(erl.sharedCapacity)) assert.NoError(t, err) // now release the capacity and observe the items return to the correct places err = c4.Release() - assert.Equal(t, 1, len(erl.capacityByClient[clientAddr])) + assert.Equal(t, 1, len(erl.capacityByClient[string(client.RoutingAddr())])) assert.Equal(t, 1, len(erl.sharedCapacity)) assert.NoError(t, err)