mirror of
https://github.com/hanzoai/tasks.git
synced 2026-10-04 02:04:03 +00:00
Stephan BehnkeandGitHub
3073126265
Reverts temporalio/temporal#8263 which reverted https://github.com/temporalio/temporal/pull/8135 So this is https://github.com/temporalio/temporal/pull/8135 plus a fix for flaky tests. Comparing the results of the test runs (4 each) for the original (1st commit) and fix (2nd commit), I can only conclude that the issue has been addressed.
92 lines
3.0 KiB
Go
92 lines
3.0 KiB
Go
package matching
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/stretchr/testify/suite"
|
|
enumspb "go.temporal.io/api/enums/v1"
|
|
"go.temporal.io/server/common/dynamicconfig"
|
|
"go.temporal.io/server/common/tqid"
|
|
)
|
|
|
|
type RateLimitManagerSuite struct {
|
|
suite.Suite
|
|
*require.Assertions
|
|
}
|
|
|
|
func TestRateLimitManagerSuite(t *testing.T) {
|
|
suite.Run(t, new(RateLimitManagerSuite))
|
|
}
|
|
|
|
func (s *RateLimitManagerSuite) SetupTest() {
|
|
s.Assertions = require.New(s.T())
|
|
}
|
|
|
|
func (s *RateLimitManagerSuite) TestUpdatePerKeySimpleRateLimitLocked_WhenFairnessKeyRateLimitDefaultIsNil() {
|
|
mockUserDataManager := &mockUserDataManager{}
|
|
config := newTaskQueueConfig(
|
|
tqid.UnsafeTaskQueueFamily("test-namespace", "test-task-queue").TaskQueue(enumspb.TASK_QUEUE_TYPE_ACTIVITY),
|
|
NewConfig(dynamicconfig.NewNoopCollection()),
|
|
"test-namespace",
|
|
)
|
|
rateLimitManager := newRateLimitManager(mockUserDataManager, config, enumspb.TASK_QUEUE_TYPE_ACTIVITY)
|
|
rateLimitManager.mu.Lock()
|
|
// Simulate the condition where fairnessKeyRateLimitDefault is nil
|
|
rateLimitManager.fairnessKeyRateLimitDefault = nil
|
|
// Add per-key ready entries to verify they get cleared
|
|
rateLimitManager.perKeyReady.Put("key1", simpleLimiter(1000))
|
|
rateLimitManager.perKeyReady.Put("key2", simpleLimiter(2000))
|
|
// Set per-key limit to verify it gets cleared
|
|
rateLimitManager.perKeyLimit = simpleLimiterParams{
|
|
interval: time.Second,
|
|
burst: 10,
|
|
}
|
|
// Verify initial state
|
|
s.Equal(2, rateLimitManager.perKeyReady.Size())
|
|
s.True(rateLimitManager.perKeyLimit.limited())
|
|
// Update the per-key simple rate limit with fairnessKeyRateLimitDefault as nil
|
|
rateLimitManager.updatePerKeySimpleRateLimitLocked(time.Second)
|
|
// Verify that clearPerKeyRateLimitsLocked was called
|
|
// The cache should be replaced with a new empty cache
|
|
s.Equal(0, rateLimitManager.perKeyReady.Size(), "All per-key ready entries should be cleared")
|
|
s.False(rateLimitManager.perKeyLimit.limited(), "Per-key limit should be cleared")
|
|
rateLimitManager.mu.Unlock()
|
|
}
|
|
|
|
// Additions to rateLimitManager for use by other unit tests:
|
|
|
|
func (r *rateLimitManager) UpdateSimpleRateLimitForTesting(burstDuration time.Duration) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
r.updateSimpleRateLimitLocked(burstDuration)
|
|
}
|
|
|
|
func (r *rateLimitManager) UpdatePerKeySimpleRateLimitForTesting(burstDuration time.Duration) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
r.updatePerKeySimpleRateLimitLocked(burstDuration)
|
|
}
|
|
|
|
func (r *rateLimitManager) SetAdminRateForTesting(rps float64) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
r.adminNsRate = rps
|
|
r.adminTqRate = rps
|
|
}
|
|
|
|
func (r *rateLimitManager) SetEffectiveRPSAndSourceForTesting(rps float64, source enumspb.RateLimitSource) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
r.effectiveRPS = rps
|
|
r.fairnessKeyRateLimitDefault = &rps
|
|
r.rateLimitSource = source
|
|
}
|
|
|
|
func (r *rateLimitManager) SetFairnessKeyRateLimitDefaultForTesting(rps float64, source enumspb.RateLimitSource) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
r.fairnessKeyRateLimitDefault = &rps
|
|
}
|