docker/swarmkit

package taskinit

import (
    "context"
    "sort"
    "time"

    gogotypes "github.com/gogo/protobuf/types"
    "github.com/moby/swarmkit/v2/api"
    "github.com/moby/swarmkit/v2/api/defaults"
    "github.com/moby/swarmkit/v2/log"
    "github.com/moby/swarmkit/v2/manager/orchestrator"
    "github.com/moby/swarmkit/v2/manager/orchestrator/restart"
    "github.com/moby/swarmkit/v2/manager/state/store"
)

// InitHandler defines orchestrator's action to fix tasks at start.
type InitHandler interface {
    IsRelatedService(service *api.Service) bool
    FixTask(ctx context.Context, batch *store.Batch, t *api.Task)
    SlotTuple(t *api.Task) orchestrator.SlotTuple
}

// CheckTasks fixes tasks in the store before orchestrator runs. The previous leader might
// not have finished processing their updates and left them in an inconsistent state.
func CheckTasks(ctx context.Context, s *store.MemoryStore, readTx store.ReadTx, initHandler InitHandler, startSupervisor restart.SupervisorInterface) error {
    instances := make(map[orchestrator.SlotTuple][]*api.Task)
    err := s.Batch(func(batch *store.Batch) error {
        tasks, err := store.FindTasks(readTx, store.All)
        if err != nil {
            return err
        }
        for _, t := range tasks {
            if t.ServiceID == "" {
                continue
            }

            // TODO(aluzzardi): We should NOT retrieve the service here.
            service := store.GetService(readTx, t.ServiceID)
            if service == nil {
                // Service was deleted
                err := batch.Update(func(tx store.Tx) error {
                    return store.DeleteTask(tx, t.ID)
                })
                if err != nil {
                    log.G(ctx).WithError(err).Error("failed to delete task")
                }
                continue
            }
            if !initHandler.IsRelatedService(service) {
                continue
            }

            tuple := initHandler.SlotTuple(t)
            instances[tuple] = append(instances[tuple], t)

            // handle task updates from agent which should have been triggered by task update events
            initHandler.FixTask(ctx, batch, t)

            // desired state ready is a transient state that it should be started.
            // however previous leader may not have started it, retry start here
            if t.DesiredState != api.TaskStateReady || t.Status.State > api.TaskStateCompleted {
                continue
            }
            restartDelay, _ := gogotypes.DurationFromProto(defaults.Service.Task.Restart.Delay)
            if t.Spec.Restart != nil && t.Spec.Restart.Delay != nil {
                var err error
                restartDelay, err = gogotypes.DurationFromProto(t.Spec.Restart.Delay)
                if err != nil {
                    log.G(ctx).WithError(err).Error("invalid restart delay")
                    restartDelay, _ = gogotypes.DurationFromProto(defaults.Service.Task.Restart.Delay)
                }
            }
            if restartDelay != 0 {
                var timestamp time.Time
                if t.Status.AppliedAt != nil {
                    timestamp, err = gogotypes.TimestampFromProto(t.Status.AppliedAt)
                } else {
                    timestamp, err = gogotypes.TimestampFromProto(t.Status.Timestamp)
                }
                if err == nil {
                    restartTime := timestamp.Add(restartDelay)
                    calculatedRestartDelay := time.Until(restartTime)
                    if calculatedRestartDelay < restartDelay {
                        restartDelay = calculatedRestartDelay
                    }
                    if restartDelay > 0 {
                        _ = batch.Update(func(tx store.Tx) error {
                            t := store.GetTask(tx, t.ID)
                            // TODO(aluzzardi): This is shady as well. We should have a more generic condition.
                            if t == nil || t.DesiredState != api.TaskStateReady {
                                return nil
                            }
                            startSupervisor.DelayStart(ctx, tx, nil, t.ID, restartDelay, true)
                            return nil
                        })
                        continue
                    }
                } else {
                    log.G(ctx).WithError(err).Error("invalid status timestamp")
                }
            }

            // Start now
            err := batch.Update(func(tx store.Tx) error {
                return startSupervisor.StartNow(tx, t.ID)
            })
            if err != nil {
                log.G(ctx).WithError(err).WithField("task.id", t.ID).Error("moving task out of delayed state failed")
            }
        }
        return nil
    })
    if err != nil {
        return err
    }

    for tuple, instance := range instances {
        // Find the most current spec version. That's the only one
        // we care about for the purpose of reconstructing restart
        // history.
        maxVersion := uint64(0)
        for _, t := range instance {
            if t.SpecVersion != nil && t.SpecVersion.Index > maxVersion {
                maxVersion = t.SpecVersion.Index
            }
        }

        // Create a new slice with just the current spec version tasks.
        var upToDate []*api.Task
        for _, t := range instance {
            if t.SpecVersion != nil && t.SpecVersion.Index == maxVersion {
                upToDate = append(upToDate, t)
            }
        }

        // Sort by creation timestamp
        sort.Sort(tasksByCreationTimestamp(upToDate))

        // All up-to-date tasks in this instance except the first one
        // should be considered restarted.
        if len(upToDate) < 2 {
            continue
        }
        for _, t := range upToDate[1:] {
            startSupervisor.RecordRestartHistory(tuple, t)
        }
    }
    return nil
}

type tasksByCreationTimestamp []*api.Task

func (t tasksByCreationTimestamp) Len() int {
    return len(t)
}
func (t tasksByCreationTimestamp) Swap(i, j int) {
    t[i], t[j] = t[j], t[i]
}
func (t tasksByCreationTimestamp) Less(i, j int) bool {
    if t[i].Meta.CreatedAt == nil {
        return true
    }
    if t[j].Meta.CreatedAt == nil {
        return false
    }
    if t[i].Meta.CreatedAt.Seconds < t[j].Meta.CreatedAt.Seconds {
        return true
    }
    if t[i].Meta.CreatedAt.Seconds > t[j].Meta.CreatedAt.Seconds {
        return false
    }
    return t[i].Meta.CreatedAt.Nanos < t[j].Meta.CreatedAt.Nanos
}