85421b7f97
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>
85 lines
2.4 KiB
Go
85 lines
2.4 KiB
Go
package importer_test
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/base64"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"sync"
|
|
"testing"
|
|
|
|
castorev1pb "code.tvl.fyi/tvix/castore-go"
|
|
"code.tvl.fyi/tvix/nar-bridge/pkg/importer"
|
|
storev1pb "code.tvl.fyi/tvix/store-go"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestRoundtrip(t *testing.T) {
|
|
// We pipe nar_1094wph9z4nwlgvsd53abfz8i117ykiv5dwnq9nnhz846s7xqd7d.nar to
|
|
// storev1pb.Export, and store all the file contents and directory objects
|
|
// received in two hashmaps.
|
|
// We then feed it to the writer, and test we come up with the same NAR file.
|
|
|
|
f, err := os.Open("../../testdata/nar_1094wph9z4nwlgvsd53abfz8i117ykiv5dwnq9nnhz846s7xqd7d.nar")
|
|
require.NoError(t, err)
|
|
|
|
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)
|
|
|
|
rootNode, _, _, err := importer.Import(
|
|
context.Background(),
|
|
bytes.NewBuffer(narContents),
|
|
func(blobReader io.Reader) ([]byte, error) {
|
|
// read in contents, we need to put it into filesMap later.
|
|
contents, err := io.ReadAll(blobReader)
|
|
require.NoError(t, err)
|
|
|
|
dgst := mustBlobDigest(bytes.NewReader(contents))
|
|
|
|
// put it in filesMap
|
|
mu.Lock()
|
|
blobsMap[base64.StdEncoding.EncodeToString(dgst)] = contents
|
|
mu.Unlock()
|
|
|
|
return dgst, nil
|
|
},
|
|
func(directory *castorev1pb.Directory) ([]byte, error) {
|
|
dgst := mustDirectoryDigest(directory)
|
|
|
|
directoriesMap[base64.StdEncoding.EncodeToString(dgst)] = directory
|
|
return dgst, nil
|
|
},
|
|
)
|
|
|
|
require.NoError(t, err)
|
|
|
|
// done populating everything, now actually test the export :-)
|
|
var narBuf bytes.Buffer
|
|
err = storev1pb.Export(
|
|
&narBuf,
|
|
rootNode,
|
|
func(directoryDgst []byte) (*castorev1pb.Directory, error) {
|
|
d, found := directoriesMap[base64.StdEncoding.EncodeToString(directoryDgst)]
|
|
if !found {
|
|
panic(fmt.Sprintf("directory %v not found", base64.StdEncoding.EncodeToString(directoryDgst)))
|
|
}
|
|
return d, nil
|
|
},
|
|
func(blobDgst []byte) (io.ReadCloser, error) {
|
|
blobContents, found := blobsMap[base64.StdEncoding.EncodeToString(blobDgst)]
|
|
if !found {
|
|
panic(fmt.Sprintf("blob %v not found", base64.StdEncoding.EncodeToString(blobDgst)))
|
|
}
|
|
return io.NopCloser(bytes.NewReader(blobContents)), nil
|
|
},
|
|
)
|
|
|
|
require.NoError(t, err, "exporter shouldn't fail")
|
|
require.Equal(t, narContents, narBuf.Bytes())
|
|
}
|