X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/e7d49ab46462902a98d812b69d7257d837f85e90..2f66d4cc05e9442a9bb69969744d0750a02a1ed4:/services/crunch-dispatch-local/crunch-dispatch-local.go

diff --git a/services/crunch-dispatch-local/crunch-dispatch-local.go b/services/crunch-dispatch-local/crunch-dispatch-local.go
index 936a9088ed..2922817b55 100644
--- a/services/crunch-dispatch-local/crunch-dispatch-local.go
+++ b/services/crunch-dispatch-local/crunch-dispatch-local.go
@@ -1,23 +1,34 @@
+// Copyright (C) The Arvados Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
 package main
 
 // Dispatcher service for Crunch that runs containers locally.
 
 import (
+	"context"
 	"flag"
-	"git.curoverse.com/arvados.git/sdk/go/arvados"
-	"git.curoverse.com/arvados.git/sdk/go/arvadosclient"
-	"git.curoverse.com/arvados.git/sdk/go/dispatch"
-	"log"
+	"fmt"
 	"os"
 	"os/exec"
+	"os/signal"
 	"sync"
+	"syscall"
 	"time"
+
+	"git.arvados.org/arvados.git/sdk/go/arvados"
+	"git.arvados.org/arvados.git/sdk/go/arvadosclient"
+	"git.arvados.org/arvados.git/sdk/go/dispatch"
+	"github.com/sirupsen/logrus"
 )
 
+var version = "dev"
+
 func main() {
 	err := doMain()
 	if err != nil {
-		log.Fatalf("%q", err)
+		logrus.Fatalf("%q", err)
 	}
 }
 
