@@ -25,8 +25,12 @@ import (
2525
2626 ocispec "github.com/opencontainers/image-spec/specs-go/v1"
2727
28+ containerd "github.com/containerd/containerd/v2/client"
29+ "github.com/containerd/containerd/v2/core/images"
2830 "github.com/containerd/containerd/v2/core/transfer"
2931 "github.com/containerd/containerd/v2/pkg/progress"
32+
33+ "github.com/containerd/nerdctl/v2/pkg/api/types"
3034)
3135
3236// From https://github.com/containerd/containerd/blob/v2.2.0-rc.0/cmd/ctr/commands/image/pull.go#L240-L473
@@ -156,6 +160,127 @@ func ProgressHandler(ctx context.Context, out io.Writer) (transfer.ProgressFunc,
156160 return progressFn , done
157161}
158162
163+ func ProgressHandlerLoadImage (ctx context.Context , client * containerd.Client , beforeSet map [string ]bool , options types.ImageLoadOptions ) (transfer.ProgressFunc , func (), * []images.Image ) {
164+ ctx , cancel := context .WithCancel (ctx )
165+ var (
166+ fw = progress .NewWriter (options .Stdout )
167+ start = time .Now ()
168+ statuses = map [string ]* progressNode {}
169+ roots = []* progressNode {}
170+ pc = make (chan transfer.Progress , 5 )
171+ status string
172+ closeC = make (chan struct {})
173+ loadedImages []images.Image
174+ imagesDisplay []string
175+ )
176+
177+ result := & loadedImages
178+ progressFn := func (p transfer.Progress ) {
179+ select {
180+ case pc <- p :
181+ case <- ctx .Done ():
182+ }
183+ }
184+
185+ done := func () {
186+ cancel ()
187+ <- closeC
188+ if ! options .Quiet {
189+ for _ , img := range imagesDisplay {
190+ fmt .Fprintf (options .Stdout , "Loaded image: %s\n " , img )
191+ }
192+ }
193+ }
194+
195+ go func () {
196+ defer close (closeC )
197+ for {
198+ select {
199+ case p := <- pc :
200+ if p .Name == "" {
201+ status = p .Event
202+ continue
203+ }
204+ if p .Event == "saved" {
205+ if img , err := client .ImageService ().Get (ctx , p .Name ); err == nil {
206+ if ! beforeSet [img .Name ] {
207+ loadedImages = append (loadedImages , img )
208+ }
209+ imagesDisplay = append (imagesDisplay , img .Name )
210+ }
211+ }
212+ if node , ok := statuses [p .Name ]; ! ok {
213+ node = & progressNode {
214+ Progress : p ,
215+ root : true ,
216+ }
217+ if len (p .Parents ) == 0 {
218+ roots = append (roots , node )
219+ } else {
220+ var parents []string
221+ for _ , parent := range p .Parents {
222+ pStatus , ok := statuses [parent ]
223+ if ok {
224+ parents = append (parents , parent )
225+ pStatus .children = append (pStatus .children , node )
226+ node .root = false
227+ }
228+ }
229+ node .Progress .Parents = parents
230+ if node .root {
231+ roots = append (roots , node )
232+ }
233+ }
234+ statuses [p .Name ] = node
235+ } else {
236+ if len (node .Progress .Parents ) != len (p .Parents ) {
237+ var parents []string
238+ var removeRoot bool
239+ for _ , parent := range p .Parents {
240+ pStatus , ok := statuses [parent ]
241+ if ok {
242+ parents = append (parents , parent )
243+ var found bool
244+ for _ , child := range pStatus .children {
245+ if child .Progress .Name == p .Name {
246+ found = true
247+ break
248+ }
249+ }
250+ if ! found {
251+ pStatus .children = append (pStatus .children , node )
252+ }
253+ if node .root {
254+ removeRoot = true
255+ }
256+ node .root = false
257+ }
258+ }
259+ p .Parents = parents
260+ // Check if needs to remove from root
261+ if removeRoot {
262+ for i := range roots {
263+ if roots [i ] == node {
264+ roots = append (roots [:i ], roots [i + 1 :]... )
265+ break
266+ }
267+ }
268+ }
269+ }
270+ node .Progress = p
271+ }
272+
273+ displayHierarchy (fw , status , roots , start )
274+ fw .Flush ()
275+
276+ case <- ctx .Done ():
277+ return
278+ }
279+ }
280+ }()
281+ return progressFn , done , result
282+ }
283+
159284func displayHierarchy (w io.Writer , status string , roots []* progressNode , start time.Time ) {
160285 total := displayNode (w , "" , roots )
161286 for _ , r := range roots {
0 commit comments