Welcome to mirror list, hosted at ThFree Co, Russian Federation.

context.go « featureflag « metadata « internal - gitlab.com/gitlab-org/gitaly.git - Unnamed repository; edit this file 'description' to name the repository.
summaryrefslogtreecommitdiff
blob: 25f269d60d4a7a2ff06acea494818128e85b22cb (plain)
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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
package featureflag

import (
	"context"
	"fmt"
	"strconv"
	"strings"

	"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)
}

// OutgoingCtxWithRubyFeatureFlag sets the Ruby feature flag for an outgoing context.
func OutgoingCtxWithRubyFeatureFlag(ctx context.Context, flag FeatureFlag, enabled bool) context.Context {
	return outgoingCtxWithFeatureFlag(ctx, rubyHeaderKey(flag.Name), 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)
}

// IncomingCtxWithRubyFeatureFlag sets the Ruby feature flag for an incoming context. This is NOT
// meant for use in clients that transfer the context across process boundaries.
func IncomingCtxWithRubyFeatureFlag(ctx context.Context, flag FeatureFlag, enabled bool) context.Context {
	return incomingCtxWithFeatureFlag(ctx, rubyHeaderKey(flag.Name), 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 {
	rawFlags := RawFromContext(ctx)

	flags := map[FeatureFlag]bool{}
	for rawName, value := range rawFlags {
		flag, err := FromMetadataKey(rawName)
		if err != nil {
			continue
		}

		flags[flag] = value == "true"
	}

	return flags
}

// Raw contains feature flags and their values in their raw form with header prefix in place
// and values unparsed.
type Raw map[string]string

// RawFromContext returns a map that contains all feature flags with their values. The feature
// flags are in their raw format with the header prefix in place. If multiple values are set a
// flag, the first occurrence is used.
//
// This is mostly intended for propagating the feature flags by other means than the metadata,
// for example into the hooks through the environment.
func RawFromContext(ctx context.Context) Raw {
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		return nil
	}

	featureFlags := map[string]string{}
	for key, values := range md {
		if !strings.HasPrefix(key, ffPrefix) || len(values) == 0 {
			continue
		}

		featureFlags[key] = values[0]
	}

	return featureFlags
}

// OutgoingWithRaw returns a new context with raw flags appended into the outgoing
// metadata.
func OutgoingWithRaw(ctx context.Context, flags Raw) context.Context {
	for key, value := range flags {
		ctx = metadata.AppendToOutgoingContext(ctx, key, value)
	}

	return ctx
}

func rubyHeaderKey(flag string) string {
	return fmt.Sprintf("gitaly-feature-ruby-%s", strings.ReplaceAll(flag, "_", "-"))
}