From ace42554a5fad48ddf7dd1debf3c2728b26054d4 Mon Sep 17 00:00:00 2001 From: Billy Zha Date: Mon, 23 Oct 2023 14:27:14 +0800 Subject: [PATCH] feat: support progress output for attach, push and copy commands (#1138) Signed-off-by: Billy Zha --- cmd/oras/internal/display/print.go | 13 ++- cmd/oras/internal/display/track/target.go | 114 ++++++++++++++++++++++ cmd/oras/root/attach.go | 17 +++- cmd/oras/root/cp.go | 97 +++++++++++------- cmd/oras/root/cp_test.go | 77 +++++++++++++++ cmd/oras/root/push.go | 62 ++++++++++-- 6 files changed, 328 insertions(+), 52 deletions(-) create mode 100644 cmd/oras/internal/display/track/target.go create mode 100644 cmd/oras/root/cp_test.go diff --git a/cmd/oras/internal/display/print.go b/cmd/oras/internal/display/print.go index 48374a327..470410fc9 100644 --- a/cmd/oras/internal/display/print.go +++ b/cmd/oras/internal/display/print.go @@ -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() @@ -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) } } @@ -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 @@ -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 } } diff --git a/cmd/oras/internal/display/track/target.go b/cmd/oras/internal/display/track/target.go new file mode 100644 index 000000000..5b640062f --- /dev/null +++ b/cmd/oras/internal/display/track/target.go @@ -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 +} diff --git a/cmd/oras/root/attach.go b/cmd/oras/root/attach.go index 41a063d13..e9c1ba7c7 100644 --- a/cmd/oras/root/attach.go +++ b/cmd/oras/root/attach.go @@ -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" @@ -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], @@ -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) { @@ -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) diff --git a/cmd/oras/root/cp.go b/cmd/oras/root/cp.go index 7de38b054..4437c3e35 100644 --- a/cmd/oras/root/cp.go +++ b/cmd/oras/root/cp.go @@ -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" @@ -120,6 +121,31 @@ 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 @@ -127,33 +153,58 @@ func runCopy(ctx context.Context, opts copyOptions) error { 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 { @@ -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. diff --git a/cmd/oras/root/cp_test.go b/cmd/oras/root/cp_test.go new file mode 100644 index 000000000..53161772c --- /dev/null +++ b/cmd/oras/root/cp_test.go @@ -0,0 +1,77 @@ +/* +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 root + +import ( + "bytes" + "context" + "fmt" + "os" + "testing" + + "github.com/opencontainers/go-digest" + ocispec "github.com/opencontainers/image-spec/specs-go/v1" + "oras.land/oras-go/v2/content/memory" + "oras.land/oras/cmd/oras/internal/display/console/testutils" +) + +var ( + src *memory.Store + desc ocispec.Descriptor +) + +func TestMain(m *testing.M) { + src = memory.New() + content := []byte("test") + r := bytes.NewReader(content) + desc = ocispec.Descriptor{ + MediaType: "application/octet-stream", + Digest: digest.FromBytes(content), + Size: int64(len(content)), + } + if err := src.Push(context.Background(), desc, r); err != nil { + fmt.Println("Setup failed:", err) + os.Exit(1) + } + if err := src.Tag(context.Background(), desc, desc.Digest.String()); err != nil { + fmt.Println("Setup failed:", err) + os.Exit(1) + } + m.Run() +} + +func Test_doCopy(t *testing.T) { + // prepare + pty, slave, err := testutils.NewPty() + if err != nil { + t.Fatal(err) + } + defer slave.Close() + var opts copyOptions + opts.TTY = slave + opts.Verbose = true + opts.From.Reference = desc.Digest.String() + dst := memory.New() + // test + _, err = doCopy(context.Background(), src, dst, opts) + if err != nil { + t.Fatal(err) + } + // validate + if err = testutils.MatchPty(pty, slave, "Copied", desc.MediaType, "100.00%", desc.Digest.String()); err != nil { + t.Fatal(err) + } +} diff --git a/cmd/oras/root/push.go b/cmd/oras/root/push.go index 071efb818..f45714570 100644 --- a/cmd/oras/root/push.go +++ b/cmd/oras/root/push.go @@ -19,6 +19,7 @@ import ( "context" "errors" "fmt" + "os" "strings" "sync" @@ -30,6 +31,7 @@ import ( "oras.land/oras-go/v2/content/memory" "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/fileref" "oras.land/oras/cmd/oras/internal/option" "oras.land/oras/internal/contentutil" @@ -182,10 +184,15 @@ func runPush(ctx context.Context, opts pushOptions) error { if err != nil { return err } + var tracked track.GraphTarget + dst, tracked, err = getTrackedTarget(dst, opts.TTY) + if err != nil { + return err + } copyOptions := oras.DefaultCopyOptions copyOptions.Concurrency = opts.concurrency union := contentutil.MultiReadOnlyTarget(memoryStore, store) - updateDisplayOption(©Options.CopyGraphOptions, union, opts.Verbose) + updateDisplayOption(©Options.CopyGraphOptions, union, opts.Verbose, tracked) copy := func(root ocispec.Descriptor) error { // add both pull and push scope hints for dst repository // to save potential push-scope token requests during copy @@ -200,7 +207,7 @@ func runPush(ctx context.Context, opts pushOptions) error { } // Push - root, err := pushArtifact(dst, pack, copy) + root, err := doPush(dst, pack, copy) if err != nil { return err } @@ -224,25 +231,62 @@ func runPush(ctx context.Context, opts pushOptions) error { return opts.ExportManifest(ctx, memoryStore, root) } -func updateDisplayOption(opts *oras.CopyGraphOptions, fetcher content.Fetcher, verbose bool) { +func doPush(dst oras.Target, pack packFunc, copy copyFunc) (ocispec.Descriptor, error) { + if tracked, ok := dst.(track.GraphTarget); ok { + defer tracked.Close() + } + // Push + return pushArtifact(dst, pack, copy) +} + +func updateDisplayOption(opts *oras.CopyGraphOptions, fetcher content.Fetcher, verbose bool, tracked track.GraphTarget) { committed := &sync.Map{} - opts.PreCopy = display.StatusPrinter("Uploading", verbose) + + if tracked == nil { + // non TTY + opts.OnCopySkipped = func(ctx context.Context, desc ocispec.Descriptor) error { + committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle]) + return display.PrintStatus(desc, "Exists ", verbose) + } + opts.PreCopy = func(ctx context.Context, desc ocispec.Descriptor) error { + return display.PrintStatus(desc, "Uploading", verbose) + } + opts.PostCopy = func(ctx context.Context, desc ocispec.Descriptor) error { + committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle]) + if err := display.PrintSuccessorStatus(ctx, desc, fetcher, committed, display.StatusPrinter("Skipped ", verbose)); err != nil { + return err + } + return display.PrintStatus(desc, "Uploaded ", verbose) + } + return + } + // TTY opts.OnCopySkipped = func(ctx context.Context, desc ocispec.Descriptor) error { committed.Store(desc.Digest.String(), desc.Annotations[ocispec.AnnotationTitle]) - return display.PrintStatus(desc, "Exists ", verbose) + return tracked.Prompt(desc, "Exists ", verbose) } opts.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 ", fetcher, committed, verbose); err != nil { - return err - } - return display.PrintStatus(desc, "Uploaded ", verbose) + return display.PrintSuccessorStatus(ctx, desc, fetcher, committed, func(d ocispec.Descriptor) error { + return tracked.Prompt(d, "Skipped ", verbose) + }) } } type packFunc func() (ocispec.Descriptor, error) type copyFunc func(desc ocispec.Descriptor) error +func getTrackedTarget(gt oras.GraphTarget, tty *os.File) (oras.GraphTarget, track.GraphTarget, error) { + if tty == nil { + return gt, nil, nil + } + tracked, err := track.NewTarget(gt, "Uploading", "Uploaded ", tty) + if err != nil { + return nil, nil, err + } + return tracked, tracked, nil +} + func pushArtifact(dst oras.Target, pack packFunc, copy copyFunc) (ocispec.Descriptor, error) { root, err := pack() if err != nil {