lib/dispatchcloud/test/queue.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package test
   6
   7 import (
   8         "fmt"
   9         "sync"
  10         "time"
  11
  12         "git.curoverse.com/arvados.git/lib/dispatchcloud/container"
  13         "git.curoverse.com/arvados.git/sdk/go/arvados"
  14 )
  15
  16 // Queue is a test stub for container.Queue. The caller specifies the
  17 // initial queue state.
  18 type Queue struct {
  19         // Containers represent the API server database contents.
  20         Containers []arvados.Container
  21
  22         // ChooseType will be called for each entry in Containers. It
  23         // must not be nil.
  24         ChooseType func(*arvados.Container) (arvados.InstanceType, error)
  25
  26         entries     map[string]container.QueueEnt
  27         updTime     time.Time
  28         subscribers map[<-chan struct{}]chan struct{}
  29
  30         mtx sync.Mutex
  31 }
  32
  33 // Entries returns the containers that were queued when Update was
  34 // last called.
  35 func (q *Queue) Entries() (map[string]container.QueueEnt, time.Time) {
  36         q.mtx.Lock()
  37         defer q.mtx.Unlock()
  38         updTime := q.updTime
  39         r := map[string]container.QueueEnt{}
  40         for uuid, ent := range q.entries {
  41                 r[uuid] = ent
  42         }
  43         return r, updTime
  44 }
  45
  46 // Get returns the container from the cached queue, i.e., as it was
  47 // when Update was last called -- just like a container.Queue does. If
  48 // the state has been changed (via Lock, Unlock, or Cancel) since the
  49 // last Update, the updated state is returned.
  50 func (q *Queue) Get(uuid string) (arvados.Container, bool) {
  51         q.mtx.Lock()
  52         defer q.mtx.Unlock()
  53         ent, ok := q.entries[uuid]
  54         return ent.Container, ok
  55 }
  56
  57 func (q *Queue) Forget(uuid string) {
  58         q.mtx.Lock()
  59         defer q.mtx.Unlock()
  60         delete(q.entries, uuid)
  61 }
  62
  63 func (q *Queue) Lock(uuid string) error {
  64         q.mtx.Lock()
  65         defer q.mtx.Unlock()
  66         return q.changeState(uuid, arvados.ContainerStateQueued, arvados.ContainerStateLocked)
  67 }
  68
  69 func (q *Queue) Unlock(uuid string) error {
  70         q.mtx.Lock()
  71         defer q.mtx.Unlock()
  72         return q.changeState(uuid, arvados.ContainerStateLocked, arvados.ContainerStateQueued)
  73 }
  74
  75 func (q *Queue) Cancel(uuid string) error {
  76         q.mtx.Lock()
  77         defer q.mtx.Unlock()
  78         return q.changeState(uuid, q.entries[uuid].Container.State, arvados.ContainerStateCancelled)
  79 }
  80
  81 func (q *Queue) Subscribe() <-chan struct{} {
  82         q.mtx.Lock()
  83         defer q.mtx.Unlock()
  84         if q.subscribers == nil {
  85                 q.subscribers = map[<-chan struct{}]chan struct{}{}
  86         }
  87         ch := make(chan struct{}, 1)
  88         q.subscribers[ch] = ch
  89         return ch
  90 }
  91
  92 func (q *Queue) Unsubscribe(ch <-chan struct{}) {
  93         q.mtx.Lock()
  94         defer q.mtx.Unlock()
  95         delete(q.subscribers, ch)
  96 }
  97
  98 func (q *Queue) notify() {
  99         for _, ch := range q.subscribers {
 100                 select {
 101                 case ch <- struct{}{}:
 102                 default:
 103                 }
 104         }
 105 }
 106
 107 // caller must have lock.
 108 func (q *Queue) changeState(uuid string, from, to arvados.ContainerState) error {
 109         ent := q.entries[uuid]
 110         if ent.Container.State != from {
 111                 return fmt.Errorf("lock failed: state=%q", ent.Container.State)
 112         }
 113         ent.Container.State = to
 114         q.entries[uuid] = ent
 115         for i, ctr := range q.Containers {
 116                 if ctr.UUID == uuid {
 117                         q.Containers[i].State = to
 118                         break
 119                 }
 120         }
 121         q.notify()
 122         return nil
 123 }
 124
 125 // Update rebuilds the current entries from the Containers slice.
 126 func (q *Queue) Update() error {
 127         q.mtx.Lock()
 128         defer q.mtx.Unlock()
 129         updTime := time.Now()
 130         upd := map[string]container.QueueEnt{}
 131         for _, ctr := range q.Containers {
 132                 _, exists := q.entries[ctr.UUID]
 133                 if !exists && (ctr.State == arvados.ContainerStateComplete || ctr.State == arvados.ContainerStateCancelled) {
 134                         continue
 135                 }
 136                 it, _ := q.ChooseType(&ctr)
 137                 upd[ctr.UUID] = container.QueueEnt{
 138                         Container:    ctr,
 139                         InstanceType: it,
 140                 }
 141         }
 142         q.entries = upd
 143         q.updTime = updTime
 144         q.notify()
 145         return nil
 146 }
 147
 148 // Notify adds/updates an entry in the Containers slice.  This
 149 // simulates the effect of an API update from someone other than the
 150 // dispatcher -- e.g., crunch-run updating state to "Complete" when a
 151 // container exits.
 152 //
 153 // The resulting changes are not exposed through Get() or Entries()
 154 // until the next call to Update().
 155 func (q *Queue) Notify(upd arvados.Container) {
 156         q.mtx.Lock()
 157         defer q.mtx.Unlock()
 158         for i, ctr := range q.Containers {
 159                 if ctr.UUID == upd.UUID {
 160                         q.Containers[i] = upd
 161                         return
 162                 }
 163         }
 164         q.Containers = append(q.Containers, upd)
 165 }