151 lines
4 KiB
Go
151 lines
4 KiB
Go
package device
|
|
|
|
import (
|
|
"context"
|
|
"io"
|
|
"time"
|
|
|
|
"github.com/LK4D4/joincontext"
|
|
"github.com/golang/protobuf/ptypes"
|
|
"github.com/hashicorp/nomad/helper/pluginutils/grpcutils"
|
|
"github.com/hashicorp/nomad/plugins/base"
|
|
"github.com/hashicorp/nomad/plugins/device/proto"
|
|
)
|
|
|
|
// devicePluginClient implements the client side of a remote device plugin, using
|
|
// gRPC to communicate to the remote plugin.
|
|
type devicePluginClient struct {
|
|
// basePluginClient is embedded to give access to the base plugin methods.
|
|
*base.BasePluginClient
|
|
|
|
client proto.DevicePluginClient
|
|
|
|
// doneCtx is closed when the plugin exits
|
|
doneCtx context.Context
|
|
}
|
|
|
|
// Fingerprint is used to retrieve the set of devices and their health from the
|
|
// device plugin. An error may be immediately returned if the fingerprint call
|
|
// could not be made or as part of the streaming response. If the context is
|
|
// cancelled, the error will be propogated.
|
|
func (d *devicePluginClient) Fingerprint(ctx context.Context) (<-chan *FingerprintResponse, error) {
|
|
// Join the passed context and the shutdown context
|
|
joinedCtx, _ := joincontext.Join(ctx, d.doneCtx)
|
|
|
|
var req proto.FingerprintRequest
|
|
stream, err := d.client.Fingerprint(joinedCtx, &req)
|
|
if err != nil {
|
|
return nil, grpcutils.HandleReqCtxGrpcErr(err, ctx, d.doneCtx)
|
|
}
|
|
|
|
out := make(chan *FingerprintResponse, 1)
|
|
go d.handleFingerprint(ctx, stream, out)
|
|
return out, nil
|
|
}
|
|
|
|
// handleFingerprint should be launched in a goroutine and handles converting
|
|
// the gRPC stream to a channel. Exits either when context is cancelled or the
|
|
// stream has an error.
|
|
func (d *devicePluginClient) handleFingerprint(
|
|
reqCtx context.Context,
|
|
stream proto.DevicePlugin_FingerprintClient,
|
|
out chan *FingerprintResponse) {
|
|
|
|
defer close(out)
|
|
for {
|
|
resp, err := stream.Recv()
|
|
if err != nil {
|
|
if err != io.EOF {
|
|
out <- &FingerprintResponse{
|
|
Error: grpcutils.HandleReqCtxGrpcErr(err, reqCtx, d.doneCtx),
|
|
}
|
|
}
|
|
|
|
// End the stream
|
|
return
|
|
}
|
|
|
|
// Send the response
|
|
f := &FingerprintResponse{
|
|
Devices: convertProtoDeviceGroups(resp.GetDeviceGroup()),
|
|
}
|
|
select {
|
|
case <-reqCtx.Done():
|
|
return
|
|
case out <- f:
|
|
}
|
|
}
|
|
}
|
|
|
|
func (d *devicePluginClient) Reserve(deviceIDs []string) (*ContainerReservation, error) {
|
|
// Build the request
|
|
req := &proto.ReserveRequest{
|
|
DeviceIds: deviceIDs,
|
|
}
|
|
|
|
// Make the request
|
|
resp, err := d.client.Reserve(d.doneCtx, req)
|
|
if err != nil {
|
|
return nil, grpcutils.HandleGrpcErr(err, d.doneCtx)
|
|
}
|
|
|
|
// Convert the response
|
|
out := convertProtoContainerReservation(resp.GetContainerRes())
|
|
return out, nil
|
|
}
|
|
|
|
// Stats is used to retrieve device statistics from the device plugin. An error
|
|
// may be immediately returned if the stats call could not be made or as part of
|
|
// the streaming response. If the context is cancelled, the error will be
|
|
// propogated.
|
|
func (d *devicePluginClient) Stats(ctx context.Context, interval time.Duration) (<-chan *StatsResponse, error) {
|
|
// Join the passed context and the shutdown context
|
|
joinedCtx, _ := joincontext.Join(ctx, d.doneCtx)
|
|
|
|
req := proto.StatsRequest{
|
|
CollectionInterval: ptypes.DurationProto(interval),
|
|
}
|
|
stream, err := d.client.Stats(joinedCtx, &req)
|
|
if err != nil {
|
|
return nil, grpcutils.HandleReqCtxGrpcErr(err, ctx, d.doneCtx)
|
|
}
|
|
|
|
out := make(chan *StatsResponse, 1)
|
|
go d.handleStats(ctx, stream, out)
|
|
return out, nil
|
|
}
|
|
|
|
// handleStats should be launched in a goroutine and handles converting
|
|
// the gRPC stream to a channel. Exits either when context is cancelled or the
|
|
// stream has an error.
|
|
func (d *devicePluginClient) handleStats(
|
|
reqCtx context.Context,
|
|
stream proto.DevicePlugin_StatsClient,
|
|
out chan *StatsResponse) {
|
|
|
|
defer close(out)
|
|
for {
|
|
resp, err := stream.Recv()
|
|
if err != nil {
|
|
if err != io.EOF {
|
|
out <- &StatsResponse{
|
|
Error: grpcutils.HandleReqCtxGrpcErr(err, reqCtx, d.doneCtx),
|
|
}
|
|
}
|
|
|
|
// End the stream
|
|
return
|
|
}
|
|
|
|
// Send the response
|
|
s := &StatsResponse{
|
|
Groups: convertProtoDeviceGroupsStats(resp.GetGroups()),
|
|
}
|
|
select {
|
|
case <-reqCtx.Done():
|
|
return
|
|
case out <- s:
|
|
}
|
|
}
|
|
}
|