12891: Don't use SIGKILL when telling crunch-run to cancel.

[arvados.git] / services / crunch-dispatch-slurm / squeue.go

diff --git a/services/crunch-dispatch-slurm/squeue.go b/services/crunch-dispatch-slurm/squeue.go
index 3f73b3ca53473f1f04a9f8ddf80fe797302f2bed..5ecfe8ff2fc049201eb27c8e58c7a02eb84cbbb4 100644 (file)
--- a/services/crunch-dispatch-slurm/squeue.go
+++ b/services/crunch-dispatch-slurm/squeue.go
@@ -8,28 +8,26 @@ import (
        "bytes"
        "fmt"
        "log"
-       "os/exec"
        "strings"
        "sync"
        "time"
 )
 
+type jobPriority struct {
+       niceness        int
+       currentPriority int
+}
+
 // Squeue implements asynchronous polling monitor of the SLURM queue using the
 // command 'squeue'.
 type SqueueChecker struct {
        Period    time.Duration
-       uuids     map[string]int
+       uuids     map[string]jobPriority
        startOnce sync.Once
        done      chan struct{}
        sync.Cond
 }
 
-func squeueFunc() *exec.Cmd {
-       return exec.Command("squeue", "--all", "--format=%j %y")
-}
-
-var squeueCmd = squeueFunc
-
 // HasUUID checks if a given container UUID is in the slurm queue.
 // This does not run squeue directly, but instead blocks until woken
 // up by next successful update of squeue.
@@ -54,7 +52,7 @@ func (sqc *SqueueChecker) GetNiceness(uuid string) int {
 
        n, exists := sqc.uuids[uuid]
        if exists {
-               return n
+               return n.niceness
        } else {
                return -1
        }
@@ -79,7 +77,7 @@ func (sqc *SqueueChecker) check() {
        sqc.L.Lock()
        defer sqc.L.Unlock()
 
-       cmd := squeueCmd()
+       cmd := theConfig.slurm.QueueCommand([]string{"--all", "--format=%j %y %Q"})
        stdout, stderr := &bytes.Buffer{}, &bytes.Buffer{}
        cmd.Stdout, cmd.Stderr = stdout, stderr
        if err := cmd.Run(); err != nil {
@@ -88,13 +86,14 @@ func (sqc *SqueueChecker) check() {
        }
 
        lines := strings.Split(stdout.String(), "\n")
-       sqc.uuids = make(map[string]int, len(lines))
+       sqc.uuids = make(map[string]jobPriority, len(lines))
        for _, line := range lines {
                var uuid string
                var nice int
-               fmt.Sscan(line, &uuid, &nice)
+               var prio int
+               fmt.Sscan(line, &uuid, &nice, &prio)
                if uuid != "" {
-                       sqc.uuids[uuid] = nice
+                       sqc.uuids[uuid] = jobPriority{nice, prio}
                }
        }
        sqc.Broadcast()