github.com/Ilhicas/nomad@v1.0.4-0.20210304152020-e86851182bc3/plugins/drivers/client.go (about)

     1  package drivers
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"io"
     7  	"time"
     8  
     9  	"github.com/LK4D4/joincontext"
    10  	"github.com/golang/protobuf/ptypes"
    11  	hclog "github.com/hashicorp/go-hclog"
    12  	cstructs "github.com/hashicorp/nomad/client/structs"
    13  	"github.com/hashicorp/nomad/helper/pluginutils/grpcutils"
    14  	"github.com/hashicorp/nomad/nomad/structs"
    15  	"github.com/hashicorp/nomad/plugins/base"
    16  	"github.com/hashicorp/nomad/plugins/drivers/proto"
    17  	"github.com/hashicorp/nomad/plugins/shared/hclspec"
    18  	pstructs "github.com/hashicorp/nomad/plugins/shared/structs"
    19  	sproto "github.com/hashicorp/nomad/plugins/shared/structs/proto"
    20  	"google.golang.org/grpc/status"
    21  )
    22  
    23  var _ DriverPlugin = &driverPluginClient{}
    24  
    25  type driverPluginClient struct {
    26  	*base.BasePluginClient
    27  
    28  	client proto.DriverClient
    29  	logger hclog.Logger
    30  
    31  	// doneCtx is closed when the plugin exits
    32  	doneCtx context.Context
    33  }
    34  
    35  func (d *driverPluginClient) TaskConfigSchema() (*hclspec.Spec, error) {
    36  	req := &proto.TaskConfigSchemaRequest{}
    37  
    38  	resp, err := d.client.TaskConfigSchema(d.doneCtx, req)
    39  	if err != nil {
    40  		return nil, grpcutils.HandleGrpcErr(err, d.doneCtx)
    41  	}
    42  
    43  	return resp.Spec, nil
    44  }
    45  
    46  func (d *driverPluginClient) Capabilities() (*Capabilities, error) {
    47  	req := &proto.CapabilitiesRequest{}
    48  
    49  	resp, err := d.client.Capabilities(d.doneCtx, req)
    50  	if err != nil {
    51  		return nil, grpcutils.HandleGrpcErr(err, d.doneCtx)
    52  	}
    53  
    54  	caps := &Capabilities{}
    55  	if resp.Capabilities != nil {
    56  		caps.SendSignals = resp.Capabilities.SendSignals
    57  		caps.Exec = resp.Capabilities.Exec
    58  		caps.MustInitiateNetwork = resp.Capabilities.MustCreateNetwork
    59  
    60  		for _, mode := range resp.Capabilities.NetworkIsolationModes {
    61  			caps.NetIsolationModes = append(caps.NetIsolationModes, netIsolationModeFromProto(mode))
    62  		}
    63  
    64  		switch resp.Capabilities.FsIsolation {
    65  		case proto.DriverCapabilities_NONE:
    66  			caps.FSIsolation = FSIsolationNone
    67  		case proto.DriverCapabilities_CHROOT:
    68  			caps.FSIsolation = FSIsolationChroot
    69  		case proto.DriverCapabilities_IMAGE:
    70  			caps.FSIsolation = FSIsolationImage
    71  		default:
    72  			caps.FSIsolation = FSIsolationNone
    73  		}
    74  
    75  		caps.MountConfigs = MountConfigSupport(resp.Capabilities.MountConfigs)
    76  	}
    77  
    78  	return caps, nil
    79  }
    80  
    81  // Fingerprint the driver, return a chan that will be pushed to periodically and on changes to health
    82  func (d *driverPluginClient) Fingerprint(ctx context.Context) (<-chan *Fingerprint, error) {
    83  	req := &proto.FingerprintRequest{}
    84  
    85  	// Join the passed context and the shutdown context
    86  	joinedCtx, _ := joincontext.Join(ctx, d.doneCtx)
    87  
    88  	stream, err := d.client.Fingerprint(joinedCtx, req)
    89  	if err != nil {
    90  		return nil, grpcutils.HandleReqCtxGrpcErr(err, ctx, d.doneCtx)
    91  	}
    92  
    93  	ch := make(chan *Fingerprint, 1)
    94  	go d.handleFingerprint(ctx, ch, stream)
    95  
    96  	return ch, nil
    97  }
    98  
    99  func (d *driverPluginClient) handleFingerprint(reqCtx context.Context, ch chan *Fingerprint, stream proto.Driver_FingerprintClient) {
   100  	defer close(ch)
   101  	for {
   102  		pb, err := stream.Recv()
   103  		if err != nil {
   104  			if err != io.EOF {
   105  				ch <- &Fingerprint{
   106  					Err: grpcutils.HandleReqCtxGrpcErr(err, reqCtx, d.doneCtx),
   107  				}
   108  			}
   109  
   110  			// End the stream
   111  			return
   112  		}
   113  
   114  		f := &Fingerprint{
   115  			Attributes:        pstructs.ConvertProtoAttributeMap(pb.Attributes),
   116  			Health:            healthStateFromProto(pb.Health),
   117  			HealthDescription: pb.HealthDescription,
   118  		}
   119  
   120  		select {
   121  		case <-reqCtx.Done():
   122  			return
   123  		case ch <- f:
   124  		}
   125  	}
   126  }
   127  
   128  // RecoverTask does internal state recovery to be able to control the task of
   129  // the given TaskHandle
   130  func (d *driverPluginClient) RecoverTask(h *TaskHandle) error {
   131  	req := &proto.RecoverTaskRequest{Handle: taskHandleToProto(h)}
   132  
   133  	_, err := d.client.RecoverTask(d.doneCtx, req)
   134  	return grpcutils.HandleGrpcErr(err, d.doneCtx)
   135  }
   136  
   137  // StartTask starts execution of a task with the given TaskConfig. A TaskHandle
   138  // is returned to the caller that can be used to recover state of the task,
   139  // should the driver crash or exit prematurely.
   140  func (d *driverPluginClient) StartTask(c *TaskConfig) (*TaskHandle, *DriverNetwork, error) {
   141  	req := &proto.StartTaskRequest{
   142  		Task: taskConfigToProto(c),
   143  	}
   144  
   145  	resp, err := d.client.StartTask(d.doneCtx, req)
   146  	if err != nil {
   147  		st := status.Convert(err)
   148  		if len(st.Details()) > 0 {
   149  			if rec, ok := st.Details()[0].(*sproto.RecoverableError); ok {
   150  				return nil, nil, structs.NewRecoverableError(err, rec.Recoverable)
   151  			}
   152  		}
   153  		return nil, nil, grpcutils.HandleGrpcErr(err, d.doneCtx)
   154  	}
   155  
   156  	var net *DriverNetwork
   157  	if resp.NetworkOverride != nil {
   158  		net = &DriverNetwork{
   159  			PortMap:       map[string]int{},
   160  			IP:            resp.NetworkOverride.Addr,
   161  			AutoAdvertise: resp.NetworkOverride.AutoAdvertise,
   162  		}
   163  		for k, v := range resp.NetworkOverride.PortMap {
   164  			net.PortMap[k] = int(v)
   165  		}
   166  	}
   167  
   168  	return taskHandleFromProto(resp.Handle), net, nil
   169  }
   170  
   171  // WaitTask returns a channel that will have an ExitResult pushed to it once when the task
   172  // exits on its own or is killed. If WaitTask is called after the task has exited, the channel
   173  // will immedialy return the ExitResult. WaitTask can be called multiple times for
   174  // the same task without issue.
   175  func (d *driverPluginClient) WaitTask(ctx context.Context, id string) (<-chan *ExitResult, error) {
   176  	ch := make(chan *ExitResult)
   177  	go d.handleWaitTask(ctx, id, ch)
   178  	return ch, nil
   179  }
   180  
   181  func (d *driverPluginClient) handleWaitTask(ctx context.Context, id string, ch chan *ExitResult) {
   182  	defer close(ch)
   183  	var result ExitResult
   184  	req := &proto.WaitTaskRequest{
   185  		TaskId: id,
   186  	}
   187  
   188  	// Join the passed context and the shutdown context
   189  	joinedCtx, _ := joincontext.Join(ctx, d.doneCtx)
   190  
   191  	resp, err := d.client.WaitTask(joinedCtx, req)
   192  	if err != nil {
   193  		result.Err = grpcutils.HandleReqCtxGrpcErr(err, ctx, d.doneCtx)
   194  	} else {
   195  		result.ExitCode = int(resp.Result.ExitCode)
   196  		result.Signal = int(resp.Result.Signal)
   197  		result.OOMKilled = resp.Result.OomKilled
   198  		if len(resp.Err) > 0 {
   199  			result.Err = errors.New(resp.Err)
   200  		}
   201  	}
   202  	ch <- &result
   203  }
   204  
   205  // StopTask stops the task with the given taskID. A timeout and signal can be
   206  // given to control a graceful termination of the task. The driver will send the
   207  // given signal to the task and wait for the given timeout for it to exit. If the
   208  // task does not exit within the timeout it will be forcefully killed.
   209  func (d *driverPluginClient) StopTask(taskID string, timeout time.Duration, signal string) error {
   210  	req := &proto.StopTaskRequest{
   211  		TaskId:  taskID,
   212  		Timeout: ptypes.DurationProto(timeout),
   213  		Signal:  signal,
   214  	}
   215  
   216  	_, err := d.client.StopTask(d.doneCtx, req)
   217  	return grpcutils.HandleGrpcErr(err, d.doneCtx)
   218  }
   219  
   220  // DestroyTask removes the task from the driver's in memory state. The task
   221  // cannot be running unless force is set to true. If force is set to true the
   222  // driver will forcefully terminate the task before removing it.
   223  func (d *driverPluginClient) DestroyTask(taskID string, force bool) error {
   224  	req := &proto.DestroyTaskRequest{
   225  		TaskId: taskID,
   226  		Force:  force,
   227  	}
   228  
   229  	_, err := d.client.DestroyTask(d.doneCtx, req)
   230  	return grpcutils.HandleGrpcErr(err, d.doneCtx)
   231  }
   232  
   233  // InspectTask returns status information for a task
   234  func (d *driverPluginClient) InspectTask(taskID string) (*TaskStatus, error) {
   235  	req := &proto.InspectTaskRequest{TaskId: taskID}
   236  
   237  	resp, err := d.client.InspectTask(d.doneCtx, req)
   238  	if err != nil {
   239  		return nil, grpcutils.HandleGrpcErr(err, d.doneCtx)
   240  	}
   241  
   242  	status, err := taskStatusFromProto(resp.Task)
   243  	if err != nil {
   244  		return nil, err
   245  	}
   246  
   247  	if resp.Driver != nil {
   248  		status.DriverAttributes = resp.Driver.Attributes
   249  	}
   250  	if resp.NetworkOverride != nil {
   251  		status.NetworkOverride = &DriverNetwork{
   252  			PortMap:       map[string]int{},
   253  			IP:            resp.NetworkOverride.Addr,
   254  			AutoAdvertise: resp.NetworkOverride.AutoAdvertise,
   255  		}
   256  		for k, v := range resp.NetworkOverride.PortMap {
   257  			status.NetworkOverride.PortMap[k] = int(v)
   258  		}
   259  	}
   260  
   261  	return status, nil
   262  }
   263  
   264  // TaskStats returns resource usage statistics for the task
   265  func (d *driverPluginClient) TaskStats(ctx context.Context, taskID string, interval time.Duration) (<-chan *cstructs.TaskResourceUsage, error) {
   266  	req := &proto.TaskStatsRequest{
   267  		TaskId:             taskID,
   268  		CollectionInterval: ptypes.DurationProto(interval),
   269  	}
   270  	ctx, _ = joincontext.Join(ctx, d.doneCtx)
   271  	stream, err := d.client.TaskStats(ctx, req)
   272  	if err != nil {
   273  		st := status.Convert(err)
   274  		if len(st.Details()) > 0 {
   275  			if rec, ok := st.Details()[0].(*sproto.RecoverableError); ok {
   276  				return nil, structs.NewRecoverableError(err, rec.Recoverable)
   277  			}
   278  		}
   279  		return nil, grpcutils.HandleGrpcErr(err, d.doneCtx)
   280  	}
   281  
   282  	ch := make(chan *cstructs.TaskResourceUsage, 1)
   283  	go d.handleStats(ctx, ch, stream)
   284  
   285  	return ch, nil
   286  }
   287  
   288  func (d *driverPluginClient) handleStats(ctx context.Context, ch chan<- *cstructs.TaskResourceUsage, stream proto.Driver_TaskStatsClient) {
   289  	defer close(ch)
   290  	for {
   291  		resp, err := stream.Recv()
   292  		if ctx.Err() != nil {
   293  			// Context canceled; exit gracefully
   294  			return
   295  		}
   296  
   297  		if err != nil {
   298  			if err != io.EOF {
   299  				d.logger.Error("error receiving stream from TaskStats driver RPC, closing stream", "error", err)
   300  			}
   301  
   302  			// End of stream
   303  			return
   304  		}
   305  
   306  		stats, err := TaskStatsFromProto(resp.Stats)
   307  		if err != nil {
   308  			d.logger.Error("failed to decode stats from RPC", "error", err, "stats", resp.Stats)
   309  			continue
   310  		}
   311  
   312  		select {
   313  		case ch <- stats:
   314  		case <-ctx.Done():
   315  		}
   316  	}
   317  }
   318  
   319  // TaskEvents returns a channel that will receive events from the driver about all
   320  // tasks such as lifecycle events, terminal errors, etc.
   321  func (d *driverPluginClient) TaskEvents(ctx context.Context) (<-chan *TaskEvent, error) {
   322  	req := &proto.TaskEventsRequest{}
   323  
   324  	// Join the passed context and the shutdown context
   325  	joinedCtx, _ := joincontext.Join(ctx, d.doneCtx)
   326  
   327  	stream, err := d.client.TaskEvents(joinedCtx, req)
   328  	if err != nil {
   329  		return nil, grpcutils.HandleReqCtxGrpcErr(err, ctx, d.doneCtx)
   330  	}
   331  
   332  	ch := make(chan *TaskEvent, 1)
   333  	go d.handleTaskEvents(ctx, ch, stream)
   334  	return ch, nil
   335  }
   336  
   337  func (d *driverPluginClient) handleTaskEvents(reqCtx context.Context, ch chan *TaskEvent, stream proto.Driver_TaskEventsClient) {
   338  	defer close(ch)
   339  	for {
   340  		ev, err := stream.Recv()
   341  		if err != nil {
   342  			if err != io.EOF {
   343  				ch <- &TaskEvent{
   344  					Err: grpcutils.HandleReqCtxGrpcErr(err, reqCtx, d.doneCtx),
   345  				}
   346  			}
   347  
   348  			// End the stream
   349  			return
   350  		}
   351  
   352  		timestamp, _ := ptypes.Timestamp(ev.Timestamp)
   353  		event := &TaskEvent{
   354  			TaskID:      ev.TaskId,
   355  			AllocID:     ev.AllocId,
   356  			TaskName:    ev.TaskName,
   357  			Annotations: ev.Annotations,
   358  			Message:     ev.Message,
   359  			Timestamp:   timestamp,
   360  		}
   361  		select {
   362  		case <-reqCtx.Done():
   363  			return
   364  		case ch <- event:
   365  		}
   366  	}
   367  }
   368  
   369  // SignalTask will send the given signal to the specified task
   370  func (d *driverPluginClient) SignalTask(taskID string, signal string) error {
   371  	req := &proto.SignalTaskRequest{
   372  		TaskId: taskID,
   373  		Signal: signal,
   374  	}
   375  	_, err := d.client.SignalTask(d.doneCtx, req)
   376  	return grpcutils.HandleGrpcErr(err, d.doneCtx)
   377  }
   378  
   379  // ExecTask will run the given command within the execution context of the task.
   380  // The driver will wait for the given timeout for the command to complete before
   381  // terminating it. The stdout and stderr of the command will be return to the caller,
   382  // along with other exit information such as exit code.
   383  func (d *driverPluginClient) ExecTask(taskID string, cmd []string, timeout time.Duration) (*ExecTaskResult, error) {
   384  	req := &proto.ExecTaskRequest{
   385  		TaskId:  taskID,
   386  		Command: cmd,
   387  		Timeout: ptypes.DurationProto(timeout),
   388  	}
   389  
   390  	resp, err := d.client.ExecTask(d.doneCtx, req)
   391  	if err != nil {
   392  		return nil, grpcutils.HandleGrpcErr(err, d.doneCtx)
   393  	}
   394  
   395  	result := &ExecTaskResult{
   396  		Stdout:     resp.Stdout,
   397  		Stderr:     resp.Stderr,
   398  		ExitResult: exitResultFromProto(resp.Result),
   399  	}
   400  
   401  	return result, nil
   402  }
   403  
   404  var _ ExecTaskStreamingRawDriver = (*driverPluginClient)(nil)
   405  
   406  func (d *driverPluginClient) ExecTaskStreamingRaw(ctx context.Context,
   407  	taskID string,
   408  	command []string,
   409  	tty bool,
   410  	execStream ExecTaskStream) error {
   411  
   412  	stream, err := d.client.ExecTaskStreaming(ctx)
   413  	if err != nil {
   414  		return grpcutils.HandleGrpcErr(err, d.doneCtx)
   415  	}
   416  
   417  	err = stream.Send(&proto.ExecTaskStreamingRequest{
   418  		Setup: &proto.ExecTaskStreamingRequest_Setup{
   419  			TaskId:  taskID,
   420  			Command: command,
   421  			Tty:     tty,
   422  		},
   423  	})
   424  	if err != nil {
   425  		return grpcutils.HandleGrpcErr(err, d.doneCtx)
   426  	}
   427  
   428  	errCh := make(chan error, 1)
   429  
   430  	go func() {
   431  		for {
   432  			m, err := execStream.Recv()
   433  			if err == io.EOF {
   434  				return
   435  			} else if err != nil {
   436  				errCh <- err
   437  				return
   438  			}
   439  
   440  			if err := stream.Send(m); err != nil {
   441  				errCh <- err
   442  				return
   443  			}
   444  
   445  		}
   446  	}()
   447  
   448  	for {
   449  		select {
   450  		case err := <-errCh:
   451  			return err
   452  		default:
   453  		}
   454  
   455  		m, err := stream.Recv()
   456  		if err == io.EOF {
   457  			// Once we get to the end of stream successfully, we can ignore errCh:
   458  			// e.g. input write failures after process terminates shouldn't cause method to fail
   459  			return nil
   460  		} else if err != nil {
   461  			return err
   462  		}
   463  
   464  		if err := execStream.Send(m); err != nil {
   465  			return err
   466  		}
   467  	}
   468  }
   469  
   470  func (d *driverPluginClient) CreateNetwork(allocID string) (*NetworkIsolationSpec, bool, error) {
   471  	req := &proto.CreateNetworkRequest{
   472  		AllocId: allocID,
   473  	}
   474  
   475  	resp, err := d.client.CreateNetwork(d.doneCtx, req)
   476  	if err != nil {
   477  		return nil, false, grpcutils.HandleGrpcErr(err, d.doneCtx)
   478  	}
   479  
   480  	return NetworkIsolationSpecFromProto(resp.IsolationSpec), resp.Created, nil
   481  }
   482  
   483  func (d *driverPluginClient) DestroyNetwork(allocID string, spec *NetworkIsolationSpec) error {
   484  	req := &proto.DestroyNetworkRequest{
   485  		AllocId:       allocID,
   486  		IsolationSpec: NetworkIsolationSpecToProto(spec),
   487  	}
   488  
   489  	_, err := d.client.DestroyNetwork(d.doneCtx, req)
   490  	if err != nil {
   491  		return grpcutils.HandleGrpcErr(err, d.doneCtx)
   492  	}
   493  
   494  	return nil
   495  }