Merge pull request #17 from Vanilla-OS/feat/async

Feat/async

Author: Mirko Brombin

client.go

@@ -17,6 +17,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"time"
 
 	"github.com/containers/buildah/define"
 	"github.com/containers/buildah/imagebuildah"
@@ -65,61 +66,106 @@ func NewPrometheus(root, graphDriverName string, maxParallelDownloads uint) (*Pr
 	}, nil
 }
 
-/* PullImage pulls an image from a remote registry and stores it in the
- * Prometheus store. It returns the manifest of the pulled image and an
- * error if any. Note that the 'docker://' prefix is automatically added
- * to the imageName to make it compatible with the alltransports.ParseImageName
- * method. */
-func (p *Prometheus) PullImage(imageName string, dstName string) (*OciManifest, error) {
-	srcRef, err := alltransports.ParseImageName(fmt.Sprintf("docker://%s", imageName))
+// PullImage pulls an image from a remote registry and stores it in the
+// Prometheus store. It returns the manifest of the pulled image and an
+// error if any. Note that the 'docker://' prefix is automatically added
+// to the imageName to make it compatible with the alltransports.ParseImageName
+// method.
+func (p *Prometheus) PullImage(imageName, dstName string) (*OciManifest, error) {
+	progressCh := make(chan types.ProgressProperties)
+	manifestCh := make(chan OciManifest)
+
+	defer close(progressCh)
+	defer close(manifestCh)
+
+	err := p.pullImage(imageName, dstName, progressCh, manifestCh)
 	if err != nil {
 		return nil, err
 	}
+	for {
+		select {
+		case report := <-progressCh:
+			fmt.Printf("%s: %v/%v\n", report.Artifact.Digest.Encoded()[:12], report.Offset, report.Artifact.Size)
+		case manifest := <-manifestCh:
+			return &manifest, nil
+		}
+	}
+}
+
+// PullImageAsync does the same thing as PullImage, but returns right
+// after starting the pull process. The user can track progress in the
+// background by reading from the `progressCh` channel, which contains
+// information about the current blob and its progress. When the pull
+// process is done, the image's manifest will be sent via the `manifestCh`
+// channel, which indicates the process is done.
+//
+// NOTE: The user is responsible for closing both channels once the operation
+// completes.
+func (p *Prometheus) PullImageAsync(imageName, dstName string, progressCh chan types.ProgressProperties, manifestCh chan OciManifest) error {
+	err := p.pullImage(imageName, dstName, progressCh, manifestCh)
+	return err
+}
+
+func (p *Prometheus) pullImage(imageName, dstName string, progressCh chan types.ProgressProperties, manifestCh chan OciManifest) error {
+	srcRef, err := alltransports.ParseImageName(fmt.Sprintf("docker://%s", imageName))
+	if err != nil {
+		return err
+	}
 
 	destRef, err := storage.Transport.ParseStoreReference(p.Store, dstName)
 	if err != nil {
-		return nil, err
+		return err
 	}
 
 	systemCtx := &types.SystemContext{}
 	policy, err := signature.DefaultPolicy(systemCtx)
 	if err != nil {
-		return nil, err
+		return err
 	}
 
 	policyCtx, err := signature.NewPolicyContext(policy)
 	if err != nil {
-		return nil, err
+		return err
 	}
 
-	pulledManifestBytes, err := copy.Image(
-		context.Background(),
-		policyCtx,
-		destRef,
-		srcRef,
-		&copy.Options{
-			ReportWriter:         os.Stdout,
-			MaxParallelDownloads: p.Config.MaxParallelDownloads,
-		},
-	)
+	duration, err := time.ParseDuration("100ms")
 	if err != nil {
-		return nil, err
-	}
+		return err
+	}
+
+	go func() {
+		pulledManifestBytes, err := copy.Image(
+			context.Background(),
+			policyCtx,
+			destRef,
+			srcRef,
+			&copy.Options{
+				MaxParallelDownloads: p.Config.MaxParallelDownloads,
+				ProgressInterval:     duration,
+				Progress:             progressCh,
+			},
+		)
+		if err != nil {
+			return
+		}
 
-	var manifest OciManifest
-	err = json.Unmarshal(pulledManifestBytes, &manifest)
-	if err != nil {
-		return nil, err
-	}
+		var manifest OciManifest
+		err = json.Unmarshal(pulledManifestBytes, &manifest)
+		if err != nil {
+			return
+		}
 
-	// here we remove the 'sha256:' prefix from the digest, so we don't have
-	// to deal with it later
-	manifest.Config.Digest = manifest.Config.Digest[7:]
-	for i := range manifest.Layers {
-		manifest.Layers[i].Digest = manifest.Layers[i].Digest[7:]
-	}
+		// here we remove the 'sha256:' prefix from the digest, so we don't have
+		// to deal with it later
+		manifest.Config.Digest = manifest.Config.Digest[7:]
+		for i := range manifest.Layers {
+			manifest.Layers[i].Digest = manifest.Layers[i].Digest[7:]
+		}
+
+		manifestCh <- manifest
+	}()
 
-	return &manifest, nil
+	return nil
 }
 
 /* GetImageByDigest returns an image from the Prometheus store by its digest. */
