11369: Fix tests. Tweak test fixture to include a mounts entry with capacity.

[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 3bebe561c380d12a3ca8b5724a7a705f27bff3ef..ee87a44fff5866632b1a27aa8910b332a0ba3b8d 100644 (file)
--- a/services/crunch-dispatch-slurm/squeue.go
+++ b/services/crunch-dispatch-slurm/squeue.go
@@ -13,7 +13,7 @@ import (
 // command 'squeue'.
 type SqueueChecker struct {
        Period    time.Duration
-       hasUUID   map[string]bool
+       uuids     map[string]bool
        startOnce sync.Once
        done      chan struct{}
        sync.Cond
@@ -36,7 +36,7 @@ func (sqc *SqueueChecker) HasUUID(uuid string) bool {
 
        // block until next squeue broadcast signaling an update.
        sqc.Wait()
-       return sqc.hasUUID[uuid]
+       return sqc.uuids[uuid]
 }
 
 // Stop stops the squeue monitoring goroutine. Do not call HasUUID
@@ -48,8 +48,8 @@ func (sqc *SqueueChecker) Stop() {
 }
 
 // check gets the names of jobs in the SLURM queue (running and
-// queued). If it succeeds, it updates squeue.hasUUID and wakes up any
-// goroutines that are waiting in HasUUID().
+// queued). If it succeeds, it updates squeue.uuids and wakes up any
+// goroutines that are waiting in HasUUID() or All().
 func (sqc *SqueueChecker) check() {
        // Mutex between squeue sync and running sbatch or scancel.  This
        // establishes a sequence so that squeue doesn't run concurrently with
@@ -67,9 +67,9 @@ func (sqc *SqueueChecker) check() {
        }
 
        uuids := strings.Split(stdout.String(), "\n")
-       sqc.hasUUID = make(map[string]bool, len(uuids))
+       sqc.uuids = make(map[string]bool, len(uuids))
        for _, uuid := range uuids {
-               sqc.hasUUID[uuid] = true
+               sqc.uuids[uuid] = true
        }
        sqc.Broadcast()
 }
@@ -92,3 +92,17 @@ func (sqc *SqueueChecker) start() {
                }
        }()
 }
+
+// All waits for the next squeue invocation, and returns all job
+// names reported by squeue.
+func (sqc *SqueueChecker) All() []string {
+       sqc.startOnce.Do(sqc.start)
+       sqc.L.Lock()
+       defer sqc.L.Unlock()
+       sqc.Wait()
+       var uuids []string
+       for uuid := range sqc.uuids {
+               uuids = append(uuids, uuid)
+       }
+       return uuids
+}