Core persistence priority rate limiting

common/headers/caller_info.go

@@ -0,0 +1,79 @@
+// The MIT License
+//
+// Copyright (c) 2020 Temporal Technologies Inc.  All rights reserved.
+//
+// Copyright (c) 2020 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package headers
+
+import (
+	"context"
+
+	"google.golang.org/grpc/metadata"
+)
+
+const (
+	CallerTypeAPI        = "api"
+	CallerTypeBackground = "background"
+)
+
+type CallerInfo struct {
+	CallerType string
+
+	// TODO: add fields for CallerName and CallerInitiation
+}
+
+func NewCallerInfo(
+	callerType string,
+) CallerInfo {
+	return CallerInfo{
+		CallerType: callerType,
+	}
+}
+
+// SetCallerInfo sets callerName and callerType value in incoming context
+// if not already exists.
+// TODO: consider only set the caller info to golang context instead of grpc metadata
+// and propagate to grpc outgoing context upon making an rpc call
+func SetCallerInfo(
+	ctx context.Context,
+	info CallerInfo,
+) context.Context {
+	mdIncoming, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		mdIncoming = metadata.MD{}
+	}
+
+	if len(mdIncoming.Get(callerTypeHeaderName)) == 0 {
+		mdIncoming.Set(callerTypeHeaderName, string(info.CallerType))
+	}
+
+	return metadata.NewIncomingContext(ctx, mdIncoming)
+}
+
+func GetCallerInfo(
+	ctx context.Context,
+) CallerInfo {
+	values := GetValues(ctx, callerTypeHeaderName)
+	return CallerInfo{
+		CallerType: values[0],
+	}
+}

common/headers/caller_info_test.go

@@ -0,0 +1,98 @@
+// The MIT License
+//
+// Copyright (c) 2020 Temporal Technologies Inc.  All rights reserved.
+//
+// Copyright (c) 2020 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package headers
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/stretchr/testify/suite"
+	"google.golang.org/grpc/metadata"
+)
+
+type (
+	callerInfoSuite struct {
+		*require.Assertions
+		suite.Suite
+	}
+)
+
+func TestCallerInfoSuite(t *testing.T) {
+	suite.Run(t, &callerInfoSuite{})
+}
+
+func (s *callerInfoSuite) SetupTest() {
+	s.Assertions = require.New(s.T())
+}
+
+func (s *callerInfoSuite) TestSetCallerInfo_PreserveOtherValues() {
+	existingKey := "key"
+	existingValue := "value"
+	callerType := CallerTypeAPI
+
+	ctx := metadata.NewIncomingContext(
+		context.Background(),
+		metadata.Pairs(existingKey, existingValue),
+	)
+
+	ctx = SetCallerInfo(ctx, NewCallerInfo(callerType))
+
+	md, ok := metadata.FromIncomingContext(ctx)
+	s.True(ok)
+	s.Equal(existingValue, md.Get(existingKey)[0])
+	s.Equal(callerType, md.Get(callerTypeHeaderName)[0])
+	s.Len(md, 2)
+}
+
+func (s *callerInfoSuite) TestSetCallerInfo_NoExistingCallerInfo() {
+	callerType := CallerTypeAPI
+
+	ctx := SetCallerInfo(context.Background(), CallerInfo{
+		CallerType: callerType,
+	})
+
+	md, ok := metadata.FromIncomingContext(ctx)
+	s.True(ok)
+	s.Equal(callerType, md.Get(callerTypeHeaderName)[0])
+	s.Len(md, 1)
+}
+
+func (s *callerInfoSuite) TestSetCallerInfo_WithExistingCallerInfo() {
+	callerType := CallerTypeAPI
+
+	ctx := SetCallerInfo(context.Background(), CallerInfo{
+		CallerType: callerType,
+	})
+
+	ctx = SetCallerInfo(ctx, CallerInfo{
+		CallerType: CallerTypeBackground,
+	})
+
+	md, ok := metadata.FromIncomingContext(ctx)
+	s.True(ok)
+	s.Equal(callerType, md.Get(callerTypeHeaderName)[0])
+	s.Len(md, 1)
+}

common/headers/headers.go

@@ -36,6 +36,8 @@ const (
 	SupportedServerVersionsHeaderName = "supported-server-versions"
 	SupportedFeaturesHeaderName       = "supported-features"
 	SupportedFeaturesHeaderDelim      = ","
+
+	callerTypeHeaderName = "caller-type"
 )
 
 var (
@@ -45,14 +47,8 @@ var (
 		ClientVersionHeaderName,
 		SupportedServerVersionsHeaderName,
 		SupportedFeaturesHeaderName,
+		callerTypeHeaderName,
 	}
-
-	internalVersionHeaders = metadata.New(map[string]string{
-		ClientNameHeaderName:              ClientNameServer,
-		ClientVersionHeaderName:           ServerVersion,
-		SupportedServerVersionsHeaderName: SupportedServerVersions,
-		SupportedFeaturesHeaderName:       AllFeatures,
-	})
 )
 
 // GetValues returns header values for passed header names.
@@ -70,7 +66,7 @@ func GetValues(ctx context.Context, headerNames ...string) []string {
 }
 
 // Propagate propagates version headers from incoming context to outgoing context.
