1 // Copyright (C) The Arvados Authors. All rights reserved.
3 // SPDX-License-Identifier: AGPL-3.0
12 "git.curoverse.com/arvados.git/sdk/go/arvados"
13 "github.com/Sirupsen/logrus"
14 "github.com/prometheus/client_golang/prometheus"
17 type typeChooser func(*arvados.Container) (arvados.InstanceType, error)
19 // An APIClient performs Arvados API requests. It is typically an
21 type APIClient interface {
22 RequestAndDecode(dst interface{}, method, path string, body io.Reader, params interface{}) error
25 // A QueueEnt is an entry in the queue, consisting of a container
26 // record and the instance type that should be used to run it.
27 type QueueEnt struct {
28 // The container to run. Only the UUID, State, Priority, and
29 // RuntimeConstraints fields are populated.
30 Container arvados.Container
31 InstanceType arvados.InstanceType
34 // String implements fmt.Stringer by returning the queued container's
36 func (c *QueueEnt) String() string {
37 return c.Container.UUID
40 // A Queue is an interface to an Arvados cluster's container
41 // database. It presents only the containers that are eligible to be
42 // run by, are already being run by, or have recently been run by the
43 // present dispatcher.
45 // The Entries, Get, and Forget methods do not block: they return
46 // immediately, using cached data.
48 // The updating methods (Cancel, Lock, Unlock, Update) do block: they
49 // return only after the operation has completed.
51 // A Queue's Update method should be called periodically to keep the
54 logger logrus.FieldLogger
55 reg *prometheus.Registry
56 chooseType typeChooser
59 auth *arvados.APIClientAuthorization
60 current map[string]QueueEnt
63 keeplocal map[string]struct{}
66 // NewQueue returns a new Queue. When a new container appears in the
67 // Arvados cluster's queue during Update, chooseType will be called to
68 // assign an appropriate arvados.InstanceType for the queue entry.
69 func NewQueue(logger logrus.FieldLogger, reg *prometheus.Registry, chooseType typeChooser, client APIClient) *Queue {
73 chooseType: chooseType,
75 current: map[string]QueueEnt{},
79 // Forget drops the specified container from the cache. It should be
80 // called on finalized containers to avoid leaking memory over
81 // time. It is a no-op if the indicated container is not in a
83 func (cq *Queue) Forget(uuid string) {
86 ctr := cq.current[uuid].Container
87 if ctr.State == arvados.ContainerStateComplete || ctr.State == arvados.ContainerStateCancelled {
88 delete(cq.current, uuid)
92 // Get returns the (partial) Container record for the specified
93 // container. Like a map lookup, its second return value is false if
94 // the specified container is not in the Queue.
95 func (cq *Queue) Get(uuid string) (arvados.Container, bool) {
98 if ctr, ok := cq.current[uuid]; !ok {
99 return arvados.Container{}, false
101 return ctr.Container, true
105 // Entries returns all cache entries, keyed by container UUID.
107 // The returned threshold indicates the maximum age of any cached data
108 // returned in the map. This makes it possible for a scheduler to
109 // determine correctly the outcome of a remote process that updates
110 // container state. It must first wait for the remote process to exit,
111 // then wait for the Queue to start and finish its next Update --
112 // i.e., it must wait until threshold > timeProcessExited.
113 func (cq *Queue) Entries() (entries map[string]QueueEnt, threshold time.Time) {
115 defer cq.mtx.Unlock()
116 entries = make(map[string]QueueEnt, len(cq.current))
117 for uuid, ctr := range cq.current {
120 threshold = cq.updated
124 // Update refreshes the cache from the Arvados API. It adds newly
125 // queued containers, and updates the state of previously queued
127 func (cq *Queue) Update() error {
129 cq.keeplocal = map[string]struct{}{}
130 updateStarted := time.Now()
133 next, err := cq.poll()
139 defer cq.mtx.Unlock()
140 for uuid, ctr := range next {
141 if _, keep := cq.keeplocal[uuid]; keep {
144 if cur, ok := cq.current[uuid]; !ok {
145 cq.addEnt(uuid, *ctr)
148 cq.current[uuid] = cur
151 for uuid := range cq.current {
152 if _, keep := cq.keeplocal[uuid]; keep {
154 } else if _, keep = next[uuid]; keep {
157 delete(cq.current, uuid)
161 cq.updated = updateStarted
165 func (cq *Queue) addEnt(uuid string, ctr arvados.Container) {
166 it, err := cq.chooseType(&ctr)
168 // FIXME: throttle warnings, cancel after timeout
169 cq.logger.Warnf("cannot run %s", &ctr)
172 cq.current[uuid] = QueueEnt{Container: ctr, InstanceType: it}
175 // Lock acquires the dispatch lock for the given container.
176 func (cq *Queue) Lock(uuid string) error {
177 return cq.apiUpdate(uuid, "lock")
180 // Unlock releases the dispatch lock for the given container.
181 func (cq *Queue) Unlock(uuid string) error {
182 return cq.apiUpdate(uuid, "unlock")
185 // Cancel cancels the given container.
186 func (cq *Queue) Cancel(uuid string) error {
187 return cq.client.RequestAndDecode(nil, "PUT", "arvados/v1/containers/"+uuid, nil, map[string]map[string]interface{}{
188 "container": {"state": arvados.ContainerStateCancelled},
192 func (cq *Queue) apiUpdate(uuid, action string) error {
193 var resp arvados.Container
194 err := cq.client.RequestAndDecode(&resp, "POST", "arvados/v1/containers/"+uuid+"/"+action, nil, nil)
200 defer cq.mtx.Unlock()
201 if cq.keeplocal != nil {
202 cq.keeplocal[uuid] = struct{}{}
204 if ent, ok := cq.current[uuid]; !ok {
205 cq.addEnt(uuid, resp)
207 ent.Container.State, ent.Container.Priority, ent.Container.LockedByUUID = resp.State, resp.Priority, resp.LockedByUUID
208 cq.current[uuid] = ent
213 func (cq *Queue) poll() (map[string]*arvados.Container, error) {
215 size := len(cq.current)
220 auth = &arvados.APIClientAuthorization{}
221 err := cq.client.RequestAndDecode(auth, "GET", "arvados/v1/api_client_authorizations/current", nil, nil)
230 next := make(map[string]*arvados.Container, size)
231 apply := func(updates []arvados.Container) {
232 for _, upd := range updates {
233 if next[upd.UUID] == nil {
234 next[upd.UUID] = &arvados.Container{}
236 *next[upd.UUID] = upd
239 selectParam := []string{"uuid", "state", "priority", "runtime_constraints"}
242 mine, err := cq.fetchAll(arvados.ResourceListParams{
247 Filters: []arvados.Filter{{"locked_by_uuid", "=", auth.UUID}},
254 avail, err := cq.fetchAll(arvados.ResourceListParams{
259 Filters: []arvados.Filter{{"state", "=", arvados.ContainerStateQueued}, {"priority", ">", "0"}},
268 for uuid, ent := range cq.current {
269 if next[uuid] == nil &&
270 ent.Container.State != arvados.ContainerStateCancelled &&
271 ent.Container.State != arvados.ContainerStateComplete {
272 missing = append(missing, uuid)
277 for i, page := 0, 20; i < len(missing); i += page {
279 if len(batch) > page {
282 ended, err := cq.fetchAll(arvados.ResourceListParams{
286 Filters: []arvados.Filter{{"uuid", "in", batch}},
296 func (cq *Queue) fetchAll(initialParams arvados.ResourceListParams) ([]arvados.Container, error) {
297 var results []arvados.Container
298 params := initialParams
301 // This list variable must be a new one declared
302 // inside the loop: otherwise, items in the API
303 // response would get deep-merged into the items
304 // loaded in previous iterations.
305 var list arvados.ContainerList
307 err := cq.client.RequestAndDecode(&list, "GET", "arvados/v1/containers", nil, params)
311 if len(list.Items) == 0 {
315 results = append(results, list.Items...)
316 if len(params.Order) == 1 && params.Order == "uuid" {
317 params.Filters = append(initialParams.Filters, arvados.Filter{"uuid", ">", list.Items[len(list.Items)-1].UUID})
319 params.Offset += len(list.Items)