snix_castore/proto/
grpc_directoryservice_wrapper.rs

1use crate::directoryservice::DirectoryService;
2use crate::directoryservice::order_validator::{LeavesToRoot, OrderValidator};
3use crate::{B3Digest, DirectoryError, proto};
4use futures::stream::BoxStream;
5use futures::{StreamExt, TryStreamExt};
6use tonic::{Request, Response, Status, Streaming, async_trait};
7use tracing::{instrument, warn};
8
9pub struct GRPCDirectoryServiceWrapper<T> {
10    directory_service: T,
11}
12
13impl<T> GRPCDirectoryServiceWrapper<T> {
14    pub fn new(directory_service: T) -> Self {
15        Self { directory_service }
16    }
17}
18
19#[async_trait]
20impl<T> proto::directory_service_server::DirectoryService for GRPCDirectoryServiceWrapper<T>
21where
22    T: DirectoryService + Clone + Send + Sync + 'static,
23{
24    type GetStream = BoxStream<'static, tonic::Result<proto::Directory>>;
25
26    #[instrument(skip_all)]
27    async fn get(
28        &self,
29        request: Request<proto::GetDirectoryRequest>,
30    ) -> Result<Response<Self::GetStream>, Status> {
31        let req_inner = request.into_inner();
32
33        match &req_inner
34            .by_what
35            .ok_or_else(|| Status::invalid_argument("invalid by_what"))?
36        {
37            proto::get_directory_request::ByWhat::Digest(digest) => {
38                let digest: B3Digest = digest
39                    .clone()
40                    .try_into()
41                    .map_err(|_e| Status::invalid_argument("invalid digest length"))?;
42
43                let directory_service = self.directory_service.clone();
44
45                Ok(tonic::Response::new({
46                    async_stream::try_stream! {
47                        if !req_inner.recursive {
48                            let directory = directory_service
49                                .get(&digest)
50                                .await
51                                .map_err(|e| {
52                                    warn!(err = %e, directory.digest=%digest, "failed to get directory");
53                                    tonic::Status::new(tonic::Code::Internal, e.to_string())
54                                })?
55                                .ok_or_else(|| {
56                                    Status::not_found(format!("directory {digest} not found"))
57                                })?;
58
59                            yield directory.into();
60                        } else {
61                            // If recursive was requested, traverse via get_recursive.
62                            // We need to use some type acrobatics as prost wants streams with 'static lifetimes.
63                            let mut s = get_recursive_owned(std::sync::Arc::new(directory_service),digest)
64                                .map_ok(proto::Directory::from)
65                                .map_err(|e| tonic::Status::new(tonic::Code::Internal, e.to_string()));
66
67                            while let Some(directory) = s.try_next().await? {
68                                yield directory
69                            }
70                        }
71
72                    }.boxed()
73                }))
74            }
75        }
76    }
77
78    #[instrument(skip_all)]
79    async fn put(
80        &self,
81        request: Request<Streaming<proto::Directory>>,
82    ) -> Result<Response<proto::PutDirectoryResponse>, Status> {
83        let mut req_inner = request.into_inner();
84
85        // Validate all received directories.
86        let mut validator = LeavesToRoot::new();
87
88        // Insert into the backing DirectoryService as we receive.
89        let mut directory_putter = self.directory_service.put_multiple_start();
90
91        while let Some(directory) = req_inner.message().await? {
92            let directory: crate::Directory =
93                directory.try_into().map_err(|e: DirectoryError| {
94                    tonic::Status::new(tonic::Code::Internal, e.to_string())
95                })?;
96            validator
97                .try_accept(&directory)
98                .map_err(|e| tonic::Status::new(tonic::Code::Internal, e.to_string()))?;
99
100            directory_putter
101                .put(directory)
102                .await
103                .map_err(|e| tonic::Status::new(tonic::Code::Internal, e.to_string()))?;
104        }
105
106        // Finalize validator, checks connectivity.
107        validator
108            .finalize()
109            .map_err(|e| tonic::Status::new(tonic::Code::Internal, e.to_string()))?;
110
111        Ok(Response::new(proto::PutDirectoryResponse {
112            // Properly close the directory putter, returning any potential errors.
113            root_digest: directory_putter
114                .close()
115                .await
116                .map_err(|e| tonic::Status::new(tonic::Code::Internal, e.to_string()))?
117                .into(),
118        }))
119    }
120}
121
122/// The same as [DirectoryService::get_recursive], but returning a stream with a static lifetime.
123/// It's only used for the gRPC server wrapper, which requires static lifetimes.
124fn get_recursive_owned<S>(
125    svc: std::sync::Arc<S>,
126    root_directory_digest: B3Digest,
127) -> BoxStream<'static, Result<crate::Directory, crate::directoryservice::Error>>
128where
129    S: DirectoryService + 'static,
130{
131    async_stream::try_stream! {
132        let mut  directories = svc.get_recursive(&root_directory_digest);
133        while let Some(directory) = directories.try_next().await? {
134            yield directory;
135        }
136    }
137    .boxed()
138}
snix_castore/proto/grpc_directoryservice_wrapper.rs

snix_castore/proto/
grpc_directoryservice_wrapper.rs