-// It copies all version headers to outgoing context only if they are exist in incoming context
+// It copies all headers to outgoing context only if they are exist in incoming context
 // and doesn't exist in outgoing context already.
 func Propagate(ctx context.Context) context.Context {
 	if mdIncoming, ok := metadata.FromIncomingContext(ctx); ok {
@@ -97,22 +93,6 @@ func Propagate(ctx context.Context) context.Context {
 	return ctx
 }
 
-// SetVersions sets headers for internal communications.
-func SetVersions(ctx context.Context) context.Context {
-	return metadata.NewOutgoingContext(ctx, internalVersionHeaders)
-}
-
-// SetVersionsForTests sets headers as they would be received from the client.
-// Must be used in tests only.
-func SetVersionsForTests(ctx context.Context, clientVersion, clientName, supportedServerVersions, supportedFeatures string) context.Context {
-	return metadata.NewIncomingContext(ctx, metadata.New(map[string]string{
-		ClientNameHeaderName:              clientName,
-		ClientVersionHeaderName:           clientVersion,
-		SupportedServerVersionsHeaderName: supportedServerVersions,
-		SupportedFeaturesHeaderName:       supportedFeatures,
-	}))
-}
-
 func getSingleHeaderValue(md metadata.MD, headerName string) string {
 	values := md.Get(headerName)
 	if len(values) == 0 {

common/headers/versionChecker.go

@@ -31,6 +31,7 @@ import (
 
 	"github.com/blang/semver/v4"
 	"golang.org/x/exp/slices"
+	"google.golang.org/grpc/metadata"
 
 	"go.temporal.io/api/serviceerror"
 )
@@ -71,6 +72,13 @@ var (
 		ClientNameServer:        "<2.0.0",
 		ClientNameUI:            "<3.0.0",
 	}
+
+	internalVersionHeaderPairs = []string{
+		ClientNameHeaderName, ClientNameServer,
+		ClientVersionHeaderName, ServerVersion,
+		SupportedServerVersionsHeaderName, SupportedServerVersions,
+		SupportedFeaturesHeaderName, AllFeatures,
+	}
 )
 
 type (
@@ -109,6 +117,22 @@ func GetClientNameAndVersion(ctx context.Context) (string, string) {
 	return clientName, clientVersion
 }
 
+// SetVersions sets headers for internal communications.
+func SetVersions(ctx context.Context) context.Context {
+	return metadata.AppendToOutgoingContext(ctx, internalVersionHeaderPairs...)
+}
+
+// SetVersionsForTests sets headers as they would be received from the client.
+// Must be used in tests only.
+func SetVersionsForTests(ctx context.Context, clientVersion, clientName, supportedServerVersions, supportedFeatures string) context.Context {
+	return metadata.NewIncomingContext(ctx, metadata.New(map[string]string{
+		ClientNameHeaderName:              clientName,
+		ClientVersionHeaderName:           clientVersion,
+		SupportedServerVersionsHeaderName: supportedServerVersions,
+		SupportedFeaturesHeaderName:       supportedFeatures,
+	}))
+}
+
 // ClientSupported returns an error if client is unsupported, nil otherwise.
 func (vc *versionChecker) ClientSupported(ctx context.Context, enableClientVersionCheck bool) error {
 	if !enableClientVersionCheck {

common/namespace/registry.go

@@ -40,6 +40,7 @@ import (
 	"go.temporal.io/server/common/cache"
 	"go.temporal.io/server/common/clock"
 	"go.temporal.io/server/common/dynamicconfig"
+	"go.temporal.io/server/common/headers"
 	"go.temporal.io/server/common/log"
 	"go.temporal.io/server/common/log/tag"
 	"go.temporal.io/server/common/metrics"
@@ -219,11 +220,13 @@ func (r *registry) Start() {
 	defer atomic.StoreInt32(&r.status, running)
 
 	// initialize the cache by initial scan
-	err := r.refreshNamespaces(context.Background())
+	ctx := headers.SetCallerInfo(context.Background(), headers.NewCallerInfo(headers.CallerTypeBackground))
+
+	err := r.refreshNamespaces(ctx)
 	if err != nil {
 		r.logger.Fatal("Unable to initialize namespace cache", tag.Error(err))
 	}
-	r.refresher = goro.NewHandle(context.Background()).Go(r.refreshLoop)
+	r.refresher = goro.NewHandle(ctx).Go(r.refreshLoop)
 }
 
 // Stop the background refresh of Namespace data

common/persistence/client/factory.go

@@ -61,7 +61,7 @@ type (
 		metricsClient    metrics.Client
 		logger           log.Logger
 		clusterName      string
-		ratelimiter      quotas.RateLimiter
+		ratelimiter      quotas.RequestRateLimiter
 	}
 )
 
@@ -75,7 +75,7 @@ type (
 func NewFactory(
 	dataStoreFactory DataStoreFactory,
 	cfg *config.Persistence,
-	ratelimiter quotas.RateLimiter,
+	ratelimiter quotas.RequestRateLimiter,
 	serializer serialization.Serializer,
 	clusterName string,
 	metricsClient metrics.Client,

common/persistence/client/fx.go

@@ -67,16 +67,17 @@ func ClusterNameProvider(config *cluster.Config) ClusterName {
 func FactoryProvider(
 	params NewFactoryParams,
 ) Factory {
-	var ratelimiter quotas.RateLimiter
+	var requestRatelimiter quotas.RequestRateLimiter
 	if params.PersistenceMaxQPS != nil && params.PersistenceMaxQPS() > 0 {
-		ratelimiter = quotas.NewDefaultOutgoingRateLimiter(
+		requestRatelimiter = NewPriorityRateLimiter(
 			func() float64 { return float64(params.PersistenceMaxQPS()) },
 		)
 	}
+
 	return NewFactory(
 		params.DataStoreFactory,
 		params.Cfg,
-		ratelimiter,
+		requestRatelimiter,
 		serialization.NewSerializer(),
 		string(params.ClusterName),
 		params.MetricsClient,