Merge branch 'master' of git.curoverse.com:arvados into 11876-r-sdk

[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 d22e02597b75d5d5d7a0b74599cc4d142d002345..819c2d2510f992ab75c18249d8b46d6c13828d6d 100644 (file)
--- a/services/crunch-dispatch-slurm/squeue.go
+++ b/services/crunch-dispatch-slurm/squeue.go
@@ -6,6 +6,7 @@ package main
 
 import (
        "bytes"
+       "fmt"
        "log"
        "os/exec"
        "strings"
@@ -13,18 +14,23 @@ import (
        "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]bool
+       uuids     map[string]jobPriority
        startOnce sync.Once
        done      chan struct{}
        sync.Cond
 }
 
 func squeueFunc() *exec.Cmd {
-       return exec.Command("squeue", "--all", "--format=%j")
+       return exec.Command("squeue", "--all", "--format=%j %y %Q")
 }
 
 var squeueCmd = squeueFunc
@@ -40,7 +46,23 @@ func (sqc *SqueueChecker) HasUUID(uuid string) bool {
 
        // block until next squeue broadcast signaling an update.
        sqc.Wait()
-       return sqc.uuids[uuid]
+       _, exists := sqc.uuids[uuid]
+       return exists
+}
+
+// GetNiceness returns the niceness of a given uuid, or -1 if it doesn't exist.
+func (sqc *SqueueChecker) GetNiceness(uuid string) int {
+       sqc.startOnce.Do(sqc.start)
+
+       sqc.L.Lock()
+       defer sqc.L.Unlock()
+
+       n, exists := sqc.uuids[uuid]
+       if exists {
+               return n.niceness
+       } else {
+               return -1
+       }
 }
 
 // Stop stops the squeue monitoring goroutine. Do not call HasUUID
@@ -70,10 +92,16 @@ func (sqc *SqueueChecker) check() {
                return
        }
 
-       uuids := strings.Split(stdout.String(), "\n")
-       sqc.uuids = make(map[string]bool, len(uuids))
-       for _, uuid := range uuids {
-               sqc.uuids[uuid] = true
+       lines := strings.Split(stdout.String(), "\n")
+       sqc.uuids = make(map[string]jobPriority, len(lines))
+       for _, line := range lines {
+               var uuid string
+               var nice int
+               var prio int
+               fmt.Sscan(line, &uuid, &nice, &prio)
+               if uuid != "" {
+                       sqc.uuids[uuid] = jobPriority{nice, prio}
+               }
        }
        sqc.Broadcast()
 }