Merge branch '16723-kill-vs-requeue'

[arvados.git] / lib / dispatchcloud / scheduler / run_queue_test.go

diff --git a/lib/dispatchcloud/scheduler/run_queue_test.go b/lib/dispatchcloud/scheduler/run_queue_test.go
index 8ab1cd9ba7337a0a2c33d9628b33f257a46fa6e9..fbc73ef50c4695a7d3605db8180f1a3d263fa869 100644 (file)
--- a/lib/dispatchcloud/scheduler/run_queue_test.go
+++ b/lib/dispatchcloud/scheduler/run_queue_test.go
@@ -36,6 +36,7 @@ type stubPool struct {
        notify    <-chan struct{}
        unalloc   map[arvados.InstanceType]int // idle+booting+unknown
        idle      map[arvados.InstanceType]int
+       unknown   map[arvados.InstanceType]int
        running   map[string]time.Time
        atQuota   bool
        canCreate int
@@ -62,7 +63,7 @@ func (p *stubPool) Unallocated() map[arvados.InstanceType]int {
        defer p.Unlock()
        r := map[arvados.InstanceType]int{}
        for it, n := range p.unalloc {
-               r[it] = n
+               r[it] = n - p.unknown[it]
        }
        return r
 }
@@ -82,8 +83,9 @@ func (p *stubPool) ForgetContainer(uuid string) {
 func (p *stubPool) KillContainer(uuid, reason string) bool {
        p.Lock()
        defer p.Unlock()
-       delete(p.running, uuid)
-       return true
+       defer delete(p.running, uuid)
+       t, ok := p.running[uuid]
+       return ok && t.IsZero()
 }
 func (p *stubPool) Shutdown(arvados.InstanceType) bool {
        p.shutdowns++
@@ -96,6 +98,7 @@ func (p *stubPool) CountWorkers() map[worker.State]int {
                worker.StateBooting: len(p.unalloc) - len(p.idle),
                worker.StateIdle:    len(p.idle),
                worker.StateRunning: len(p.running),
+               worker.StateUnknown: len(p.unknown),
        }
 }
 func (p *stubPool) StartContainer(it arvados.InstanceType, ctr arvados.Container) bool {
@@ -349,7 +352,7 @@ func (*SchedulerSuite) TestKillNonexistentContainer(c *check.C) {
                        test.InstanceType(2): 0,
                },
                running: map[string]time.Time{
-                       test.ContainerUUID(2): time.Time{},
+                       test.ContainerUUID(2): {},
                },
        }
        queue := test.Queue{