/build/source/nativelink-worker/src/worker_api_client_wrapper.rs
Line | Count | Source (jump to first uncovered line) |
1 | | // Copyright 2024 The NativeLink Authors. All rights reserved. |
2 | | // |
3 | | // Licensed under the Apache License, Version 2.0 (the "License"); |
4 | | // you may not use this file except in compliance with the License. |
5 | | // You may obtain a copy of the License at |
6 | | // |
7 | | // http://www.apache.org/licenses/LICENSE-2.0 |
8 | | // |
9 | | // Unless required by applicable law or agreed to in writing, software |
10 | | // distributed under the License is distributed on an "AS IS" BASIS, |
11 | | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | | // See the License for the specific language governing permissions and |
13 | | // limitations under the License. |
14 | | |
15 | | use std::future::Future; |
16 | | |
17 | | use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::worker_api_client::WorkerApiClient; |
18 | | use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{ |
19 | | ExecuteResult, GoingAwayRequest, KeepAliveRequest, SupportedProperties, UpdateForWorker, |
20 | | }; |
21 | | use tonic::codec::Streaming; |
22 | | use tonic::transport::Channel; |
23 | | use tonic::{Response, Status}; |
24 | | |
25 | | /// This is used in order to allow unit tests to intercept these calls. This should always match |
26 | | /// the API of WorkerApiClient defined in the worker_api.proto file. |
27 | | pub trait WorkerApiClientTrait: Clone + Sync + Send + Sized + Unpin { |
28 | | fn connect_worker( |
29 | | &mut self, |
30 | | request: SupportedProperties, |
31 | | ) -> impl Future<Output = Result<Response<Streaming<UpdateForWorker>>, Status>> + Send; |
32 | | |
33 | | fn keep_alive( |
34 | | &mut self, |
35 | | request: KeepAliveRequest, |
36 | | ) -> impl Future<Output = Result<Response<()>, Status>> + Send; |
37 | | |
38 | | fn going_away( |
39 | | &mut self, |
40 | | request: GoingAwayRequest, |
41 | | ) -> impl Future<Output = Result<Response<()>, Status>> + Send; |
42 | | |
43 | | fn execution_response( |
44 | | &mut self, |
45 | | request: ExecuteResult, |
46 | | ) -> impl Future<Output = Result<Response<()>, Status>> + Send; |
47 | | } |
48 | | |
49 | | #[derive(Clone)] |
50 | | pub struct WorkerApiClientWrapper { |
51 | | inner: WorkerApiClient<Channel>, |
52 | | } |
53 | | |
54 | | impl From<WorkerApiClient<Channel>> for WorkerApiClientWrapper { |
55 | 0 | fn from(other: WorkerApiClient<Channel>) -> Self { |
56 | 0 | Self { inner: other } |
57 | 0 | } |
58 | | } |
59 | | |
60 | | impl WorkerApiClientTrait for WorkerApiClientWrapper { |
61 | 0 | async fn connect_worker( |
62 | 0 | &mut self, |
63 | 0 | request: SupportedProperties, |
64 | 0 | ) -> Result<Response<Streaming<UpdateForWorker>>, Status> { |
65 | 0 | self.inner.connect_worker(request).await |
66 | 0 | } |
67 | | |
68 | 0 | async fn keep_alive(&mut self, request: KeepAliveRequest) -> Result<Response<()>, Status> { |
69 | 0 | self.inner.keep_alive(request).await |
70 | 0 | } |
71 | | |
72 | 0 | async fn going_away(&mut self, request: GoingAwayRequest) -> Result<Response<()>, Status> { |
73 | 0 | self.inner.going_away(request).await |
74 | 0 | } |
75 | | |
76 | 0 | async fn execution_response(&mut self, request: ExecuteResult) -> Result<Response<()>, Status> { |
77 | 0 | self.inner.execution_response(request).await |
78 | 0 | } |
79 | | } |