feat: implement pull-state for properly keeping track of layer progress and pull events

Author: silenium-dev

examples/pull/main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"fmt"
 	"github.com/distribution/reference"
 	client2 "github.com/docker/docker/client"
 	"github.com/silenium-dev/docker-wrapper/pkg/client"
@@ -17,21 +18,30 @@ func main() {
 	cli, err := client.NewWithOpts(
 		client.WithAuthProvider(authProvider),
 		client.FromEnv,
-		client.WithDockerOpts(client2.WithTimeout(time.Second*10)),
+		client.WithDockerOpts(client2.WithTimeout(time.Hour*1)),
 	)
 	if err != nil {
 		panic(err)
 	}
 	//ref, err := reference.ParseDockerRef("quay.io/prometheus/node-exporter@sha256:a25fbdaa3e4d03e0d735fd03f231e9a48332ecf40ca209b2f103b1f970d1cde0")
-	ref, err := reference.ParseDockerRef("confluentinc/cp-kafka:latest")
+	ref, err := reference.ParseDockerRef("localstack/localstack:latest")
 	if err != nil {
 		panic(err)
 	}
-	eventChan, err := cli.Pull(context.Background(), ref)
+	stateChan, err := cli.PullWithState(context.Background(), ref)
 	if err != nil {
 		panic(err)
 	}
-	for event := range eventChan {
-		println(event.String())
+	for state := range stateChan {
+		print("\033[2J")
+		fmt.Printf("%s\n", state.Status())
+		for idx, l := range state.Layers() {
+			fmt.Printf("%02d [%s]: %s\n", idx, l.Id(), l.Status())
+		}
 	}
+	digest, err := cli.Pull(context.Background(), ref)
+	if err != nil {
+		panic(err)
+	}
+	fmt.Printf("Digest: %s\n", digest)
 }

pkg/client/pull.go

@@ -2,14 +2,25 @@ package client
 
 import (
 	"context"
+	"fmt"
 	"github.com/distribution/reference"
 	"github.com/docker/docker/api/types"
 	"github.com/docker/docker/api/types/registry"
+	"github.com/opencontainers/go-digest"
 	"github.com/silenium-dev/docker-wrapper/pkg/client/pull"
 	"github.com/silenium-dev/docker-wrapper/pkg/client/pull/events"
+	"github.com/silenium-dev/docker-wrapper/pkg/client/pull/state"
 )
 
-func (c *Client) Pull(ctx context.Context, ref reference.Named) (chan events.PullEvent, error) {
+func (c *Client) PullWithState(ctx context.Context, ref reference.Named) (chan state.Pull, error) {
+	eventChan, err := c.PullWithEvents(ctx, ref)
+	if err != nil {
+		return nil, err
+	}
+	return pull.StateFromStream(ctx, ref, eventChan), nil
+}
+
+func (c *Client) PullWithEvents(ctx context.Context, ref reference.Named) (chan events.PullEvent, error) {
 	var encodedAuth string
 	var err error
 	if c.authProvider != nil {
@@ -24,3 +35,22 @@ func (c *Client) Pull(ctx context.Context, ref reference.Named) (chan events.Pul
 	}
 	return pull.ParseStream(ctx, reader), nil
 }
+
+func (c *Client) Pull(ctx context.Context, ref reference.Named) (digest.Digest, error) {
+	eventChan, err := c.PullWithEvents(ctx, ref)
+	if err != nil {
+		return "", err
+	}
+
+	var digestEvent *events.Digest
+	for event := range eventChan {
+		if _, ok := event.(*events.Digest); digestEvent == nil && ok {
+			digestEvent = event.(*events.Digest)
+		}
+	}
+	if digestEvent == nil {
+		return "", fmt.Errorf("no digest event received")
+	}
+
+	return digestEvent.Digest, nil
+}

pkg/client/pull/events/error.go

@@ -0,0 +1,14 @@
+package events
+
+type LayerError struct {
+	LayerBase
+	PullError
+}
+
+type PullError struct {
+	Error string
+}
+
+func (e *PullError) String() string {
+	return e.Error
+}

pkg/client/pull/events/event.go

@@ -13,6 +13,7 @@ type PullEvent interface {
 }
 
 type LayerEvent interface {
+	PullEvent
 	LayerId() string
 }
 
@@ -47,6 +48,7 @@ func Parse(event base.PullProgressEvent) (PullEvent, error) {
 		Hide:    event.ProgressDetail.HideCounts,
 	}
 	progressBase := ProgressBase{layer, progress}
+	errorEvent := PullError{Error: event.Error}
 
 	switch event.Status {
 	case AlreadyExistsStatus:
@@ -76,6 +78,11 @@ func Parse(event base.PullProgressEvent) (PullEvent, error) {
 	case PullCompleteStatus:
 		return &PullComplete{layer}, nil
 	default:
+		if event.Error != "" && event.ID != "" {
+			return &LayerError{layer, errorEvent}, nil
+		} else if event.Error != "" {
+			return &errorEvent, nil
+		}
 		if strings.HasPrefix(event.Status, "Pulling from") {
 			return &PullStarted{}, nil
 		}

pkg/client/pull/events/final.go

@@ -1,6 +1,8 @@
 package events
 
-import "fmt"
+import (
+	"fmt"
+)
 
 type DownloadedNewerImage struct {
 	Final

pkg/client/pull/state.go

@@ -0,0 +1,43 @@
+package pull
+
+import (
+	"context"
+	"github.com/distribution/reference"
+	"github.com/silenium-dev/docker-wrapper/pkg/client/pull/events"
+	"github.com/silenium-dev/docker-wrapper/pkg/client/pull/state"
+)
+
+func StateFromStream(ctx context.Context, ref reference.Named, ch chan events.PullEvent) chan state.Pull {
+	out := make(chan state.Pull)
+
+	go processEvents(ctx, ref, ch, out)
+
+	return out
+}
+
+func processEvents(ctx context.Context, ref reference.Named, ch chan events.PullEvent, out chan state.Pull) {
+	defer close(out)
+	var current state.Pull
+	var err error
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case event, ok := <-ch:
+			if !ok {
+				return
+			}
+			var next state.Pull
+			if current == nil {
+				next, err = state.NewPullState(ref, event)
+			} else {
+				next, err = current.Next(event)
+			}
+			if err != nil {
+				panic(err)
+			}
+			current = next
+			out <- current
+		}
+	}
+}