github.com/bigcommerce/nomad@v0.9.3-bc/drivers/shared/executor/client.go (about)

     1  package executor
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"os"
     8  	"syscall"
     9  	"time"
    10  
    11  	"github.com/LK4D4/joincontext"
    12  	"github.com/golang/protobuf/ptypes"
    13  	hclog "github.com/hashicorp/go-hclog"
    14  	cstructs "github.com/hashicorp/nomad/client/structs"
    15  	"github.com/hashicorp/nomad/drivers/shared/executor/proto"
    16  	"github.com/hashicorp/nomad/helper/pluginutils/grpcutils"
    17  	"github.com/hashicorp/nomad/plugins/drivers"
    18  	dproto "github.com/hashicorp/nomad/plugins/drivers/proto"
    19  )
    20  
    21  var _ Executor = (*grpcExecutorClient)(nil)
    22  
    23  type grpcExecutorClient struct {
    24  	client proto.ExecutorClient
    25  	logger hclog.Logger
    26  
    27  	// doneCtx is close when the plugin exits
    28  	doneCtx context.Context
    29  }
    30  
    31  func (c *grpcExecutorClient) Launch(cmd *ExecCommand) (*ProcessState, error) {
    32  	ctx := context.Background()
    33  	req := &proto.LaunchRequest{
    34  		Cmd:                cmd.Cmd,
    35  		Args:               cmd.Args,
    36  		Resources:          drivers.ResourcesToProto(cmd.Resources),
    37  		StdoutPath:         cmd.StdoutPath,
    38  		StderrPath:         cmd.StderrPath,
    39  		Env:                cmd.Env,
    40  		User:               cmd.User,
    41  		TaskDir:            cmd.TaskDir,
    42  		ResourceLimits:     cmd.ResourceLimits,
    43  		BasicProcessCgroup: cmd.BasicProcessCgroup,
    44  		Mounts:             drivers.MountsToProto(cmd.Mounts),
    45  		Devices:            drivers.DevicesToProto(cmd.Devices),
    46  	}
    47  	resp, err := c.client.Launch(ctx, req)
    48  	if err != nil {
    49  		return nil, err
    50  	}
    51  
    52  	ps, err := processStateFromProto(resp.Process)
    53  	if err != nil {
    54  		return nil, err
    55  	}
    56  	return ps, nil
    57  }
    58  
    59  func (c *grpcExecutorClient) Wait(ctx context.Context) (*ProcessState, error) {
    60  	// Join the passed context and the shutdown context
    61  	ctx, _ = joincontext.Join(ctx, c.doneCtx)
    62  
    63  	resp, err := c.client.Wait(ctx, &proto.WaitRequest{})
    64  	if err != nil {
    65  		return nil, err
    66  	}
    67  
    68  	ps, err := processStateFromProto(resp.Process)
    69  	if err != nil {
    70  		return nil, err
    71  	}
    72  
    73  	return ps, nil
    74  }
    75  
    76  func (c *grpcExecutorClient) Shutdown(signal string, gracePeriod time.Duration) error {
    77  	ctx := context.Background()
    78  	req := &proto.ShutdownRequest{
    79  		Signal:      signal,
    80  		GracePeriod: gracePeriod.Nanoseconds(),
    81  	}
    82  	if _, err := c.client.Shutdown(ctx, req); err != nil {
    83  		return err
    84  	}
    85  
    86  	return nil
    87  }
    88  
    89  func (c *grpcExecutorClient) UpdateResources(r *drivers.Resources) error {
    90  	ctx := context.Background()
    91  	req := &proto.UpdateResourcesRequest{Resources: drivers.ResourcesToProto(r)}
    92  	if _, err := c.client.UpdateResources(ctx, req); err != nil {
    93  		return err
    94  	}
    95  
    96  	return nil
    97  }
    98  
    99  func (c *grpcExecutorClient) Version() (*ExecutorVersion, error) {
   100  	ctx := context.Background()
   101  	resp, err := c.client.Version(ctx, &proto.VersionRequest{})
   102  	if err != nil {
   103  		return nil, err
   104  	}
   105  	return &ExecutorVersion{Version: resp.Version}, nil
   106  }
   107  
   108  func (c *grpcExecutorClient) Stats(ctx context.Context, interval time.Duration) (<-chan *cstructs.TaskResourceUsage, error) {
   109  	stream, err := c.client.Stats(ctx, &proto.StatsRequest{})
   110  	if err != nil {
   111  		return nil, err
   112  	}
   113  
   114  	ch := make(chan *cstructs.TaskResourceUsage)
   115  	go c.handleStats(ctx, stream, ch)
   116  	return ch, nil
   117  }
   118  
   119  func (c *grpcExecutorClient) handleStats(ctx context.Context, stream proto.Executor_StatsClient, ch chan<- *cstructs.TaskResourceUsage) {
   120  	defer close(ch)
   121  	for {
   122  		resp, err := stream.Recv()
   123  		if ctx.Err() != nil {
   124  			// Context canceled; exit gracefully
   125  			return
   126  		}
   127  
   128  		if err != nil {
   129  			if err != io.EOF {
   130  				c.logger.Error("error receiving stream from Stats executor RPC, closing stream", "error", err)
   131  			}
   132  
   133  			// End stream
   134  			return
   135  		}
   136  
   137  		stats, err := drivers.TaskStatsFromProto(resp.Stats)
   138  		if err != nil {
   139  			c.logger.Error("failed to decode stats from RPC", "error", err, "stats", resp.Stats)
   140  			continue
   141  		}
   142  
   143  		select {
   144  		case ch <- stats:
   145  		case <-ctx.Done():
   146  			return
   147  		}
   148  	}
   149  }
   150  
   151  func (c *grpcExecutorClient) Signal(s os.Signal) error {
   152  	ctx := context.Background()
   153  	sig, ok := s.(syscall.Signal)
   154  	if !ok {
   155  		return fmt.Errorf("unsupported signal type: %q", s.String())
   156  	}
   157  	req := &proto.SignalRequest{
   158  		Signal: int32(sig),
   159  	}
   160  	if _, err := c.client.Signal(ctx, req); err != nil {
   161  		return err
   162  	}
   163  
   164  	return nil
   165  }
   166  
   167  func (c *grpcExecutorClient) Exec(deadline time.Time, cmd string, args []string) ([]byte, int, error) {
   168  	ctx := context.Background()
   169  	pbDeadline, err := ptypes.TimestampProto(deadline)
   170  	if err != nil {
   171  		return nil, 0, err
   172  	}
   173  	req := &proto.ExecRequest{
   174  		Deadline: pbDeadline,
   175  		Cmd:      cmd,
   176  		Args:     args,
   177  	}
   178  
   179  	resp, err := c.client.Exec(ctx, req)
   180  	if err != nil {
   181  		return nil, 0, err
   182  	}
   183  
   184  	return resp.Output, int(resp.ExitCode), nil
   185  }
   186  
   187  func (d *grpcExecutorClient) ExecStreaming(ctx context.Context,
   188  	command []string,
   189  	tty bool,
   190  	execStream drivers.ExecTaskStream) error {
   191  
   192  	err := d.execStreaming(ctx, command, tty, execStream)
   193  	if err != nil {
   194  		return grpcutils.HandleGrpcErr(err, d.doneCtx)
   195  	}
   196  	return nil
   197  }
   198  
   199  func (d *grpcExecutorClient) execStreaming(ctx context.Context,
   200  	command []string,
   201  	tty bool,
   202  	execStream drivers.ExecTaskStream) error {
   203  
   204  	stream, err := d.client.ExecStreaming(ctx)
   205  	if err != nil {
   206  		return err
   207  	}
   208  
   209  	err = stream.Send(&dproto.ExecTaskStreamingRequest{
   210  		Setup: &dproto.ExecTaskStreamingRequest_Setup{
   211  			Command: command,
   212  			Tty:     tty,
   213  		},
   214  	})
   215  	if err != nil {
   216  		return err
   217  	}
   218  
   219  	errCh := make(chan error, 1)
   220  	go func() {
   221  		for {
   222  			m, err := execStream.Recv()
   223  			if err == io.EOF {
   224  				return
   225  			} else if err != nil {
   226  				errCh <- err
   227  				return
   228  			}
   229  
   230  			if err := stream.Send(m); err != nil {
   231  				errCh <- err
   232  				return
   233  			}
   234  
   235  		}
   236  	}()
   237  
   238  	for {
   239  		select {
   240  		case err := <-errCh:
   241  			return err
   242  		default:
   243  		}
   244  
   245  		m, err := stream.Recv()
   246  		if err == io.EOF {
   247  			return nil
   248  		} else if err != nil {
   249  			return err
   250  		}
   251  
   252  		if err := execStream.Send(m); err != nil {
   253  			return err
   254  		}
   255  	}
   256  }