-
Notifications
You must be signed in to change notification settings - Fork 648
/
serve.go
77 lines (69 loc) · 1.87 KB
/
serve.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
package provisionersdk
import (
"context"
"errors"
"io"
"net"
"os"
"golang.org/x/xerrors"
"storj.io/drpc/drpcmux"
"storj.io/drpc/drpcserver"
"github.com/hashicorp/yamux"
"github.com/coder/coder/provisionersdk/proto"
)
// ServeOptions are configurations to serve a provisioner.
type ServeOptions struct {
// Conn specifies a custom transport to serve the dRPC connection.
Listener net.Listener
}
// Serve starts a dRPC connection for the provisioner and transport provided.
func Serve(ctx context.Context, server proto.DRPCProvisionerServer, options *ServeOptions) error {
if options == nil {
options = &ServeOptions{}
}
// Default to using stdio.
if options.Listener == nil {
config := yamux.DefaultConfig()
config.LogOutput = io.Discard
stdio, err := yamux.Server(&readWriteCloser{
ReadCloser: os.Stdin,
Writer: os.Stdout,
}, config)
if err != nil {
return xerrors.Errorf("create yamux: %w", err)
}
options.Listener = stdio
}
// dRPC is a drop-in replacement for gRPC with less generated code, and faster transports.
// See: https://www.storj.io/blog/introducing-drpc-our-replacement-for-grpc
mux := drpcmux.New()
err := proto.DRPCRegisterProvisioner(mux, server)
if err != nil {
return xerrors.Errorf("register provisioner: %w", err)
}
srv := drpcserver.New(mux)
// Only serve a single connection on the transport.
// Transports are not multiplexed, and provisioners are
// short-lived processes that can be executed concurrently.
err = srv.Serve(ctx, options.Listener)
if err != nil {
if errors.Is(err, io.EOF) {
return nil
}
if errors.Is(err, context.Canceled) {
return nil
}
if errors.Is(err, io.ErrClosedPipe) {
return nil
}
if errors.Is(err, yamux.ErrSessionShutdown) {
return nil
}
return xerrors.Errorf("serve transport: %w", err)
}
return nil
}
type readWriteCloser struct {
io.ReadCloser
io.Writer
}