fork of https://github.com/sourcegraph/zoekt
grpc: port messagesize interceptors and raise default client message size to 90mb (#640)
author
Geoffrey Gilmore
committer
GitHub
date
(Aug 22, 2023, 2:24 PM -0700)
commit
f75df3d8
f75df3d8a3f8de3003da19567bf13429563036cb
parent
993cfdb2
993cfdb27a57e730f9d05cba7af3c06b142e52ad
+298
-33
+44
-22
cmd/zoekt-sourcegraph-indexserver/main.go
+44
-22
cmd/zoekt-sourcegraph-indexserver/main.go
···
37
37
"github.com/prometheus/client_golang/prometheus/promauto"
38
38
sglog "github.com/sourcegraph/log"
39
39
"github.com/sourcegraph/zoekt/grpc/internalerrs"
40
+
"github.com/sourcegraph/zoekt/grpc/messagesize"
40
41
"go.uber.org/automaxprocs/maxprocs"
41
42
"golang.org/x/net/trace"
42
43
"golang.org/x/sys/unix"
···
1372
1373
}
1373
1374
1374
1375
logger := sglog.Scoped("zoektConfigurationGRPCClient", "")
1375
-
1376
-
gRPCConnectionOptions := []grpc.DialOption{
1377
-
grpc.WithTransportCredentials(insecure.NewCredentials()),
1378
-
grpc.WithChainStreamInterceptor(
1379
-
internalActorStreamInterceptor(),
1380
-
internalerrs.LoggingStreamClientInterceptor(logger),
1381
-
internalerrs.PrometheusStreamClientInterceptor,
1382
-
),
1383
-
grpc.WithChainUnaryInterceptor(
1384
-
internalActorUnaryInterceptor(),
1385
-
internalerrs.LoggingUnaryClientInterceptor(logger),
1386
-
internalerrs.PrometheusUnaryClientInterceptor,
1387
-
),
1388
-
grpc.WithDefaultServiceConfig(defaultGRPCServiceConfigurationJSON),
1389
-
}
1390
-
1391
-
// This dialer is used to connect via gRPC to the Sourcegraph instance.
1392
-
// This is done lazily, so we can provide the client to use regardless of
1393
-
// whether we enabled gRPC or not initially.
1394
-
cc, err := grpc.Dial(rootURL.Host, gRPCConnectionOptions...)
1376
+
client, err := dialGRPCClient(rootURL.Host, logger)
1395
1377
if err != nil {
1396
1378
return nil, fmt.Errorf("initializing gRPC connection to %q: %w", rootURL.Host, err)
1397
1379
}
1398
1380
1399
-
client := proto.NewZoektConfigurationServiceClient(cc)
1400
1381
opts = append(opts, WithGRPCClient(client))
1401
-
1402
1382
sg = newSourcegraphClient(rootURL, conf.hostname, opts...)
1403
1383
1404
1384
} else {
···
1470
1450
ctx = metadata.AppendToOutgoingContext(ctx, "X-Sourcegraph-Actor-UID", "internal")
1471
1451
return streamer(ctx, desc, cc, method, opts...)
1472
1452
}
1453
+
}
1454
+
1455
+
// defaultGRPCMessageReceiveSizeBytes is the default message size that gRPCs servers and clients are allowed to process.
1456
+
// This can be overridden by providing custom Server/Dial options.
1457
+
const defaultGRPCMessageReceiveSizeBytes = 90 * 1024 * 1024 // 90 MB
1458
+
1459
+
func dialGRPCClient(addr string, logger sglog.Logger, additionalOpts ...grpc.DialOption) (proto.ZoektConfigurationServiceClient, error) {
1460
+
opts := []grpc.DialOption{
1461
+
grpc.WithTransportCredentials(insecure.NewCredentials()),
1462
+
grpc.WithChainStreamInterceptor(
1463
+
internalActorStreamInterceptor(),
1464
+
internalerrs.LoggingStreamClientInterceptor(logger),
1465
+
internalerrs.PrometheusStreamClientInterceptor,
1466
+
),
1467
+
grpc.WithChainUnaryInterceptor(
1468
+
internalActorUnaryInterceptor(),
1469
+
internalerrs.LoggingUnaryClientInterceptor(logger),
1470
+
internalerrs.PrometheusUnaryClientInterceptor,
1471
+
),
1472
+
grpc.WithDefaultServiceConfig(defaultGRPCServiceConfigurationJSON),
1473
+
grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(defaultGRPCMessageReceiveSizeBytes)),
1474
+
}
1475
+
1476
+
opts = append(opts, additionalOpts...)
1477
+
1478
+
// Ensure that the message size options are set last, so they override any other
1479
+
// client-specific options that tweak the message size.
1480
+
//
1481
+
// The message size options are only provided if the environment variable is set. These options serve as an escape hatch, so they
1482
+
// take precedence over everything else with a uniform size setting that's easy to reason about.
1483
+
opts = append(opts, messagesize.MustGetClientMessageSizeFromEnv()...)
1484
+
1485
+
// This dialer is used to connect via gRPC to the Sourcegraph instance.
1486
+
// This is done lazily, so we can provide the client to use regardless of
1487
+
// whether we enabled gRPC or not initially.
1488
+
cc, err := grpc.Dial(addr, opts...)
1489
+
if err != nil {
1490
+
return nil, fmt.Errorf("dialing %q: %w", addr, err)
1491
+
}
1492
+
1493
+
client := proto.NewZoektConfigurationServiceClient(cc)
1494
+
return client, nil
1473
1495
}
1474
1496
1475
1497
// addDefaultPort adds a default port to a URL if one is not specified.
+30
-11
cmd/zoekt-webserver/main.go
+30
-11
cmd/zoekt-webserver/main.go
···
39
39
40
40
"github.com/sourcegraph/mountinfo"
41
41
"github.com/sourcegraph/zoekt/grpc/internalerrs"
42
+
"github.com/sourcegraph/zoekt/grpc/messagesize"
42
43
zoektgrpc "github.com/sourcegraph/zoekt/grpc/server"
43
44
"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
44
45
"golang.org/x/net/http2"
···
291
292
292
293
logger := sglog.Scoped("ZoektWebserverGRPCServer", "The Zoekt Webserver GRPC Server")
293
294
294
-
grpcServer := grpc.NewServer(
295
-
grpc.ChainStreamInterceptor(
296
-
otelgrpc.StreamServerInterceptor(),
297
-
internalerrs.LoggingStreamServerInterceptor(logger),
298
-
),
299
-
grpc.ChainUnaryInterceptor(
300
-
otelgrpc.UnaryServerInterceptor(),
301
-
internalerrs.LoggingUnaryServerInterceptor(logger),
302
-
),
303
-
)
304
-
v1.RegisterWebserverServiceServer(grpcServer, zoektgrpc.NewServer(web.NewTraceAwareSearcher(s.Searcher)))
295
+
streamer := web.NewTraceAwareSearcher(s.Searcher)
296
+
grpcServer := newGRPCServer(logger, streamer)
305
297
306
298
handler = multiplexGRPC(grpcServer, handler)
307
299
···
634
626
}
635
627
636
628
return sglog.TraceContext{}
629
+
}
630
+
631
+
func newGRPCServer(logger sglog.Logger, streamer zoekt.Streamer, additionalOpts ...grpc.ServerOption) *grpc.Server {
632
+
opts := []grpc.ServerOption{
633
+
grpc.ChainStreamInterceptor(
634
+
otelgrpc.StreamServerInterceptor(),
635
+
internalerrs.LoggingStreamServerInterceptor(logger),
636
+
),
637
+
grpc.ChainUnaryInterceptor(
638
+
otelgrpc.UnaryServerInterceptor(),
639
+
internalerrs.LoggingUnaryServerInterceptor(logger),
640
+
),
641
+
}
642
+
643
+
opts = append(opts, additionalOpts...)
644
+
645
+
// Ensure that the message size options are set last, so they override any other
646
+
// server-specific options that tweak the message size.
647
+
//
648
+
// The message size options are only provided if the environment variable is set. These options serve as an escape hatch, so they
649
+
// take precedence over everything else with a uniform size setting that's easy to reason about.
650
+
opts = append(opts, messagesize.MustGetServerMessageSizeFromEnv()...)
651
+
652
+
s := grpc.NewServer(opts...)
653
+
v1.RegisterWebserverServiceServer(s, zoektgrpc.NewServer(streamer))
654
+
655
+
return s
637
656
}
638
657
639
658
var (
+127
grpc/messagesize/messagesize.go
+127
grpc/messagesize/messagesize.go
···
1
+
package messagesize
2
+
3
+
import (
4
+
"fmt"
5
+
"math"
6
+
"os"
7
+
8
+
"google.golang.org/grpc"
9
+
10
+
"github.com/dustin/go-humanize"
11
+
)
12
+
13
+
var (
14
+
smallestAllowedMaxMessageSize = uint64(4 * 1024 * 1024) // 4 MB: There isn't a scenario where we'd want to dip below the default of 4MB.
15
+
largestAllowedMaxMessageSize = uint64(math.MaxInt) // This is the largest allowed value for the type accepted by the grpc.MaxSize[...] options.
16
+
17
+
envClientMessageSize = getEnv("GRPC_CLIENT_MAX_MESSAGE_SIZE", messageSizeDisabled) // set the maximum message size for gRPC clients (ex: "40MB")
18
+
envServerMessageSize = getEnv("GRPC_SERVER_MAX_MESSAGE_SIZE", messageSizeDisabled) // set the maximum message size for gRPC servers (ex: "40MB")
19
+
20
+
messageSizeDisabled = "message_size_disabled" // sentinel value for when the message size env var isn't set
21
+
)
22
+
23
+
// MustGetClientMessageSizeFromEnv returns a slice of grpc.DialOptions that set the maximum message size for gRPC clients if
24
+
// the "SRC_GRPC_CLIENT_MAX_MESSAGE_SIZE" environment variable is set to a valid size value (ex: "40 MB").
25
+
//
26
+
// If the environment variable isn't set, it returns nil.
27
+
// If the size value in the environment variable is invalid (too small, not parsable, etc.), it panics.
28
+
func MustGetClientMessageSizeFromEnv() []grpc.DialOption {
29
+
if envClientMessageSize == messageSizeDisabled {
30
+
return nil
31
+
}
32
+
33
+
messageSize, err := getMessageSizeBytesFromString(envClientMessageSize, smallestAllowedMaxMessageSize, largestAllowedMaxMessageSize)
34
+
if err != nil {
35
+
panic(fmt.Sprintf("failed to get gRPC client message size: %s", err))
36
+
}
37
+
38
+
return []grpc.DialOption{
39
+
grpc.WithDefaultCallOptions(
40
+
grpc.MaxCallRecvMsgSize(messageSize),
41
+
grpc.MaxCallSendMsgSize(messageSize),
42
+
),
43
+
}
44
+
}
45
+
46
+
// MustGetServerMessageSizeFromEnv returns a slice of grpc.ServerOption that set the maximum message size for gRPC servers if
47
+
// the "SRC_GRPC_SERVER_MAX_MESSAGE_SIZE" environment variable is set to a valid size value (ex: "40 MB").
48
+
//
49
+
// If the environment variable isn't set, it returns nil.
50
+
// If the size value in the environment variable is invalid (too small, not parsable, etc.), it panics.
51
+
func MustGetServerMessageSizeFromEnv() []grpc.ServerOption {
52
+
if envServerMessageSize == messageSizeDisabled {
53
+
return nil
54
+
}
55
+
56
+
messageSize, err := getMessageSizeBytesFromString(envServerMessageSize, smallestAllowedMaxMessageSize, largestAllowedMaxMessageSize)
57
+
if err != nil {
58
+
panic(fmt.Sprintf("failed to get gRPC server message size: %s", err))
59
+
}
60
+
61
+
return []grpc.ServerOption{
62
+
grpc.MaxRecvMsgSize(messageSize),
63
+
grpc.MaxSendMsgSize(messageSize),
64
+
}
65
+
}
66
+
67
+
// getMessageSizeBytesFromEnv parses rawSize returns the message size in bytes within the range [minSize, maxSize].
68
+
//
69
+
// If rawSize isn't a valid size is not set or the value is outside the allowed range, it returns an error.
70
+
func getMessageSizeBytesFromString(rawSize string, minSize, maxSize uint64) (size int, err error) {
71
+
sizeBytes, err := humanize.ParseBytes(rawSize)
72
+
if err != nil {
73
+
return 0, &parseError{
74
+
rawSize: rawSize,
75
+
err: err,
76
+
}
77
+
}
78
+
79
+
if sizeBytes < minSize || sizeBytes > maxSize {
80
+
return 0, &sizeOutOfRangeError{
81
+
size: humanize.IBytes(sizeBytes),
82
+
min: humanize.IBytes(minSize),
83
+
max: humanize.IBytes(maxSize),
84
+
}
85
+
}
86
+
87
+
return int(sizeBytes), nil
88
+
}
89
+
90
+
// parseError occurs when the environment variable's value cannot be parsed as a byte size.
91
+
type parseError struct {
92
+
// rawSize is the raw size string that was attempted to be parsed
93
+
rawSize string
94
+
// err is the error that occurred while parsing rawSize
95
+
err error
96
+
}
97
+
98
+
func (e *parseError) Error() string {
99
+
return fmt.Sprintf("failed to parse %q as bytes: %s", e.rawSize, e.err)
100
+
}
101
+
102
+
func (e *parseError) Unwrap() error {
103
+
return e.err
104
+
}
105
+
106
+
// sizeOutOfRangeError occurs when the environment variable's value is outside of the allowed range.
107
+
type sizeOutOfRangeError struct {
108
+
// size is the size that was out of range
109
+
size string
110
+
// min is the minimum allowed size
111
+
min string
112
+
// max is the maximum allowed size
113
+
max string
114
+
}
115
+
116
+
func (e *sizeOutOfRangeError) Error() string {
117
+
return fmt.Sprintf("size %s is outside of allowed range [%s, %s]", e.size, e.min, e.max)
118
+
}
119
+
120
+
func getEnv(key string, defaultValue string) string {
121
+
value, ok := os.LookupEnv(key)
122
+
if !ok {
123
+
return defaultValue
124
+
}
125
+
126
+
return value
127
+
}
+97
grpc/messagesize/messagesize_test.go
+97
grpc/messagesize/messagesize_test.go
···
1
+
package messagesize
2
+
3
+
import (
4
+
"errors"
5
+
"math"
6
+
"testing"
7
+
8
+
"github.com/google/go-cmp/cmp"
9
+
)
10
+
11
+
func TestGetMessageSizeBytesFromString(t *testing.T) {
12
+
13
+
t.Run("8 MB", func(t *testing.T) {
14
+
sizeString := "8MB"
15
+
16
+
size, err := getMessageSizeBytesFromString(sizeString, 0, math.MaxInt)
17
+
18
+
if err != nil {
19
+
t.Fatalf("unexpected error: %s", err)
20
+
}
21
+
22
+
expectedSize := 8 * 1000 * 1000
23
+
if diff := cmp.Diff(expectedSize, size); diff != "" {
24
+
t.Fatalf("unexpected size (-want +got):\n%s", diff)
25
+
}
26
+
})
27
+
28
+
t.Run("just small enough", func(t *testing.T) {
29
+
sizeString := "4MB" // inside large-end of range
30
+
31
+
fourMegaBytes := 4 * 1000 * 1000
32
+
size, err := getMessageSizeBytesFromString(sizeString, 0, uint64(fourMegaBytes))
33
+
if err != nil {
34
+
t.Fatalf("unexpected error: %s", err)
35
+
}
36
+
37
+
if diff := cmp.Diff(fourMegaBytes, size); diff != "" {
38
+
t.Fatalf("unexpected size (-want +got):\n%s", diff)
39
+
}
40
+
})
41
+
42
+
t.Run("just large enough", func(t *testing.T) {
43
+
sizeString := "4MB" // inside low-end of range
44
+
45
+
fourMegaBytes := 4 * 1000 * 1000
46
+
size, err := getMessageSizeBytesFromString(sizeString, uint64(fourMegaBytes), math.MaxInt)
47
+
if err != nil {
48
+
t.Fatalf("unexpected error: %s", err)
49
+
}
50
+
51
+
if diff := cmp.Diff(fourMegaBytes, size); diff != "" {
52
+
t.Fatalf("unexpected size (-want +got):\n%s", diff)
53
+
}
54
+
})
55
+
56
+
t.Run("invalid size", func(t *testing.T) {
57
+
sizeString := "this-is-not-a-size"
58
+
59
+
_, err := getMessageSizeBytesFromString(sizeString, 0, math.MaxInt)
60
+
var expectedErr *parseError
61
+
if !errors.As(err, &expectedErr) {
62
+
t.Fatalf("expected parseError, got error %q", err)
63
+
}
64
+
})
65
+
66
+
t.Run("empty", func(t *testing.T) {
67
+
sizeString := ""
68
+
69
+
_, err := getMessageSizeBytesFromString(sizeString, 0, math.MaxInt)
70
+
var expectedErr *parseError
71
+
if !errors.As(err, &expectedErr) {
72
+
t.Fatalf("expected parseError, got error %q", err)
73
+
}
74
+
})
75
+
76
+
t.Run("too large", func(t *testing.T) {
77
+
sizeString := "4MB" // above range
78
+
79
+
twoMegaBytes := 2 * 1024 * 1024
80
+
_, err := getMessageSizeBytesFromString(sizeString, 0, uint64(twoMegaBytes))
81
+
var expectedErr *sizeOutOfRangeError
82
+
if !errors.As(err, &expectedErr) {
83
+
t.Fatalf("expected sizeOutOfRangeError, got error %q", err)
84
+
}
85
+
})
86
+
87
+
t.Run("too small", func(t *testing.T) {
88
+
sizeString := "1MB" // below range
89
+
90
+
twoMegaBytes := 2 * 1024 * 1024
91
+
_, err := getMessageSizeBytesFromString(sizeString, uint64(twoMegaBytes), math.MaxInt)
92
+
var expectedErr *sizeOutOfRangeError
93
+
if !errors.As(err, &expectedErr) {
94
+
t.Fatalf("expected sizeOutOfRangeError, got error %q", err)
95
+
}
96
+
})
97
+
}