@@ -29,6 +40,14 @@ var (
 )
 
 func doMain() error {
+	logger := logrus.StandardLogger()
+	if os.Getenv("DEBUG") != "" {
+		logger.SetLevel(logrus.DebugLevel)
+	}
+	logger.Formatter = &logrus.JSONFormatter{
+		TimestampFormat: "2006-01-02T15:04:05.000000000Z07:00",
+	}
+
 	flags := flag.NewFlagSet("crunch-dispatch-local", flag.ExitOnError)
 
 	pollInterval := flags.Int(
@@ -41,29 +60,53 @@ func doMain() error {
 		"/usr/bin/crunch-run",
 		"Crunch command to run container")
 
+	getVersion := flags.Bool(
+		"version",
+		false,
+		"Print version information and exit.")
+
 	// Parse args; omit the first arg which is the command name
 	flags.Parse(os.Args[1:])
 
+	// Print version information if requested
+	if *getVersion {
+		fmt.Printf("crunch-dispatch-local %s\n", version)
+		return nil
+	}
+
+	logger.Printf("crunch-dispatch-local %s started", version)
+
 	runningCmds = make(map[string]*exec.Cmd)
 
 	arv, err := arvadosclient.MakeArvadosClient()
 	if err != nil {
-		log.Printf("Error making Arvados client: %v", err)
+		logger.Errorf("error making Arvados client: %v", err)
 		return err
 	}
 	arv.Retries = 25
 
+	ctx, cancel := context.WithCancel(context.Background())
+
 	dispatcher := dispatch.Dispatcher{
-		Arv:            arv,
-		RunContainer:   run,
-		PollInterval:   time.Duration(*pollInterval) * time.Second,
-		DoneProcessing: make(chan struct{})}
+		Logger:       logger,
+		Arv:          arv,
+		RunContainer: (&LocalRun{startFunc, make(chan bool, 8), ctx}).run,
+		PollPeriod:   time.Duration(*pollInterval) * time.Second,
+	}
 
-	err = dispatcher.RunDispatcher()
+	err = dispatcher.Run(ctx)
 	if err != nil {
 		return err
 	}
 
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, os.Interrupt, syscall.SIGTERM, syscall.SIGQUIT)
+	sig := <-c
+	logger.Printf("Received %s, shutting down", sig)
+	signal.Stop(c)
+
+	cancel()
+
 	runningCmdsMutex.Lock()
 	// Finished dispatching; interrupt any crunch jobs that are still running
 	for _, cmd := range runningCmds {
@@ -81,7 +124,11 @@ func startFunc(container arvados.Container, cmd *exec.Cmd) error {
 	return cmd.Start()
 }
 
-var startCmd = startFunc
+type LocalRun struct {
+	startCmd         func(container arvados.Container, cmd *exec.Cmd) error
+	concurrencyLimit chan bool
+	ctx              context.Context
+}
 
 // Run a container.
 //
@@ -91,29 +138,51 @@ var startCmd = startFunc
 //
 // If the container is in any other state, or is not Complete/Cancelled after
 // crunch-run terminates, mark the container as Cancelled.
-func run(dispatcher *dispatch.Dispatcher,
+func (lr *LocalRun) run(dispatcher *dispatch.Dispatcher,
 	container arvados.Container,
-	status chan arvados.Container) {
+	status <-chan arvados.Container) {
 
 	uuid := container.UUID
 
 	if container.State == dispatch.Locked {
+
+		select {
+		case lr.concurrencyLimit <- true:
+			break
+		case <-lr.ctx.Done():
+			return
+		}
+
+		defer func() { <-lr.concurrencyLimit }()
+
+		select {
+		case c := <-status:
+			// Check for state updates after possibly
+			// waiting to be ready-to-run
+			if c.Priority == 0 {
+				goto Finish
+			}
+		default:
+			break
+		}
+
 		waitGroup.Add(1)
+		defer waitGroup.Done()
 
 		cmd := exec.Command(*crunchRunCommand, uuid)
 		cmd.Stdin = nil
 		cmd.Stderr = os.Stderr
 		cmd.Stdout = os.Stderr
 
-		log.Printf("Starting container %v", uuid)
+		dispatcher.Logger.Printf("starting container %v", uuid)
 
 		// Add this crunch job to the list of runningCmds only if we
 		// succeed in starting crunch-run.
 
 		runningCmdsMutex.Lock()
-		if err := startCmd(container, cmd); err != nil {
+		if err := lr.startCmd(container, cmd); err != nil {
 			runningCmdsMutex.Unlock()
-			log.Printf("Error starting %v for %v: %q", *crunchRunCommand, uuid, err)
+			dispatcher.Logger.Warnf("error starting %q for %s: %s", *crunchRunCommand, uuid, err)
 			dispatcher.UpdateState(uuid, dispatch.Cancelled)
 		} else {
 			runningCmds[uuid] = cmd
@@ -124,9 +193,9 @@ func run(dispatcher *dispatch.Dispatcher,
 
 			go func() {
 				if _, err := cmd.Process.Wait(); err != nil {
-					log.Printf("Error while waiting for crunch job to finish for %v: %q", uuid, err)
+					dispatcher.Logger.Warnf("error while waiting for crunch job to finish for %v: %q", uuid, err)
 				}
-				log.Printf("sending done")
+				dispatcher.Logger.Debugf("sending done")
 				done <- struct{}{}
 			}()
 
@@ -138,38 +207,38 @@ func run(dispatcher *dispatch.Dispatcher,
 				case c := <-status:
 					// Interrupt the child process if priority changes to 0
 					if (c.State == dispatch.Locked || c.State == dispatch.Running) && c.Priority == 0 {
-						log.Printf("Sending SIGINT to pid %d to cancel container %v", cmd.Process.Pid, uuid)
+						dispatcher.Logger.Printf("sending SIGINT to pid %d to cancel container %v", cmd.Process.Pid, uuid)
 						cmd.Process.Signal(os.Interrupt)
 					}
 				}
 			}
 			close(done)
 
-			log.Printf("Finished container run for %v", uuid)
+			dispatcher.Logger.Printf("finished container run for %v", uuid)
 
 			// Remove the crunch job from runningCmds
 			runningCmdsMutex.Lock()
 			delete(runningCmds, uuid)
 			runningCmdsMutex.Unlock()
 		}
-		waitGroup.Done()
 	}
 
+Finish:
+
 	// If the container is not finalized, then change it to "Cancelled".
 	err := dispatcher.Arv.Get("containers", uuid, nil, &container)
 	if err != nil {
-		log.Printf("Error getting final container state: %v", err)
+		dispatcher.Logger.Warnf("error getting final container state: %v", err)
 	}
-	if container.LockedByUUID == dispatcher.Auth.UUID &&
-		(container.State == dispatch.Locked || container.State == dispatch.Running) {
-		log.Printf("After %s process termination, container state for %v is %q.  Updating it to %q",
-			*crunchRunCommand, container.State, uuid, dispatch.Cancelled)
+	if container.State == dispatch.Locked || container.State == dispatch.Running {
+		dispatcher.Logger.Warnf("after %q process termination, container state for %v is %q; updating it to %q",
+			*crunchRunCommand, uuid, container.State, dispatch.Cancelled)
 		dispatcher.UpdateState(uuid, dispatch.Cancelled)
 	}
 
 	// drain any subsequent status changes
-	for _ = range status {
+	for range status {
 	}
 
-	log.Printf("Finalized container %v", uuid)
+	dispatcher.Logger.Printf("finalized container %v", uuid)
 }