@@ -185,8 +231,8 @@ func (p *Prometheus) BuildContainerFile(dockerfilePath string, imageName string)
 		context.Background(),
 		p.Store,
 		define.BuildOptions{
-            Output: imageName,
-        },
+			Output: imageName,
+		},
 		dockerfilePath,
 	)
 	if err != nil {

main.go

@@ -1,5 +1,5 @@
 package prometheus
 
-var version = "0.1.6"
+var version = "0.2.0"
 
 func main() {}

tests/build_containerfile_test.go

@@ -1,7 +1,7 @@
 package tests
 
 import (
-	"io/ioutil"
+	"os"
 	"testing"
 
 	"github.com/containers/storage/pkg/reexec"
@@ -22,8 +22,8 @@ func TestBuildContainerfile(t *testing.T) {
 		t.Fatal("prometheus instance is nil")
 	}
 
-	containerfile := []byte("FROM alpine:latest")
-	err = ioutil.WriteFile("Containerfile", containerfile, 0644)
+	containerfile := []byte("FROM docker.io/library/alpine:latest")
+	err = os.WriteFile("Containerfile", containerfile, 0644)
 	if err != nil {
 		t.Fatalf("error creating Containerfile: %v", err)
 	}

tests/pull_image_test.go

@@ -1,26 +1,37 @@
 package tests
 
 import (
+	"fmt"
+	"os"
 	"testing"
 
+	"github.com/containers/image/v5/types"
 	"github.com/containers/storage/pkg/reexec"
 	"github.com/vanilla-os/prometheus"
 )
 
-func TestPullImage(t *testing.T) {
+var pmt *prometheus.Prometheus
+
+func TestMain(m *testing.M) {
 	if reexec.Init() {
 		return
 	}
 
-	pmt, err := prometheus.NewPrometheus("storage", "vfs", 5)
+	var err error
+	pmt, err = prometheus.NewPrometheus("storage", "vfs", 5)
 	if err != nil {
-		t.Fatalf("error creating Prometheus instance: %v", err)
+		panic("error creating Prometheus instance: " + err.Error())
 	}
 
 	if pmt == nil {
-		t.Fatal("prometheus instance is nil")
+		panic("prometheus instance is nil")
 	}
 
+	status := m.Run()
+	os.Exit(status)
+}
+
+func TestPullImage(t *testing.T) {
 	image, err := pmt.PullImage("docker.io/library/alpine:latest", "my-alpine")
 	if err != nil {
 		t.Fatalf("error pulling image: %v", err)
@@ -60,3 +71,26 @@ func TestPullImage(t *testing.T) {
 		}
 	}
 }
+
+func TestPullImageAsync(t *testing.T) {
+	progressCh := make(chan types.ProgressProperties)
+	manifestCh := make(chan prometheus.OciManifest)
+
+	defer close(progressCh)
+	defer close(manifestCh)
+
+	err := pmt.PullImageAsync("docker.io/library/alpine:latest", "my-alpine", progressCh, manifestCh)
+	if err != nil {
+		t.Fatalf("error pulling image: %v", err)
+	}
+
+	for {
+		select {
+		case report := <-progressCh:
+			fmt.Printf("%s: %v/%v\n", report.Artifact.Digest.Encoded()[:12], report.Offset, report.Artifact.Size)
+		case manifest := <-manifestCh:
+			fmt.Printf("Got manifest: %v\n", manifest)
+			return
+		}
+	}
+}