Merge pull request #31 from resin-os/pull-combined-progress

Print combined progress when pulling

distribution/pull_v1.go

@@ -370,3 +370,7 @@ func (ld *v1LayerDescriptor) Registered(diffID layer.DiffID) {
  // Cache mapping from this layer's DiffID to the blobsum
  ld.v1IDService.Set(ld.v1LayerID, ld.indexName, diffID)
 }
+
+func (ld *v1LayerDescriptor) Size() int64 {
+ return ld.layerSize
+}

distribution/pull_v2.go

@@ -237,6 +237,10 @@ func (ld *v2LayerDescriptor) Registered(diffID layer.DiffID) {
  ld.V2MetadataService.Add(diffID, metadata.V2Metadata{Digest: ld.digest, SourceRepository: ld.repoInfo.Name.Name()})
 }
 
+func (ld *v2LayerDescriptor) Size() int64 {
+ return ld.src.Size
+}
+
 func (p *v2Puller) pullV2Tag(ctx context.Context, ref reference.Named) (tagUpdated bool, err error) {
  manSvc, err := p.repo.Manifests(ctx)
  if err != nil {

distribution/xfer/download.go

@@ -72,6 +72,7 @@ type DownloadDescriptor interface {
  // if it is unknown (for example, if it has not been downloaded
  // before).
  DiffID() (layer.DiffID, error)
+ Size() int64
  // Download is called to perform the download.
  Download(ctx context.Context, progressOutput progress.Output) (io.ReadCloser, int64, error)
  // Return the DeltaBase if any
@@ -114,6 +115,8 @@ func (ldm *LayerDownloadManager) Download(ctx context.Context, initialRootFS ima
  platform = layer.Platform(runtime.GOOS)
  }
 
+ totalProgress := progress.NewProgressSink(progressOutput, 0, "Total", "")
+
  rootFS := initialRootFS
  for _, descriptor := range layers {
  key := descriptor.Key()
@@ -159,13 +162,14 @@ func (ldm *LayerDownloadManager) Download(ctx context.Context, initialRootFS ima
 
  // Layer is not known to exist - download and register it.
  progress.Update(progressOutput, descriptor.ID(), "Pulling fs layer")
+ totalProgress.Size += descriptor.Size()
 
  var xferFunc DoFunc
  if topDownload != nil {
- xferFunc = ldm.makeDownloadFunc(descriptor, "", topDownload, platform)
+ xferFunc = ldm.makeDownloadFunc(descriptor, "", topDownload, platform, totalProgress)
  defer topDownload.Transfer.Release(watcher)
  } else {
- xferFunc = ldm.makeDownloadFunc(descriptor, rootFS.ChainID(), nil, platform)
+ xferFunc = ldm.makeDownloadFunc(descriptor, rootFS.ChainID(), nil, platform, totalProgress)
  }
  topDownloadUncasted, watcher = ldm.tm.Transfer(transferKey, xferFunc, progressOutput)
  topDownload = topDownloadUncasted.(*downloadTransfer)
@@ -222,7 +226,7 @@ func (ldm *LayerDownloadManager) Download(ctx context.Context, initialRootFS ima
 // complete before the registration step, and registers the downloaded data
 // on top of parentDownload's resulting layer. Otherwise, it registers the
 // layer on top of the ChainID given by parentLayer.
-func (ldm *LayerDownloadManager) makeDownloadFunc(descriptor DownloadDescriptor, parentLayer layer.ChainID, parentDownload *downloadTransfer, platform layer.Platform) DoFunc {
+func (ldm *LayerDownloadManager) makeDownloadFunc(descriptor DownloadDescriptor, parentLayer layer.ChainID, parentDownload *downloadTransfer, platform layer.Platform, totalProgress io.Writer) DoFunc {
  return func(progressChan chan<- progress.Progress, start <-chan struct{}, inactive chan<- struct{}) Transfer {
  d := &downloadTransfer{
  Transfer: NewTransfer(),
@@ -334,7 +338,7 @@ func (ldm *LayerDownloadManager) makeDownloadFunc(descriptor DownloadDescriptor,
  reader := progress.NewProgressReader(ioutils.NewCancelReadCloser(d.Transfer.Context(), downloadReader), progressOutput, size, descriptor.ID(), "Extracting")
  defer reader.Close()
 
- inflatedLayerData, err := archive.DecompressStream(reader)
+ inflatedLayerData, err := archive.DecompressStream(io.TeeReader(reader, totalProgress))
  if err != nil {
  d.err = fmt.Errorf("could not get decompression stream: %v", err)
  return

pkg/progress/progresssink.go

@@ -0,0 +1,53 @@
+package progress
+
+import (
+ "time"
+ "sync"
+
+ "golang.org/x/time/rate"
+)
+
+type Sink struct {
+ sync.Mutex
+ out Output // Where to send progress bar to
+ Size int64
+ current int64
+ lastUpdate int64
+ id string
+ action string
+ rateLimiter *rate.Limiter
+}
+
+// NewProgressSink creates a new ProgressSink.
+func NewProgressSink(out Output, size int64, id, action string) *Sink {
+ return &Sink{
+ out: out,
+ Size: size,
+ id: id,
+ action: action,
+ rateLimiter: rate.NewLimiter(rate.Every(100*time.Millisecond), 1),
+ }
+}
+
+func (p *Sink) Write(buf []byte) (int, error) {
+ p.Lock()
+ defer p.Unlock()
+
+ n := len(buf)
+ p.current += int64(n)
+ updateEvery := int64(1024 * 512) //512kB
+ if p.Size > 0 {
+ // Update progress for every 1% if 1% < 512kB
+ if increment := int64(0.01 * float64(p.Size)); increment < updateEvery {
+ updateEvery = increment
+ }
+ }
+ if p.current-p.lastUpdate > updateEvery || p.current == p.Size {
+ if p.current == p.Size || p.rateLimiter.Allow() {
+ p.out.WriteProgress(Progress{ID: p.id, Action: p.action, Current: p.current, Total: p.Size, LastUpdate: p.current == p.Size})
+ }
+ p.lastUpdate = p.current
+ }
+
+ return n, nil
+}