feat(tvix/store/nar-info): Support async blob upload
Right now all blob uploads are performed synchronously, this means if a NAR contains many small files, the import time is dominated by round trip time to the blob service. For small files, we can buffer them in memory and upload them asynchronously to the blob service. Before returning we make sure to join all the uploads to make sure they complete successfully before responding OK. This reduces time to import a bash-interactive closure on my machine from 1m19s to 7s. Change-Id: Ica3695c159e6c8ad8769281ac20d037e3143e856 Reviewed-on: https://cl.tvl.fyi/c/depot/+/10679 Tested-by: BuildkiteCI Reviewed-by: flokli <flokli@flokli.de> Autosubmit: Connor Brewster <cbrewster@hey.com>
This commit is contained in:
parent
d056329412
commit
85421b7f97
5 changed files with 85 additions and 11 deletions
|
@ -6,5 +6,5 @@ pkgs.buildGoModule {
|
|||
name = "nar-bridge";
|
||||
src = depot.third_party.gitignoreSource ./.;
|
||||
|
||||
vendorHash = "sha256-9tEpICef6xCuRQwyXCg15KugcvknnP53WHotohcbeis=";
|
||||
vendorHash = "sha256-7jugbC5sEGhppjiZgnoLP5A6kQSaHK9vE6cXVZBG22s=";
|
||||
}
|
||||
|
|
|
@ -18,6 +18,7 @@ require (
|
|||
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.22.0
|
||||
go.opentelemetry.io/otel/sdk v1.22.0
|
||||
go.opentelemetry.io/otel/sdk/metric v1.22.0
|
||||
golang.org/x/sync v0.4.0
|
||||
google.golang.org/grpc v1.60.1
|
||||
google.golang.org/protobuf v1.32.0
|
||||
lukechampine.com/blake3 v1.2.1
|
||||
|
|
|
@ -90,6 +90,8 @@ golang.org/x/crypto v0.18.0/go.mod h1:R0j02AL6hcrfOiy9T4ZYp/rcWeMxM3L6QYxlOuEG1m
|
|||
golang.org/x/net v0.20.0 h1:aCL9BSgETF1k+blQaYUBx9hJ9LOGP3gAVemcZlf1Kpo=
|
||||
golang.org/x/net v0.20.0/go.mod h1:z8BVo6PvndSri0LbOE3hAn0apkU+1YvI6E70E9jsnvY=
|
||||
golang.org/x/oauth2 v0.13.0 h1:jDDenyj+WgFtmV3zYVoi8aE2BwtXFLWOA67ZfNWftiY=
|
||||
golang.org/x/sync v0.4.0 h1:zxkM55ReGkDlKSM+Fu41A+zmbZuaPVbGMzvvdUPznYQ=
|
||||
golang.org/x/sync v0.4.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y=
|
||||
golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.16.0 h1:xWw16ngr6ZMtmxDyKyIgsE93KNKz5HKmMa3b8ALHidU=
|
||||
|
|
|
@ -1,6 +1,7 @@
|
|||
package importer
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"errors"
|
||||
|
@ -11,6 +12,17 @@ import (
|
|||
|
||||
castorev1pb "code.tvl.fyi/tvix/castore-go"
|
||||
"github.com/nix-community/go-nix/pkg/nar"
|
||||
"golang.org/x/sync/errgroup"
|
||||
"lukechampine.com/blake3"
|
||||
)
|
||||
|
||||
const (
|
||||
// asyncUploadThreshold controls when a file is buffered into memory and uploaded
|
||||
// asynchronously. Files must be smaller than the threshold to be uploaded asynchronously.
|
||||
asyncUploadThreshold = 1024 * 1024 // 1 MiB
|
||||
// maxAsyncUploadBufferBytes is the maximum number of async blob uploads allowed to be
|
||||
// running concurrently at any given time for a simple import operation.
|
||||
maxConcurrentAsyncUploads = 128
|
||||
)
|
||||
|
||||
// An item on the directories stack
|
||||
|
@ -51,6 +63,11 @@ func Import(
|
|||
var rootFile *castorev1pb.FileNode
|
||||
var stackDirectory *castorev1pb.Directory
|
||||
|
||||
// Keep track of all asynch blob uploads so we can make sure they all succeed
|
||||
// before returning.
|
||||
var asyncBlobWg errgroup.Group
|
||||
asyncBlobWg.SetLimit(maxConcurrentAsyncUploads)
|
||||
|
||||
var stack = []stackItem{}
|
||||
|
||||
// popFromStack is used when we transition to a different directory or
|
||||
|
@ -124,6 +141,12 @@ func Import(
|
|||
}
|
||||
}
|
||||
|
||||
// Wait for any pending blob uploads.
|
||||
err := asyncBlobWg.Wait()
|
||||
if err != nil {
|
||||
return nil, 0, nil, fmt.Errorf("async blob upload: %w", err)
|
||||
}
|
||||
|
||||
// Stack is empty.
|
||||
// Now either root{File,Symlink,Directory} is not nil,
|
||||
// and we can return the root node.
|
||||
|
@ -192,19 +215,63 @@ func Import(
|
|||
|
||||
}
|
||||
if hdr.Type == nar.TypeRegular {
|
||||
// wrap reader with a reader counting the number of bytes read
|
||||
blobCountW := &CountingWriter{}
|
||||
blobReader := io.TeeReader(narReader, blobCountW)
|
||||
uploadBlob := func(r io.Reader) ([]byte, error) {
|
||||
// wrap reader with a reader counting the number of bytes read
|
||||
blobCountW := &CountingWriter{}
|
||||
blobReader := io.TeeReader(r, blobCountW)
|
||||
|
||||
blobDigest, err := blobCb(blobReader)
|
||||
if err != nil {
|
||||
return nil, 0, nil, fmt.Errorf("failure from blobCb: %w", err)
|
||||
blobDigest, err := blobCb(blobReader)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failure from blobCb: %w", err)
|
||||
}
|
||||
|
||||
// ensure blobCb did read all the way to the end.
|
||||
// If it didn't, the blobCb function is wrong and we should bail out.
|
||||
if blobCountW.BytesWritten() != uint64(hdr.Size) {
|
||||
return nil, fmt.Errorf("blobCb did not read all: %d/%d bytes", blobCountW.BytesWritten(), hdr.Size)
|
||||
}
|
||||
|
||||
return blobDigest, nil
|
||||
}
|
||||
|
||||
// ensure blobCb did read all the way to the end.
|
||||
// If it didn't, the blobCb function is wrong and we should bail out.
|
||||
if blobCountW.BytesWritten() != uint64(hdr.Size) {
|
||||
panic("blobCB did not read to end")
|
||||
h := blake3.New(32, nil)
|
||||
blobReader := io.TeeReader(narReader, io.MultiWriter(h))
|
||||
var blobDigest []byte
|
||||
|
||||
// If this file is small enough, read it off the wire immediately and
|
||||
// upload to the blob service asynchronously. This helps reduce the
|
||||
// RTT on blob uploads for NARs with many small files.
|
||||
doAsync := hdr.Size < asyncUploadThreshold
|
||||
if doAsync {
|
||||
blobContents, err := io.ReadAll(blobReader)
|
||||
if err != nil {
|
||||
return nil, 0, nil, fmt.Errorf("read blob: %w", err)
|
||||
}
|
||||
|
||||
blobDigest = h.Sum(nil)
|
||||
|
||||
asyncBlobWg.Go(func() error {
|
||||
blobDigestFromCb, err := uploadBlob(bytes.NewReader(blobContents))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if !bytes.Equal(blobDigest, blobDigestFromCb) {
|
||||
return fmt.Errorf("unexpected digest (got %x, expected %x)", blobDigestFromCb, blobDigest)
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
} else {
|
||||
blobDigestFromCb, err := uploadBlob(blobReader)
|
||||
if err != nil {
|
||||
return nil, 0, nil, fmt.Errorf("upload blob: %w", err)
|
||||
}
|
||||
|
||||
blobDigest = h.Sum(nil)
|
||||
if !bytes.Equal(blobDigest, blobDigestFromCb) {
|
||||
return nil, 0, nil, fmt.Errorf("unexpected digest (got %x, expected %x)", blobDigestFromCb, blobDigest)
|
||||
}
|
||||
}
|
||||
|
||||
fileNode := &castorev1pb.FileNode{
|
||||
|
|
|
@ -7,6 +7,7 @@ import (
|
|||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
castorev1pb "code.tvl.fyi/tvix/castore-go"
|
||||
|
@ -27,6 +28,7 @@ func TestRoundtrip(t *testing.T) {
|
|||
narContents, err := io.ReadAll(f)
|
||||
require.NoError(t, err)
|
||||
|
||||
var mu sync.Mutex
|
||||
blobsMap := make(map[string][]byte, 0)
|
||||
directoriesMap := make(map[string]*castorev1pb.Directory)
|
||||
|
||||
|
@ -41,7 +43,9 @@ func TestRoundtrip(t *testing.T) {
|
|||
dgst := mustBlobDigest(bytes.NewReader(contents))
|
||||
|
||||
// put it in filesMap
|
||||
mu.Lock()
|
||||
blobsMap[base64.StdEncoding.EncodeToString(dgst)] = contents
|
||||
mu.Unlock()
|
||||
|
||||
return dgst, nil
|
||||
},
|
||||
|
|
Loading…
Reference in a new issue