refactor(tvix/nar-bridge): update to new separated protos

Change-Id: I5eaadc837a4d3a7f635574437127a22de88f556b
Reviewed-on: https://cl.tvl.fyi/c/depot/+/9407
Tested-by: BuildkiteCI
Reviewed-by: Connor Brewster <cbrewster@hey.com>
Autosubmit: flokli <flokli@flokli.de>
This commit is contained in:
Florian Klink 2023-09-22 16:38:10 +03:00 committed by flokli
parent e13c86d916
commit 732dc68727
16 changed files with 170 additions and 158 deletions

View file

@ -8,7 +8,7 @@ import (
"fmt"
"io"
storev1pb "code.tvl.fyi/tvix/store/protos"
castorev1pb "code.tvl.fyi/tvix/castore/protos"
log "github.com/sirupsen/logrus"
)
@ -17,7 +17,7 @@ const chunkSize = 1024 * 1024
// this produces a callback function that can be used as blobCb for the
// reader.Import function call
func genBlobServiceWriteCb(ctx context.Context, blobServiceClient storev1pb.BlobServiceClient) func(io.Reader) error {
func genBlobServiceWriteCb(ctx context.Context, blobServiceClient castorev1pb.BlobServiceClient) func(io.Reader) error {
return func(blobReader io.Reader) error {
// Ensure the blobReader is buffered to at least the chunk size.
blobReader = bufio.NewReaderSize(blobReader, chunkSize)
@ -42,7 +42,7 @@ func genBlobServiceWriteCb(ctx context.Context, blobServiceClient storev1pb.Blob
blobSize += n
// send the blob chunk to the server. The err is only valid in the inner scope
if err := putter.Send(&storev1pb.BlobChunk{
if err := putter.Send(&castorev1pb.BlobChunk{
Data: chunk[:n],
}); err != nil {
return fmt.Errorf("sending blob chunk: %w", err)

View file

@ -5,17 +5,17 @@ import (
"encoding/base64"
"fmt"
storev1pb "code.tvl.fyi/tvix/store/protos"
castorev1pb "code.tvl.fyi/tvix/castore/protos"
log "github.com/sirupsen/logrus"
)
type DirectoriesUploader struct {
ctx context.Context
directoryServiceClient storev1pb.DirectoryServiceClient
directoryServicePutStream storev1pb.DirectoryService_PutClient
directoryServiceClient castorev1pb.DirectoryServiceClient
directoryServicePutStream castorev1pb.DirectoryService_PutClient
}
func NewDirectoriesUploader(ctx context.Context, directoryServiceClient storev1pb.DirectoryServiceClient) *DirectoriesUploader {
func NewDirectoriesUploader(ctx context.Context, directoryServiceClient castorev1pb.DirectoryServiceClient) *DirectoriesUploader {
return &DirectoriesUploader{
ctx: ctx,
directoryServiceClient: directoryServiceClient,
@ -23,7 +23,7 @@ func NewDirectoriesUploader(ctx context.Context, directoryServiceClient storev1p
}
}
func (du *DirectoriesUploader) Put(directory *storev1pb.Directory) error {
func (du *DirectoriesUploader) Put(directory *castorev1pb.Directory) error {
directoryDgst, err := directory.Digest()
if err != nil {
return fmt.Errorf("failed calculating directory digest: %w", err)
@ -50,7 +50,7 @@ func (du *DirectoriesUploader) Put(directory *storev1pb.Directory) error {
}
// Done is called whenever we're
func (du *DirectoriesUploader) Done() (*storev1pb.PutDirectoryResponse, error) {
func (du *DirectoriesUploader) Done() (*castorev1pb.PutDirectoryResponse, error) {
// only close once, and only if we opened.
if du.directoryServicePutStream == nil {
return nil, nil

View file

@ -12,6 +12,7 @@ import (
"net/http"
"sync"
castorev1pb "code.tvl.fyi/tvix/castore/protos"
"code.tvl.fyi/tvix/nar-bridge/pkg/writer"
storev1pb "code.tvl.fyi/tvix/store/protos"
"github.com/go-chi/chi/v5"
@ -27,8 +28,8 @@ const (
func renderNar(
ctx context.Context,
log *log.Entry,
directoryServiceClient storev1pb.DirectoryServiceClient,
blobServiceClient storev1pb.BlobServiceClient,
directoryServiceClient castorev1pb.DirectoryServiceClient,
blobServiceClient castorev1pb.BlobServiceClient,
narHashToPathInfoMu *sync.Mutex,
narHashToPathInfo map[string]*storev1pb.PathInfo,
w io.Writer,
@ -50,15 +51,15 @@ func renderNar(
return nil
}
directories := make(map[string]*storev1pb.Directory)
directories := make(map[string]*castorev1pb.Directory)
// If the root node is a directory, ask the directory service for all directories
if pathInfoDirectory := pathInfo.GetNode().GetDirectory(); pathInfoDirectory != nil {
rootDirectoryDigest := pathInfoDirectory.GetDigest()
log = log.WithField("root_directory", base64.StdEncoding.EncodeToString(rootDirectoryDigest))
directoryStream, err := directoryServiceClient.Get(ctx, &storev1pb.GetDirectoryRequest{
ByWhat: &storev1pb.GetDirectoryRequest_Digest{
directoryStream, err := directoryServiceClient.Get(ctx, &castorev1pb.GetDirectoryRequest{
ByWhat: &castorev1pb.GetDirectoryRequest_Digest{
Digest: rootDirectoryDigest,
},
Recursive: true,
@ -96,7 +97,7 @@ func renderNar(
err := writer.Export(
w,
pathInfo,
func(directoryDigest []byte) (*storev1pb.Directory, error) {
func(directoryDigest []byte) (*castorev1pb.Directory, error) {
log.WithField("directory", base64.StdEncoding.EncodeToString(directoryDigest)).Debug("Get directory")
directoryRefStr := hex.EncodeToString(directoryDigest)
directory, found := directories[directoryRefStr]
@ -112,7 +113,7 @@ func renderNar(
},
func(blobDigest []byte) (io.ReadCloser, error) {
log.WithField("blob", base64.StdEncoding.EncodeToString(blobDigest)).Debug("Get blob")
resp, err := blobServiceClient.Read(ctx, &storev1pb.ReadBlobRequest{
resp, err := blobServiceClient.Read(ctx, &castorev1pb.ReadBlobRequest{
Digest: blobDigest,
})
if err != nil {

View file

@ -6,8 +6,8 @@ import (
"fmt"
"net/http"
castorev1pb "code.tvl.fyi/tvix/castore/protos"
"code.tvl.fyi/tvix/nar-bridge/pkg/reader"
storev1pb "code.tvl.fyi/tvix/store/protos"
"github.com/go-chi/chi/v5"
nixhash "github.com/nix-community/go-nix/pkg/hash"
"github.com/nix-community/go-nix/pkg/nixbase32"
@ -44,7 +44,7 @@ func registerNarPut(s *Server) {
pathInfo, err := rd.Import(
ctx,
genBlobServiceWriteCb(ctx, s.blobServiceClient),
func(directory *storev1pb.Directory) error {
func(directory *castorev1pb.Directory) error {
return directoriesUploader.Put(directory)
},
)

View file

@ -12,6 +12,7 @@ import (
"strings"
"sync"
castorev1pb "code.tvl.fyi/tvix/castore/protos"
storev1pb "code.tvl.fyi/tvix/store/protos"
"github.com/go-chi/chi/v5"
nixhash "github.com/nix-community/go-nix/pkg/hash"
@ -82,11 +83,11 @@ func renderNarinfo(
// extract the name of the node in the pathInfo structure, which will become the output path
var nodeName []byte
switch v := (pathInfo.GetNode().GetNode()).(type) {
case *storev1pb.Node_File:
case *castorev1pb.Node_File:
nodeName = v.File.GetName()
case *storev1pb.Node_Symlink:
case *castorev1pb.Node_Symlink:
nodeName = v.Symlink.GetName()
case *storev1pb.Node_Directory:
case *castorev1pb.Node_Directory:
nodeName = v.Directory.GetName()
}

View file

@ -4,6 +4,7 @@ import (
"net/http"
"path"
castorev1pb "code.tvl.fyi/tvix/castore/protos"
storev1pb "code.tvl.fyi/tvix/store/protos"
"github.com/go-chi/chi/v5"
"github.com/nix-community/go-nix/pkg/narinfo"
@ -123,10 +124,10 @@ func registerNarinfoPut(s *Server) {
// We need to add the basename of the storepath from the .narinfo
// to the pathInfo to be sent.
switch v := (pathInfo.GetNode().GetNode()).(type) {
case *storev1pb.Node_File:
pathInfoToUpload.Node = &storev1pb.Node{
Node: &storev1pb.Node_File{
File: &storev1pb.FileNode{
case *castorev1pb.Node_File:
pathInfoToUpload.Node = &castorev1pb.Node{
Node: &castorev1pb.Node_File{
File: &castorev1pb.FileNode{
Name: []byte(path.Base(narInfo.StorePath)),
Digest: v.File.Digest,
Size: v.File.Size,
@ -134,19 +135,19 @@ func registerNarinfoPut(s *Server) {
},
},
}
case *storev1pb.Node_Symlink:
pathInfoToUpload.Node = &storev1pb.Node{
Node: &storev1pb.Node_Symlink{
Symlink: &storev1pb.SymlinkNode{
case *castorev1pb.Node_Symlink:
pathInfoToUpload.Node = &castorev1pb.Node{
Node: &castorev1pb.Node_Symlink{
Symlink: &castorev1pb.SymlinkNode{
Name: []byte(path.Base(narInfo.StorePath)),
Target: v.Symlink.Target,
},
},
}
case *storev1pb.Node_Directory:
pathInfoToUpload.Node = &storev1pb.Node{
Node: &storev1pb.Node_Directory{
Directory: &storev1pb.DirectoryNode{
case *castorev1pb.Node_Directory:
pathInfoToUpload.Node = &castorev1pb.Node{
Node: &castorev1pb.Node_Directory{
Directory: &castorev1pb.DirectoryNode{
Name: []byte(path.Base(narInfo.StorePath)),
Digest: v.Directory.Digest,
Size: v.Directory.Size,

View file

@ -7,6 +7,7 @@ import (
"sync"
"time"
castorev1pb "code.tvl.fyi/tvix/castore/protos"
storev1pb "code.tvl.fyi/tvix/store/protos"
"github.com/go-chi/chi/middleware"
"github.com/go-chi/chi/v5"
@ -17,8 +18,8 @@ type Server struct {
srv *http.Server
handler chi.Router
directoryServiceClient storev1pb.DirectoryServiceClient
blobServiceClient storev1pb.BlobServiceClient
directoryServiceClient castorev1pb.DirectoryServiceClient
blobServiceClient castorev1pb.BlobServiceClient
pathInfoServiceClient storev1pb.PathInfoServiceClient
// When uploading NAR files to a HTTP binary cache, the .nar
@ -32,8 +33,8 @@ type Server struct {
}
func New(
directoryServiceClient storev1pb.DirectoryServiceClient,
blobServiceClient storev1pb.BlobServiceClient,
directoryServiceClient castorev1pb.DirectoryServiceClient,
blobServiceClient castorev1pb.BlobServiceClient,
pathInfoServiceClient storev1pb.PathInfoServiceClient,
enableAccessLog bool,
priority int,