Use streaming compression for files.

cmd/llama/objects.go

@@ -49,7 +49,7 @@ func (c *StoreCommand) Execute(ctx context.Context, flag *flag.FlagSet, _ ...int
 			return subcommands.ExitFailure
 		}
 
-		id, err := global.MustStore().Store(ctx, bytes)
+		id, err := global.MustStore().StoreBytes(ctx, bytes)
 		if err != nil {
 			log.Printf("storing %q: %v\n", arg, err)
 			return subcommands.ExitFailure

files/list.go

@@ -17,7 +17,6 @@ package files
 import (
 	"context"
 	"fmt"
-	"io/ioutil"
 	"os"
 	"path"
 	"strings"
@@ -76,35 +75,38 @@ func (f List) Append(mapped ...Mapped) List {
 
 func uploadWorker(ctx context.Context, store store.Store, jobs <-chan Mapped, out chan<- *protocol.FileAndPath) {
 	for file := range jobs {
-		data, mode, err := func() ([]byte, os.FileMode, error) {
-			if file.Local.Bytes != nil {
-				if file.Local.Path != "" {
-					panic("MappedFile: got both Path and Bytes")
-				}
-				return file.Local.Bytes, file.Local.Mode, nil
-			} else {
-				data, err := ioutil.ReadFile(file.Local.Path)
-				if err != nil {
-					return nil, 0, fmt.Errorf("reading file %q: %w", file.Local.Path, err)
-				}
-				st, err := os.Stat(file.Local.Path)
-				if err != nil {
-					return nil, 0, fmt.Errorf("stat %q: %w", file.Local.Path, err)
-				}
-				return data, st.Mode(), nil
-			}
-		}()
-		var blob *protocol.Blob
-		if err == nil {
-			blob, err = files.NewBlob(ctx, store, data)
-		}
-		if err != nil {
-			blob = &protocol.Blob{Err: err.Error()}
+		if file.Local.Bytes != nil && file.Local.Path != "" {
+			panic("MappedFile: got both Path and Bytes")
 		}
-		out <- &protocol.FileAndPath{
-			File: protocol.File{Blob: *blob, Mode: mode},
+
+		result := protocol.FileAndPath{
+			File: protocol.File{
+				Mode: file.Local.Mode,
+			},
 			Path: file.Remote,
 		}
+
+		if file.Local.Path != "" {
+			pfile, err := files.ReadFile(ctx, store, file.Local.Path)
+			switch err {
+			case nil:
+				result.File = *pfile
+			default:
+				result.File.Blob = protocol.Blob{Err: err.Error()}
+			}
+		}
+
+		if file.Local.Bytes != nil {
+			blob, err := files.NewBlob(ctx, store, file.Local.Bytes)
+			switch err {
+			case nil:
+				result.File.Blob = *blob
+			default:
+				result.File.Blob = protocol.Blob{Err: err.Error()}
+			}
+		}
+
+		out <- &result
 	}
 }
 

protocol/files/blobs.go

@@ -79,7 +79,7 @@ func NewBlob(ctx context.Context, store store.Store, bytes []byte) (*protocol.Bl
 	if base64.StdEncoding.EncodedLen(len(bytes)) < protocol.MaxInlineBlob {
 		return &protocol.Blob{Bytes: bytes}, nil
 	}
-	id, err := store.Store(ctx, bytes)
+	id, err := store.StoreBytes(ctx, bytes)
 	if err != nil {
 		return nil, err
 	}
@@ -100,16 +100,14 @@ func ReadFile(ctx context.Context, store store.Store, path string) (*protocol.Fi
 	if fi.Mode().IsDir() {
 		return nil, errors.New("ReadFile: got directory")
 	}
-	bytes, err := ioutil.ReadAll(fh)
-	if err != nil {
-		return nil, err
-	}
-	blob, err := NewBlob(ctx, store, bytes)
+
+	id, err := store.Store(ctx, fh)
 	if err != nil {
 		return nil, err
 	}
+
 	return &protocol.File{
-		Blob: *blob,
+		Blob: protocol.Blob{Ref: id},
 		Mode: fi.Mode(),
 	}, nil
 }

store/memory.go

@@ -17,6 +17,8 @@ package store
 import (
 	"context"
 	"encoding/hex"
+	"io"
+	"io/ioutil"
 
 	"github.com/nelhage/llama/protocol"
 	"golang.org/x/crypto/blake2b"
@@ -26,7 +28,16 @@ type inMemory struct {
 	objects map[string][]byte
 }
 
-func (s *inMemory) Store(ctx context.Context, obj []byte) (string, error) {
+func (s *inMemory) Store(ctx context.Context, obj io.Reader) (string, error) {
+	buf, err := ioutil.ReadAll(obj)
+	if err != nil {
+		return "", err
+	}
+
+	return s.StoreBytes(ctx, buf)
+}
+
+func (s *inMemory) StoreBytes(ctx context.Context, obj []byte) (string, error) {
 	sha := blake2b.Sum256(obj)
 	id := hex.EncodeToString(sha[:])
 	s.objects[id] = append([]byte(nil), obj...)

store/s3store/s3.go

@@ -18,6 +18,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"log"
 	"net/url"
@@ -133,9 +134,13 @@ func FromSessionAndOptions(s *session.Session, address string, opts Options) (*S
 	}, nil
 }
 
-func (s *Store) Store(ctx context.Context, obj []byte) (string, error) {
+// StoreCompressedObj takes a compressed data object and stores it in S3,
+// returning an object ID.
+func (s *Store) StoreCompressedObj(ctx context.Context, obj []byte) (string, error) {
 	ctx, span := tracing.StartSpan(ctx, "s3.store")
 	defer span.End()
+
+	// The object ID is the hash of the compressed object.
 	id := storeutil.HashObject(obj) + ":zstd"
 
 	span.AddField("object_id", id)
@@ -170,23 +175,48 @@ func (s *Store) Store(ctx context.Context, obj []byte) (string, error) {
 		}
 	}
 
-	compressed := encode.EncodeAll(obj, nil)
-	span.AddField("s3.write_bytes", len(compressed))
+	span.AddField("s3.write_bytes", len(obj))
 
 	usage.WriteRequests += 1
 	_, err = s.s3.PutObjectWithContext(ctx, &s3.PutObjectInput{
-		Body:   bytes.NewReader(compressed),
+		Body:   bytes.NewReader(obj),
 		Bucket: &s.url.Host,
 		Key:    key,
 	})
 	if err != nil {
 		return "", err
 	}
+
 	s.metrics.XferIn += uint64(len(obj))
 	upload.Complete()
 	return id, nil
 }
 
+func (s *Store) Store(ctx context.Context, obj io.Reader) (string, error) {
+	compressed := bytes.Buffer{}
+
+	encoder, err := zstd.NewWriter(&compressed)
+	if err != nil {
+		return "", err
+	}
+
+	if _, err := io.Copy(encoder, obj); err != nil {
+		encoder.Close()
+		return "", err
+	}
+
+	if err := encoder.Close(); err != nil {
+		return "", err
+	}
+
+	return s.StoreCompressedObj(ctx, compressed.Bytes())
+}
+
+func (s *Store) StoreBytes(ctx context.Context, obj []byte) (string, error) {
+	compressed := encode.EncodeAll(obj, nil)
+	return s.StoreCompressedObj(ctx, compressed)
+}
+
 const getConcurrency = 32
 
 func (s *Store) getFromS3(ctx context.Context, id string, usage *usageMetrics) ([]byte, error) {
@@ -226,9 +256,11 @@ func (s *Store) decompress(id string, body []byte) (string, []byte, error) {
 			return expectHash, nil, fmt.Errorf("%q: unknown compression %s", id, coding)
 		}
 		var err error
+		inbytes := len(body)
 		body, err = decode.DecodeAll(body, nil)
 		if err != nil {
-			return expectHash, nil, fmt.Errorf("%q: decoding:  %w", id, err)
+			return expectHash, nil,
+				fmt.Errorf("%q: decoding %d %s bytes: %w", id, inbytes, coding, err)
 		}
 	}
 	return expectHash, body, nil
@@ -247,15 +279,18 @@ func (s *Store) getOne(ctx context.Context, id string, usage *usageMetrics) ([]b
 		}
 	}
 
+	// The ID is generated from the compressed representation.
+	gotHash := storeutil.HashObject(body)
+
 	hash, body, err := s.decompress(id, body)
 	if err != nil {
 		return nil, err
 	}
 
-	gotHash := storeutil.HashObject(body)
 	if gotHash != hash {
 		return nil, fmt.Errorf("object store mismatch: got csum=%s expected %s", gotHash, id)
 	}
+
 	u := s.seen.StartUpload(id)
 	u.Complete()
 

store/store.go

@@ -17,6 +17,7 @@ package store
 import (
 	"context"
 	"errors"
+	"io"
 
 	"github.com/nelhage/llama/protocol"
 )
@@ -30,7 +31,9 @@ type GetRequest struct {
 var ErrNotExists = errors.New("Requested object does not exist")
 
 type Store interface {
-	Store(ctx context.Context, obj []byte) (string, error)
+	StoreBytes(ctx context.Context, obj []byte) (string, error)
+	Store(ctx context.Context, obj io.Reader) (string, error)
+
 	GetObjects(ctx context.Context, gets []GetRequest)
 	FetchAWSUsage(u *protocol.UsageMetrics)
 }