From 4e8147427f05bf50902e416f6ec2d108d7779a18 Mon Sep 17 00:00:00 2001
From: Xiaodong Ye <xiaodong.ye@mthreads.com>
Date: Mon, 26 Aug 2024 19:03:33 +0800
Subject: [PATCH] Support storing Ollama layers

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
---
 drivers/driver.go          |  1 +
 drivers/overlay/overlay.go | 47 +++++++++++++++++++++++++++++---------
 layers.go                  | 25 +++++++++++++++-----
 store.go                   |  2 ++
 4 files changed, 58 insertions(+), 17 deletions(-)

diff --git a/drivers/driver.go b/drivers/driver.go
index b62234e57e..2bd5825502 100644
--- a/drivers/driver.go
+++ b/drivers/driver.go
@@ -72,6 +72,7 @@ type ApplyDiffOpts struct {
 	MountLabel        string
 	IgnoreChownErrors bool
 	ForceMask         *os.FileMode
+	LayerFilename     *string
 }
 
 // ApplyDiffWithDifferOpts contains optional arguments for ApplyDiffWithDiffer methods.
diff --git a/drivers/overlay/overlay.go b/drivers/overlay/overlay.go
index 82a9086b68..2781ddfe44 100644
--- a/drivers/overlay/overlay.go
+++ b/drivers/overlay/overlay.go
@@ -790,6 +790,23 @@ func supportsOverlay(home string, homeMagic graphdriver.FsMagic, rootUID, rootGI
 	return supportsDType, fmt.Errorf("'overlay' not found as a supported filesystem on this host. Please ensure kernel is new enough and has overlay support loaded.: %w", graphdriver.ErrNotSupported)
 }
 
+func cp(r io.Reader, dest string, filename string) error {
+	if seeker, ok := r.(io.Seeker); ok {
+		if _, err := seeker.Seek(0, io.SeekStart); err != nil {
+			return err
+		}
+	}
+
+	f, err := os.Create(filepath.Join(dest, filename))
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	_, err = io.Copy(f, r)
+	return err
+}
+
 func (d *Driver) useNaiveDiff() bool {
 	if d.usingComposefs {
 		return true
@@ -2329,17 +2346,25 @@ func (d *Driver) ApplyDiff(id, parent string, options graphdriver.ApplyDiffOpts)
 		return 0, err
 	}
 
-	logrus.Debugf("Applying tar in %s", applyDir)
-	// Overlay doesn't need the parent id to apply the diff
-	if err := untar(options.Diff, applyDir, &archive.TarOptions{
-		UIDMaps:           idMappings.UIDs(),
-		GIDMaps:           idMappings.GIDs(),
-		IgnoreChownErrors: d.options.ignoreChownErrors,
-		ForceMask:         d.options.forceMask,
-		WhiteoutFormat:    d.getWhiteoutFormat(),
-		InUserNS:          unshare.IsRootless(),
-	}); err != nil {
-		return 0, err
+	if options.LayerFilename != nil {
+		logrus.Debugf("Applying file in %s", applyDir)
+		err := cp(options.Diff, applyDir, *options.LayerFilename)
+		if err != nil {
+			return 0, err
+		}
+	} else {
+		logrus.Debugf("Applying tar in %s", applyDir)
+		// Overlay doesn't need the parent id to apply the diff
+		if err := untar(options.Diff, applyDir, &archive.TarOptions{
+			UIDMaps:           idMappings.UIDs(),
+			GIDMaps:           idMappings.GIDs(),
+			IgnoreChownErrors: d.options.ignoreChownErrors,
+			ForceMask:         d.options.forceMask,
+			WhiteoutFormat:    d.getWhiteoutFormat(),
+			InUserNS:          unshare.IsRootless(),
+		}); err != nil {
+			return 0, err
+		}
 	}
 
 	return directory.Size(applyDir)
diff --git a/layers.go b/layers.go
index 6caf28ab71..2c0afe5ef9 100644
--- a/layers.go
+++ b/layers.go
@@ -2422,14 +2422,21 @@ func (r *layerStore) applyDiffWithOptions(to string, layerOptions *LayerOptions,
 		if uncompressedDigester != nil {
 			uncompressedWriter = io.MultiWriter(uncompressedWriter, uncompressedDigester.Hash())
 		}
-		payload, err := asm.NewInputTarStream(io.TeeReader(uncompressed, uncompressedWriter), metadata, storage.NewDiscardFilePutter())
-		if err != nil {
-			return -1, err
+
+		var payload io.Reader
+		if layerOptions != nil && layerOptions.LayerFilename != nil {
+			payload = diff
+		} else {
+			payload, err = asm.NewInputTarStream(io.TeeReader(uncompressed, uncompressedWriter), metadata, storage.NewDiscardFilePutter())
+			if err != nil {
+				return -1, err
+			}
 		}
 		options := drivers.ApplyDiffOpts{
-			Diff:       payload,
-			Mappings:   r.layerMappings(layer),
-			MountLabel: layer.MountLabel,
+			Diff:          payload,
+			Mappings:      r.layerMappings(layer),
+			MountLabel:    layer.MountLabel,
+			LayerFilename: layerOptions.LayerFilename,
 		}
 		size, err := r.driver.ApplyDiff(layer.ID, layer.Parent, options)
 		if err != nil {
@@ -2468,6 +2475,12 @@ func (r *layerStore) applyDiffWithOptions(to string, layerOptions *LayerOptions,
 	layer.UncompressedDigest = uncompressedDigest
 	layer.UncompressedSize = uncompressedCounter.Count
 	layer.CompressionType = compression
+
+	if layerOptions != nil && layerOptions.LayerFilename != nil {
+		layer.CompressedSize = size
+		layer.UncompressedSize = size
+	}
+
 	layer.UIDs = make([]uint32, 0, len(uidLog))
 	for uid := range uidLog {
 		layer.UIDs = append(layer.UIDs, uid)
diff --git a/store.go b/store.go
index 8510d033c6..2114da0f73 100644
--- a/store.go
+++ b/store.go
@@ -641,6 +641,8 @@ type LayerOptions struct {
 	// Currently these can only be set when the layer record is created, but that
 	// could change in the future.
 	Flags map[string]interface{}
+	// LayerFilename is the target filename of the layer blob.
+	LayerFilename *string
 }
 
 type LayerBigDataOption struct {