Merge pull request #588 from oasisprotocol/andrej/feature/oci-progressbar

Show progress bar when uploading to OCI repository

Author: abukosek

.golangci.yml

@@ -82,6 +82,7 @@ linters-settings:
           - github.com/wI2L/jsondiff
           - github.com/google/uuid
           - golang.org/x/net/idna
+          - golang.org/x/term
   exhaustive:
     # Switch statements are to be considered exhaustive if a 'default' case is
     # present, even if all enum members aren't listed in the switch.

build/rofl/oci.go

@@ -3,12 +3,16 @@ package rofl
 import (
 	"context"
 	"fmt"
+	"io"
 	"maps"
 	"net/http"
 	"os"
 	"path/filepath"
 	"slices"
 	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
 
 	v1 "github.com/opencontainers/image-spec/specs-go/v1"
 	oras "oras.land/oras-go/v2"
@@ -20,6 +24,8 @@ import (
 
 	"github.com/oasisprotocol/oasis-core/go/common/crypto/hash"
 	"github.com/oasisprotocol/oasis-core/go/runtime/bundle"
+
+	"github.com/oasisprotocol/cli/cmd/common/progress"
 )
 
 const (
@@ -31,6 +37,118 @@ const (
 // DefaultOCIRegistry is the default OCI registry.
 const DefaultOCIRegistry = "rofl.sh"
 
+// progressBarUpdateInterval specifies how often the progress bar should be updated.
+const progressBarUpdateInterval = 1 * time.Second
+
+// TargetWithProgress wraps oras.Target and provides updates via a progress bar.
+type TargetWithProgress struct {
+	oras.Target
+
+	Message        string
+	UpdateInterval time.Duration
+	BytesRead      *atomic.Uint64
+	BytesTotal     uint64
+
+	stopUpdate chan struct{}
+	wg         sync.WaitGroup
+}
+
+// NewTargetWithProgress creates a new TargetWithProgress.
+// bytesTotal is the size to use for the 100% value (use 0 if unknown).
+// msg is the message to display in front of the progress bar.
+func NewTargetWithProgress(target oras.Target, bytesTotal uint64, msg string) *TargetWithProgress {
+	return &TargetWithProgress{
+		Target:         target,
+		Message:        msg,
+		UpdateInterval: progressBarUpdateInterval,
+		BytesRead:      &atomic.Uint64{},
+		BytesTotal:     bytesTotal,
+		stopUpdate:     make(chan struct{}),
+	}
+}
+
+// Push wraps the oras.Target Push method with progress bar updates.
+func (t *TargetWithProgress) Push(ctx context.Context, desc v1.Descriptor, content io.Reader) error {
+	var progReader io.Reader
+
+	// Wrap the reader with our variant.
+	pReader := &progressReader{
+		reader:    content,
+		bytesRead: t.BytesRead,
+	}
+	progReader = pReader
+
+	// If the reader also has a WriteTo method, wrap it appropriately.
+	if wt, ok := content.(io.WriterTo); ok {
+		progReader = &progressWriterToReader{
+			progressReader: pReader,
+			writerTo:       wt,
+		}
+	}
+
+	// Do the actual push using our wrappers.
+	return t.Target.Push(ctx, desc, progReader)
+}
+
+// StartProgress starts updating the progress bar.
+func (t *TargetWithProgress) StartProgress() {
+	t.wg.Add(1)
+	go func() {
+		defer t.wg.Done()
+
+		ticker := time.NewTicker(t.UpdateInterval)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-ticker.C:
+				progress.PrintProgressBar(os.Stderr, t.Message, t.BytesRead.Load(), t.BytesTotal, false)
+			case <-t.stopUpdate:
+				return
+			}
+		}
+	}()
+}
+
+// StopProgress stops the progress bar updates.
+func (t *TargetWithProgress) StopProgress() {
+	if t.stopUpdate != nil {
+		// Print the final stage of the progress bar.
+		progress.PrintProgressBar(os.Stderr, t.Message, t.BytesRead.Load(), t.BytesTotal, true)
+
+		close(t.stopUpdate)
+		t.stopUpdate = nil
+	}
+	t.wg.Wait()
+}
+
+type progressReader struct {
+	reader    io.Reader
+	bytesRead *atomic.Uint64
+}
+
+func (pr *progressReader) Read(b []byte) (int, error) {
+	n, err := pr.reader.Read(b)
+	if n > 0 {
+		pr.bytesRead.Add(uint64(n))
+	}
+	return n, err
+}
+
+type progressWriterToReader struct {
+	*progressReader
+
+	writerTo io.WriterTo
+}
+
+func (pwr *progressWriterToReader) WriteTo(w io.Writer) (int64, error) {
+	written, err := pwr.writerTo.WriteTo(w)
+	if written > 0 {
+		pwr.bytesRead.Add(uint64(written))
+	}
+	return written, err
+}
+
 // PushBundleToOciRepository pushes an ORC bundle to the given remote OCI repository.
 //
 // Returns the OCI manifest digest and the ORC manifest hash.
@@ -70,12 +188,26 @@ func PushBundleToOciRepository(bundleFn, dst string) (string, hash.Hash, error)
 		return "", hash.Hash{}, fmt.Errorf("failed to explode bundle: %w", err)
 	}
 
