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 }