-
Notifications
You must be signed in to change notification settings - Fork 1
/
grpc.go
83 lines (68 loc) · 1.94 KB
/
grpc.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
package sentry
import (
"context"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
)
type options struct {
Repanic bool
ReportOn ReportOn
}
func BuildOptions(ff ...Option) options {
opts := options{
ReportOn: ReportAlways,
}
for _, f := range ff {
f(&opts)
}
return opts
}
// Option configures reporting behavior.
type Option func(*options)
// WithRepanic configures whether to panic again after recovering from
// a panic. Use this option if you have other panic handlers.
func WithRepanic(b bool) Option {
return func(o *options) {
o.Repanic = b
}
}
// WithReportOn configures whether to report on errors.
func WithReportOn(r ReportOn) Option {
return func(o *options) {
o.ReportOn = r
}
}
// ReportOn decides error should be reported to sentry.
type ReportOn func(error) bool
// ReportAlways returns true if err is non-nil.
func ReportAlways(err error) bool {
return err != nil
}
// ReportOnCodes returns true if error code matches on of the given codes.
func ReportOnCodes(cc ...codes.Code) ReportOn {
cm := make(map[codes.Code]bool)
for _, c := range cc {
cm[c] = true
}
return func(err error) bool {
return cm[status.Code(err)]
}
}
// WrappedServerStream is a thin wrapper around grpc.ServerStream that allows modifying context.
type WrappedServerStream struct {
grpc.ServerStream
// WrappedContext is the wrapper's own Context. You can assign it.
WrappedContext context.Context
}
// Context returns the wrapper's WrappedContext, overwriting the nested grpc.ServerStream.Context()
func (w *WrappedServerStream) Context() context.Context {
return w.WrappedContext
}
// WrapServerStream returns a ServerStream that has the ability to overwrite context.
func WrapServerStream(stream grpc.ServerStream) *WrappedServerStream {
if existing, ok := stream.(*WrappedServerStream); ok {
return existing
}
return &WrappedServerStream{ServerStream: stream, WrappedContext: stream.Context()}
}