fork of https://github.com/sourcegraph/zoekt
1package messagesize
2
3import (
4 "fmt"
5 "math"
6 "os"
7
8 "google.golang.org/grpc"
9
10 "github.com/dustin/go-humanize"
11)
12
13var (
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.
28func 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.
51func 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.
70func 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.
91type 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
98func (e *parseError) Error() string {
99 return fmt.Sprintf("failed to parse %q as bytes: %s", e.rawSize, e.err)
100}
101
102func (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.
107type 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
116func (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
120func getEnv(key string, defaultValue string) string {
121 value, ok := os.LookupEnv(key)
122 if !ok {
123 return defaultValue
124 }
125
126 return value
127}