1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
|
package featureflag
import (
"context"
"strconv"
"google.golang.org/grpc/metadata"
)
const (
// explicitFeatureFlagKey is used by ContextWithExplicitFeatureFlags to mark a context as
// requiring all feature flags to have been explicitly defined.
explicitFeatureFlagKey = "require_explicit_feature_flag_checks"
)
// ContextWithExplicitFeatureFlags marks the context such that all feature flags which are checked
// must have been explicitly set in that context. If a feature flag wasn't set to an explicit value,
// then checking this feature flag will panic. This is not for use in production systems, but is
// intended for tests to verify that we test each feature flag properly.
func ContextWithExplicitFeatureFlags(ctx context.Context) context.Context {
return injectIntoIncomingAndOutgoingContext(ctx, explicitFeatureFlagKey, true)
}
// ContextWithFeatureFlag sets the feature flag in both the incoming and outgoing context.
func ContextWithFeatureFlag(ctx context.Context, flag FeatureFlag, enabled bool) context.Context {
return injectIntoIncomingAndOutgoingContext(ctx, flag.MetadataKey(), enabled)
}
// OutgoingCtxWithFeatureFlag sets the feature flag for an outgoing context.
func OutgoingCtxWithFeatureFlag(ctx context.Context, flag FeatureFlag, enabled bool) context.Context {
return outgoingCtxWithFeatureFlag(ctx, flag.MetadataKey(), enabled)
}
func outgoingCtxWithFeatureFlag(ctx context.Context, key string, enabled bool) context.Context {
md, ok := metadata.FromOutgoingContext(ctx)
if !ok {
md = metadata.New(map[string]string{})
}
md = md.Copy()
md.Set(key, strconv.FormatBool(enabled))
return metadata.NewOutgoingContext(ctx, md)
}
// IncomingCtxWithFeatureFlag sets the feature flag for an incoming context. This is NOT meant for
// use in clients that transfer the context across process boundaries.
func IncomingCtxWithFeatureFlag(ctx context.Context, flag FeatureFlag, enabled bool) context.Context {
return incomingCtxWithFeatureFlag(ctx, flag.MetadataKey(), enabled)
}
func incomingCtxWithFeatureFlag(ctx context.Context, key string, enabled bool) context.Context {
md, ok := metadata.FromIncomingContext(ctx)
if !ok {
md = metadata.New(map[string]string{})
}
md = md.Copy()
md.Set(key, strconv.FormatBool(enabled))
return metadata.NewIncomingContext(ctx, md)
}
func injectIntoIncomingAndOutgoingContext(ctx context.Context, key string, enabled bool) context.Context {
incomingMD, ok := metadata.FromIncomingContext(ctx)
if !ok {
incomingMD = metadata.New(map[string]string{})
}
incomingMD.Set(key, strconv.FormatBool(enabled))
ctx = metadata.NewIncomingContext(ctx, incomingMD)
return metadata.AppendToOutgoingContext(ctx, key, strconv.FormatBool(enabled))
}
// FromContext returns the set of all feature flags defined in the context. This returns both
// feature flags that are currently defined by Gitaly, but may also return some that aren't defined
// by us in case they match the feature flag prefix but don't have a definition. This function also
// returns the state of the feature flag *as defined in the context*. This value may be overridden.
func FromContext(ctx context.Context) map[FeatureFlag]bool {
md, ok := metadata.FromIncomingContext(ctx)
if !ok {
return map[FeatureFlag]bool{}
}
flags := map[FeatureFlag]bool{}
for metadataName, values := range md {
if len(values) == 0 {
continue
}
flag, err := FromMetadataKey(metadataName)
if err != nil {
continue
}
flags[flag] = values[0] == "true"
}
return flags
}
|
