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: 1,
  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: 1,
  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                 { // specify 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], "433"}},
  42                 },
  43                 { // ignore fake2 because SetPriority() not called
  44                         spread: 1,
  45                         squeue: uuids[0] + " 10000 4294000000\n" + uuids[1] + " 10000 4294000111\n" + uuids[2] + " 10000 4294000222\n",
  46                         want:   map[string]int64{uuids[0]: 999, uuids[1]: 1},
  47                         expect: [][]string{{uuids[0], "0"}, {uuids[1], "112"}},
  48                 },
  49         } {
  50                 c.Logf("spread=%d squeue=%q want=%v -> expect=%v", test.spread, test.squeue, test.want, test.expect)
  51                 slurm := &slurmFake{
  52                         queue: test.squeue,
  53                 }
  54                 sqc := &SqueueChecker{
  55                         Slurm:          slurm,
  56                         PrioritySpread: test.spread,
  57                         Period:         time.Hour,
  58                 }
  59                 sqc.startOnce.Do(sqc.start)
  60                 sqc.check()
  61                 for uuid, pri := range test.want {
  62                         sqc.SetPriority(uuid, pri)
  63                 }
  64                 sqc.reniceAll()
  65                 c.Check(slurm.didRenice, DeepEquals, test.expect)
  66                 sqc.Stop()
  67         }
  68 }
  69
  70 // If the given UUID isn't in the slurm queue yet, SetPriority()
  71 // should wait for it to appear on the very next poll, then give up.
  72 func (s *SqueueSuite) TestSetPriorityBeforeQueued(c *C) {
  73         uuidGood := "zzzzz-dz642-fake0fake0fake0"
  74         uuidBad := "zzzzz-dz642-fake1fake1fake1"
  75
  76         slurm := &slurmFake{}
  77         sqc := &SqueueChecker{
  78                 Slurm:  slurm,
  79                 Period: time.Hour,
  80         }
  81         sqc.startOnce.Do(sqc.start)
  82         sqc.Stop()
  83         sqc.check()
  84
  85         done := make(chan struct{})
  86         go func() {
  87                 sqc.SetPriority(uuidGood, 123)
  88                 sqc.SetPriority(uuidBad, 345)
  89                 close(done)
  90         }()
  91         c.Check(sqc.queue[uuidGood], IsNil)
  92         c.Check(sqc.queue[uuidBad], IsNil)
  93         timeout := time.NewTimer(time.Second)
  94         defer timeout.Stop()
  95         tick := time.NewTicker(time.Millisecond)
  96         defer tick.Stop()
  97         for {
  98                 select {
  99                 case <-tick.C:
 100                         slurm.queue = uuidGood + " 0 12345\n"
 101                         sqc.check()
 102                 case <-timeout.C:
 103                         c.Fatal("timed out")
 104                 case <-done:
 105                         c.Assert(sqc.queue[uuidGood], NotNil)
 106                         c.Check(sqc.queue[uuidGood].wantPriority, Equals, int64(123))
 107                         c.Check(sqc.queue[uuidBad], IsNil)
 108                         return
 109                 }
 110         }
 111 }