fix: k3s && trace collector

client/client.go

@@ -337,7 +337,7 @@ func (c *Client) updateManifest(ctr *containerd.Client, d string, chainIds []dig
 	if err != nil {
 		return errors.Wrapf(err, "failed to mark manifest as completed")
 	}
-	log.G(c.ctx).WithField("digest", info.Digest).Info("download completed")
+	log.G(c.ctx).WithField("m", info.Digest).Info("download completed")
 	return nil
 }
 
@@ -511,12 +511,6 @@ func (c *Client) PullImage(
 ) {
 	// init vars
 	is := ctr.ImageService()
-	closedReady := false
-	defer func() {
-		if !closedReady {
-			close(*ready)
-		}
-	}()
 	localCtx := context.Background()
 
 	// check local image
@@ -530,9 +524,14 @@ func (c *Client) PullImage(
 	if img != nil {
 		labels := img.Labels()
 		if _, has := labels[util.ContentLabelCompletion]; has {
+			if _, err = c.LoadImage(ctr, img.Target().Digest); err != nil {
+				*ready <- PullFinishedMessage{nil, "", errors.Wrapf(err, "failed to load image %s", ref)}
+				return
+			}
 			meta := img.Metadata()
 			*ready <- PullFinishedMessage{&meta, "", fmt.Errorf("requested image %s already exists", ref)}
 			return
+
 		}
 		log.G(c.ctx).
 			WithField("image", ref).
@@ -601,7 +600,6 @@ func (c *Client) PullImage(
 	}
 
 	// check if the image is in containerd's image pool
-
 	var (
 		buf         *bytes.Buffer
 		man, con    []byte
@@ -739,7 +737,6 @@ func (c *Client) PullImage(
 	// Image is ready (content is still on the way)
 	// close(*ready)
 	*ready <- PullFinishedMessage{&ctrImg, baseRef, nil}
-	closedReady = true
 
 	// 6. Extract file content
 	// download content

client/fs/tracer.go

@@ -151,8 +151,8 @@ func NewTracer(ctx context.Context, optimizeGroup, digest, outputDir string) (*T
 // OptimizedTraceItem with ranking
 type OptimizedTraceItem struct {
 	TraceItem
-	Rank        int `json:"r"`
-	SourceImage int `json:"s"`
+	Rank        int `json:"rank"`
+	SourceImage int `json:"img"`
 }
 
 func (oti OptimizedTraceItem) Key() string {
@@ -304,14 +304,16 @@ func NewTraceCollection(ctx context.Context, p string) (*TraceCollection, error)
 				idx := lookupMap[trace.Image]
 
 				for _, t := range trace.Seq {
+					// remove duplicated items
 					if !visited[t.FileName] {
 						g.History = append(g.History, &OptimizedTraceItem{
 							TraceItem: TraceItem{
 								FileName: t.FileName,
+								Stack:    t.Stack,
 								Access:   t.Access + traceOffset,
 								Wait:     t.Wait,
 							},
-							Rank:        0,
+							Rank:        0, // will be updated later
 							SourceImage: idx,
 						})
 						visited[t.FileName] = true

client/fs/tracer_test.go

@@ -41,3 +41,17 @@ func TestLoadTraces(t *testing.T) {
 	}
 	_ = ioutil.WriteFile(path.Join(os.TempDir(), "group-optimize.json"), buf, 0644)
 }
+
+func TestLoadTraces2(t *testing.T) {
+	ctx := util.ConfigLogger()
+	tc, err := NewTraceCollection(ctx, "/root/traces2")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	buf, err := json.MarshalIndent(tc, "", "\t")
+	if err != nil {
+		t.Fatal(err)
+	}
+	_ = ioutil.WriteFile(path.Join(os.TempDir(), "group-optimize.json"), buf, 0644)
+}

client/snapshotter/operator_test.go

@@ -5,9 +5,100 @@
 
 package snapshotter
 
-import "testing"
+import (
+	"context"
+	"fmt"
+	"github.com/containerd/containerd"
+	"github.com/containerd/containerd/snapshots"
+	"github.com/containerd/containerd/snapshots/storage"
+	"github.com/mc256/starlight/util"
+	"testing"
+)
 
-func TestNewOperator(t *testing.T) {
-	//c.client.SnapshotService("starlight")
+func TestSnapshotList(t *testing.T) {
+	socket := "/run/k3s/containerd/containerd.sock"
+	ns := "k8s.io"
 
-}
+	client, err := containerd.New(socket, containerd.WithDefaultNamespace(ns))
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	defer client.Close()
+
+	ctx := context.Background()
+
+	sns := client.SnapshotService("starlight")
+
+	sns.Walk(ctx, func(ctx context.Context, info snapshots.Info) error {
+		fmt.Println(info.Name, info.Parent)
+		fmt.Println("-->", info.Labels[util.SnapshotLabelRefImage])
+		ssid, inf, p, err := storage.GetInfo(ctx, info.Name)
+		if err != nil {
+			fmt.Println(ssid, inf, p)
+		} else {
+			fmt.Println(err)
+		}
+		return nil
+	})
+
+	/*
+				36c584644174b103238ecd17cf43487b38485bf321f64c8cf525fafc5580ab86 sha256:1021ef88c7974bfff89c5a0ec4fd3160daac6c48a075f74cff721f85dd104e68
+				-->
+				40b6e9770e3ec864efd193670e72640ed5a9364f2d1287234fa4b8d7731a9cba sha256:68ec11adbdc694363fbd1026d772664049569ee10fc3109b0fc44cc890b60c2c
+				-->
+				47aac9e355fe94ed399ea2dcbbac7e4c912eb936eb80941d8bf08537f2956c5e sha256:1021ef88c7974bfff89c5a0ec4fd3160daac6c48a075f74cff721f85dd104e68
+				-->
+				9823e5be2c32d3b2547b9a5344d49628e38189a4beafbb6c7dee3858545b98f1 sha256:ef04ecfb1d007266224a5b5b67992c74afdf12e984433a6978d061c7b852ee10
+				-->
+				c7f1c4db6e6fdb52efd09abeb5e51e7af6dc6fc753e32c34e4f9f04544dcd6cc sha256:ef04ecfb1d007266224a5b5b67992c74afdf12e984433a6978d061c7b852ee10
+				-->
+				sha256:1021ef88c7974bfff89c5a0ec4fd3160daac6c48a075f74cff721f85dd104e68
+				-->
+				sha256:1ad27bdd166b922492031b1938a4fb2f775e3d98c8f1b72051dad0570a4dd1b5
+				-->
+				sha256:40cf597a9181e86497f4121c604f9f0ab208950a98ca21db883f26b0a548a2eb
+				-->
+				sha256:4ac94bd63114d70c68e73d408d559bbd681b7c5094715b1a38bd7b361312555f sha256:1ad27bdd166b922492031b1938a4fb2f775e3d98c8f1b72051dad0570a4dd1b5
+				-->
+				sha256:68ec11adbdc694363fbd1026d772664049569ee10fc3109b0fc44cc890b60c2c sha256:7ad6b790083aeda0ad3e5e2888d9cf1a64d6c5ff6b41487a3404bc6f5684a2d3
+				--> sha256:c20f5ef500bc93f26fec480dfea1d4cbda2d39791d611031a3aaf3c096f35c73
+				sha256:7ad6b790083aeda0ad3e5e2888d9cf1a64d6c5ff6b41487a3404bc6f5684a2d3 sha256:b11040115c1a7bad75c256bfad158a1bf2aec69418ffc1f49326ba59e33f69b2
+				--> sha256:c20f5ef500bc93f26fec480dfea1d4cbda2d39791d611031a3aaf3c096f35c73
+				sha256:a5878da390fcf0e6324ce5593f03ae9609caaa9bce0a522b504c00b91e579b46
+				--> sha256:c20f5ef500bc93f26fec480dfea1d4cbda2d39791d611031a3aaf3c096f35c73
+				sha256:b11040115c1a7bad75c256bfad158a1bf2aec69418ffc1f49326ba59e33f69b2 sha256:d44df6880f5a4e309d25a92dc413458a1f964493868c99000519dbc365abe1d7
+				--> sha256:c20f5ef500bc93f26fec480dfea1d4cbda2d39791d611031a3aaf3c096f35c73
+				sha256:d44df6880f5a4e309d25a92dc413458a1f964493868c99000519dbc365abe1d7 sha256:f8521e5ffbea025f78787bf0d407a4295dff5fcc3cc5c8936126fd17fa1819a8
+				--> sha256:c20f5ef500bc93f26fec480dfea1d4cbda2d39791d611031a3aaf3c096f35c73
+				sha256:ef04ecfb1d007266224a5b5b67992c74afdf12e984433a6978d061c7b852ee10 sha256:4ac94bd63114d70c68e73d408d559bbd681b7c5094715b1a38bd7b361312555f
+				-->
+				sha256:f8521e5ffbea025f78787bf0d407a4295dff5fcc3cc5c8936126fd17fa1819a8 sha256:a5878da390fcf0e6324ce5593f03ae9609caaa9bce0a522b504c00b91e579b46
+				--> sha256:c20f5ef500bc93f26fec480dfea1d4cbda2d39791d611031a3aaf3c096f35c73
+
+
+		sha256:c20f5ef500bc93f26fec480dfea1d4cbda2d39791d611031a3aaf3c096f35c73
+		manifest
+
+		complete.starlight.mc256.dev=2022-11-27T18:34:46Z,
+		puller.containerd.io=starlight,
+		mediaType.starlight.mc256.dev=manifest,
+		containerd.io/gc.ref.snapshot.starlight/0=sha256:a5878da390fcf0e6324ce5593f03ae9609caaa9bce0a522b504c00b91e579b46,
+		containerd.io/gc.ref.content.starlight=sha256:5c409db91624a306224d9a98efcfded2ae94da0a483cf2f1d40b8cb1ea776385,
+		starlight.mc256.dev/distribution.source.in-cluster=starlight-registry.default.svc.cluster.local:5000/starlight/redis:6.2.1,
+		containerd.io/gc.ref.content.config=sha256:982983f3c76aeb79bddb66239c878f011c147b15cac9d7c71abedf76e772c7c6,
+		containerd.io/gc.ref.snapshot.starlight/5=sha256:68ec11adbdc694363fbd1026d772664049569ee10fc3109b0fc44cc890b60c2c,
+		containerd.io/gc.ref.snapshot.starlight/4=sha256:7ad6b790083aeda0ad3e5e2888d9cf1a64d6c5ff6b41487a3404bc6f5684a2d3,
+		containerd.io/gc.ref.snapshot.starlight/2=sha256:d44df6880f5a4e309d25a92dc413458a1f964493868c99000519dbc365abe1d7,
+		containerd.io/gc.ref.snapshot.starlight/1=sha256:f8521e5ffbea025f78787bf0d407a4295dff5fcc3cc5c8936126fd17fa1819a8,
+		containerd.io/gc.ref.snapshot.starlight/3=sha256:b11040115c1a7bad75c256bfad158a1bf2aec69418ffc1f49326ba59e33f69b2
+
+		sha256:e813af18bfe9565a5a4b67e4e80be064c488c4630770951c97f240fa337192e8 945B    53 minutes
+		containerd.io/gc.ref.content.config=sha256:4b167e69a056cad9f179b344d7208fefcbc99fee9d48c8988098d4af45cbc9ed,
+		containerd.io/distribution.source.ghcr.io=mc256/starlight/cli,
+		containerd.io/gc.ref.content.l.2=sha256:4f4fb700ef54461cfa02571ae0db9a0dc1e0cdb5577484a6d75e68dc38e8acc1,
+		containerd.io/gc.ref.content.l.1=sha256:76fc34d44084250f3a5e66218596e2105e9e0537b69456999dccf1740be1795e,
+		containerd.io/gc.ref.content.l.0=sha256:1b7ca6aea1ddfe716f3694edb811ab35114db9e93f3ce38d7dab6b4d9270cb0c
+
+	*/
+}

client/snapshotter/plugin.go

@@ -150,6 +150,7 @@ func (s *Plugin) Update(ctx context.Context, info snapshots.Info, fieldpaths ...
 	log.G(s.ctx).WithFields(logrus.Fields{
 		"name":  info.Name,
 		"usage": fieldpaths,
+		"info":  info,
 	}).Debug("sn: updated")
 	return info, nil
 }

demo/k3s-reset.sh

@@ -0,0 +1,11 @@
+#!/usr/bin/env bash
+
+# remove containerd filesystems
+rm -rf /var/lib/rancher/k3s/agent/containerd
+
+# remove starlight layer cache
+rm -rf /var/lib/starlight/layers
+
+
+# systemctl restart k3s-agent
+# systemctl stop containerd

proxy/api_test.go

@@ -17,7 +17,7 @@ func TestStarlightProxy_Ping(t *testing.T) {
 	proxy := NewStarlightProxy(context.TODO(), "https", "test.yuri.moe")
 	proxy.auth = *url.UserPassword("username", "password")
 
-	if err := proxy.Ping(); err != nil {
+	if _, _, _, err := proxy.Ping(); err != nil {
 		t.Error(err)
 	}
 }

proxy/database.go

@@ -468,52 +468,57 @@ func (d *Database) GetUniqueFiles(layers []*send.ImageLayer) ([]*send.File, erro
 	return fl, nil
 }
 
-func (d *Database) UpdateFileRanks(collection *fs.TraceCollection) (fs []int64, err error) {
-	var layersMap []int64
-
-	for _, img := range collection.Groups {
-		// get image serial id
-		i := img.Images[0]
-		var imageSerial, nlayer int64
-		if err = d.db.QueryRow(`
-			SELECT id, nlayer FROM image
-			WHERE ready IS NOT NULL AND hash=$1 LIMIT 1`,
-			i).Scan(&imageSerial, &nlayer); err != nil {
-			return nil, err
-		}
+func (d *Database) UpdateFileRanks(collection *fs.TraceCollection) (fs [][][]int64, err error) {
+	res := make([][][]int64, 0, len(collection.Groups))
+	for _, group := range collection.Groups {
+		layersImageMap := make([][]int64, len(group.Images))
+		for idx, img := range group.Images {
+			// get image serial id
+			var imageSerial, nlayer int64
+			if err = d.db.QueryRow(`
+					SELECT id, nlayer FROM image
+					WHERE ready IS NOT NULL AND hash=$1 LIMIT 1`,
+				img).Scan(&imageSerial, &nlayer); err != nil {
+				return nil, err
+			}
 
-		// get layers of the image
-		var rows *sql.Rows
-		rows, err = d.db.Query(`
+			// get layers of the image
+			var rows *sql.Rows
+			rows, err = d.db.Query(`
 				SELECT 
 					L.layer
 				FROM layer AS L
 				WHERE L.image=$1
 				ORDER BY "stackIndex" ASC`, imageSerial)
-		if err != nil {
-			return nil, err
-		}
-		defer rows.Close()
-		layersMap = make([]int64, nlayer)
-		idx := 0
-		for rows.Next() {
-			if err = rows.Scan(&layersMap[idx]); err != nil {
+			if err != nil {
 				return nil, err
 			}
-			idx += 1
+			defer rows.Close()
+
+			// update mapping
+			layersImageMap[idx] = make([]int64, nlayer)
+			layerIdx := 0
+			for rows.Next() {
+				if err = rows.Scan(&layersImageMap[idx][layerIdx]); err != nil {
+					return nil, err
+				}
+				layerIdx += 1
+			}
+
 		}
 
 		stmt, _ := d.db.Prepare(`
 			UPDATE file SET "order" = array_append("order",$1) WHERE fs=$2 and file=$3
 		`)
-		for _, f := range img.History {
-			_, err = stmt.Exec(f.Rank, layersMap[f.Stack], f.FileName)
+		for _, f := range group.History {
+			_, err = stmt.Exec(f.Rank, layersImageMap[f.SourceImage][f.Stack], f.FileName)
 			if err != nil {
 				return nil, err
 			}
 		}
+		res = append(res, layersImageMap)
 	}
-	return layersMap, nil
+	return res, nil
 }
 
 func (d *Database) GetFilesWithRanks(imageSerial int64) ([]*send.RankedFile, error) {

proxy/database_test.go

@@ -32,7 +32,7 @@ func TestMain(m *testing.M) {
 
 func TestDatabase_Init(t *testing.T) {
 	ctx := context.Background()
-	db, err := NewDatabase(ctx, "postgres://postgres:[email protected]:5432/postgres?sslmode=disable")
+	db, err := NewDatabase(ctx, "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable")
 	if err != nil {
 		t.Error(err)
 		return
@@ -70,7 +70,7 @@ func TestDatabase_GetFilesWithRanks(t *testing.T) {
 }
 
 func TestDatabase_UpdateFileRanks(t *testing.T) {
-	p := "../sandbox/group-optimize.json"
+	p := "/tmp/group-optimize.json"
 	b, err := ioutil.ReadFile(p)
 	if err != nil {
 		t.Error(err)
@@ -79,7 +79,7 @@ func TestDatabase_UpdateFileRanks(t *testing.T) {
 	var col *fs.TraceCollection
 	err = json.Unmarshal(b, &col)
 
-	var arr []int64
+	var arr [][][]int64
 	arr, err = db.UpdateFileRanks(col)
 	if err != nil {
 		t.Error(err)

util/config.go

@@ -55,9 +55,9 @@ const (
 	// Snapshot labels have a prefix of "containerd.io/snapshot/"
 	// or are the "containerd.io/snapshot.ref" label.
 	// SnapshotLabelRefImage is the digest of the image manifest
-	SnapshotLabelRefImage        = "containerd.io/snapshot/starlight/ref.image"
-	SnapshotLabelRefUncompressed = "containerd.io/snapshot/starlight/ref.uncompressed"
-	SnapshotLabelRefLayer        = "containerd.io/snapshot/starlight/ref.layer"
+	SnapshotLabelRefImage        = "containerd.io/snapshot/starlight.ref.image"
+	SnapshotLabelRefUncompressed = "containerd.io/snapshot/starlight.ref.uncompressed"
+	SnapshotLabelRefLayer        = "containerd.io/snapshot/starlight.ref.layer"
 
 	// ---------------------------------------------------------------------------------
 	// Switch to false in `Makefile` when build for production environment