X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/ab3d39a8fdbc76425f48988c27187f3b54967ff9..54e5ad11af9ee6804c908e49249edf87de7b35dd:/services/crunch-dispatch-slurm/squeue.go?ds=inline

diff --git a/services/crunch-dispatch-slurm/squeue.go b/services/crunch-dispatch-slurm/squeue.go
index 3f73b3ca53..5ecfe8ff2f 100644
--- 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()