go.temporal.io/server@v1.23.0/common/persistence/persistence_rate_limited_clients.go (about) 1 // The MIT License 2 // 3 // Copyright (c) 2020 Temporal Technologies Inc. All rights reserved. 4 // 5 // Copyright (c) 2020 Uber Technologies, Inc. 6 // 7 // Permission is hereby granted, free of charge, to any person obtaining a copy 8 // of this software and associated documentation files (the "Software"), to deal 9 // in the Software without restriction, including without limitation the rights 10 // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 // copies of the Software, and to permit persons to whom the Software is 12 // furnished to do so, subject to the following conditions: 13 // 14 // The above copyright notice and this permission notice shall be included in 15 // all copies or substantial portions of the Software. 16 // 17 // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 23 // THE SOFTWARE. 24 25 package persistence 26 27 import ( 28 "context" 29 "time" 30 31 commonpb "go.temporal.io/api/common/v1" 32 enumspb "go.temporal.io/api/enums/v1" 33 "go.temporal.io/api/serviceerror" 34 "go.temporal.io/server/common/headers" 35 "go.temporal.io/server/common/log" 36 "go.temporal.io/server/common/quotas" 37 "go.temporal.io/server/service/history/tasks" 38 ) 39 40 const ( 41 RateLimitDefaultToken = 1 42 CallerSegmentMissing = -1 43 ) 44 45 var ( 46 // ErrPersistenceLimitExceeded is the error indicating QPS limit reached. 47 ErrPersistenceLimitExceeded = serviceerror.NewResourceExhausted(enumspb.RESOURCE_EXHAUSTED_CAUSE_PERSISTENCE_LIMIT, "Persistence Max QPS Reached.") 48 ) 49 50 type ( 51 shardRateLimitedPersistenceClient struct { 52 rateLimiter quotas.RequestRateLimiter 53 persistence ShardManager 54 logger log.Logger 55 } 56 57 executionRateLimitedPersistenceClient struct { 58 rateLimiter quotas.RequestRateLimiter 59 persistence ExecutionManager 60 logger log.Logger 61 } 62 63 taskRateLimitedPersistenceClient struct { 64 rateLimiter quotas.RequestRateLimiter 65 persistence TaskManager 66 logger log.Logger 67 } 68 69 metadataRateLimitedPersistenceClient struct { 70 rateLimiter quotas.RequestRateLimiter 71 persistence MetadataManager 72 logger log.Logger 73 } 74 75 clusterMetadataRateLimitedPersistenceClient struct { 76 rateLimiter quotas.RequestRateLimiter 77 persistence ClusterMetadataManager 78 logger log.Logger 79 } 80 81 queueRateLimitedPersistenceClient struct { 82 rateLimiter quotas.RequestRateLimiter 83 persistence Queue 84 logger log.Logger 85 } 86 ) 87 88 var _ ShardManager = (*shardRateLimitedPersistenceClient)(nil) 89 var _ ExecutionManager = (*executionRateLimitedPersistenceClient)(nil) 90 var _ TaskManager = (*taskRateLimitedPersistenceClient)(nil) 91 var _ MetadataManager = (*metadataRateLimitedPersistenceClient)(nil) 92 var _ ClusterMetadataManager = (*clusterMetadataRateLimitedPersistenceClient)(nil) 93 var _ Queue = (*queueRateLimitedPersistenceClient)(nil) 94 95 // NewShardPersistenceRateLimitedClient creates a client to manage shards 96 func NewShardPersistenceRateLimitedClient(persistence ShardManager, rateLimiter quotas.RequestRateLimiter, logger log.Logger) ShardManager { 97 return &shardRateLimitedPersistenceClient{ 98 persistence: persistence, 99 rateLimiter: rateLimiter, 100 logger: logger, 101 } 102 } 103 104 // NewExecutionPersistenceRateLimitedClient creates a client to manage executions 105 func NewExecutionPersistenceRateLimitedClient(persistence ExecutionManager, rateLimiter quotas.RequestRateLimiter, logger log.Logger) ExecutionManager { 106 return &executionRateLimitedPersistenceClient{ 107 persistence: persistence, 108 rateLimiter: rateLimiter, 109 logger: logger, 110 } 111 } 112 113 // NewTaskPersistenceRateLimitedClient creates a client to manage tasks 114 func NewTaskPersistenceRateLimitedClient(persistence TaskManager, rateLimiter quotas.RequestRateLimiter, logger log.Logger) TaskManager { 115 return &taskRateLimitedPersistenceClient{ 116 persistence: persistence, 117 rateLimiter: rateLimiter, 118 logger: logger, 119 } 120 } 121 122 // NewMetadataPersistenceRateLimitedClient creates a MetadataManager client to manage metadata 123 func NewMetadataPersistenceRateLimitedClient(persistence MetadataManager, rateLimiter quotas.RequestRateLimiter, logger log.Logger) MetadataManager { 124 return &metadataRateLimitedPersistenceClient{ 125 persistence: persistence, 126 rateLimiter: rateLimiter, 127 logger: logger, 128 } 129 } 130 131 // NewClusterMetadataPersistenceRateLimitedClient creates a MetadataManager client to manage metadata 132 func NewClusterMetadataPersistenceRateLimitedClient(persistence ClusterMetadataManager, rateLimiter quotas.RequestRateLimiter, logger log.Logger) ClusterMetadataManager { 133 return &clusterMetadataRateLimitedPersistenceClient{ 134 persistence: persistence, 135 rateLimiter: rateLimiter, 136 logger: logger, 137 } 138 } 139 140 // NewQueuePersistenceRateLimitedClient creates a client to manage queue 141 func NewQueuePersistenceRateLimitedClient(persistence Queue, rateLimiter quotas.RequestRateLimiter, logger log.Logger) Queue { 142 return &queueRateLimitedPersistenceClient{ 143 persistence: persistence, 144 rateLimiter: rateLimiter, 145 logger: logger, 146 } 147 } 148 149 func (p *shardRateLimitedPersistenceClient) GetName() string { 150 return p.persistence.GetName() 151 } 152 153 func (p *shardRateLimitedPersistenceClient) GetOrCreateShard( 154 ctx context.Context, 155 request *GetOrCreateShardRequest, 156 ) (*GetOrCreateShardResponse, error) { 157 if ok := allow(ctx, "GetOrCreateShard", request.ShardID, p.rateLimiter); !ok { 158 return nil, ErrPersistenceLimitExceeded 159 } 160 161 response, err := p.persistence.GetOrCreateShard(ctx, request) 162 return response, err 163 } 164 165 func (p *shardRateLimitedPersistenceClient) UpdateShard( 166 ctx context.Context, 167 request *UpdateShardRequest, 168 ) error { 169 if ok := allow(ctx, "UpdateShard", request.ShardInfo.ShardId, p.rateLimiter); !ok { 170 return ErrPersistenceLimitExceeded 171 } 172 173 return p.persistence.UpdateShard(ctx, request) 174 } 175 176 func (p *shardRateLimitedPersistenceClient) AssertShardOwnership( 177 ctx context.Context, 178 request *AssertShardOwnershipRequest, 179 ) error { 180 if ok := allow(ctx, "AssertShardOwnership", request.ShardID, p.rateLimiter); !ok { 181 return ErrPersistenceLimitExceeded 182 } 183 184 return p.persistence.AssertShardOwnership(ctx, request) 185 } 186 187 func (p *shardRateLimitedPersistenceClient) Close() { 188 p.persistence.Close() 189 } 190 191 func (p *executionRateLimitedPersistenceClient) GetName() string { 192 return p.persistence.GetName() 193 } 194 195 func (p *executionRateLimitedPersistenceClient) GetHistoryBranchUtil() HistoryBranchUtil { 196 return p.persistence.GetHistoryBranchUtil() 197 } 198 199 func (p *executionRateLimitedPersistenceClient) CreateWorkflowExecution( 200 ctx context.Context, 201 request *CreateWorkflowExecutionRequest, 202 ) (*CreateWorkflowExecutionResponse, error) { 203 if ok := allow(ctx, "CreateWorkflowExecution", request.ShardID, p.rateLimiter); !ok { 204 return nil, ErrPersistenceLimitExceeded 205 } 206 207 response, err := p.persistence.CreateWorkflowExecution(ctx, request) 208 return response, err 209 } 210 211 func (p *executionRateLimitedPersistenceClient) GetWorkflowExecution( 212 ctx context.Context, 213 request *GetWorkflowExecutionRequest, 214 ) (*GetWorkflowExecutionResponse, error) { 215 if ok := allow(ctx, "GetWorkflowExecution", request.ShardID, p.rateLimiter); !ok { 216 return nil, ErrPersistenceLimitExceeded 217 } 218 219 response, err := p.persistence.GetWorkflowExecution(ctx, request) 220 return response, err 221 } 222 223 func (p *executionRateLimitedPersistenceClient) SetWorkflowExecution( 224 ctx context.Context, 225 request *SetWorkflowExecutionRequest, 226 ) (*SetWorkflowExecutionResponse, error) { 227 if ok := allow(ctx, "SetWorkflowExecution", request.ShardID, p.rateLimiter); !ok { 228 return nil, ErrPersistenceLimitExceeded 229 } 230 231 response, err := p.persistence.SetWorkflowExecution(ctx, request) 232 return response, err 233 } 234 235 func (p *executionRateLimitedPersistenceClient) UpdateWorkflowExecution( 236 ctx context.Context, 237 request *UpdateWorkflowExecutionRequest, 238 ) (*UpdateWorkflowExecutionResponse, error) { 239 if ok := allow(ctx, "UpdateWorkflowExecution", request.ShardID, p.rateLimiter); !ok { 240 return nil, ErrPersistenceLimitExceeded 241 } 242 243 resp, err := p.persistence.UpdateWorkflowExecution(ctx, request) 244 return resp, err 245 } 246 247 func (p *executionRateLimitedPersistenceClient) ConflictResolveWorkflowExecution( 248 ctx context.Context, 249 request *ConflictResolveWorkflowExecutionRequest, 250 ) (*ConflictResolveWorkflowExecutionResponse, error) { 251 if ok := allow(ctx, "ConflictResolveWorkflowExecution", request.ShardID, p.rateLimiter); !ok { 252 return nil, ErrPersistenceLimitExceeded 253 } 254 255 response, err := p.persistence.ConflictResolveWorkflowExecution(ctx, request) 256 return response, err 257 } 258 259 func (p *executionRateLimitedPersistenceClient) DeleteWorkflowExecution( 260 ctx context.Context, 261 request *DeleteWorkflowExecutionRequest, 262 ) error { 263 if ok := allow(ctx, "DeleteWorkflowExecution", request.ShardID, p.rateLimiter); !ok { 264 return ErrPersistenceLimitExceeded 265 } 266 267 return p.persistence.DeleteWorkflowExecution(ctx, request) 268 } 269 270 func (p *executionRateLimitedPersistenceClient) DeleteCurrentWorkflowExecution( 271 ctx context.Context, 272 request *DeleteCurrentWorkflowExecutionRequest, 273 ) error { 274 if ok := allow(ctx, "DeleteCurrentWorkflowExecution", request.ShardID, p.rateLimiter); !ok { 275 return ErrPersistenceLimitExceeded 276 } 277 278 return p.persistence.DeleteCurrentWorkflowExecution(ctx, request) 279 } 280 281 func (p *executionRateLimitedPersistenceClient) GetCurrentExecution( 282 ctx context.Context, 283 request *GetCurrentExecutionRequest, 284 ) (*GetCurrentExecutionResponse, error) { 285 if ok := allow(ctx, "GetCurrentExecution", request.ShardID, p.rateLimiter); !ok { 286 return nil, ErrPersistenceLimitExceeded 287 } 288 289 response, err := p.persistence.GetCurrentExecution(ctx, request) 290 return response, err 291 } 292 293 func (p *executionRateLimitedPersistenceClient) ListConcreteExecutions( 294 ctx context.Context, 295 request *ListConcreteExecutionsRequest, 296 ) (*ListConcreteExecutionsResponse, error) { 297 if ok := allow(ctx, "ListConcreteExecutions", request.ShardID, p.rateLimiter); !ok { 298 return nil, ErrPersistenceLimitExceeded 299 } 300 301 response, err := p.persistence.ListConcreteExecutions(ctx, request) 302 return response, err 303 } 304 305 func (p *executionRateLimitedPersistenceClient) RegisterHistoryTaskReader( 306 ctx context.Context, 307 request *RegisterHistoryTaskReaderRequest, 308 ) error { 309 // hint methods don't actually hint DB, so don't go through persistence rate limiter 310 return p.persistence.RegisterHistoryTaskReader(ctx, request) 311 } 312 313 func (p *executionRateLimitedPersistenceClient) UnregisterHistoryTaskReader( 314 ctx context.Context, 315 request *UnregisterHistoryTaskReaderRequest, 316 ) { 317 // hint methods don't actually hint DB, so don't go through persistence rate limiter 318 p.persistence.UnregisterHistoryTaskReader(ctx, request) 319 } 320 321 func (p *executionRateLimitedPersistenceClient) UpdateHistoryTaskReaderProgress( 322 ctx context.Context, 323 request *UpdateHistoryTaskReaderProgressRequest, 324 ) { 325 // hint methods don't actually hint DB, so don't go through persistence rate limiter 326 p.persistence.UpdateHistoryTaskReaderProgress(ctx, request) 327 } 328 329 func (p *executionRateLimitedPersistenceClient) AddHistoryTasks( 330 ctx context.Context, 331 request *AddHistoryTasksRequest, 332 ) error { 333 if ok := allow(ctx, "AddHistoryTasks", request.ShardID, p.rateLimiter); !ok { 334 return ErrPersistenceLimitExceeded 335 } 336 337 return p.persistence.AddHistoryTasks(ctx, request) 338 } 339 340 func (p *executionRateLimitedPersistenceClient) GetHistoryTasks( 341 ctx context.Context, 342 request *GetHistoryTasksRequest, 343 ) (*GetHistoryTasksResponse, error) { 344 if ok := allow( 345 ctx, 346 ConstructHistoryTaskAPI("GetHistoryTasks", request.TaskCategory), 347 request.ShardID, 348 p.rateLimiter, 349 ); !ok { 350 return nil, ErrPersistenceLimitExceeded 351 } 352 353 response, err := p.persistence.GetHistoryTasks(ctx, request) 354 return response, err 355 } 356 357 func (p *executionRateLimitedPersistenceClient) CompleteHistoryTask( 358 ctx context.Context, 359 request *CompleteHistoryTaskRequest, 360 ) error { 361 if ok := allow( 362 ctx, 363 ConstructHistoryTaskAPI("CompleteHistoryTask", request.TaskCategory), 364 request.ShardID, 365 p.rateLimiter, 366 ); !ok { 367 return ErrPersistenceLimitExceeded 368 } 369 370 return p.persistence.CompleteHistoryTask(ctx, request) 371 } 372 373 func (p *executionRateLimitedPersistenceClient) RangeCompleteHistoryTasks( 374 ctx context.Context, 375 request *RangeCompleteHistoryTasksRequest, 376 ) error { 377 if ok := allow( 378 ctx, 379 ConstructHistoryTaskAPI("RangeCompleteHistoryTasks", request.TaskCategory), 380 request.ShardID, 381 p.rateLimiter, 382 ); !ok { 383 return ErrPersistenceLimitExceeded 384 } 385 386 return p.persistence.RangeCompleteHistoryTasks(ctx, request) 387 } 388 389 func (p *executionRateLimitedPersistenceClient) PutReplicationTaskToDLQ( 390 ctx context.Context, 391 request *PutReplicationTaskToDLQRequest, 392 ) error { 393 if ok := allow(ctx, "PutReplicationTaskToDLQ", request.ShardID, p.rateLimiter); !ok { 394 return ErrPersistenceLimitExceeded 395 } 396 397 return p.persistence.PutReplicationTaskToDLQ(ctx, request) 398 } 399 400 func (p *executionRateLimitedPersistenceClient) GetReplicationTasksFromDLQ( 401 ctx context.Context, 402 request *GetReplicationTasksFromDLQRequest, 403 ) (*GetHistoryTasksResponse, error) { 404 if ok := allow(ctx, "GetReplicationTasksFromDLQ", request.ShardID, p.rateLimiter); !ok { 405 return nil, ErrPersistenceLimitExceeded 406 } 407 408 return p.persistence.GetReplicationTasksFromDLQ(ctx, request) 409 } 410 411 func (p *executionRateLimitedPersistenceClient) DeleteReplicationTaskFromDLQ( 412 ctx context.Context, 413 request *DeleteReplicationTaskFromDLQRequest, 414 ) error { 415 if ok := allow(ctx, "DeleteReplicationTaskFromDLQ", request.ShardID, p.rateLimiter); !ok { 416 return ErrPersistenceLimitExceeded 417 } 418 419 return p.persistence.DeleteReplicationTaskFromDLQ(ctx, request) 420 } 421 422 func (p *executionRateLimitedPersistenceClient) RangeDeleteReplicationTaskFromDLQ( 423 ctx context.Context, 424 request *RangeDeleteReplicationTaskFromDLQRequest, 425 ) error { 426 if ok := allow(ctx, "RangeDeleteReplicationTaskFromDLQ", request.ShardID, p.rateLimiter); !ok { 427 return ErrPersistenceLimitExceeded 428 } 429 430 return p.persistence.RangeDeleteReplicationTaskFromDLQ(ctx, request) 431 } 432 433 func (p *executionRateLimitedPersistenceClient) IsReplicationDLQEmpty( 434 ctx context.Context, 435 request *GetReplicationTasksFromDLQRequest, 436 ) (bool, error) { 437 if ok := allow(ctx, "IsReplicationDLQEmpty", request.ShardID, p.rateLimiter); !ok { 438 return true, ErrPersistenceLimitExceeded 439 } 440 441 return p.persistence.IsReplicationDLQEmpty(ctx, request) 442 } 443 444 func (p *executionRateLimitedPersistenceClient) Close() { 445 p.persistence.Close() 446 } 447 448 func (p *taskRateLimitedPersistenceClient) GetName() string { 449 return p.persistence.GetName() 450 } 451 452 func (p *taskRateLimitedPersistenceClient) CreateTasks( 453 ctx context.Context, 454 request *CreateTasksRequest, 455 ) (*CreateTasksResponse, error) { 456 if ok := allow(ctx, "CreateTasks", CallerSegmentMissing, p.rateLimiter); !ok { 457 return nil, ErrPersistenceLimitExceeded 458 } 459 460 response, err := p.persistence.CreateTasks(ctx, request) 461 return response, err 462 } 463 464 func (p *taskRateLimitedPersistenceClient) GetTasks( 465 ctx context.Context, 466 request *GetTasksRequest, 467 ) (*GetTasksResponse, error) { 468 if ok := allow(ctx, "GetTasks", CallerSegmentMissing, p.rateLimiter); !ok { 469 return nil, ErrPersistenceLimitExceeded 470 } 471 472 response, err := p.persistence.GetTasks(ctx, request) 473 return response, err 474 } 475 476 func (p *taskRateLimitedPersistenceClient) CompleteTask( 477 ctx context.Context, 478 request *CompleteTaskRequest, 479 ) error { 480 if ok := allow(ctx, "CompleteTask", CallerSegmentMissing, p.rateLimiter); !ok { 481 return ErrPersistenceLimitExceeded 482 } 483 484 return p.persistence.CompleteTask(ctx, request) 485 } 486 487 func (p *taskRateLimitedPersistenceClient) CompleteTasksLessThan( 488 ctx context.Context, 489 request *CompleteTasksLessThanRequest, 490 ) (int, error) { 491 if ok := allow(ctx, "CompleteTasksLessThan", CallerSegmentMissing, p.rateLimiter); !ok { 492 return 0, ErrPersistenceLimitExceeded 493 } 494 return p.persistence.CompleteTasksLessThan(ctx, request) 495 } 496 497 func (p *taskRateLimitedPersistenceClient) CreateTaskQueue( 498 ctx context.Context, 499 request *CreateTaskQueueRequest, 500 ) (*CreateTaskQueueResponse, error) { 501 if ok := allow(ctx, "CreateTaskQueue", CallerSegmentMissing, p.rateLimiter); !ok { 502 return nil, ErrPersistenceLimitExceeded 503 } 504 return p.persistence.CreateTaskQueue(ctx, request) 505 } 506 507 func (p *taskRateLimitedPersistenceClient) UpdateTaskQueue( 508 ctx context.Context, 509 request *UpdateTaskQueueRequest, 510 ) (*UpdateTaskQueueResponse, error) { 511 if ok := allow(ctx, "UpdateTaskQueue", CallerSegmentMissing, p.rateLimiter); !ok { 512 return nil, ErrPersistenceLimitExceeded 513 } 514 return p.persistence.UpdateTaskQueue(ctx, request) 515 } 516 517 func (p *taskRateLimitedPersistenceClient) GetTaskQueue( 518 ctx context.Context, 519 request *GetTaskQueueRequest, 520 ) (*GetTaskQueueResponse, error) { 521 if ok := allow(ctx, "GetTaskQueue", CallerSegmentMissing, p.rateLimiter); !ok { 522 return nil, ErrPersistenceLimitExceeded 523 } 524 return p.persistence.GetTaskQueue(ctx, request) 525 } 526 527 func (p *taskRateLimitedPersistenceClient) ListTaskQueue( 528 ctx context.Context, 529 request *ListTaskQueueRequest, 530 ) (*ListTaskQueueResponse, error) { 531 if ok := allow(ctx, "ListTaskQueue", CallerSegmentMissing, p.rateLimiter); !ok { 532 return nil, ErrPersistenceLimitExceeded 533 } 534 return p.persistence.ListTaskQueue(ctx, request) 535 } 536 537 func (p *taskRateLimitedPersistenceClient) DeleteTaskQueue( 538 ctx context.Context, 539 request *DeleteTaskQueueRequest, 540 ) error { 541 if ok := allow(ctx, "DeleteTaskQueue", CallerSegmentMissing, p.rateLimiter); !ok { 542 return ErrPersistenceLimitExceeded 543 } 544 return p.persistence.DeleteTaskQueue(ctx, request) 545 } 546 547 func (p taskRateLimitedPersistenceClient) GetTaskQueueUserData( 548 ctx context.Context, 549 request *GetTaskQueueUserDataRequest, 550 ) (*GetTaskQueueUserDataResponse, error) { 551 if ok := allow(ctx, "GetTaskQueueUserData", CallerSegmentMissing, p.rateLimiter); !ok { 552 return nil, ErrPersistenceLimitExceeded 553 } 554 return p.persistence.GetTaskQueueUserData(ctx, request) 555 } 556 557 func (p taskRateLimitedPersistenceClient) UpdateTaskQueueUserData( 558 ctx context.Context, 559 request *UpdateTaskQueueUserDataRequest, 560 ) error { 561 if ok := allow(ctx, "UpdateTaskQueueUserData", CallerSegmentMissing, p.rateLimiter); !ok { 562 return ErrPersistenceLimitExceeded 563 } 564 return p.persistence.UpdateTaskQueueUserData(ctx, request) 565 } 566 567 func (p taskRateLimitedPersistenceClient) ListTaskQueueUserDataEntries( 568 ctx context.Context, 569 request *ListTaskQueueUserDataEntriesRequest, 570 ) (*ListTaskQueueUserDataEntriesResponse, error) { 571 if ok := allow(ctx, "ListTaskQueueUserDataEntries", CallerSegmentMissing, p.rateLimiter); !ok { 572 return nil, ErrPersistenceLimitExceeded 573 } 574 return p.persistence.ListTaskQueueUserDataEntries(ctx, request) 575 } 576 577 func (p taskRateLimitedPersistenceClient) GetTaskQueuesByBuildId(ctx context.Context, request *GetTaskQueuesByBuildIdRequest) ([]string, error) { 578 if ok := allow(ctx, "GetTaskQueuesByBuildId", CallerSegmentMissing, p.rateLimiter); !ok { 579 return nil, ErrPersistenceLimitExceeded 580 } 581 return p.persistence.GetTaskQueuesByBuildId(ctx, request) 582 } 583 584 func (p taskRateLimitedPersistenceClient) CountTaskQueuesByBuildId(ctx context.Context, request *CountTaskQueuesByBuildIdRequest) (int, error) { 585 if ok := allow(ctx, "CountTaskQueuesByBuildId", CallerSegmentMissing, p.rateLimiter); !ok { 586 return 0, ErrPersistenceLimitExceeded 587 } 588 return p.persistence.CountTaskQueuesByBuildId(ctx, request) 589 } 590 591 func (p *taskRateLimitedPersistenceClient) Close() { 592 p.persistence.Close() 593 } 594 595 func (p *metadataRateLimitedPersistenceClient) GetName() string { 596 return p.persistence.GetName() 597 } 598 599 func (p *metadataRateLimitedPersistenceClient) CreateNamespace( 600 ctx context.Context, 601 request *CreateNamespaceRequest, 602 ) (*CreateNamespaceResponse, error) { 603 if ok := allow(ctx, "CreateNamespace", CallerSegmentMissing, p.rateLimiter); !ok { 604 return nil, ErrPersistenceLimitExceeded 605 } 606 607 response, err := p.persistence.CreateNamespace(ctx, request) 608 return response, err 609 } 610 611 func (p *metadataRateLimitedPersistenceClient) GetNamespace( 612 ctx context.Context, 613 request *GetNamespaceRequest, 614 ) (*GetNamespaceResponse, error) { 615 if ok := allow(ctx, "GetNamespace", CallerSegmentMissing, p.rateLimiter); !ok { 616 return nil, ErrPersistenceLimitExceeded 617 } 618 619 response, err := p.persistence.GetNamespace(ctx, request) 620 return response, err 621 } 622 623 func (p *metadataRateLimitedPersistenceClient) UpdateNamespace( 624 ctx context.Context, 625 request *UpdateNamespaceRequest, 626 ) error { 627 if ok := allow(ctx, "UpdateNamespace", CallerSegmentMissing, p.rateLimiter); !ok { 628 return ErrPersistenceLimitExceeded 629 } 630 631 return p.persistence.UpdateNamespace(ctx, request) 632 } 633 634 func (p *metadataRateLimitedPersistenceClient) RenameNamespace( 635 ctx context.Context, 636 request *RenameNamespaceRequest, 637 ) error { 638 if ok := allow(ctx, "RenameNamespace", CallerSegmentMissing, p.rateLimiter); !ok { 639 return ErrPersistenceLimitExceeded 640 } 641 642 return p.persistence.RenameNamespace(ctx, request) 643 } 644 645 func (p *metadataRateLimitedPersistenceClient) DeleteNamespace( 646 ctx context.Context, 647 request *DeleteNamespaceRequest, 648 ) error { 649 if ok := allow(ctx, "DeleteNamespace", CallerSegmentMissing, p.rateLimiter); !ok { 650 return ErrPersistenceLimitExceeded 651 } 652 653 return p.persistence.DeleteNamespace(ctx, request) 654 } 655 656 func (p *metadataRateLimitedPersistenceClient) DeleteNamespaceByName( 657 ctx context.Context, 658 request *DeleteNamespaceByNameRequest, 659 ) error { 660 if ok := allow(ctx, "DeleteNamespaceByName", CallerSegmentMissing, p.rateLimiter); !ok { 661 return ErrPersistenceLimitExceeded 662 } 663 664 return p.persistence.DeleteNamespaceByName(ctx, request) 665 } 666 667 func (p *metadataRateLimitedPersistenceClient) ListNamespaces( 668 ctx context.Context, 669 request *ListNamespacesRequest, 670 ) (*ListNamespacesResponse, error) { 671 if ok := allow(ctx, "ListNamespaces", CallerSegmentMissing, p.rateLimiter); !ok { 672 return nil, ErrPersistenceLimitExceeded 673 } 674 675 response, err := p.persistence.ListNamespaces(ctx, request) 676 return response, err 677 } 678 679 func (p *metadataRateLimitedPersistenceClient) GetMetadata( 680 ctx context.Context, 681 ) (*GetMetadataResponse, error) { 682 if ok := allow(ctx, "GetMetadata", CallerSegmentMissing, p.rateLimiter); !ok { 683 return nil, ErrPersistenceLimitExceeded 684 } 685 686 response, err := p.persistence.GetMetadata(ctx) 687 return response, err 688 } 689 690 func (p *metadataRateLimitedPersistenceClient) InitializeSystemNamespaces( 691 ctx context.Context, 692 currentClusterName string, 693 ) error { 694 if ok := allow(ctx, "InitializeSystemNamespaces", CallerSegmentMissing, p.rateLimiter); !ok { 695 return ErrPersistenceLimitExceeded 696 } 697 return p.persistence.InitializeSystemNamespaces(ctx, currentClusterName) 698 } 699 700 func (p *metadataRateLimitedPersistenceClient) Close() { 701 p.persistence.Close() 702 } 703 704 // AppendHistoryNodes add a node to history node table 705 func (p *executionRateLimitedPersistenceClient) AppendHistoryNodes( 706 ctx context.Context, 707 request *AppendHistoryNodesRequest, 708 ) (*AppendHistoryNodesResponse, error) { 709 if ok := allow(ctx, "AppendHistoryNodes", request.ShardID, p.rateLimiter); !ok { 710 return nil, ErrPersistenceLimitExceeded 711 } 712 return p.persistence.AppendHistoryNodes(ctx, request) 713 } 714 715 // AppendRawHistoryNodes add a node to history node table 716 func (p *executionRateLimitedPersistenceClient) AppendRawHistoryNodes( 717 ctx context.Context, 718 request *AppendRawHistoryNodesRequest, 719 ) (*AppendHistoryNodesResponse, error) { 720 if ok := allow(ctx, "AppendRawHistoryNodes", request.ShardID, p.rateLimiter); !ok { 721 return nil, ErrPersistenceLimitExceeded 722 } 723 return p.persistence.AppendRawHistoryNodes(ctx, request) 724 } 725 726 // ReadHistoryBranch returns history node data for a branch 727 func (p *executionRateLimitedPersistenceClient) ReadHistoryBranch( 728 ctx context.Context, 729 request *ReadHistoryBranchRequest, 730 ) (*ReadHistoryBranchResponse, error) { 731 if ok := allow(ctx, "ReadHistoryBranch", request.ShardID, p.rateLimiter); !ok { 732 return nil, ErrPersistenceLimitExceeded 733 } 734 response, err := p.persistence.ReadHistoryBranch(ctx, request) 735 return response, err 736 } 737 738 // ReadHistoryBranchReverse returns history node data for a branch 739 func (p *executionRateLimitedPersistenceClient) ReadHistoryBranchReverse( 740 ctx context.Context, 741 request *ReadHistoryBranchReverseRequest, 742 ) (*ReadHistoryBranchReverseResponse, error) { 743 if ok := allow(ctx, "ReadHistoryBranchReverse", request.ShardID, p.rateLimiter); !ok { 744 return nil, ErrPersistenceLimitExceeded 745 } 746 response, err := p.persistence.ReadHistoryBranchReverse(ctx, request) 747 return response, err 748 } 749 750 // ReadHistoryBranchByBatch returns history node data for a branch 751 func (p *executionRateLimitedPersistenceClient) ReadHistoryBranchByBatch( 752 ctx context.Context, 753 request *ReadHistoryBranchRequest, 754 ) (*ReadHistoryBranchByBatchResponse, error) { 755 if ok := allow(ctx, "ReadHistoryBranchByBatch", request.ShardID, p.rateLimiter); !ok { 756 return nil, ErrPersistenceLimitExceeded 757 } 758 response, err := p.persistence.ReadHistoryBranchByBatch(ctx, request) 759 return response, err 760 } 761 762 // ReadHistoryBranchByBatch returns history node data for a branch 763 func (p *executionRateLimitedPersistenceClient) ReadRawHistoryBranch( 764 ctx context.Context, 765 request *ReadHistoryBranchRequest, 766 ) (*ReadRawHistoryBranchResponse, error) { 767 if ok := allow(ctx, "ReadRawHistoryBranch", request.ShardID, p.rateLimiter); !ok { 768 return nil, ErrPersistenceLimitExceeded 769 } 770 response, err := p.persistence.ReadRawHistoryBranch(ctx, request) 771 return response, err 772 } 773 774 // ForkHistoryBranch forks a new branch from a old branch 775 func (p *executionRateLimitedPersistenceClient) ForkHistoryBranch( 776 ctx context.Context, 777 request *ForkHistoryBranchRequest, 778 ) (*ForkHistoryBranchResponse, error) { 779 if ok := allow(ctx, "ForkHistoryBranch", request.ShardID, p.rateLimiter); !ok { 780 return nil, ErrPersistenceLimitExceeded 781 } 782 response, err := p.persistence.ForkHistoryBranch(ctx, request) 783 return response, err 784 } 785 786 // DeleteHistoryBranch removes a branch 787 func (p *executionRateLimitedPersistenceClient) DeleteHistoryBranch( 788 ctx context.Context, 789 request *DeleteHistoryBranchRequest, 790 ) error { 791 if ok := allow(ctx, "DeleteHistoryBranch", request.ShardID, p.rateLimiter); !ok { 792 return ErrPersistenceLimitExceeded 793 } 794 return p.persistence.DeleteHistoryBranch(ctx, request) 795 } 796 797 // TrimHistoryBranch trims a branch 798 func (p *executionRateLimitedPersistenceClient) TrimHistoryBranch( 799 ctx context.Context, 800 request *TrimHistoryBranchRequest, 801 ) (*TrimHistoryBranchResponse, error) { 802 if ok := allow(ctx, "TrimHistoryBranch", request.ShardID, p.rateLimiter); !ok { 803 return nil, ErrPersistenceLimitExceeded 804 } 805 resp, err := p.persistence.TrimHistoryBranch(ctx, request) 806 return resp, err 807 } 808 809 // GetHistoryTree returns all branch information of a tree 810 func (p *executionRateLimitedPersistenceClient) GetHistoryTree( 811 ctx context.Context, 812 request *GetHistoryTreeRequest, 813 ) (*GetHistoryTreeResponse, error) { 814 if ok := allow(ctx, "GetHistoryTree", request.ShardID, p.rateLimiter); !ok { 815 return nil, ErrPersistenceLimitExceeded 816 } 817 response, err := p.persistence.GetHistoryTree(ctx, request) 818 return response, err 819 } 820 821 func (p *executionRateLimitedPersistenceClient) GetAllHistoryTreeBranches( 822 ctx context.Context, 823 request *GetAllHistoryTreeBranchesRequest, 824 ) (*GetAllHistoryTreeBranchesResponse, error) { 825 if ok := allow(ctx, "GetAllHistoryTreeBranches", CallerSegmentMissing, p.rateLimiter); !ok { 826 return nil, ErrPersistenceLimitExceeded 827 } 828 response, err := p.persistence.GetAllHistoryTreeBranches(ctx, request) 829 return response, err 830 } 831 832 func (p *queueRateLimitedPersistenceClient) EnqueueMessage( 833 ctx context.Context, 834 blob *commonpb.DataBlob, 835 ) error { 836 if ok := allow(ctx, "EnqueueMessage", CallerSegmentMissing, p.rateLimiter); !ok { 837 return ErrPersistenceLimitExceeded 838 } 839 840 return p.persistence.EnqueueMessage(ctx, blob) 841 } 842 843 func (p *queueRateLimitedPersistenceClient) ReadMessages( 844 ctx context.Context, 845 lastMessageID int64, 846 maxCount int, 847 ) ([]*QueueMessage, error) { 848 if ok := allow(ctx, "ReadMessages", CallerSegmentMissing, p.rateLimiter); !ok { 849 return nil, ErrPersistenceLimitExceeded 850 } 851 852 return p.persistence.ReadMessages(ctx, lastMessageID, maxCount) 853 } 854 855 func (p *queueRateLimitedPersistenceClient) UpdateAckLevel( 856 ctx context.Context, 857 metadata *InternalQueueMetadata, 858 ) error { 859 if ok := allow(ctx, "UpdateAckLevel", CallerSegmentMissing, p.rateLimiter); !ok { 860 return ErrPersistenceLimitExceeded 861 } 862 863 return p.persistence.UpdateAckLevel(ctx, metadata) 864 } 865 866 func (p *queueRateLimitedPersistenceClient) GetAckLevels( 867 ctx context.Context, 868 ) (*InternalQueueMetadata, error) { 869 if ok := allow(ctx, "GetAckLevels", CallerSegmentMissing, p.rateLimiter); !ok { 870 return nil, ErrPersistenceLimitExceeded 871 } 872 873 return p.persistence.GetAckLevels(ctx) 874 } 875 876 func (p *queueRateLimitedPersistenceClient) DeleteMessagesBefore( 877 ctx context.Context, 878 messageID int64, 879 ) error { 880 if ok := allow(ctx, "DeleteMessagesBefore", CallerSegmentMissing, p.rateLimiter); !ok { 881 return ErrPersistenceLimitExceeded 882 } 883 884 return p.persistence.DeleteMessagesBefore(ctx, messageID) 885 } 886 887 func (p *queueRateLimitedPersistenceClient) EnqueueMessageToDLQ( 888 ctx context.Context, 889 blob *commonpb.DataBlob, 890 ) (int64, error) { 891 if ok := allow(ctx, "EnqueueMessageToDLQ", CallerSegmentMissing, p.rateLimiter); !ok { 892 return EmptyQueueMessageID, ErrPersistenceLimitExceeded 893 } 894 895 return p.persistence.EnqueueMessageToDLQ(ctx, blob) 896 } 897 898 func (p *queueRateLimitedPersistenceClient) ReadMessagesFromDLQ( 899 ctx context.Context, 900 firstMessageID int64, 901 lastMessageID int64, 902 pageSize int, 903 pageToken []byte, 904 ) ([]*QueueMessage, []byte, error) { 905 if ok := allow(ctx, "ReadMessagesFromDLQ", CallerSegmentMissing, p.rateLimiter); !ok { 906 return nil, nil, ErrPersistenceLimitExceeded 907 } 908 909 return p.persistence.ReadMessagesFromDLQ(ctx, firstMessageID, lastMessageID, pageSize, pageToken) 910 } 911 912 func (p *queueRateLimitedPersistenceClient) RangeDeleteMessagesFromDLQ( 913 ctx context.Context, 914 firstMessageID int64, 915 lastMessageID int64, 916 ) error { 917 if ok := allow(ctx, "RangeDeleteMessagesFromDLQ", CallerSegmentMissing, p.rateLimiter); !ok { 918 return ErrPersistenceLimitExceeded 919 } 920 921 return p.persistence.RangeDeleteMessagesFromDLQ(ctx, firstMessageID, lastMessageID) 922 } 923 func (p *queueRateLimitedPersistenceClient) UpdateDLQAckLevel( 924 ctx context.Context, 925 metadata *InternalQueueMetadata, 926 ) error { 927 if ok := allow(ctx, "UpdateDLQAckLevel", CallerSegmentMissing, p.rateLimiter); !ok { 928 return ErrPersistenceLimitExceeded 929 } 930 931 return p.persistence.UpdateDLQAckLevel(ctx, metadata) 932 } 933 934 func (p *queueRateLimitedPersistenceClient) GetDLQAckLevels( 935 ctx context.Context, 936 ) (*InternalQueueMetadata, error) { 937 if ok := allow(ctx, "GetDLQAckLevels", CallerSegmentMissing, p.rateLimiter); !ok { 938 return nil, ErrPersistenceLimitExceeded 939 } 940 941 return p.persistence.GetDLQAckLevels(ctx) 942 } 943 944 func (p *queueRateLimitedPersistenceClient) DeleteMessageFromDLQ( 945 ctx context.Context, 946 messageID int64, 947 ) error { 948 if ok := allow(ctx, "DeleteMessageFromDLQ", CallerSegmentMissing, p.rateLimiter); !ok { 949 return ErrPersistenceLimitExceeded 950 } 951 952 return p.persistence.DeleteMessageFromDLQ(ctx, messageID) 953 } 954 955 func (p *queueRateLimitedPersistenceClient) Close() { 956 p.persistence.Close() 957 } 958 959 func (p *queueRateLimitedPersistenceClient) Init( 960 ctx context.Context, 961 blob *commonpb.DataBlob, 962 ) error { 963 return p.persistence.Init(ctx, blob) 964 } 965 966 func (c *clusterMetadataRateLimitedPersistenceClient) Close() { 967 c.persistence.Close() 968 } 969 970 func (c *clusterMetadataRateLimitedPersistenceClient) GetName() string { 971 return c.persistence.GetName() 972 } 973 974 func (c *clusterMetadataRateLimitedPersistenceClient) GetClusterMembers( 975 ctx context.Context, 976 request *GetClusterMembersRequest, 977 ) (*GetClusterMembersResponse, error) { 978 if ok := allow(ctx, "GetClusterMembers", CallerSegmentMissing, c.rateLimiter); !ok { 979 return nil, ErrPersistenceLimitExceeded 980 } 981 return c.persistence.GetClusterMembers(ctx, request) 982 } 983 984 func (c *clusterMetadataRateLimitedPersistenceClient) UpsertClusterMembership( 985 ctx context.Context, 986 request *UpsertClusterMembershipRequest, 987 ) error { 988 if ok := allow(ctx, "UpsertClusterMembership", CallerSegmentMissing, c.rateLimiter); !ok { 989 return ErrPersistenceLimitExceeded 990 } 991 return c.persistence.UpsertClusterMembership(ctx, request) 992 } 993 994 func (c *clusterMetadataRateLimitedPersistenceClient) PruneClusterMembership( 995 ctx context.Context, 996 request *PruneClusterMembershipRequest, 997 ) error { 998 if ok := allow(ctx, "PruneClusterMembership", CallerSegmentMissing, c.rateLimiter); !ok { 999 return ErrPersistenceLimitExceeded 1000 } 1001 return c.persistence.PruneClusterMembership(ctx, request) 1002 } 1003 1004 func (c *clusterMetadataRateLimitedPersistenceClient) ListClusterMetadata( 1005 ctx context.Context, 1006 request *ListClusterMetadataRequest, 1007 ) (*ListClusterMetadataResponse, error) { 1008 if ok := allow(ctx, "ListClusterMetadata", CallerSegmentMissing, c.rateLimiter); !ok { 1009 return nil, ErrPersistenceLimitExceeded 1010 } 1011 return c.persistence.ListClusterMetadata(ctx, request) 1012 } 1013 1014 func (c *clusterMetadataRateLimitedPersistenceClient) GetCurrentClusterMetadata( 1015 ctx context.Context, 1016 ) (*GetClusterMetadataResponse, error) { 1017 if ok := allow(ctx, "GetCurrentClusterMetadata", CallerSegmentMissing, c.rateLimiter); !ok { 1018 return nil, ErrPersistenceLimitExceeded 1019 } 1020 return c.persistence.GetCurrentClusterMetadata(ctx) 1021 } 1022 1023 func (c *clusterMetadataRateLimitedPersistenceClient) GetClusterMetadata( 1024 ctx context.Context, 1025 request *GetClusterMetadataRequest, 1026 ) (*GetClusterMetadataResponse, error) { 1027 if ok := allow(ctx, "GetClusterMetadata", CallerSegmentMissing, c.rateLimiter); !ok { 1028 return nil, ErrPersistenceLimitExceeded 1029 } 1030 return c.persistence.GetClusterMetadata(ctx, request) 1031 } 1032 1033 func (c *clusterMetadataRateLimitedPersistenceClient) SaveClusterMetadata( 1034 ctx context.Context, 1035 request *SaveClusterMetadataRequest, 1036 ) (bool, error) { 1037 if ok := allow(ctx, "SaveClusterMetadata", CallerSegmentMissing, c.rateLimiter); !ok { 1038 return false, ErrPersistenceLimitExceeded 1039 } 1040 return c.persistence.SaveClusterMetadata(ctx, request) 1041 } 1042 1043 func (c *clusterMetadataRateLimitedPersistenceClient) DeleteClusterMetadata( 1044 ctx context.Context, 1045 request *DeleteClusterMetadataRequest, 1046 ) error { 1047 if ok := allow(ctx, "DeleteClusterMetadata", CallerSegmentMissing, c.rateLimiter); !ok { 1048 return ErrPersistenceLimitExceeded 1049 } 1050 return c.persistence.DeleteClusterMetadata(ctx, request) 1051 } 1052 1053 func allow( 1054 ctx context.Context, 1055 api string, 1056 shardID int32, 1057 rateLimiter quotas.RequestRateLimiter, 1058 ) bool { 1059 callerInfo := headers.GetCallerInfo(ctx) 1060 return rateLimiter.Allow(time.Now().UTC(), quotas.NewRequest( 1061 api, 1062 RateLimitDefaultToken, 1063 callerInfo.CallerName, 1064 callerInfo.CallerType, 1065 shardID, 1066 callerInfo.CallOrigin, 1067 )) 1068 } 1069 1070 // TODO: change the value returned so it can also be used by 1071 // persistence metrics client. For now, it's only used by rate 1072 // limit client, and we don't really care about the actual value 1073 // returned, as long as they are different from each task category. 1074 func ConstructHistoryTaskAPI( 1075 baseAPI string, 1076 taskCategory tasks.Category, 1077 ) string { 1078 return baseAPI + taskCategory.Name() 1079 }