-
Notifications
You must be signed in to change notification settings - Fork 568
/
server.go
118 lines (102 loc) · 3.61 KB
/
server.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
package grpcutil
import (
"context"
gotls "crypto/tls"
"fmt"
"math"
"net"
"time"
"go.uber.org/zap"
"golang.org/x/sync/errgroup"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials"
// Import registers the grpc GZIP decoder
_ "google.golang.org/grpc/encoding/gzip"
"google.golang.org/grpc/keepalive"
grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
"github.com/pachyderm/pachyderm/v2/src/internal/errors"
"github.com/pachyderm/pachyderm/v2/src/internal/log"
"github.com/pachyderm/pachyderm/v2/src/internal/middleware/logging"
"github.com/pachyderm/pachyderm/v2/src/internal/pctx"
"github.com/pachyderm/pachyderm/v2/src/internal/tls"
)
// Interceptor can be used to configure Unary and Stream interceptors
type Interceptor struct {
UnaryServerInterceptor grpc.UnaryServerInterceptor
StreamServerInterceptor grpc.StreamServerInterceptor
}
// Server is a convenience wrapper to gRPC servers that simplifies their
// setup and execution
type Server struct {
Server *grpc.Server
eg *errgroup.Group
}
// NewServer creates a new gRPC server, but does not start serving yet.
//
// If 'publicPortTLSAllowed' is set, grpcutil may enable TLS. This should be
// set for public ports that serve GRPC services to 3rd party clients. If set,
// the criterion for actually serving over TLS is: if a signed TLS cert and
// corresponding private key in 'TLSVolumePath', this will serve GRPC traffic
// over TLS. If either are missing this will serve GRPC traffic over
// unencrypted HTTP,
func NewServer(ctx context.Context, publicPortTLSAllowed bool, options ...grpc.ServerOption) (*Server, error) {
baseInterceptor := logging.NewBaseContextInterceptor(pctx.Child(ctx, "grpc"))
opts := append([]grpc.ServerOption{
grpc.MaxConcurrentStreams(math.MaxUint32),
grpc.MaxRecvMsgSize(MaxMsgSize),
grpc.MaxSendMsgSize(MaxMsgSize),
grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{
MinTime: 5 * time.Second,
PermitWithoutStream: true,
}),
grpc.ChainUnaryInterceptor(baseInterceptor.UnaryServerInterceptor, grpc_prometheus.UnaryServerInterceptor),
grpc.ChainStreamInterceptor(baseInterceptor.StreamServerInterceptor, grpc_prometheus.StreamServerInterceptor),
}, options...)
var cLoader *tls.CertLoader
if publicPortTLSAllowed {
// Validate environment
certPath, keyPath, err := tls.GetCertPaths()
if err != nil {
log.Info(ctx, "TLS disabled", zap.Error(err))
} else {
cLoader = tls.NewCertLoader(certPath, keyPath, tls.CertCheckFrequency)
// Read TLS cert and key
err := cLoader.LoadAndStart()
if err != nil {
return nil, errors.Wrapf(err, "couldn't build transport creds: %v", err)
}
transportCreds := credentials.NewTLS(&gotls.Config{GetCertificate: cLoader.GetCertificate})
opts = append(opts, grpc.Creds(transportCreds))
}
}
server := grpc.NewServer(opts...)
eg, ctx := errgroup.WithContext(ctx)
eg.Go(func() error {
<-ctx.Done()
server.GracefulStop() // This also closes the listeners
if cLoader != nil {
cLoader.Stop()
}
return nil
})
return &Server{
Server: server,
eg: eg,
}, nil
}
// ListenTCP causes the gRPC server to listen on a given TCP host and port
func (s *Server) ListenTCP(host string, port uint16) (net.Listener, error) {
listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", host, port))
if err != nil {
return nil, errors.EnsureStack(err)
}
s.eg.Go(func() error {
return errors.EnsureStack(s.Server.Serve(listener))
})
return listener, nil
}
// Wait causes the gRPC server to wait until it finishes, returning any errors
// that happened
func (s *Server) Wait() error {
return errors.EnsureStack(s.eg.Wait())
}