X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/925a1526383299a1ade38a18e616564ab8c38da4..5a625a53394fa99ef7fb44ab63104b8195d963df:/services/crunch-dispatch-slurm/squeue.go

diff --git a/services/crunch-dispatch-slurm/squeue.go b/services/crunch-dispatch-slurm/squeue.go
index 3bebe561c3..ee87a44fff 100644
--- 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
+}