Files
tasks/common/persistence/client/persistence_rate_limited_clients_test.go
Alex ShtinandGitHub 8e8f1b0f7c Remove license header from every file (#7689)
## What changed?
<!-- Describe what has changed in this PR -->
Remove license header from every file. Because it is really hard to
follow in this PR here is the summary:
1. License header is removed from all `.go` and `.proto` files
:fireworks::fireworks:🎆.
2. `LICENSE` file in the root directory has only Temporal and Uber
copyrights.
3. 5 other `LICENSE` files added to the packages which have copyrights
different from Temporal and Uber: Datadog, Xargin, "Mat Ryer, Tyler
Bunnell and contributors".
4. `license_file` flag is removed from all code generation tools.
5. `copyright_file` flag is removed from `go:generate mockgen`
directive.
6. All copyright related targets are removed from `Makefile`.
7. Updated Temporal copyright year to 2025 everywhere.

## Why?
<!-- Tell your future self why have you made these changes -->
I double checked with legal department that it is not needed to have
license header in every file. One file per repo is enough. I put all
copyrights to the root `LICENSE` file and removed header from all other
files. Also updated tools and `Makefile`.
2025-05-01 18:50:21 -07:00

250 lines
8.1 KiB
Go

package client_test
import (
"context"
"errors"
"fmt"
"testing"
"time"
"github.com/stretchr/testify/assert"
enumspb "go.temporal.io/api/enums/v1"
"go.temporal.io/api/serviceerror"
"go.temporal.io/server/common/config"
"go.temporal.io/server/common/persistence"
"go.temporal.io/server/common/persistence/client"
"go.temporal.io/server/common/persistence/mock"
"go.temporal.io/server/common/persistence/serialization"
"go.temporal.io/server/common/quotas"
"go.uber.org/mock/gomock"
)
func TestRateLimitedPersistenceClients(t *testing.T) {
t.Parallel()
for _, tc := range []struct {
name string
err error
numRequests int
namespaceRPS int
systemRPS int
namespaceShardRPS int
expectRateLimit bool
expectedScope enumspb.ResourceExhaustedScope
expectedMessage string
}{
{
name: "Namespace limit allow",
err: nil,
numRequests: 10,
namespaceRPS: 10,
namespaceShardRPS: 100,
systemRPS: 100,
expectRateLimit: false,
},
{
name: "Namespace limit hit",
err: nil,
numRequests: 11,
namespaceRPS: 10,
namespaceShardRPS: 100,
systemRPS: 100,
expectRateLimit: true,
expectedScope: enumspb.RESOURCE_EXHAUSTED_SCOPE_NAMESPACE,
expectedMessage: "Namespace Persistence Max QPS Reached.",
},
{
name: "System limit allow",
err: nil,
numRequests: 10,
namespaceRPS: 100,
namespaceShardRPS: 100,
systemRPS: 10,
expectRateLimit: false,
},
{
name: "System limit hit",
err: nil,
numRequests: 11,
namespaceRPS: 100,
namespaceShardRPS: 100,
systemRPS: 10,
expectRateLimit: true,
expectedScope: enumspb.RESOURCE_EXHAUSTED_SCOPE_SYSTEM,
expectedMessage: "System Persistence Max QPS Reached.",
},
{
name: "Shard limit allow",
err: nil,
numRequests: 10,
namespaceRPS: 100,
namespaceShardRPS: 10,
systemRPS: 100,
expectRateLimit: false,
},
{
name: "Shard limit hit",
err: nil,
numRequests: 11,
namespaceRPS: 100,
namespaceShardRPS: 10,
systemRPS: 100,
expectRateLimit: true,
expectedScope: enumspb.RESOURCE_EXHAUSTED_SCOPE_NAMESPACE,
expectedMessage: "Namespace Per-Shard Persistence Max QPS Reached.",
},
} {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctr := gomock.NewController(t)
dataStoreFactory := mock.NewMockDataStoreFactory(ctr)
shardStore := mock.NewMockShardStore(ctr)
shardStore.EXPECT().AssertShardOwnership(gomock.Any(), gomock.Any()).AnyTimes().Return(nil)
dataStoreFactory.EXPECT().NewShardStore().AnyTimes().Return(shardStore, nil)
queue := mock.NewMockQueue(ctr)
queue.EXPECT().Init(gomock.Any(), gomock.Any()).Return(nil)
queue.EXPECT().DeleteMessageFromDLQ(gomock.Any(), gomock.Any()).AnyTimes().Return(nil)
dataStoreFactory.EXPECT().NewQueue(gomock.Any()).AnyTimes().Return(queue, nil)
executionStore := mock.NewMockExecutionStore(ctr)
executionStore.EXPECT().DeleteWorkflowExecution(gomock.Any(), gomock.Any()).AnyTimes().Return(nil)
dataStoreFactory.EXPECT().NewExecutionStore().AnyTimes().Return(executionStore, nil)
metadataStore := mock.NewMockMetadataStore(ctr)
metadataStore.EXPECT().DeleteNamespace(gomock.Any(), gomock.Any()).AnyTimes().Return(nil)
metadataStore.EXPECT().GetNamespace(gomock.Any(), gomock.Any()).AnyTimes().Return(nil, nil)
dataStoreFactory.EXPECT().NewMetadataStore().AnyTimes().Return(metadataStore, nil)
clusterMetadataStore := mock.NewMockClusterMetadataStore(ctr)
clusterMetadataStore.EXPECT().DeleteClusterMetadata(gomock.Any(), gomock.Any()).AnyTimes().Return(nil)
clusterMetadataStore.EXPECT().GetClusterMetadata(gomock.Any(), gomock.Any()).AnyTimes().Return(nil, nil)
dataStoreFactory.EXPECT().NewClusterMetadataStore().AnyTimes().Return(clusterMetadataStore, nil)
nexusStore := mock.NewMockNexusEndpointStore(ctr)
nexusStore.EXPECT().DeleteNexusEndpoint(gomock.Any(), gomock.Any()).AnyTimes().Return(nil)
nexusStore.EXPECT().ListNexusEndpoints(gomock.Any(), gomock.Any()).AnyTimes().Return(nil, nil)
dataStoreFactory.EXPECT().NewNexusEndpointStore().AnyTimes().Return(nexusStore, nil)
taskStore := mock.NewMockTaskStore(ctr)
taskStore.EXPECT().GetTasks(gomock.Any(), gomock.Any()).AnyTimes().Return(nil, nil)
dataStoreFactory.EXPECT().NewTaskStore().AnyTimes().Return(taskStore, nil)
burstRatioFn := func() float64 {
return 1.0
}
systemRequestRateLimiter := quotas.NewRequestRateLimiterAdapter(
quotas.NewDefaultRateLimiter(
func() float64 { return float64(tc.systemRPS) },
burstRatioFn,
),
)
namespaceRequestRateLimiter := quotas.NewRequestRateLimiterAdapter(
quotas.NewDefaultRateLimiter(
func() float64 { return float64(tc.namespaceRPS) },
burstRatioFn,
),
)
shardRequestRateLimiter := quotas.NewRequestRateLimiterAdapter(
quotas.NewDefaultRateLimiter(
func() float64 { return float64(tc.namespaceShardRPS) },
burstRatioFn,
),
)
factory := client.NewFactory(
dataStoreFactory,
&config.Persistence{
NumHistoryShards: 1,
},
systemRequestRateLimiter,
namespaceRequestRateLimiter,
shardRequestRateLimiter,
serialization.NewSerializer(),
nil,
"",
nil,
nil,
nil,
)
shardManager, _ := factory.NewShardManager()
executionManager, _ := factory.NewExecutionManager()
metadataManager, _ := factory.NewMetadataManager()
clusterMetadataManager, _ := factory.NewClusterMetadataManager()
nexusManager, _ := factory.NewNexusEndpointManager()
namespaceQueue, _ := factory.NewNamespaceReplicationQueue()
// Make calls to different manager objects to verify that RPS is enforced.
persistenceCalls := []struct {
name string
call func() error
}{
{
name: "AssertShardOwnership",
call: func() error {
return shardManager.AssertShardOwnership(context.Background(), &persistence.AssertShardOwnershipRequest{ShardID: 0})
},
},
{
name: "DeleteWorkflowExecution",
call: func() error {
return executionManager.DeleteWorkflowExecution(context.Background(), &persistence.DeleteWorkflowExecutionRequest{})
},
},
{
name: "DeleteNamespace",
call: func() error {
return metadataManager.DeleteNamespace(context.Background(), &persistence.DeleteNamespaceRequest{})
},
},
{
name: "DeleteClusterMetadata",
call: func() error {
return clusterMetadataManager.DeleteClusterMetadata(context.Background(), &persistence.DeleteClusterMetadataRequest{ClusterName: "test"})
},
},
{
name: "DeleteNexusEndpoint",
call: func() error {
return nexusManager.DeleteNexusEndpoint(context.Background(), &persistence.DeleteNexusEndpointRequest{})
},
},
{
name: "DeleteMessageFromDLQ",
call: func() error {
return namespaceQueue.DeleteMessageFromDLQ(context.Background(), 0)
},
},
}
var err error
for _, persistenceCall := range persistenceCalls {
t.Run(fmt.Sprintf("%s %s", tc.name, persistenceCall.name), func(t *testing.T) {
// Generate load by sending a number of requests.
for i := 0; i < tc.numRequests; i++ {
err = persistenceCall.call()
if err != nil {
// Assert resource exhausted at last request.
assert.Equal(t, tc.numRequests-1, i)
break
}
}
// Check if the rate limit is hit.
if tc.expectRateLimit {
var resourceExhausted *serviceerror.ResourceExhausted
errors.As(err, &resourceExhausted)
assert.Equal(t, enumspb.RESOURCE_EXHAUSTED_CAUSE_PERSISTENCE_LIMIT, resourceExhausted.Cause)
assert.Equal(t, tc.expectedScope, resourceExhausted.Scope)
assert.Equal(t, tc.expectedMessage, resourceExhausted.Message)
} else {
assert.NoError(t, err)
}
// Sleep for 1 second for the rate limits to reset.
time.Sleep(1 * time.Second)
})
}
})
}
}