-
Notifications
You must be signed in to change notification settings - Fork 3.8k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
grpcutil: new+fast FromIncomingContext variant
If we want to properly leverage gRPC metadata, we need to find a way around the fact that the default `metatada.FromIncomingContext` is slow and expensive. This patch fixes it. ``` BenchmarkFromIncomingContext/stdlib-32 10987958 327.5 ns/op 432 B/op 3 allocs/op BenchmarkFromIncomingContext/fast-32 698772889 5.152 ns/op 0 B/op 0 allocs/op ``` Release note: None
- Loading branch information
Showing
8 changed files
with
155 additions
and
11 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,104 @@ | ||
| // Copyright 2023 The Cockroach Authors. | ||
| // | ||
| // Use of this software is governed by the Business Source License | ||
| // included in the file licenses/BSL.txt. | ||
| // | ||
| // As of the Change Date specified in that file, in accordance with | ||
| // the Business Source License, use of this software will be governed | ||
| // by the Apache License, Version 2.0, included in the file | ||
| // licenses/APL.txt. | ||
|
|
||
| package grpcutil | ||
|
|
||
| import ( | ||
| "context" | ||
| "strings" | ||
| "time" | ||
|
|
||
| "google.golang.org/grpc/metadata" | ||
| ) | ||
|
|
||
| // FastFromIncomingContext is a specialization of | ||
| // metadata.FromIncomingContext() which extracts the metadata.MD from | ||
| // the context, if any, by reference. Main differences: | ||
| // | ||
| // - This variant does not guarantee that all the MD keys are | ||
| // lowercase. This happens to be true when the MD is populated by | ||
| // gRPC itself on an incoming RPC call, but it may not be true for | ||
| // MD populated elsewhere. | ||
| // - The caller promises to not modify the returned MD -- the gRPC | ||
| // APIs assume that the map in the context remains constant. | ||
| func FastFromIncomingContext(ctx context.Context) (metadata.MD, bool) { | ||
| md, ok := ctx.Value(grpcIncomingKeyObj).(metadata.MD) | ||
| return md, ok | ||
| } | ||
|
|
||
| // FastFirstValueFromIncomingContext is a specialization of | ||
| // metadata.ValueFromIncomingContext() which extracts the first string | ||
| // from the given metadata key, if it exists. No extra objects are | ||
| // allocated. The key is assumed to contain only ASCII characters. | ||
| func FastFirstValueFromIncomingContext(ctx context.Context, key string) (string, bool) { | ||
| md, ok := ctx.Value(grpcIncomingKeyObj).(metadata.MD) | ||
| if !ok { | ||
| return "", false | ||
| } | ||
| if v, ok := md[key]; ok { | ||
| if len(v) > 0 { | ||
| return v[0], true | ||
| } | ||
| return "", false | ||
| } | ||
| for k, v := range md { | ||
| // The letter casing may not have been set properly when MD was | ||
| // attached to the context. So we need to normalize it here. | ||
| // | ||
| // We add len(k) == len(key) to avoid the overhead of | ||
| // strings.ToLower when the keys of different length, because then | ||
| // they are guaranteed to not match anyway. This is the | ||
| // optimization that requires the key to be all ASCII, as | ||
| // generally ToLower() on non-ascii unicode can change the length | ||
| // of the string. | ||
| if len(k) == len(key) && strings.ToLower(k) == key { | ||
| if len(v) > 0 { | ||
| return v[0], true | ||
| } | ||
| return "", false | ||
| } | ||
| } | ||
| return "", false | ||
| } | ||
|
|
||
| // grpcIncomingKeyObj is a copy of a value with the Go type | ||
| // `metadata.incomingKey{}` (from the grpc metadata package). We | ||
| // cannot construct an object of that type directly, but we can | ||
| // "steal" it by forcing the metadata package to give it to us: | ||
| // `metadata.FromIncomingContext` gives an instance of this object as | ||
| // parameter to the `Value` method of the context you give it as | ||
| // argument. We use a custom implementation of that to "steal" the | ||
| // argument of type `incomingKey{}` given to us that way. | ||
| var grpcIncomingKeyObj = func() interface{} { | ||
| var f fakeContext | ||
| _, _ = metadata.FromIncomingContext(&f) | ||
| if f.recordedKey == nil { | ||
| panic("ValueFromIncomingContext did not request a key") | ||
| } | ||
| return f.recordedKey | ||
| }() | ||
|
|
||
| type fakeContext struct { | ||
| recordedKey interface{} | ||
| } | ||
|
|
||
| var _ context.Context = (*fakeContext)(nil) | ||
|
|
||
| // Value implements the context.Context interface and is our helper | ||
| // that "steals" an instance of the private type `incomingKey` in the | ||
| // grpc metadata package. | ||
| func (f *fakeContext) Value(keyObj interface{}) interface{} { | ||
| f.recordedKey = keyObj | ||
| return nil | ||
| } | ||
|
|
||
| func (*fakeContext) Deadline() (time.Time, bool) { panic("unused") } | ||
| func (*fakeContext) Done() <-chan struct{} { panic("unused") } | ||
| func (*fakeContext) Err() error { panic("unused") } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,35 @@ | ||
| // Copyright 2023 The Cockroach Authors. | ||
| // | ||
| // Use of this software is governed by the Business Source License | ||
| // included in the file licenses/BSL.txt. | ||
| // | ||
| // As of the Change Date specified in that file, in accordance with | ||
| // the Business Source License, use of this software will be governed | ||
| // by the Apache License, Version 2.0, included in the file | ||
| // licenses/APL.txt. | ||
|
|
||
| package grpcutil | ||
|
|
||
| import ( | ||
| "context" | ||
| "testing" | ||
|
|
||
| "github.com/cockroachdb/cockroach/pkg/util/leaktest" | ||
| "github.com/stretchr/testify/require" | ||
| "google.golang.org/grpc/metadata" | ||
| ) | ||
|
|
||
| func TestFastFromIncomingContext(t *testing.T) { | ||
| defer leaktest.AfterTest(t)() | ||
|
|
||
| md := metadata.MD{"hello": []string{"world", "universe"}} | ||
|
|
||
| ctx := metadata.NewIncomingContext(context.Background(), md) | ||
| md2, ok := FastFromIncomingContext(ctx) | ||
| require.True(t, ok) | ||
| require.Equal(t, md2, md) | ||
|
|
||
| v, ok := FastFirstValueFromIncomingContext(ctx, "hello") | ||
| require.True(t, ok) | ||
| require.Equal(t, v, "world") | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters