1 // Copyright (C) The Arvados Authors. All rights reserved.
3 // SPDX-License-Identifier: AGPL-3.0
16 type jobPriority struct {
21 // Squeue implements asynchronous polling monitor of the SLURM queue using the
23 type SqueueChecker struct {
25 uuids map[string]jobPriority
31 // HasUUID checks if a given container UUID is in the slurm queue.
32 // This does not run squeue directly, but instead blocks until woken
33 // up by next successful update of squeue.
34 func (sqc *SqueueChecker) HasUUID(uuid string) bool {
35 sqc.startOnce.Do(sqc.start)
40 // block until next squeue broadcast signaling an update.
42 _, exists := sqc.uuids[uuid]
46 // GetNiceness returns the niceness of a given uuid, or -1 if it doesn't exist.
47 func (sqc *SqueueChecker) GetNiceness(uuid string) int {
48 sqc.startOnce.Do(sqc.start)
53 n, exists := sqc.uuids[uuid]
61 // Stop stops the squeue monitoring goroutine. Do not call HasUUID
62 // after calling Stop.
63 func (sqc *SqueueChecker) Stop() {
69 // check gets the names of jobs in the SLURM queue (running and
70 // queued). If it succeeds, it updates squeue.uuids and wakes up any
71 // goroutines that are waiting in HasUUID() or All().
72 func (sqc *SqueueChecker) check() {
73 // Mutex between squeue sync and running sbatch or scancel. This
74 // establishes a sequence so that squeue doesn't run concurrently with
75 // sbatch or scancel; the next update of squeue will occur only after
76 // sbatch or scancel has completed.
80 cmd := theConfig.slurm.QueueCommand([]string{"--all", "--format=%j %y %Q"})
81 stdout, stderr := &bytes.Buffer{}, &bytes.Buffer{}
82 cmd.Stdout, cmd.Stderr = stdout, stderr
83 if err := cmd.Run(); err != nil {
84 log.Printf("Error running %q %q: %s %q", cmd.Path, cmd.Args, err, stderr.String())
88 lines := strings.Split(stdout.String(), "\n")
89 sqc.uuids = make(map[string]jobPriority, len(lines))
90 for _, line := range lines {
94 fmt.Sscan(line, &uuid, &nice, &prio)
96 sqc.uuids[uuid] = jobPriority{nice, prio}
102 // Initialize, and start a goroutine to call check() once per
103 // squeue.Period until terminated by calling Stop().
104 func (sqc *SqueueChecker) start() {
105 sqc.L = &sync.Mutex{}
106 sqc.done = make(chan struct{})
108 ticker := time.NewTicker(sqc.Period)
121 // All waits for the next squeue invocation, and returns all job
122 // names reported by squeue.
123 func (sqc *SqueueChecker) All() []string {
124 sqc.startOnce.Do(sqc.start)
129 for uuid := range sqc.uuids {
130 uuids = append(uuids, uuid)