open-nomad/client/driver/docker_progress.go

package driver

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"sync"
	"time"

	"github.com/docker/docker/pkg/jsonmessage"
	units "github.com/docker/go-units"
)

const (
	// defaultPullActivityDeadline is the default value set in the imageProgressManager
	// when newImageProgressManager is called
	defaultPullActivityDeadline = 2 * time.Minute

	// defaultImageProgressReportInterval is the default value set in the
	// imageProgressManager when newImageProgressManager is called
	defaultImageProgressReportInterval = 10 * time.Second
)

// layerProgress tracks the state and downloaded bytes of a single layer within
// a docker image
type layerProgress struct {
	id           string
	status       layerProgressStatus
	currentBytes int64
	totalBytes   int64
}

type layerProgressStatus int

const (
	layerProgressStatusUnknown layerProgressStatus = iota
	layerProgressStatusStarting
	layerProgressStatusWaiting
	layerProgressStatusDownloading
	layerProgressStatusVerifying
	layerProgressStatusDownloaded
	layerProgressStatusExtracting
	layerProgressStatusComplete
	layerProgressStatusExists
)

func lpsFromString(status string) layerProgressStatus {
	switch status {
	case "Pulling fs layer":
		return layerProgressStatusStarting
	case "Waiting":
		return layerProgressStatusWaiting
	case "Downloading":
		return layerProgressStatusDownloading
	case "Verifying Checksum":
		return layerProgressStatusVerifying
	case "Download complete":
		return layerProgressStatusDownloaded
	case "Extracting":
		return layerProgressStatusExtracting
	case "Pull complete":
		return layerProgressStatusComplete
	case "Already exists":
		return layerProgressStatusExists
	default:
		return layerProgressStatusUnknown
	}
}

// imageProgress tracks the status of each child layer as its pulled from a
// docker image repo
type imageProgress struct {
	sync.RWMutex
	lastMessage *jsonmessage.JSONMessage
	timestamp   time.Time
	layers      map[string]*layerProgress
	pullStart   time.Time
}

// get returns a status message and the timestamp of the last status update
func (p *imageProgress) get() (string, time.Time) {
	p.RLock()
	defer p.RUnlock()

	if p.lastMessage == nil {
		return "No progress", p.timestamp
	}

	var pulled, pulling int
	for _, l := range p.layers {
		if l.status == layerProgressStatusDownloading {
			pulling++
		} else if l.status > layerProgressStatusVerifying {
			pulled++
		}
	}

	elapsed := time.Now().Sub(p.pullStart)
	cur := p.currentBytes()
	total := p.totalBytes()
	var est int64
	if cur != 0 {
		est = (elapsed.Nanoseconds() / cur * total) - elapsed.Nanoseconds()
	}

	return fmt.Sprintf("Pulled %d/%d (%s/%s) pulling %d layers - est %.1fs remaining",
		pulled, len(p.layers), units.BytesSize(float64(cur)), units.BytesSize(float64(total)), pulling,
		time.Duration(est).Seconds()), p.timestamp
}

// set takes a status message received from the docker engine api during an image
// pull and updates the status of the coorisponding layer
func (p *imageProgress) set(msg *jsonmessage.JSONMessage) {
	p.Lock()
	defer p.Unlock()

	p.lastMessage = msg
	p.timestamp = time.Now()

	lps := lpsFromString(msg.Status)
	if lps == layerProgressStatusUnknown {
		return
	}

	layer, ok := p.layers[msg.ID]
	if !ok {
		layer = &layerProgress{id: msg.ID}
		p.layers[msg.ID] = layer
	}
	layer.status = lps
	if msg.Progress != nil && lps == layerProgressStatusDownloading {
		layer.currentBytes = msg.Progress.Current
		layer.totalBytes = msg.Progress.Total
	} else if lps == layerProgressStatusDownloaded {
		layer.currentBytes = layer.totalBytes
	}
}

// currentBytes iterates through all image layers and sums the total of
// current bytes. The caller is responsible for acquiring a read lock on the
// imageProgress struct
func (p *imageProgress) currentBytes() int64 {
	var b int64
	for _, l := range p.layers {
		b += l.currentBytes
	}
	return b
}

// totalBytes iterates through all image layers and sums the total of
// total bytes. The caller is responsible for acquiring a read lock on the
// imageProgress struct
func (p *imageProgress) totalBytes() int64 {
	var b int64
	for _, l := range p.layers {
		b += l.totalBytes
	}
	return b
}

// progressReporterFunc defines the method for handeling inactivity and report
// events from the imageProgressManager. The image name, current status message,
// timestamp of last received status update and timestamp of when the pull started
// are passed in.
type progressReporterFunc func(image string, msg string, timestamp time.Time, pullStart time.Time)

// imageProgressManager tracks the progress of pulling a docker image from an
// image repository.
// It also implemented the io.Writer interface so as to be passed to the docker
// client pull image method in order to receive status updates from the docker
// engine api.
type imageProgressManager struct {
	imageProgress    *imageProgress
	image            string
	activityDeadline time.Duration
	inactivityFunc   progressReporterFunc
	reportInterval   time.Duration
	reporter         progressReporterFunc
	cancel           context.CancelFunc
	stopCh           chan struct{}
	buf              bytes.Buffer
}

func newImageProgressManager(
	image string, cancel context.CancelFunc,
	inactivityFunc, reporter progressReporterFunc) *imageProgressManager {

	return &imageProgressManager{
		image:            image,
		activityDeadline: defaultPullActivityDeadline,
		inactivityFunc:   inactivityFunc,
		reportInterval:   defaultImageProgressReportInterval,
		reporter:         reporter,
		imageProgress: &imageProgress{
			timestamp: time.Now(),
			layers:    make(map[string]*layerProgress),
		},
		cancel: cancel,
		stopCh: make(chan struct{}),
	}
}

// start intiates the ticker to trigger the inactivity and reporter handlers
func (pm *imageProgressManager) start() {
	pm.imageProgress.pullStart = time.Now()
	go func() {
		ticker := time.NewTicker(defaultImageProgressReportInterval)
		for {
			select {
			case <-ticker.C:
				msg, timestamp := pm.imageProgress.get()
				if time.Now().Sub(timestamp) > pm.activityDeadline {
					pm.inactivityFunc(pm.image, msg, timestamp, pm.imageProgress.pullStart)
					pm.cancel()
					return
				}
				pm.reporter(pm.image, msg, timestamp, pm.imageProgress.pullStart)
			case <-pm.stopCh:
				return
			}
		}
	}()
}

func (pm *imageProgressManager) stop() {
	close(pm.stopCh)
}

func (pm *imageProgressManager) Write(p []byte) (n int, err error) {
	n, err = pm.buf.Write(p)
	var msg jsonmessage.JSONMessage

	for {
		line, err := pm.buf.ReadBytes('\n')
		if err == io.EOF {
			// Partial write of line; push back onto buffer and break until full line
			pm.buf.Write(line)
			break
		}
		if err != nil {
			return n, err
		}
		err = json.Unmarshal(line, &msg)
		if err != nil {
			return n, err
		}

		if msg.Error != nil {
			// error received from the docker engine api
			return n, msg.Error
		}

		pm.imageProgress.set(&msg)
	}

	return
}