forked from sourcegraph/zoekt
-
Notifications
You must be signed in to change notification settings - Fork 1
/
messagesize.go
127 lines (105 loc) · 4.18 KB
/
messagesize.go
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
package messagesize
import (
"fmt"
"math"
"os"
"google.golang.org/grpc"
"github.com/dustin/go-humanize"
)
var (
smallestAllowedMaxMessageSize = uint64(4 * 1024 * 1024) // 4 MB: There isn't a scenario where we'd want to dip below the default of 4MB.
largestAllowedMaxMessageSize = uint64(math.MaxInt) // This is the largest allowed value for the type accepted by the grpc.MaxSize[...] options.
envClientMessageSize = getEnv("GRPC_CLIENT_MAX_MESSAGE_SIZE", messageSizeDisabled) // set the maximum message size for gRPC clients (ex: "40MB")
envServerMessageSize = getEnv("GRPC_SERVER_MAX_MESSAGE_SIZE", messageSizeDisabled) // set the maximum message size for gRPC servers (ex: "40MB")
messageSizeDisabled = "message_size_disabled" // sentinel value for when the message size env var isn't set
)
// MustGetClientMessageSizeFromEnv returns a slice of grpc.DialOptions that set the maximum message size for gRPC clients if
// the "SRC_GRPC_CLIENT_MAX_MESSAGE_SIZE" environment variable is set to a valid size value (ex: "40 MB").
//
// If the environment variable isn't set, it returns nil.
// If the size value in the environment variable is invalid (too small, not parsable, etc.), it panics.
func MustGetClientMessageSizeFromEnv() []grpc.DialOption {
if envClientMessageSize == messageSizeDisabled {
return nil
}
messageSize, err := getMessageSizeBytesFromString(envClientMessageSize, smallestAllowedMaxMessageSize, largestAllowedMaxMessageSize)
if err != nil {
panic(fmt.Sprintf("failed to get gRPC client message size: %s", err))
}
return []grpc.DialOption{
grpc.WithDefaultCallOptions(
grpc.MaxCallRecvMsgSize(messageSize),
grpc.MaxCallSendMsgSize(messageSize),
),
}
}
// MustGetServerMessageSizeFromEnv returns a slice of grpc.ServerOption that set the maximum message size for gRPC servers if
// the "SRC_GRPC_SERVER_MAX_MESSAGE_SIZE" environment variable is set to a valid size value (ex: "40 MB").
//
// If the environment variable isn't set, it returns nil.
// If the size value in the environment variable is invalid (too small, not parsable, etc.), it panics.
func MustGetServerMessageSizeFromEnv() []grpc.ServerOption {
if envServerMessageSize == messageSizeDisabled {
return nil
}
messageSize, err := getMessageSizeBytesFromString(envServerMessageSize, smallestAllowedMaxMessageSize, largestAllowedMaxMessageSize)
if err != nil {
panic(fmt.Sprintf("failed to get gRPC server message size: %s", err))
}
return []grpc.ServerOption{
grpc.MaxRecvMsgSize(messageSize),
grpc.MaxSendMsgSize(messageSize),
}
}
// getMessageSizeBytesFromEnv parses rawSize returns the message size in bytes within the range [minSize, maxSize].
//
// If rawSize isn't a valid size is not set or the value is outside the allowed range, it returns an error.
func getMessageSizeBytesFromString(rawSize string, minSize, maxSize uint64) (size int, err error) {
sizeBytes, err := humanize.ParseBytes(rawSize)
if err != nil {
return 0, &parseError{
rawSize: rawSize,
err: err,
}
}
if sizeBytes < minSize || sizeBytes > maxSize {
return 0, &sizeOutOfRangeError{
size: humanize.IBytes(sizeBytes),
min: humanize.IBytes(minSize),
max: humanize.IBytes(maxSize),
}
}
return int(sizeBytes), nil
}
// parseError occurs when the environment variable's value cannot be parsed as a byte size.
type parseError struct {
// rawSize is the raw size string that was attempted to be parsed
rawSize string
// err is the error that occurred while parsing rawSize
err error
}
func (e *parseError) Error() string {
return fmt.Sprintf("failed to parse %q as bytes: %s", e.rawSize, e.err)
}
func (e *parseError) Unwrap() error {
return e.err
}
// sizeOutOfRangeError occurs when the environment variable's value is outside of the allowed range.
type sizeOutOfRangeError struct {
// size is the size that was out of range
size string
// min is the minimum allowed size
min string
// max is the maximum allowed size
max string
}
func (e *sizeOutOfRangeError) Error() string {
return fmt.Sprintf("size %s is outside of allowed range [%s, %s]", e.size, e.min, e.max)
}
func getEnv(key string, defaultValue string) string {
value, ok := os.LookupEnv(key)
if !ok {
return defaultValue
}
return value
}