+	// Keep track of the total size of the bundle as we add files to it.
+	var totalSize int64
+
+	getFileSize := func(path string) int64 {
+		f, err := os.Stat(path)
+		if err != nil {
+			// This shouldn't happen, because store.Add should fail if the file doesn't exist.
+			panic(err)
+		}
+		return f.Size()
+	}
+
 	// Generate the config object from the manifest.
 	const manifestName = "META-INF/MANIFEST.MF"
-	configDsc, err := store.Add(ctx, manifestName, ociTypeOrcConfig, filepath.Join(bundleDir, manifestName))
+	manifestPath := filepath.Join(bundleDir, manifestName)
+	configDsc, err := store.Add(ctx, manifestName, ociTypeOrcConfig, manifestPath)
 	if err != nil {
 		return "", hash.Hash{}, fmt.Errorf("failed to add config object from manifest: %w", err)
 	}
+	totalSize += getFileSize(manifestPath)
 
 	// Add other files as layers.
 	layers := make([]v1.Descriptor, 0, len(bnd.Data)-1)
@@ -86,10 +218,12 @@ func PushBundleToOciRepository(bundleFn, dst string) (string, hash.Hash, error)
 		}
 
 		var layerDsc v1.Descriptor
-		layerDsc, err = store.Add(ctx, fn, ociTypeOrcLayer, filepath.Join(bundleDir, fn))
+		filePath := filepath.Join(bundleDir, fn)
+		layerDsc, err = store.Add(ctx, fn, ociTypeOrcLayer, filePath)
 		if err != nil {
 			return "", hash.Hash{}, fmt.Errorf("failed to add OCI layer: %w", err)
 		}
+		totalSize += getFileSize(filePath)
 
 		layers = append(layers, layerDsc)
 	}
@@ -134,10 +268,20 @@ func PushBundleToOciRepository(bundleFn, dst string) (string, hash.Hash, error)
 	}
 	repo.Client = client
 
+	repoSize := uint64(totalSize) //nolint:gosec
+	repoWithProgress := NewTargetWithProgress(repo, repoSize, "Pushing...")
+	repoWithProgress.StartProgress()
+	defer repoWithProgress.StopProgress()
+
 	// Push to remote repository.
-	if _, err = oras.Copy(ctx, store, tag, repo, tag, oras.DefaultCopyOptions); err != nil {
+	if _, err = oras.Copy(ctx, store, tag, repoWithProgress, tag, oras.DefaultCopyOptions); err != nil {
 		return "", hash.Hash{}, fmt.Errorf("failed to push to remote OCI repository: %w", err)
 	}
 
+	// Force progress to 100% in case the ORC already exists on the remote.
+	// This is necessary so we get a full progressbar instead of an empty one
+	// when we're done.
+	repoWithProgress.BytesRead.Store(repoWithProgress.BytesTotal)
+
 	return manifestDescriptor.Digest.String(), bnd.Manifest.Hash(), nil
 }