go.temporal.io/server@v1.23.0/common/persistence/visibility/visibility_manager_rate_limited.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 visibility
    26  
    27  import (
    28  	"context"
    29  	"time"
    30  
    31  	"go.temporal.io/server/common/dynamicconfig"
    32  	"go.temporal.io/server/common/headers"
    33  	"go.temporal.io/server/common/namespace"
    34  	"go.temporal.io/server/common/persistence"
    35  	"go.temporal.io/server/common/persistence/visibility/manager"
    36  	"go.temporal.io/server/common/quotas"
    37  )
    38  
    39  const (
    40  	RateLimitDefaultToken = 1
    41  )
    42  
    43  var _ manager.VisibilityManager = (*visibilityManagerRateLimited)(nil)
    44  
    45  type visibilityManagerRateLimited struct {
    46  	delegate         manager.VisibilityManager
    47  	readRateLimiter  quotas.RequestRateLimiter
    48  	writeRateLimiter quotas.RequestRateLimiter
    49  }
    50  
    51  func NewVisibilityManagerRateLimited(
    52  	delegate manager.VisibilityManager,
    53  	readMaxQPS dynamicconfig.IntPropertyFn,
    54  	writeMaxQPS dynamicconfig.IntPropertyFn,
    55  	operatorRPSRatio dynamicconfig.FloatPropertyFn,
    56  ) *visibilityManagerRateLimited {
    57  	return &visibilityManagerRateLimited{
    58  		delegate:         delegate,
    59  		readRateLimiter:  newPriorityRateLimiter(readMaxQPS, operatorRPSRatio),
    60  		writeRateLimiter: newPriorityRateLimiter(writeMaxQPS, operatorRPSRatio),
    61  	}
    62  }
    63  
    64  func (m *visibilityManagerRateLimited) Close() {
    65  	m.delegate.Close()
    66  }
    67  
    68  func (m *visibilityManagerRateLimited) GetReadStoreName(nsName namespace.Name) string {
    69  	return m.delegate.GetReadStoreName(nsName)
    70  }
    71  
    72  func (m *visibilityManagerRateLimited) GetStoreNames() []string {
    73  	return m.delegate.GetStoreNames()
    74  }
    75  
    76  func (m *visibilityManagerRateLimited) HasStoreName(stName string) bool {
    77  	return m.delegate.HasStoreName(stName)
    78  }
    79  
    80  func (m *visibilityManagerRateLimited) GetIndexName() string {
    81  	return m.delegate.GetIndexName()
    82  }
    83  
    84  func (m *visibilityManagerRateLimited) ValidateCustomSearchAttributes(
    85  	searchAttributes map[string]any,
    86  ) (map[string]any, error) {
    87  	return m.delegate.ValidateCustomSearchAttributes(searchAttributes)
    88  }
    89  
    90  // Below are write APIs.
    91  
    92  func (m *visibilityManagerRateLimited) RecordWorkflowExecutionStarted(
    93  	ctx context.Context,
    94  	request *manager.RecordWorkflowExecutionStartedRequest,
    95  ) error {
    96  	if ok := allow(ctx, "RecordWorkflowExecutionStarted", m.writeRateLimiter); !ok {
    97  		return persistence.ErrPersistenceLimitExceeded
    98  	}
    99  	return m.delegate.RecordWorkflowExecutionStarted(ctx, request)
   100  }
   101  
   102  func (m *visibilityManagerRateLimited) RecordWorkflowExecutionClosed(
   103  	ctx context.Context,
   104  	request *manager.RecordWorkflowExecutionClosedRequest,
   105  ) error {
   106  	if ok := allow(ctx, "RecordWorkflowExecutionClosed", m.writeRateLimiter); !ok {
   107  		return persistence.ErrPersistenceLimitExceeded
   108  	}
   109  	return m.delegate.RecordWorkflowExecutionClosed(ctx, request)
   110  }
   111  
   112  func (m *visibilityManagerRateLimited) UpsertWorkflowExecution(
   113  	ctx context.Context,
   114  	request *manager.UpsertWorkflowExecutionRequest,
   115  ) error {
   116  	if ok := allow(ctx, "UpsertWorkflowExecution", m.writeRateLimiter); !ok {
   117  		return persistence.ErrPersistenceLimitExceeded
   118  	}
   119  	return m.delegate.UpsertWorkflowExecution(ctx, request)
   120  }
   121  
   122  func (m *visibilityManagerRateLimited) DeleteWorkflowExecution(
   123  	ctx context.Context,
   124  	request *manager.VisibilityDeleteWorkflowExecutionRequest,
   125  ) error {
   126  	if ok := allow(ctx, "DeleteWorkflowExecution", m.writeRateLimiter); !ok {
   127  		return persistence.ErrPersistenceLimitExceeded
   128  	}
   129  	return m.delegate.DeleteWorkflowExecution(ctx, request)
   130  }
   131  
   132  // Below are read APIs.
   133  
   134  func (m *visibilityManagerRateLimited) ListOpenWorkflowExecutions(
   135  	ctx context.Context,
   136  	request *manager.ListWorkflowExecutionsRequest,
   137  ) (*manager.ListWorkflowExecutionsResponse, error) {
   138  	if ok := allow(ctx, "ListOpenWorkflowExecutions", m.readRateLimiter); !ok {
   139  		return nil, persistence.ErrPersistenceLimitExceeded
   140  	}
   141  	return m.delegate.ListOpenWorkflowExecutions(ctx, request)
   142  }
   143  
   144  func (m *visibilityManagerRateLimited) ListClosedWorkflowExecutions(
   145  	ctx context.Context,
   146  	request *manager.ListWorkflowExecutionsRequest,
   147  ) (*manager.ListWorkflowExecutionsResponse, error) {
   148  	if ok := allow(ctx, "ListClosedWorkflowExecutions", m.readRateLimiter); !ok {
   149  		return nil, persistence.ErrPersistenceLimitExceeded
   150  	}
   151  	return m.delegate.ListClosedWorkflowExecutions(ctx, request)
   152  }
   153  
   154  func (m *visibilityManagerRateLimited) ListOpenWorkflowExecutionsByType(
   155  	ctx context.Context,
   156  	request *manager.ListWorkflowExecutionsByTypeRequest,
   157  ) (*manager.ListWorkflowExecutionsResponse, error) {
   158  	if ok := allow(ctx, "ListOpenWorkflowExecutionsByType", m.readRateLimiter); !ok {
   159  		return nil, persistence.ErrPersistenceLimitExceeded
   160  	}
   161  	return m.delegate.ListOpenWorkflowExecutionsByType(ctx, request)
   162  }
   163  
   164  func (m *visibilityManagerRateLimited) ListClosedWorkflowExecutionsByType(
   165  	ctx context.Context,
   166  	request *manager.ListWorkflowExecutionsByTypeRequest,
   167  ) (*manager.ListWorkflowExecutionsResponse, error) {
   168  	if ok := allow(ctx, "ListClosedWorkflowExecutionsByType", m.readRateLimiter); !ok {
   169  		return nil, persistence.ErrPersistenceLimitExceeded
   170  	}
   171  	return m.delegate.ListClosedWorkflowExecutionsByType(ctx, request)
   172  }
   173  
   174  func (m *visibilityManagerRateLimited) ListOpenWorkflowExecutionsByWorkflowID(
   175  	ctx context.Context,
   176  	request *manager.ListWorkflowExecutionsByWorkflowIDRequest,
   177  ) (*manager.ListWorkflowExecutionsResponse, error) {
   178  	if ok := allow(ctx, "ListOpenWorkflowExecutionsByWorkflowID", m.readRateLimiter); !ok {
   179  		return nil, persistence.ErrPersistenceLimitExceeded
   180  	}
   181  	return m.delegate.ListOpenWorkflowExecutionsByWorkflowID(ctx, request)
   182  }
   183  
   184  func (m *visibilityManagerRateLimited) ListClosedWorkflowExecutionsByWorkflowID(
   185  	ctx context.Context,
   186  	request *manager.ListWorkflowExecutionsByWorkflowIDRequest,
   187  ) (*manager.ListWorkflowExecutionsResponse, error) {
   188  	if ok := allow(ctx, "ListClosedWorkflowExecutionsByWorkflowID", m.readRateLimiter); !ok {
   189  		return nil, persistence.ErrPersistenceLimitExceeded
   190  	}
   191  	return m.delegate.ListClosedWorkflowExecutionsByWorkflowID(ctx, request)
   192  }
   193  
   194  func (m *visibilityManagerRateLimited) ListClosedWorkflowExecutionsByStatus(
   195  	ctx context.Context,
   196  	request *manager.ListClosedWorkflowExecutionsByStatusRequest,
   197  ) (*manager.ListWorkflowExecutionsResponse, error) {
   198  	if ok := allow(ctx, "ListClosedWorkflowExecutionsByStatus", m.readRateLimiter); !ok {
   199  		return nil, persistence.ErrPersistenceLimitExceeded
   200  	}
   201  	return m.delegate.ListClosedWorkflowExecutionsByStatus(ctx, request)
   202  }
   203  
   204  func (m *visibilityManagerRateLimited) ListWorkflowExecutions(
   205  	ctx context.Context,
   206  	request *manager.ListWorkflowExecutionsRequestV2,
   207  ) (*manager.ListWorkflowExecutionsResponse, error) {
   208  	if ok := allow(ctx, "ListWorkflowExecutions", m.readRateLimiter); !ok {
   209  		return nil, persistence.ErrPersistenceLimitExceeded
   210  	}
   211  	return m.delegate.ListWorkflowExecutions(ctx, request)
   212  }
   213  
   214  func (m *visibilityManagerRateLimited) ScanWorkflowExecutions(
   215  	ctx context.Context,
   216  	request *manager.ListWorkflowExecutionsRequestV2,
   217  ) (*manager.ListWorkflowExecutionsResponse, error) {
   218  	if ok := allow(ctx, "ScanWorkflowExecutions", m.readRateLimiter); !ok {
   219  		return nil, persistence.ErrPersistenceLimitExceeded
   220  	}
   221  	return m.delegate.ScanWorkflowExecutions(ctx, request)
   222  }
   223  
   224  func (m *visibilityManagerRateLimited) CountWorkflowExecutions(
   225  	ctx context.Context,
   226  	request *manager.CountWorkflowExecutionsRequest,
   227  ) (*manager.CountWorkflowExecutionsResponse, error) {
   228  	if ok := allow(ctx, "CountWorkflowExecutions", m.readRateLimiter); !ok {
   229  		return nil, persistence.ErrPersistenceLimitExceeded
   230  	}
   231  	return m.delegate.CountWorkflowExecutions(ctx, request)
   232  }
   233  
   234  func (m *visibilityManagerRateLimited) GetWorkflowExecution(
   235  	ctx context.Context,
   236  	request *manager.GetWorkflowExecutionRequest,
   237  ) (*manager.GetWorkflowExecutionResponse, error) {
   238  	if ok := allow(ctx, "GetWorkflowExecution", m.readRateLimiter); !ok {
   239  		return nil, persistence.ErrPersistenceLimitExceeded
   240  	}
   241  	return m.delegate.GetWorkflowExecution(ctx, request)
   242  }
   243  
   244  func allow(
   245  	ctx context.Context,
   246  	api string,
   247  	rateLimiter quotas.RequestRateLimiter,
   248  ) bool {
   249  	callerInfo := headers.GetCallerInfo(ctx)
   250  	// Currently only CallerType is used. See common/persistence/visibility/quotas.go for rate limiter details.
   251  	return rateLimiter.Allow(time.Now().UTC(), quotas.NewRequest(
   252  		api,
   253  		RateLimitDefaultToken,
   254  		callerInfo.CallerName,
   255  		callerInfo.CallerType,
   256  		-1,
   257  		callerInfo.CallOrigin,
   258  	))
   259  }