github.com/zoomfoo/nomad@v0.8.5-0.20180907175415-f28fd3a1a056/plugins/device/server.go (about)

     1  package device
     2  
     3  import (
     4  	context "golang.org/x/net/context"
     5  
     6  	plugin "github.com/hashicorp/go-plugin"
     7  	"github.com/hashicorp/nomad/plugins/device/proto"
     8  )
     9  
    10  // devicePluginServer wraps a device plugin and exposes it via gRPC.
    11  type devicePluginServer struct {
    12  	broker *plugin.GRPCBroker
    13  	impl   DevicePlugin
    14  }
    15  
    16  func (d *devicePluginServer) Fingerprint(req *proto.FingerprintRequest, stream proto.DevicePlugin_FingerprintServer) error {
    17  	ctx := stream.Context()
    18  	outCh, err := d.impl.Fingerprint(ctx)
    19  	if err != nil {
    20  		return err
    21  	}
    22  
    23  	for {
    24  		select {
    25  		case <-ctx.Done():
    26  			return nil
    27  		case resp, ok := <-outCh:
    28  			// The output channel has been closed, end the stream
    29  			if !ok {
    30  				return nil
    31  			}
    32  
    33  			// Handle any error
    34  			if resp.Error != nil {
    35  				return resp.Error
    36  			}
    37  
    38  			// Convert the devices
    39  			out := convertStructDeviceGroups(resp.Devices)
    40  
    41  			// Build the response
    42  			presp := &proto.FingerprintResponse{
    43  				DeviceGroup: out,
    44  			}
    45  
    46  			// Send the devices
    47  			if err := stream.Send(presp); err != nil {
    48  				return err
    49  			}
    50  		}
    51  	}
    52  }
    53  
    54  func (d *devicePluginServer) Reserve(ctx context.Context, req *proto.ReserveRequest) (*proto.ReserveResponse, error) {
    55  	resp, err := d.impl.Reserve(req.GetDeviceIds())
    56  	if err != nil {
    57  		return nil, err
    58  	}
    59  
    60  	// Make the response
    61  	presp := &proto.ReserveResponse{
    62  		ContainerRes: convertStructContainerReservation(resp),
    63  	}
    64  
    65  	return presp, nil
    66  }
    67  
    68  func (d *devicePluginServer) Stats(req *proto.StatsRequest, stream proto.DevicePlugin_StatsServer) error {
    69  	ctx := stream.Context()
    70  	outCh, err := d.impl.Stats(ctx)
    71  	if err != nil {
    72  		return err
    73  	}
    74  
    75  	for {
    76  		select {
    77  		case <-ctx.Done():
    78  			return nil
    79  		case resp, ok := <-outCh:
    80  			// The output channel has been closed, end the stream
    81  			if !ok {
    82  				return nil
    83  			}
    84  
    85  			// Handle any error
    86  			if resp.Error != nil {
    87  				return resp.Error
    88  			}
    89  
    90  			// Convert the devices
    91  			out := convertStructDeviceGroupsStats(resp.Groups)
    92  
    93  			// Build the response
    94  			presp := &proto.StatsResponse{
    95  				Groups: out,
    96  			}
    97  
    98  			// Send the devices
    99  			if err := stream.Send(presp); err != nil {
   100  				return err
   101  			}
   102  		}
   103  	}
   104  }