Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ERL: Use IP address instead of IP:port pair #6176

Draft
wants to merge 2 commits into
base: master
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 10 additions & 7 deletions util/rateLimit.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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,
}
Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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
Expand All @@ -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() {
Expand All @@ -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)
Copy link
Contributor

@cce cce Nov 20, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's say you have N conns per IP, ERL sets up the closeReservation(wsPeer) is called by the wsPeer's OnClose() handler for each of N conns (not per IP) so you will end up calling delete multiple times for each conn (not per IP) and also, the capacity that is reserved for this IP (CapacityPerReservation) will be returned N times

Copy link
Contributor

@cce cce Nov 20, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

actually, since N conns are sharing one erl.capacityByClient[IP] entry, ERL will only call wsPeer.OnClose() when it sees the IP missing from the map. So that means if the first conn goes away, it will return the CapacityPerReservation back and delete itself from the map, even though the other conns are still running.

I think there needs to be an implementation of the ErlClient interface, like a meta-client, that manages setting OnClose() handlers for all the wsPeers sharing the same IP, something like that, and function that takes a wsPeer and looks it up by IP, and returns a IPBasedERLClient impl to pass to ConsumeCapacity(c ErlClient). Then when the IPBasedErlClient notices all the child wsPeers have called OnClose, it calls its own OnClose, which in turns calls ERL.closeReservation()

// 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++ {
Expand Down
26 changes: 15 additions & 11 deletions util/rateLimit_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -57,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[client]))
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[client]))
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[client]))
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[client]))
assert.Equal(t, 0, len(erl.capacityByClient[string(client.RoutingAddr())]))
assert.Equal(t, 0, len(erl.sharedCapacity))
assert.NoError(t, err)
}
Expand Down Expand Up @@ -138,40 +142,40 @@ func TestConsumeReleaseCapacity(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, 2, len(erl.capacityByClient[client]))
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[client]))
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[client]))
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[client]))
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[client]))
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[client]))
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[client]))
assert.Equal(t, 1, len(erl.capacityByClient[string(client.RoutingAddr())]))
assert.Equal(t, 1, len(erl.sharedCapacity))
assert.NoError(t, err)

Expand Down
Loading