1 // Copyright (C) The Arvados Authors. All rights reserved.
3 // SPDX-License-Identifier: Apache-2.0
5 // Package dispatch is a helper library for building Arvados container
16 "git.curoverse.com/arvados.git/sdk/go/arvados"
17 "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
21 Queued = arvados.ContainerStateQueued
22 Locked = arvados.ContainerStateLocked
23 Running = arvados.ContainerStateRunning
24 Complete = arvados.ContainerStateComplete
25 Cancelled = arvados.ContainerStateCancelled
29 type Dispatcher struct {
30 Arv *arvadosclient.ArvadosClient
32 // Batch size for container queries
35 // Queue polling frequency
36 PollPeriod time.Duration
38 // Time to wait between successive attempts to run the same container
39 MinRetryPeriod time.Duration
41 // Func that implements the container lifecycle. Must be set
42 // to a non-nil DispatchFunc before calling Run().
43 RunContainer DispatchFunc
45 auth arvados.APIClientAuthorization
47 trackers map[string]*runTracker
51 // A DispatchFunc executes a container (if the container record is
52 // Locked) or resume monitoring an already-running container, and wait
53 // until that container exits.
55 // While the container runs, the DispatchFunc should listen for
56 // updated container records on the provided channel. When the channel
57 // closes, the DispatchFunc should stop the container if it's still
58 // running, and return.
60 // The DispatchFunc should not return until the container is finished.
61 type DispatchFunc func(*Dispatcher, arvados.Container, <-chan arvados.Container)
63 // Run watches the API server's queue for containers that are either
64 // ready to run and available to lock, or are already locked by this
65 // dispatcher's token. When a new one appears, Run calls RunContainer
66 // in a new goroutine.
67 func (d *Dispatcher) Run(ctx context.Context) error {
68 err := d.Arv.Call("GET", "api_client_authorizations", "", "current", nil, &d.auth)
70 return fmt.Errorf("error getting my token UUID: %v", err)
73 d.throttle.hold = d.MinRetryPeriod
75 poll := time.NewTicker(d.PollPeriod)
90 todo := make(map[string]*runTracker)
92 // Make a copy of trackers
93 for uuid, tracker := range d.trackers {
98 // Containers I currently own (Locked/Running)
99 querySuccess := d.checkForUpdates([][]interface{}{
100 {"locked_by_uuid", "=", d.auth.UUID}}, todo)
102 // Containers I should try to dispatch
103 querySuccess = d.checkForUpdates([][]interface{}{
104 {"state", "=", Queued},
105 {"priority", ">", "0"}}, todo) && querySuccess
108 // There was an error in one of the previous queries,
109 // we probably didn't get updates for all the
110 // containers we should have. Don't check them
111 // individually because it may be expensive.
115 // Containers I know about but didn't fall into the
116 // above two categories (probably Complete/Cancelled)
118 for uuid := range todo {
119 missed = append(missed, uuid)
122 for len(missed) > 0 {
124 if len(missed) > 20 {
131 querySuccess = d.checkForUpdates([][]interface{}{
132 {"uuid", "in", batch}}, todo) && querySuccess
136 // There was an error in one of the previous queries, we probably
137 // didn't see all the containers we should have, so don't shut down
138 // the missed containers.
142 // Containers that I know about that didn't show up in any
143 // query should be let go.
144 for uuid, tracker := range todo {
145 log.Printf("Container %q not returned by any query, stopping tracking.", uuid)
152 // Start a runner in a new goroutine, and send the initial container
153 // record to its updates channel.
154 func (d *Dispatcher) start(c arvados.Container) *runTracker {
155 tracker := &runTracker{updates: make(chan arvados.Container, 1)}
158 d.RunContainer(d, c, tracker.updates)
159 // RunContainer blocks for the lifetime of the container. When
160 // it returns, the tracker should delete itself.
162 delete(d.trackers, c.UUID)
168 func (d *Dispatcher) checkForUpdates(filters [][]interface{}, todo map[string]*runTracker) bool {
169 var countList arvados.ContainerList
170 params := arvadosclient.Dict{
174 "order": []string{"priority desc"}}
175 err := d.Arv.List("containers", params, &countList)
177 log.Printf("error getting count of containers: %q", err)
180 itemsAvailable := countList.ItemsAvailable
181 params = arvadosclient.Dict{
184 "limit": d.BatchSize,
185 "order": []string{"priority desc"}}
188 params["offset"] = offset
190 // This list variable must be a new one declared
191 // inside the loop: otherwise, items in the API
192 // response would get deep-merged into the items
193 // loaded in previous iterations.
194 var list arvados.ContainerList
196 err := d.Arv.List("containers", params, &list)
198 log.Printf("Error getting list of containers: %q", err)
201 d.checkListForUpdates(list.Items, todo)
202 offset += len(list.Items)
203 if len(list.Items) == 0 || itemsAvailable <= offset {
209 func (d *Dispatcher) checkListForUpdates(containers []arvados.Container, todo map[string]*runTracker) {
212 if d.trackers == nil {
213 d.trackers = make(map[string]*runTracker)
216 for _, c := range containers {
217 tracker, alreadyTracking := d.trackers[c.UUID]
220 if c.LockedByUUID != "" && c.LockedByUUID != d.auth.UUID {
221 log.Printf("debug: ignoring %s locked by %s", c.UUID, c.LockedByUUID)
222 } else if alreadyTracking {
226 case Locked, Running:
228 case Cancelled, Complete:
234 if !d.throttle.Check(c.UUID) {
237 err := d.lock(c.UUID)
239 log.Printf("debug: error locking container %s: %s", c.UUID, err)
243 d.trackers[c.UUID] = d.start(c)
244 case Locked, Running:
245 if !d.throttle.Check(c.UUID) {
248 d.trackers[c.UUID] = d.start(c)
249 case Cancelled, Complete:
250 // no-op (we already stopped monitoring)
256 // UpdateState makes an API call to change the state of a container.
257 func (d *Dispatcher) UpdateState(uuid string, state arvados.ContainerState) error {
258 err := d.Arv.Update("containers", uuid,
260 "container": arvadosclient.Dict{"state": state},
263 log.Printf("Error updating container %s to state %q: %s", uuid, state, err)
268 // Lock makes the lock API call which updates the state of a container to Locked.
269 func (d *Dispatcher) lock(uuid string) error {
270 return d.Arv.Call("POST", "containers", uuid, "lock", nil, nil)
273 // Unlock makes the unlock API call which updates the state of a container to Queued.
274 func (d *Dispatcher) Unlock(uuid string) error {
275 return d.Arv.Call("POST", "containers", uuid, "unlock", nil, nil)
278 // TrackContainer ensures a tracker is running for the given UUID,
279 // regardless of the current state of the container (except: if the
280 // container is locked by a different dispatcher, a tracker will not
281 // be started). If the container is not in Locked or Running state,
282 // the new tracker will close down immediately.
284 // This allows the dispatcher to put its own RunContainer func into a
285 // cleanup phase (for example, to kill local processes created by a
286 // prevous dispatch process that are still running even though the
287 // container state is final) without the risk of having multiple
288 // goroutines monitoring the same UUID.
289 func (d *Dispatcher) TrackContainer(uuid string) error {
290 var cntr arvados.Container
291 err := d.Arv.Call("GET", "containers", uuid, "", nil, &cntr)
295 if cntr.LockedByUUID != "" && cntr.LockedByUUID != d.auth.UUID {
301 if _, alreadyTracking := d.trackers[uuid]; alreadyTracking {
304 if d.trackers == nil {
305 d.trackers = make(map[string]*runTracker)
307 d.trackers[uuid] = d.start(cntr)
309 case Queued, Cancelled, Complete:
310 d.trackers[uuid].close()
315 type runTracker struct {
317 updates chan arvados.Container
320 func (tracker *runTracker) close() {
321 if !tracker.closing {
322 close(tracker.updates)
324 tracker.closing = true
327 func (tracker *runTracker) update(c arvados.Container) {
332 case <-tracker.updates:
333 log.Printf("debug: runner is handling updates slowly, discarded previous update for %s", c.UUID)