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 }