Merge pull request #109 from jmorganca/fix-create-memory

fix memory leak in create

server/images.go

@@ -3,7 +3,6 @@ package server
 import (
 	"bytes"
 	"crypto/sha256"
-	"encoding/hex"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -42,10 +41,9 @@ type Layer struct {
 	Size      int    `json:"size"`
 }
 
-type LayerWithBuffer struct {
+type LayerReader struct {
 	Layer
-
-	Buffer *bytes.Buffer
+	io.Reader
 }
 
 type ConfigV2 struct {
@@ -161,7 +159,7 @@ func CreateModel(name string, mf io.Reader, fn func(status string)) error {
 		return err
 	}
 
-	var layers []*LayerWithBuffer
+	var layers []*LayerReader
 	params := make(map[string]string)
 
 	for _, c := range commands {
@@ -274,7 +272,7 @@ func CreateModel(name string, mf io.Reader, fn func(status string)) error {
 	return nil
 }
 
-func removeLayerFromLayers(layers []*LayerWithBuffer, mediaType string) []*LayerWithBuffer {
+func removeLayerFromLayers(layers []*LayerReader, mediaType string) []*LayerReader {
 	j := 0
 	for _, l := range layers {
 		if l.MediaType != mediaType {
@@ -285,7 +283,7 @@ func removeLayerFromLayers(layers []*LayerWithBuffer, mediaType string) []*Layer
 	return layers[:j]
 }
 
-func SaveLayers(layers []*LayerWithBuffer, fn func(status string), force bool) error {
+func SaveLayers(layers []*LayerReader, fn func(status string), force bool) error {
 	// Write each of the layers to disk
 	for _, layer := range layers {
 		fp, err := GetBlobsPath(layer.Digest)
@@ -303,10 +301,10 @@ func SaveLayers(layers []*LayerWithBuffer, fn func(status string), force bool) e
 			}
 			defer out.Close()
 
-			_, err = io.Copy(out, layer.Buffer)
-			if err != nil {
+			if _, err = io.Copy(out, layer.Reader); err != nil {
 				return err
 			}
+
 		} else {
 			fn(fmt.Sprintf("using already created layer %s", layer.Digest))
 		}
@@ -315,7 +313,7 @@ func SaveLayers(layers []*LayerWithBuffer, fn func(status string), force bool) e
 	return nil
 }
 
-func CreateManifest(name string, cfg *LayerWithBuffer, layers []*Layer) error {
+func CreateManifest(name string, cfg *LayerReader, layers []*Layer) error {
 	mp := ParseModelPath(name)
 
 	manifest := ManifestV2{
@@ -341,7 +339,7 @@ func CreateManifest(name string, cfg *LayerWithBuffer, layers []*Layer) error {
 	return os.WriteFile(fp, manifestJSON, 0o644)
 }
 
-func GetLayerWithBufferFromLayer(layer *Layer) (*LayerWithBuffer, error) {
+func GetLayerWithBufferFromLayer(layer *Layer) (*LayerReader, error) {
 	fp, err := GetBlobsPath(layer.Digest)
 	if err != nil {
 		return nil, err
@@ -361,7 +359,7 @@ func GetLayerWithBufferFromLayer(layer *Layer) (*LayerWithBuffer, error) {
 	return newLayer, nil
 }
 
-func paramsToReader(params map[string]string) (io.Reader, error) {
+func paramsToReader(params map[string]string) (io.ReadSeeker, error) {
 	opts := api.DefaultOptions()
 	typeOpts := reflect.TypeOf(opts)
 
@@ -419,7 +417,7 @@ func paramsToReader(params map[string]string) (io.Reader, error) {
 	return bytes.NewReader(bts), nil
 }
 
-func getLayerDigests(layers []*LayerWithBuffer) ([]string, error) {
+func getLayerDigests(layers []*LayerReader) ([]string, error) {
 	var digests []string
 	for _, l := range layers {
 		if l.Digest == "" {
@@ -431,22 +429,17 @@ func getLayerDigests(layers []*LayerWithBuffer) ([]string, error) {
 }
 
 // CreateLayer creates a Layer object from a given file
-func CreateLayer(f io.Reader) (*LayerWithBuffer, error) {
-	buf := new(bytes.Buffer)
-	_, err := io.Copy(buf, f)
-	if err != nil {
-		return nil, err
-	}
+func CreateLayer(f io.ReadSeeker) (*LayerReader, error) {
+	digest, size := GetSHA256Digest(f)
+	f.Seek(0, 0)
 
-	digest, size := GetSHA256Digest(buf)
-
-	layer := &LayerWithBuffer{
+	layer := &LayerReader{
 		Layer: Layer{
 			MediaType: "application/vnd.docker.image.rootfs.diff.tar",
 			Digest:    digest,
 			Size:      size,
 		},
-		Buffer: buf,
+		Reader: f,
 	}
 
 	return layer, nil
@@ -609,7 +602,7 @@ func pullModelManifest(mp ModelPath, username, password string) (*ManifestV2, er
 	return m, err
 }
 
-func createConfigLayer(layers []string) (*LayerWithBuffer, error) {
+func createConfigLayer(layers []string) (*LayerReader, error) {
 	// TODO change architecture and OS
 	config := ConfigV2{
 		Architecture: "arm64",
@@ -628,22 +621,26 @@ func createConfigLayer(layers []string) (*LayerWithBuffer, error) {
 	buf := bytes.NewBuffer(configJSON)
 	digest, size := GetSHA256Digest(buf)
 
-	layer := &LayerWithBuffer{
+	layer := &LayerReader{
 		Layer: Layer{
 			MediaType: "application/vnd.docker.container.image.v1+json",
 			Digest:    digest,
 			Size:      size,
 		},
-		Buffer: buf,
+		Reader: buf,
 	}
 	return layer, nil
 }
 
 // GetSHA256Digest returns the SHA256 hash of a given buffer and returns it, and the size of buffer
-func GetSHA256Digest(data *bytes.Buffer) (string, int) {
-	layerBytes := data.Bytes()
-	hash := sha256.Sum256(layerBytes)
-	return "sha256:" + hex.EncodeToString(hash[:]), len(layerBytes)
+func GetSHA256Digest(r io.Reader) (string, int) {
+	h := sha256.New()
+	n, err := io.Copy(h, r)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	return fmt.Sprintf("sha256:%x", h.Sum(nil)), int(n)
 }
 
 func startUpload(mp ModelPath, username string, password string) (string, error) {