services/crunch-dispatch-slurm/squeue_test.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package main
   6
   7 import (
   8         "time"
   9
  10         . "gopkg.in/check.v1"
  11 )
  12
  13 var _ = Suite(&SqueueSuite{})
  14
  15 type SqueueSuite struct{}
  16
  17 func (s *SqueueSuite) TestReniceAll(c *C) {
  18         uuids := []string{"zzzzz-dz642-fake0fake0fake0", "zzzzz-dz642-fake1fake1fake1", "zzzzz-dz642-fake2fake2fake2"}
  19         for _, test := range []struct {
  20                 spread int64
  21                 squeue string
  22                 want   map[string]int64
  23                 expect [][]string
  24         }{
  25                 {
  26                         spread: 0,
  27                         squeue: uuids[0] + " 10000 4294000000\n",
  28                         want:   map[string]int64{uuids[0]: 1},
  29                         expect: [][]string{{uuids[0], "0"}},
  30                 },
  31                 { // fake0 priority is too high
  32                         spread: 0,
  33                         squeue: uuids[0] + " 10000 4294000777\n" + uuids[1] + " 10000 4294000444\n",
  34                         want:   map[string]int64{uuids[0]: 1, uuids[1]: 999},
  35                         expect: [][]string{{uuids[1], "0"}, {uuids[0], "334"}},
  36                 },
  37                 { // non-zero spread
  38                         spread: 100,
  39                         squeue: uuids[0] + " 10000 4294000777\n" + uuids[1] + " 10000 4294000444\n",
  40                         want:   map[string]int64{uuids[0]: 1, uuids[1]: 999},
  41                         expect: [][]string{{uuids[1], "0"}, {uuids[0], "434"}},
  42                 },
  43                 { // ignore fake2 because SetPriority() not called
  44                         squeue: uuids[0] + " 10000 4294000000\n" + uuids[1] + " 10000 4294000111\n" + uuids[2] + " 10000 4294000222\n",
  45                         want:   map[string]int64{uuids[0]: 999, uuids[1]: 1},
  46                         expect: [][]string{{uuids[0], "0"}, {uuids[1], "112"}},
  47                 },
  48         } {
  49                 c.Logf("spread=%d squeue=%q want=%v -> expect=%v", test.spread, test.squeue, test.want, test.expect)
  50                 slurm := &slurmFake{
  51                         queue: test.squeue,
  52                 }
  53                 sqc := &SqueueChecker{
  54                         Slurm:          slurm,
  55                         PrioritySpread: test.spread,
  56                         Period:         time.Hour,
  57                 }
  58                 sqc.startOnce.Do(sqc.start)
  59                 sqc.check()
  60                 for uuid, pri := range test.want {
  61                         sqc.SetPriority(uuid, pri)
  62                 }
  63                 sqc.reniceAll()
  64                 c.Check(slurm.didRenice, DeepEquals, test.expect)
  65                 sqc.Stop()
  66         }
  67 }
  68
  69 // If the given UUID isn't in the slurm queue yet, SetPriority()
  70 // should wait for it to appear on the very next poll, then give up.
  71 func (s *SqueueSuite) TestSetPriorityBeforeQueued(c *C) {
  72         uuidGood := "zzzzz-dz642-fake0fake0fake0"
  73         uuidBad := "zzzzz-dz642-fake1fake1fake1"
  74
  75         slurm := &slurmFake{}
  76         sqc := &SqueueChecker{
  77                 Slurm:  slurm,
  78                 Period: time.Hour,
  79         }
  80         sqc.startOnce.Do(sqc.start)
  81         sqc.Stop()
  82         sqc.check()
  83
  84         done := make(chan struct{})
  85         go func() {
  86                 sqc.SetPriority(uuidGood, 123)
  87                 sqc.SetPriority(uuidBad, 345)
  88                 close(done)
  89         }()
  90         c.Check(sqc.queue[uuidGood], IsNil)
  91         c.Check(sqc.queue[uuidBad], IsNil)
  92         timeout := time.NewTimer(time.Second)
  93         defer timeout.Stop()
  94         tick := time.NewTicker(time.Millisecond)
  95         defer tick.Stop()
  96         for {
  97                 select {
  98                 case <-tick.C:
  99                         slurm.queue = uuidGood + " 0 12345\n"
 100                         sqc.check()
 101                 case <-timeout.C:
 102                         c.Fatal("timed out")
 103                 case <-done:
 104                         c.Assert(sqc.queue[uuidGood], NotNil)
 105                         c.Check(sqc.queue[uuidGood].wantPriority, Equals, int64(123))
 106                         c.Check(sqc.queue[uuidBad], IsNil)
 107                         return
 108                 }
 109         }
 110 }