cubefs

1
/*
2
 *
3
 * Copyright 2020 gRPC authors.
4
 *
5
 * Licensed under the Apache License, Version 2.0 (the "License");
6
 * you may not use this file except in compliance with the License.
7
 * You may obtain a copy of the License at
8
 *
9
 *     http://www.apache.org/licenses/LICENSE-2.0
10
 *
11
 * Unless required by applicable law or agreed to in writing, software
12
 * distributed under the License is distributed on an "AS IS" BASIS,
13
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14
 * See the License for the specific language governing permissions and
15
 * limitations under the License.
16
 *
17
 */
18

19
// Package resolver provides internal resolver-related functionality.
20
package resolver
21

22
import (
23
	"context"
24
	"sync"
25

26
	"google.golang.org/grpc/internal/serviceconfig"
27
	"google.golang.org/grpc/resolver"
28
)
29

30
// ConfigSelector controls what configuration to use for every RPC.
31
type ConfigSelector interface {
32
	// Selects the configuration for the RPC, or terminates it using the error.
33
	// This error will be converted by the gRPC library to a status error with
34
	// code UNKNOWN if it is not returned as a status error.
35
	SelectConfig(RPCInfo) (*RPCConfig, error)
36
}
37

38
// RPCInfo contains RPC information needed by a ConfigSelector.
39
type RPCInfo struct {
40
	// Context is the user's context for the RPC and contains headers and
41
	// application timeout.  It is passed for interception purposes and for
42
	// efficiency reasons.  SelectConfig should not be blocking.
43
	Context context.Context
44
	Method  string // i.e. "/Service/Method"
45
}
46

47
// RPCConfig describes the configuration to use for each RPC.
48
type RPCConfig struct {
49
	// The context to use for the remainder of the RPC; can pass info to LB
50
	// policy or affect timeout or metadata.
51
	Context      context.Context
52
	MethodConfig serviceconfig.MethodConfig // configuration to use for this RPC
53
	OnCommitted  func()                     // Called when the RPC has been committed (retries no longer possible)
54
}
55

56
type csKeyType string
57

58
const csKey = csKeyType("grpc.internal.resolver.configSelector")
59

60
// SetConfigSelector sets the config selector in state and returns the new
61
// state.
62
func SetConfigSelector(state resolver.State, cs ConfigSelector) resolver.State {
63
	state.Attributes = state.Attributes.WithValues(csKey, cs)
64
	return state
65
}
66

67
// GetConfigSelector retrieves the config selector from state, if present, and
68
// returns it or nil if absent.
69
func GetConfigSelector(state resolver.State) ConfigSelector {
70
	cs, _ := state.Attributes.Value(csKey).(ConfigSelector)
71
	return cs
72
}
73

74
// SafeConfigSelector allows for safe switching of ConfigSelector
75
// implementations such that previous values are guaranteed to not be in use
76
// when UpdateConfigSelector returns.
77
type SafeConfigSelector struct {
78
	mu sync.RWMutex
79
	cs ConfigSelector
80
}
81

82
// UpdateConfigSelector swaps to the provided ConfigSelector and blocks until
83
// all uses of the previous ConfigSelector have completed.
84
func (scs *SafeConfigSelector) UpdateConfigSelector(cs ConfigSelector) {
85
	scs.mu.Lock()
86
	defer scs.mu.Unlock()
87
	scs.cs = cs
88
}
89

90
// SelectConfig defers to the current ConfigSelector in scs.
91
func (scs *SafeConfigSelector) SelectConfig(r RPCInfo) (*RPCConfig, error) {
92
	scs.mu.RLock()
93
	defer scs.mu.RUnlock()
94
	return scs.cs.SelectConfig(r)
95
}
96