Skip to content

Commit

Permalink
feat: support progress output for attach, push and copy commands (#1138)
Browse files Browse the repository at this point in the history
Signed-off-by: Billy Zha <jinzha1@microsoft.com>
  • Loading branch information
qweeah committed Oct 23, 2023
1 parent 58811f2 commit ace4255
Show file tree
Hide file tree
Showing 6 changed files with 328 additions and 52 deletions.
13 changes: 8 additions & 5 deletions cmd/oras/internal/display/print.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,10 @@ import (

var printLock sync.Mutex

// Print objects to display concurrent-safely
// PrintFunc is the function type returned by StatusPrinter.
type PrintFunc func(ocispec.Descriptor) error

// Print objects to display concurrent-safely.
func Print(a ...any) error {
printLock.Lock()
defer printLock.Unlock()
Expand All @@ -38,8 +41,8 @@ func Print(a ...any) error {
}

// StatusPrinter returns a tracking function for transfer status.
func StatusPrinter(status string, verbose bool) func(context.Context, ocispec.Descriptor) error {
return func(ctx context.Context, desc ocispec.Descriptor) error {
func StatusPrinter(status string, verbose bool) PrintFunc {
return func(desc ocispec.Descriptor) error {
return PrintStatus(desc, status, verbose)
}
}
Expand All @@ -58,7 +61,7 @@ func PrintStatus(desc ocispec.Descriptor, status string, verbose bool) error {
}

// PrintSuccessorStatus prints transfer status of successors.
func PrintSuccessorStatus(ctx context.Context, desc ocispec.Descriptor, status string, fetcher content.Fetcher, committed *sync.Map, verbose bool) error {
func PrintSuccessorStatus(ctx context.Context, desc ocispec.Descriptor, fetcher content.Fetcher, committed *sync.Map, print PrintFunc) error {
successors, err := content.Successors(ctx, fetcher, desc)
if err != nil {
return err
Expand All @@ -67,7 +70,7 @@ func PrintSuccessorStatus(ctx context.Context, desc ocispec.Descriptor, status s
name := s.Annotations[ocispec.AnnotationTitle]
if v, ok := committed.Load(s.Digest.String()); ok && v != name {
// Reprint status for deduplicated content
if err := PrintStatus(s, status, verbose); err != nil {
if err := print(s); err != nil {
return err
}
}
Expand Down
114 changes: 114 additions & 0 deletions cmd/oras/internal/display/track/target.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
/*
Copyright The ORAS Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package track

import (
"context"
"io"
"os"

ocispec "github.com/opencontainers/image-spec/specs-go/v1"
"oras.land/oras-go/v2"
"oras.land/oras-go/v2/registry"
"oras.land/oras/cmd/oras/internal/display/progress"
)

// GraphTarget is a tracked oras.GraphTarget.
type GraphTarget interface {
oras.GraphTarget
io.Closer
Prompt(desc ocispec.Descriptor, prompt string, verbose bool) error
}

type graphTarget struct {
oras.GraphTarget
manager progress.Manager
actionPrompt string
donePrompt string
}

type referenceGraphTarget struct {
*graphTarget
}

// NewTarget creates a new tracked Target.
func NewTarget(t oras.GraphTarget, actionPrompt, donePrompt string, tty *os.File) (GraphTarget, error) {
manager, err := progress.NewManager(tty)
if err != nil {
return nil, err
}
gt := &graphTarget{
GraphTarget: t,
manager: manager,
actionPrompt: actionPrompt,
donePrompt: donePrompt,
}

if _, ok := t.(registry.ReferencePusher); ok {
return &referenceGraphTarget{
graphTarget: gt,
}, nil
}
return gt, nil
}

// Push pushes the content to the base oras.GraphTarget with tracking.
func (t *graphTarget) Push(ctx context.Context, expected ocispec.Descriptor, content io.Reader) error {
r, err := managedReader(content, expected, t.manager, t.actionPrompt, t.donePrompt)
if err != nil {
return err
}
defer r.Close()
r.Start()
if err := t.GraphTarget.Push(ctx, expected, r); err != nil {
return err
}
r.Done()
return nil
}

// PushReference pushes the content to the base oras.GraphTarget with tracking.
func (rgt *referenceGraphTarget) PushReference(ctx context.Context, expected ocispec.Descriptor, content io.Reader, reference string) error {
r, err := managedReader(content, expected, rgt.manager, rgt.actionPrompt, rgt.donePrompt)
if err != nil {
return err
}
defer r.Close()
r.Start()
err = rgt.GraphTarget.(registry.ReferencePusher).PushReference(ctx, expected, r, reference)
if err != nil {
return err
}
r.Done()
return nil
}

// Close closes the tracking manager.
func (t *graphTarget) Close() error {
return t.manager.Close()
}

// Prompt prompts the user with the provided prompt and descriptor.
func (t *graphTarget) Prompt(desc ocispec.Descriptor, prompt string, verbose bool) error {
status, err := t.manager.Add()
if err != nil {
return err
}
defer close(status)
status <- progress.NewStatus(prompt, desc, desc.Size)
status <- progress.EndTiming()
return nil
}
17 changes: 12 additions & 5 deletions cmd/oras/root/attach.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ import (
"oras.land/oras-go/v2/content"
"oras.land/oras-go/v2/content/file"
"oras.land/oras-go/v2/registry/remote/auth"
"oras.land/oras/cmd/oras/internal/display/track"
"oras.land/oras/cmd/oras/internal/option"
"oras.land/oras/internal/graph"
"oras.land/oras/internal/registryutil"
Expand Down Expand Up @@ -131,6 +132,15 @@ func runAttach(ctx context.Context, opts attachOptions) error {
}

// prepare push
var tracked track.GraphTarget
dst, tracked, err = getTrackedTarget(dst, opts.TTY)
if err != nil {
return err
}
graphCopyOptions := oras.DefaultCopyGraphOptions
graphCopyOptions.Concurrency = opts.concurrency
updateDisplayOption(&graphCopyOptions, store, opts.Verbose, tracked)

packOpts := oras.PackManifestOptions{
Subject: &subject,
ManifestAnnotations: annotations[option.AnnotationManifest],
Expand All @@ -140,9 +150,6 @@ func runAttach(ctx context.Context, opts attachOptions) error {
return oras.PackManifest(ctx, store, oras.PackManifestVersion1_1_RC4, opts.artifactType, packOpts)
}

graphCopyOptions := oras.DefaultCopyGraphOptions
graphCopyOptions.Concurrency = opts.concurrency
updateDisplayOption(&graphCopyOptions, store, opts.Verbose)
copy := func(root ocispec.Descriptor) error {
graphCopyOptions.FindSuccessors = func(ctx context.Context, fetcher content.Fetcher, node ocispec.Descriptor) ([]ocispec.Descriptor, error) {
if content.Equal(node, root) {
Expand All @@ -161,11 +168,11 @@ func runAttach(ctx context.Context, opts attachOptions) error {
return oras.CopyGraph(ctx, store, dst, root, graphCopyOptions)
}

root, err := pushArtifact(dst, pack, copy)
// Attach
root, err := doPush(dst, pack, copy)
if err != nil {
return err
}

digest := subject.Digest.String()
if !strings.HasSuffix(opts.RawReference, digest) {
opts.RawReference = fmt.Sprintf("%s@%s", opts.Path, subject.Digest)
Expand Down
97 changes: 64 additions & 33 deletions cmd/oras/root/cp.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ import (
"oras.land/oras-go/v2/content"
"oras.land/oras-go/v2/registry/remote/auth"
"oras.land/oras/cmd/oras/internal/display"
"oras.land/oras/cmd/oras/internal/display/track"
"oras.land/oras/cmd/oras/internal/option"
"oras.land/oras/internal/docker"
"oras.land/oras/internal/graph"
Expand Down Expand Up @@ -120,40 +121,90 @@ func runCopy(ctx context.Context, opts copyOptions) error {
}
ctx = registryutil.WithScopeHint(ctx, dst, auth.ActionPull, auth.ActionPush)

desc, err := doCopy(ctx, src, dst, opts)
if err != nil {
return err
}

if from, err := digest.Parse(opts.From.Reference); err == nil && from != desc.Digest {
// correct source digest
opts.From.RawReference = fmt.Sprintf("%s@%s", opts.From.Path, desc.Digest.String())
}
fmt.Println("Copied", opts.From.AnnotatedReference(), "=>", opts.To.AnnotatedReference())

if len(opts.extraRefs) != 0 {
tagNOpts := oras.DefaultTagNOptions
tagNOpts.Concurrency = opts.concurrency
if _, err = oras.TagN(ctx, display.NewTagStatusPrinter(dst), opts.To.Reference, opts.extraRefs, tagNOpts); err != nil {
return err
}
}

fmt.Println("Digest:", desc.Digest)

return nil
}

func doCopy(ctx context.Context, src oras.ReadOnlyGraphTarget, dst oras.GraphTarget, opts copyOptions) (ocispec.Descriptor, error) {
// Prepare copy options
committed := &sync.Map{}
extendedCopyOptions := oras.DefaultExtendedCopyOptions
extendedCopyOptions.Concurrency = opts.concurrency
extendedCopyOptions.FindPredecessors = func(ctx context.Context, src content.ReadOnlyGraphStorage, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
return graph.Referrers(ctx, src, desc, "")
}
extendedCopyOptions.PreCopy = display.StatusPrinter("Copying", opts.Verbose)
extendedCopyOptions.PostCopy = func(ctx context.Context, desc ocispec.Descriptor) error {
committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle])
if err := display.PrintSuccessorStatus(ctx, desc, "Skipped", dst, committed, opts.Verbose); err != nil {
return err

if opts.TTY == nil {
// none TTY output
extendedCopyOptions.OnCopySkipped = func(ctx context.Context, desc ocispec.Descriptor) error {
committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle])
return display.PrintStatus(desc, "Exists ", opts.Verbose)
}
extendedCopyOptions.PreCopy = func(ctx context.Context, desc ocispec.Descriptor) error {
return display.PrintStatus(desc, "Copying", opts.Verbose)
}
extendedCopyOptions.PostCopy = func(ctx context.Context, desc ocispec.Descriptor) error {
committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle])
if err := display.PrintSuccessorStatus(ctx, desc, dst, committed, display.StatusPrinter("Skipped", opts.Verbose)); err != nil {
return err
}
return display.PrintStatus(desc, "Copied ", opts.Verbose)
}
} else {
// TTY output
tracked, err := track.NewTarget(dst, "Copying ", "Copied ", opts.TTY)
if err != nil {
return ocispec.Descriptor{}, err
}
defer tracked.Close()
dst = tracked
extendedCopyOptions.OnCopySkipped = func(ctx context.Context, desc ocispec.Descriptor) error {
committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle])
return tracked.Prompt(desc, "Exists ", opts.Verbose)
}
extendedCopyOptions.PostCopy = func(ctx context.Context, desc ocispec.Descriptor) error {
committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle])
return display.PrintSuccessorStatus(ctx, desc, tracked, committed, func(desc ocispec.Descriptor) error {
return tracked.Prompt(desc, "Skipped", opts.Verbose)
})
}
return display.PrintStatus(desc, "Copied ", opts.Verbose)
}
extendedCopyOptions.OnCopySkipped = func(ctx context.Context, desc ocispec.Descriptor) error {
committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle])
return display.PrintStatus(desc, "Exists ", opts.Verbose)
}

var desc ocispec.Descriptor
var err error
rOpts := oras.DefaultResolveOptions
rOpts.TargetPlatform = opts.Platform.Platform
if opts.recursive {
desc, err = oras.Resolve(ctx, src, opts.From.Reference, rOpts)
if err != nil {
return fmt.Errorf("failed to resolve %s: %w", opts.From.Reference, err)
return ocispec.Descriptor{}, fmt.Errorf("failed to resolve %s: %w", opts.From.Reference, err)
}
err = recursiveCopy(ctx, src, dst, opts.To.Reference, desc, extendedCopyOptions)
} else {
if opts.To.Reference == "" {
desc, err = oras.Resolve(ctx, src, opts.From.Reference, rOpts)
if err != nil {
return fmt.Errorf("failed to resolve %s: %w", opts.From.Reference, err)
return ocispec.Descriptor{}, fmt.Errorf("failed to resolve %s: %w", opts.From.Reference, err)
}
err = oras.CopyGraph(ctx, src, dst, desc, extendedCopyOptions.CopyGraphOptions)
} else {
Expand All @@ -166,27 +217,7 @@ func runCopy(ctx context.Context, opts copyOptions) error {
desc, err = oras.Copy(ctx, src, opts.From.Reference, dst, opts.To.Reference, copyOptions)
}
}
if err != nil {
return err
}

if from, err := digest.Parse(opts.From.Reference); err == nil && from != desc.Digest {
// correct source digest
opts.From.RawReference = fmt.Sprintf("%s@%s", opts.From.Path, desc.Digest.String())
}
fmt.Println("Copied", opts.From.AnnotatedReference(), "=>", opts.To.AnnotatedReference())

if len(opts.extraRefs) != 0 {
tagNOpts := oras.DefaultTagNOptions
tagNOpts.Concurrency = opts.concurrency
if _, err = oras.TagN(ctx, display.NewTagStatusPrinter(dst), opts.To.Reference, opts.extraRefs, tagNOpts); err != nil {
return err
}
}

fmt.Println("Digest:", desc.Digest)

return nil
return desc, err
}

// recursiveCopy copies an artifact and its referrers from one target to another.
Expand Down
Loading

0 comments on commit ace4255

Please sign in to comment.