forked from hashicorp/nomad
/
plugin.go
51 lines (45 loc) · 1.38 KB
/
plugin.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
package device
import (
"context"
log "github.com/hashicorp/go-hclog"
plugin "github.com/hashicorp/go-plugin"
"github.com/hashicorp/nomad/plugins/base"
bproto "github.com/hashicorp/nomad/plugins/base/proto"
"github.com/hashicorp/nomad/plugins/device/proto"
"google.golang.org/grpc"
)
// PluginDevice is wraps a DevicePlugin and implements go-plugins GRPCPlugin
// interface to expose the interface over gRPC.
type PluginDevice struct {
plugin.NetRPCUnsupportedPlugin
Impl DevicePlugin
}
func (p *PluginDevice) GRPCServer(broker *plugin.GRPCBroker, s *grpc.Server) error {
proto.RegisterDevicePluginServer(s, &devicePluginServer{
impl: p.Impl,
broker: broker,
})
return nil
}
func (p *PluginDevice) GRPCClient(ctx context.Context, broker *plugin.GRPCBroker, c *grpc.ClientConn) (interface{}, error) {
return &devicePluginClient{
doneCtx: ctx,
client: proto.NewDevicePluginClient(c),
BasePluginClient: &base.BasePluginClient{
Client: bproto.NewBasePluginClient(c),
DoneCtx: ctx,
},
}, nil
}
// Serve is used to serve a device plugin
func Serve(dev DevicePlugin, logger log.Logger) {
plugin.Serve(&plugin.ServeConfig{
HandshakeConfig: base.Handshake,
Plugins: map[string]plugin.Plugin{
base.PluginTypeBase: &base.PluginBase{Impl: dev},
base.PluginTypeDevice: &PluginDevice{Impl: dev},
},
GRPCServer: plugin.DefaultGRPCServer,
Logger: logger,
})
}