docker/swarmkit

package agent

import (
    "context"
    "reflect"
    "sync"

    "github.com/moby/swarmkit/v2/api"
    "github.com/moby/swarmkit/v2/log"
)

// StatusReporter receives updates to task status. Method may be called
// concurrently, so implementations should be goroutine-safe.
type StatusReporter interface {
    UpdateTaskStatus(ctx context.Context, taskID string, status *api.TaskStatus) error
}

// Reporter receives update to both task and volume status.
type Reporter interface {
    StatusReporter
    ReportVolumeUnpublished(ctx context.Context, volumeID string) error
}

type statusReporterFunc func(ctx context.Context, taskID string, status *api.TaskStatus) error

func (fn statusReporterFunc) UpdateTaskStatus(ctx context.Context, taskID string, status *api.TaskStatus) error {
    return fn(ctx, taskID, status)
}

//nolint:unused // currently only used in tests.
type volumeReporterFunc func(ctx context.Context, volumeID string) error

//nolint:unused // currently only used in tests.
func (fn volumeReporterFunc) ReportVolumeUnpublished(ctx context.Context, volumeID string) error {
    return fn(ctx, volumeID)
}

//nolint:unused // currently only used in tests.
type statusReporterCombined struct {
    statusReporterFunc
    volumeReporterFunc
}

// statusReporter creates a reliable StatusReporter that will always succeed.
// It handles several tasks at once, ensuring all statuses are reported.
//
// The reporter will continue reporting the current status until it succeeds.
type statusReporter struct {
    reporter Reporter
    statuses map[string]*api.TaskStatus
    // volumes is a set of volumes which are to be reported unpublished.
    volumes map[string]struct{}
    mu      sync.Mutex
    cond    sync.Cond
    closed  bool
}

func newStatusReporter(ctx context.Context, upstream Reporter) *statusReporter {
    r := &statusReporter{
        reporter: upstream,
        statuses: make(map[string]*api.TaskStatus),
        volumes:  make(map[string]struct{}),
    }

    r.cond.L = &r.mu

    go r.run(ctx)
    return r
}

func (sr *statusReporter) UpdateTaskStatus(ctx context.Context, taskID string, status *api.TaskStatus) error {
    sr.mu.Lock()
    defer sr.mu.Unlock()

    current, ok := sr.statuses[taskID]
    if ok {
        if reflect.DeepEqual(current, status) {
            return nil
        }

        if current.State > status.State {
            return nil // ignore old updates
        }
    }
    sr.statuses[taskID] = status
    sr.cond.Signal()

    return nil
}

func (sr *statusReporter) ReportVolumeUnpublished(ctx context.Context, volumeID string) error {
    sr.mu.Lock()
    defer sr.mu.Unlock()

    sr.volumes[volumeID] = struct{}{}
    sr.cond.Signal()

    return nil
}

func (sr *statusReporter) Close() error {
    sr.mu.Lock()
    defer sr.mu.Unlock()

    sr.closed = true
    sr.cond.Signal()

    return nil
}

func (sr *statusReporter) run(ctx context.Context) {
    done := make(chan struct{})
    defer close(done)

    sr.mu.Lock() // released during wait, below.
    defer sr.mu.Unlock()

    go func() {
        select {
        case <-ctx.Done():
            sr.Close()
        case <-done:
            return
        }
    }()

    for {
        if len(sr.statuses) == 0 && len(sr.volumes) == 0 {
            sr.cond.Wait()
        }

        if sr.closed {
            // TODO(stevvooe): Add support here for waiting until all
            // statuses are flushed before shutting down.
            return
        }

        for taskID, status := range sr.statuses {
            delete(sr.statuses, taskID) // delete the entry, while trying to send.

            sr.mu.Unlock()
            err := sr.reporter.UpdateTaskStatus(ctx, taskID, status)
            sr.mu.Lock()

            // reporter might be closed during UpdateTaskStatus call
            if sr.closed {
                return
            }

            if err != nil {
                log.G(ctx).WithError(err).Error("status reporter failed to report status to agent")

                // place it back in the map, if not there, allowing us to pick
                // the value if a new one came in when we were sending the last
                // update.
                if _, ok := sr.statuses[taskID]; !ok {
                    sr.statuses[taskID] = status
                }
            }
        }

        for volumeID := range sr.volumes {
            delete(sr.volumes, volumeID)

            sr.mu.Unlock()
            err := sr.reporter.ReportVolumeUnpublished(ctx, volumeID)
            sr.mu.Lock()

            // reporter might be closed during ReportVolumeUnpublished call
            if sr.closed {
                return
            }

            if err != nil {
                log.G(ctx).WithError(err).Error("status reporter failed to report volume status to agent")
                sr.volumes[volumeID] = struct{}{}
            }
        }
    }
}