refactor(tvix/store/pathinfoservice): make more generic

We don't need Arcs in most of the cases, we're fine with some container.

Change-Id: Ic4f8acb5b9d93e2b0923bb607463fb91e9d0e4fe
Reviewed-on: https://cl.tvl.fyi/c/depot/+/10606
Autosubmit: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI
Reviewed-by: raitobezarius <tvl@lahfa.xyz>
This commit is contained in:
Florian Klink 2024-01-12 12:18:43 +02:00 committed by clbot
parent 7d51193f7d
commit b59df53774
6 changed files with 51 additions and 54 deletions

View file

@ -159,8 +159,8 @@ mod tests {
let router = server.add_service(
crate::proto::path_info_service_server::PathInfoServiceServer::new(
GRPCPathInfoServiceWrapper::new(Box::new(MemoryPathInfoService::new(
gen_blob_service().into(),
gen_directory_service().into(),
gen_blob_service(),
gen_directory_service(),
))
as Box<dyn PathInfoService>),
),

View file

@ -11,18 +11,15 @@ use tvix_castore::proto as castorepb;
use tvix_castore::Error;
use tvix_castore::{blobservice::BlobService, directoryservice::DirectoryService};
pub struct MemoryPathInfoService {
pub struct MemoryPathInfoService<BS, DS> {
db: Arc<RwLock<HashMap<[u8; 20], PathInfo>>>,
blob_service: Arc<dyn BlobService>,
directory_service: Arc<dyn DirectoryService>,
blob_service: BS,
directory_service: DS,
}
impl MemoryPathInfoService {
pub fn new(
blob_service: Arc<dyn BlobService>,
directory_service: Arc<dyn DirectoryService>,
) -> Self {
impl<BS, DS> MemoryPathInfoService<BS, DS> {
pub fn new(blob_service: BS, directory_service: DS) -> Self {
Self {
db: Default::default(),
blob_service,
@ -32,7 +29,11 @@ impl MemoryPathInfoService {
}
#[async_trait]
impl PathInfoService for MemoryPathInfoService {
impl<BS, DS> PathInfoService for MemoryPathInfoService<BS, DS>
where
BS: AsRef<dyn BlobService> + Send + Sync,
DS: AsRef<dyn DirectoryService> + Send + Sync,
{
async fn get(&self, digest: [u8; 20]) -> Result<Option<PathInfo>, Error> {
let db = self.db.read().unwrap();
@ -65,13 +66,9 @@ impl PathInfoService for MemoryPathInfoService {
&self,
root_node: &castorepb::node::Node,
) -> Result<(u64, [u8; 32]), Error> {
calculate_size_and_sha256(
root_node,
self.blob_service.clone(),
self.directory_service.clone(),
)
.await
.map_err(|e| Error::StorageError(e.to_string()))
calculate_size_and_sha256(root_node, &self.blob_service, &self.directory_service)
.await
.map_err(|e| Error::StorageError(e.to_string()))
}
fn list(&self) -> Pin<Box<dyn Stream<Item = Result<PathInfo, Error>> + Send>> {

View file

@ -1,7 +1,6 @@
use std::{
io::{self, BufRead, Read, Write},
pin::Pin,
sync::Arc,
};
use data_encoding::BASE64;
@ -38,24 +37,20 @@ use super::PathInfoService;
/// [PathInfoService::put] and [PathInfoService::calculate_nar] are not
/// implemented and return an error if called.
/// TODO: what about reading from nix-cache-info?
pub struct NixHTTPPathInfoService {
pub struct NixHTTPPathInfoService<BS, DS> {
base_url: url::Url,
http_client: reqwest::Client,
blob_service: Arc<dyn BlobService>,
directory_service: Arc<dyn DirectoryService>,
blob_service: BS,
directory_service: DS,
/// An optional list of [narinfo::PubKey].
/// If set, the .narinfo files received need to have correct signature by at least one of these.
public_keys: Option<Vec<narinfo::PubKey>>,
}
impl NixHTTPPathInfoService {
pub fn new(
base_url: url::Url,
blob_service: Arc<dyn BlobService>,
directory_service: Arc<dyn DirectoryService>,
) -> Self {
impl<BS, DS> NixHTTPPathInfoService<BS, DS> {
pub fn new(base_url: url::Url, blob_service: BS, directory_service: DS) -> Self {
Self {
base_url,
http_client: reqwest::Client::new(),
@ -73,7 +68,11 @@ impl NixHTTPPathInfoService {
}
#[async_trait]
impl PathInfoService for NixHTTPPathInfoService {
impl<BS, DS> PathInfoService for NixHTTPPathInfoService<BS, DS>
where
BS: AsRef<dyn BlobService> + Send + Sync + Clone + 'static,
DS: AsRef<dyn DirectoryService> + Send + Sync + Clone + 'static,
{
#[instrument(skip_all, err, fields(path.digest=BASE64.encode(&digest)))]
async fn get(&self, digest: [u8; 20]) -> Result<Option<PathInfo>, Error> {
let narinfo_url = self

View file

@ -3,7 +3,7 @@ use crate::nar::calculate_size_and_sha256;
use crate::proto::PathInfo;
use futures::{stream::iter, Stream};
use prost::Message;
use std::{path::Path, pin::Pin, sync::Arc};
use std::{path::Path, pin::Pin};
use tonic::async_trait;
use tracing::warn;
use tvix_castore::proto as castorepb;
@ -13,18 +13,18 @@ use tvix_castore::{blobservice::BlobService, directoryservice::DirectoryService,
///
/// The PathInfo messages are stored as encoded protos, and keyed by their output hash,
/// as that's currently the only request type available.
pub struct SledPathInfoService {
pub struct SledPathInfoService<BS, DS> {
db: sled::Db,
blob_service: Arc<dyn BlobService>,
directory_service: Arc<dyn DirectoryService>,
blob_service: BS,
directory_service: DS,
}
impl SledPathInfoService {
impl<BS, DS> SledPathInfoService<BS, DS> {
pub fn new<P: AsRef<Path>>(
p: P,
blob_service: Arc<dyn BlobService>,
directory_service: Arc<dyn DirectoryService>,
blob_service: BS,
directory_service: DS,
) -> Result<Self, sled::Error> {
let config = sled::Config::default()
.use_compression(false) // is a required parameter
@ -38,10 +38,7 @@ impl SledPathInfoService {
})
}
pub fn new_temporary(
blob_service: Arc<dyn BlobService>,
directory_service: Arc<dyn DirectoryService>,
) -> Result<Self, sled::Error> {
pub fn new_temporary(blob_service: BS, directory_service: DS) -> Result<Self, sled::Error> {
let config = sled::Config::default().temporary(true);
let db = config.open()?;
@ -54,7 +51,11 @@ impl SledPathInfoService {
}
#[async_trait]
impl PathInfoService for SledPathInfoService {
impl<BS, DS> PathInfoService for SledPathInfoService<BS, DS>
where
BS: AsRef<dyn BlobService> + Send + Sync,
DS: AsRef<dyn DirectoryService> + Send + Sync,
{
async fn get(&self, digest: [u8; 20]) -> Result<Option<PathInfo>, Error> {
match self.db.get(digest) {
Ok(None) => Ok(None),
@ -106,13 +107,9 @@ impl PathInfoService for SledPathInfoService {
&self,
root_node: &castorepb::node::Node,
) -> Result<(u64, [u8; 32]), Error> {
calculate_size_and_sha256(
root_node,
self.blob_service.clone(),
self.directory_service.clone(),
)
.await
.map_err(|e| Error::StorageError(e.to_string()))
calculate_size_and_sha256(root_node, &self.blob_service, &self.directory_service)
.await
.map_err(|e| Error::StorageError(e.to_string()))
}
fn list(&self) -> Pin<Box<dyn Stream<Item = Result<PathInfo, Error>> + Send>> {

View file

@ -22,8 +22,8 @@ fn gen_grpc_service(
let blob_service = gen_blob_service();
let directory_service = gen_directory_service();
Arc::new(GRPCPathInfoServiceWrapper::new(gen_pathinfo_service(
blob_service.into(),
directory_service.into(),
blob_service,
directory_service,
)))
}

View file

@ -4,9 +4,13 @@ use tvix_castore::{blobservice::BlobService, directoryservice::DirectoryService}
pub use tvix_castore::utils::*;
pub fn gen_pathinfo_service(
blob_service: Arc<dyn BlobService>,
directory_service: Arc<dyn DirectoryService>,
) -> Arc<dyn PathInfoService> {
pub fn gen_pathinfo_service<BS, DS>(
blob_service: BS,
directory_service: DS,
) -> Arc<dyn PathInfoService>
where
BS: AsRef<dyn BlobService> + Send + Sync + 'static,
DS: AsRef<dyn DirectoryService> + Send + Sync + 'static,
{
Arc::new(MemoryPathInfoService::new(blob_service, directory_service))
}