7816: Fixes recommended by golint.
[arvados.git] / services / crunch-exec / crunchexec.go
1 package main
2
3 import (
4         "errors"
5         "flag"
6         "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
7         "git.curoverse.com/arvados.git/sdk/go/keepclient"
8         "git.curoverse.com/arvados.git/sdk/go/manifest"
9         "github.com/curoverse/dockerclient"
10         "io"
11         "log"
12         "os"
13         "os/signal"
14         "strings"
15         "sync"
16         "syscall"
17 )
18
19 // IArvadosClient is the minimal Arvados API methods used by crunchexec.
20 type IArvadosClient interface {
21         Create(resourceType string, parameters arvadosclient.Dict, output interface{}) error
22         Get(resourceType string, uuid string, parameters arvadosclient.Dict, output interface{}) error
23         Update(resourceType string, uuid string, parameters arvadosclient.Dict, output interface{}) (err error)
24 }
25
26 // ErrCancelled is the error returned when the container is cancelled.
27 var ErrCancelled = errors.New("Cancelled")
28
29 // IKeepClient is the minimal Keep API methods used by crunchexec.
30 type IKeepClient interface {
31         PutHB(hash string, buf []byte) (string, int, error)
32         ManifestFileReader(m manifest.Manifest, filename string) (keepclient.ReadCloserWithLen, error)
33 }
34
35 // Mount describes the mount points to create inside the container.
36 type Mount struct{}
37
38 // Collection record returned by the API server.
39 type Collection struct {
40         ManifestText string `json:"manifest_text"`
41 }
42
43 // ContainerRecord is the container record returned by the API server.
44 type ContainerRecord struct {
45         UUID               string                 `json:"uuid"`
46         Command            []string               `json:"command"`
47         ContainerImage     string                 `json:"container_image"`
48         Cwd                string                 `json:"cwd"`
49         Environment        map[string]string      `json:"environment"`
50         Mounts             map[string]Mount       `json:"mounts"`
51         OutputPath         string                 `json:"output_path"`
52         Priority           int                    `json:"priority"`
53         RuntimeConstraints map[string]interface{} `json:"runtime_constraints"`
54         State              string                 `json:"state"`
55 }
56
57 // NewLogWriter is a factory function to create a new log writer.
58 type NewLogWriter func(name string) io.WriteCloser
59
60 // ThinDockerClient is the minimal Docker client interface used by crunchexec.
61 type ThinDockerClient interface {
62         StopContainer(id string, timeout int) error
63         InspectImage(id string) (*dockerclient.ImageInfo, error)
64         LoadImage(reader io.Reader) error
65         CreateContainer(config *dockerclient.ContainerConfig, name string, authConfig *dockerclient.AuthConfig) (string, error)
66         StartContainer(id string, config *dockerclient.HostConfig) error
67         ContainerLogs(id string, options *dockerclient.LogOptions) (io.ReadCloser, error)
68         Wait(id string) <-chan dockerclient.WaitResult
69         RemoveImage(name string, force bool) ([]*dockerclient.ImageDelete, error)
70 }
71
72 // ContainerRunner is the main stateful struct used for a single execution of a
73 // container.
74 type ContainerRunner struct {
75         Docker    ThinDockerClient
76         ArvClient IArvadosClient
77         Kc        IKeepClient
78         ContainerRecord
79         dockerclient.ContainerConfig
80         ContainerID string
81         ExitCode    *int
82         NewLogWriter
83         loggingDone   chan bool
84         CrunchLog     *ThrottledLogger
85         Stdout        *ThrottledLogger
86         Stderr        *ThrottledLogger
87         LogCollection *CollectionWriter
88         LogsPDH       *string
89         CancelLock    sync.Mutex
90         Cancelled     bool
91         SigChan       chan os.Signal
92         finalState    string
93 }
94
95 // SetupSignals sets up signal handling to gracefully terminate the underlying
96 // Docker container and update state when receiving a TERM, INT or QUIT signal.
97 func (runner *ContainerRunner) SetupSignals() error {
98         runner.SigChan = make(chan os.Signal, 1)
99         signal.Notify(runner.SigChan, syscall.SIGTERM)
100         signal.Notify(runner.SigChan, syscall.SIGINT)
101         signal.Notify(runner.SigChan, syscall.SIGQUIT)
102
103         go func(sig <-chan os.Signal) {
104                 for _ = range sig {
105                         if !runner.Cancelled {
106                                 runner.CancelLock.Lock()
107                                 runner.Cancelled = true
108                                 if runner.ContainerID != "" {
109                                         runner.Docker.StopContainer(runner.ContainerID, 10)
110                                 }
111                                 runner.CancelLock.Unlock()
112                         }
113                 }
114         }(runner.SigChan)
115
116         return nil
117 }
118
119 // LoadImage determines the docker image id from the container record and
120 // checks if it is available in the local Docker image store.  If not, it loads
121 // the image from Keep.
122 func (runner *ContainerRunner) LoadImage() (err error) {
123
124         runner.CrunchLog.Printf("Fetching Docker image from collection '%s'", runner.ContainerRecord.ContainerImage)
125
126         var collection Collection
127         err = runner.ArvClient.Get("collections", runner.ContainerRecord.ContainerImage, nil, &collection)
128         if err != nil {
129                 return err
130         }
131         manifest := manifest.Manifest{Text: collection.ManifestText}
132         var img, imageID string
133         for ms := range manifest.StreamIter() {
134                 img = ms.FileStreamSegments[0].Name
135                 if !strings.HasSuffix(img, ".tar") {
136                         return errors.New("First file in the collection does not end in .tar")
137                 }
138                 imageID = img[:len(img)-4]
139         }
140
141         runner.CrunchLog.Printf("Using Docker image id '%s'", imageID)
142
143         _, err = runner.Docker.InspectImage(imageID)
144         if err != nil {
145                 runner.CrunchLog.Print("Loading Docker image from keep")
146
147                 var readCloser io.ReadCloser
148                 readCloser, err = runner.Kc.ManifestFileReader(manifest, img)
149                 if err != nil {
150                         return err
151                 }
152
153                 err = runner.Docker.LoadImage(readCloser)
154                 if err != nil {
155                         return err
156                 }
157         } else {
158                 runner.CrunchLog.Print("Docker image is available")
159         }
160
161         runner.ContainerConfig.Image = imageID
162
163         return nil
164 }
165
166 // StartContainer creates the container and runs it.
167 func (runner *ContainerRunner) StartContainer() (err error) {
168         runner.CrunchLog.Print("Creating Docker container")
169
170         runner.CancelLock.Lock()
171         defer runner.CancelLock.Unlock()
172
173         if runner.Cancelled {
174                 return ErrCancelled
175         }
176
177         runner.ContainerConfig.Cmd = runner.ContainerRecord.Command
178         if runner.ContainerRecord.Cwd != "." {
179                 runner.ContainerConfig.WorkingDir = runner.ContainerRecord.Cwd
180         }
181         for k, v := range runner.ContainerRecord.Environment {
182                 runner.ContainerConfig.Env = append(runner.ContainerConfig.Env, k+"="+v)
183         }
184         runner.ContainerID, err = runner.Docker.CreateContainer(&runner.ContainerConfig, "", nil)
185         if err != nil {
186                 return
187         }
188         hostConfig := &dockerclient.HostConfig{}
189
190         runner.CrunchLog.Printf("Starting Docker container id '%s'", runner.ContainerID)
191         err = runner.Docker.StartContainer(runner.ContainerID, hostConfig)
192         if err != nil {
193                 return
194         }
195
196         return nil
197 }
198
199 // AttachLogs connects the docker container stdout and stderr logs to the
200 // Arvados logger which logs to Keep and the API server logs table.
201 func (runner *ContainerRunner) AttachLogs() (err error) {
202
203         runner.CrunchLog.Print("Attaching container logs")
204
205         var stderrReader, stdoutReader io.Reader
206         stderrReader, err = runner.Docker.ContainerLogs(runner.ContainerID, &dockerclient.LogOptions{Follow: true, Stderr: true})
207         if err != nil {
208                 return
209         }
210         stdoutReader, err = runner.Docker.ContainerLogs(runner.ContainerID, &dockerclient.LogOptions{Follow: true, Stdout: true})
211         if err != nil {
212                 return
213         }
214
215         runner.loggingDone = make(chan bool)
216
217         runner.Stdout = NewThrottledLogger(runner.NewLogWriter("stdout"))
218         runner.Stderr = NewThrottledLogger(runner.NewLogWriter("stderr"))
219         go CopyReaderToLog(stdoutReader, runner.Stdout.Logger, runner.loggingDone)
220         go CopyReaderToLog(stderrReader, runner.Stderr.Logger, runner.loggingDone)
221
222         return nil
223 }
224
225 // WaitFinish waits for the container to terminate, capture the exit code, and
226 // close the stdout/stderr logging.
227 func (runner *ContainerRunner) WaitFinish() error {
228         result := runner.Docker.Wait(runner.ContainerID)
229         wr := <-result
230         if wr.Error != nil {
231                 return wr.Error
232         }
233         runner.ExitCode = &wr.ExitCode
234
235         // drain stdout/stderr
236         <-runner.loggingDone
237         <-runner.loggingDone
238
239         runner.Stdout.Close()
240         runner.Stderr.Close()
241
242         return nil
243 }
244
245 // CommitLogs posts the collection containing the final container logs.
246 func (runner *ContainerRunner) CommitLogs() error {
247         runner.CrunchLog.Print(runner.finalState)
248         runner.CrunchLog.Close()
249
250         // Closing CrunchLog above allows it to be committed to Keep at this
251         // point, but re-open crunch log with ArvClient in case there are any
252         // other further (such as failing to write the log to Keep!) while
253         // shutting down
254         runner.CrunchLog = NewThrottledLogger(&ArvLogWriter{runner.ArvClient, runner.ContainerRecord.UUID,
255                 "crunchexec", nil})
256
257         mt, err := runner.LogCollection.ManifestText()
258         if err != nil {
259                 return err
260         }
261
262         response := make(map[string]string)
263         err = runner.ArvClient.Create("collections",
264                 arvadosclient.Dict{"name": "logs for " + runner.ContainerRecord.UUID,
265                         "manifest_text": mt},
266                 response)
267         if err != nil {
268                 return err
269         }
270
271         runner.LogsPDH = new(string)
272         *runner.LogsPDH = response["portable_data_hash"]
273
274         return nil
275 }
276
277 // UpdateContainerRecordRunning updates the container state to "Running"
278 func (runner *ContainerRunner) UpdateContainerRecordRunning() error {
279         update := arvadosclient.Dict{"state": "Running"}
280         return runner.ArvClient.Update("containers", runner.ContainerRecord.UUID, update, nil)
281 }
282
283 // UpdateContainerRecordComplete updates the container record state on API
284 // server to "Complete" or "Cancelled"
285 func (runner *ContainerRunner) UpdateContainerRecordComplete() error {
286         update := arvadosclient.Dict{}
287         if runner.LogsPDH != nil {
288                 update["log"] = *runner.LogsPDH
289         }
290         if runner.ExitCode != nil {
291                 update["exit_code"] = *runner.ExitCode
292         }
293
294         update["state"] = runner.finalState
295
296         return runner.ArvClient.Update("containers", runner.ContainerRecord.UUID, update, nil)
297 }
298
299 // NewArvLogWriter creates an ArvLogWriter
300 func (runner *ContainerRunner) NewArvLogWriter(name string) io.WriteCloser {
301         return &ArvLogWriter{runner.ArvClient, runner.ContainerRecord.UUID, name, runner.LogCollection.Open(name + ".txt")}
302 }
303
304 // Run the full container lifecycle.
305 func (runner *ContainerRunner) Run(containerUUID string) (err error) {
306         runner.CrunchLog.Printf("Executing container '%s'", containerUUID)
307
308         var runerr, waiterr error
309
310         defer func() {
311                 if err != nil {
312                         runner.CrunchLog.Print(err)
313                 }
314
315                 if runner.Cancelled {
316                         runner.finalState = "Cancelled"
317                 } else {
318                         runner.finalState = "Complete"
319                 }
320
321                 // (6) write logs
322                 logerr := runner.CommitLogs()
323                 if logerr != nil {
324                         runner.CrunchLog.Print(logerr)
325                 }
326
327                 // (7) update container record with results
328                 updateerr := runner.UpdateContainerRecordComplete()
329                 if updateerr != nil {
330                         runner.CrunchLog.Print(updateerr)
331                 }
332
333                 runner.CrunchLog.Close()
334
335                 if err == nil {
336                         if runerr != nil {
337                                 err = runerr
338                         } else if waiterr != nil {
339                                 err = runerr
340                         } else if logerr != nil {
341                                 err = logerr
342                         } else if updateerr != nil {
343                                 err = updateerr
344                         }
345                 }
346         }()
347
348         err = runner.ArvClient.Get("containers", containerUUID, nil, &runner.ContainerRecord)
349         if err != nil {
350                 return
351         }
352
353         // (0) setup signal handling
354         err = runner.SetupSignals()
355         if err != nil {
356                 return
357         }
358
359         // (1) check for and/or load image
360         err = runner.LoadImage()
361         if err != nil {
362                 return
363         }
364
365         // (2) start container
366         err = runner.StartContainer()
367         if err != nil {
368                 if err == ErrCancelled {
369                         err = nil
370                 }
371                 return
372         }
373
374         // (3) update container record state
375         err = runner.UpdateContainerRecordRunning()
376         if err != nil {
377                 runner.CrunchLog.Print(err)
378         }
379
380         // (4) attach container logs
381         runerr = runner.AttachLogs()
382         if runerr != nil {
383                 runner.CrunchLog.Print(runerr)
384         }
385
386         // (5) wait for container to finish
387         waiterr = runner.WaitFinish()
388
389         return
390 }
391
392 // NewContainerRunner creates a new container runner.
393 func NewContainerRunner(api IArvadosClient,
394         kc IKeepClient,
395         docker ThinDockerClient) *ContainerRunner {
396
397         cr := &ContainerRunner{ArvClient: api, Kc: kc, Docker: docker}
398         cr.NewLogWriter = cr.NewArvLogWriter
399         cr.LogCollection = &CollectionWriter{kc, nil}
400         cr.CrunchLog = NewThrottledLogger(cr.NewLogWriter("crunchexec"))
401         return cr
402 }
403
404 func main() {
405         flag.Parse()
406
407         api, err := arvadosclient.MakeArvadosClient()
408         if err != nil {
409                 log.Fatal(err)
410         }
411         api.Retries = 8
412
413         var kc *keepclient.KeepClient
414         kc, err = keepclient.MakeKeepClient(&api)
415         if err != nil {
416                 log.Fatal(err)
417         }
418         kc.Retries = 4
419
420         var docker *dockerclient.DockerClient
421         docker, err = dockerclient.NewDockerClient("unix:///var/run/docker.sock", nil)
422         if err != nil {
423                 log.Fatal(err)
424         }
425
426         cr := NewContainerRunner(api, kc, docker)
427
428         err = cr.Run(flag.Arg(0))
429         if err != nil {
430                 log.Fatal(err)
431         }
432
433 }