/build/source/nativelink-worker/src/worker_api_client_wrapper.rs
Line | Count | Source |
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 core::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 | | ConnectWorkerRequest, ExecuteComplete, ExecuteResult, GoingAwayRequest, KeepAliveRequest, |
20 | | UpdateForWorker, |
21 | | }; |
22 | | use tonic::codec::Streaming; |
23 | | use tonic::transport::Channel; |
24 | | use tonic::{Response, Status}; |
25 | | |
26 | | /// This is used in order to allow unit tests to intercept these calls. This should always match |
27 | | /// the API of `WorkerApiClient` defined in the `worker_api.proto` file. |
28 | | pub trait WorkerApiClientTrait: Clone + Sync + Send + Sized + Unpin { |
29 | | fn connect_worker( |
30 | | &mut self, |
31 | | request: ConnectWorkerRequest, |
32 | | ) -> impl Future<Output = Result<Response<Streaming<UpdateForWorker>>, Status>> + Send; |
33 | | |
34 | | fn keep_alive( |
35 | | &mut self, |
36 | | request: KeepAliveRequest, |
37 | | ) -> impl Future<Output = Result<Response<()>, Status>> + Send; |
38 | | |
39 | | fn going_away( |
40 | | &mut self, |
41 | | request: GoingAwayRequest, |
42 | | ) -> impl Future<Output = Result<Response<()>, Status>> + Send; |
43 | | |
44 | | fn execution_response( |
45 | | &mut self, |
46 | | request: ExecuteResult, |
47 | | ) -> impl Future<Output = Result<Response<()>, Status>> + Send; |
48 | | |
49 | | fn execution_complete( |
50 | | &mut self, |
51 | | request: ExecuteComplete, |
52 | | ) -> impl Future<Output = Result<Response<()>, Status>> + Send; |
53 | | } |
54 | | |
55 | | #[derive(Debug, Clone)] |
56 | | pub struct WorkerApiClientWrapper { |
57 | | inner: WorkerApiClient<Channel>, |
58 | | } |
59 | | |
60 | | impl From<WorkerApiClient<Channel>> for WorkerApiClientWrapper { |
61 | 0 | fn from(other: WorkerApiClient<Channel>) -> Self { |
62 | 0 | Self { inner: other } |
63 | 0 | } |
64 | | } |
65 | | |
66 | | impl WorkerApiClientTrait for WorkerApiClientWrapper { |
67 | 0 | async fn connect_worker( |
68 | 0 | &mut self, |
69 | 0 | request: ConnectWorkerRequest, |
70 | 0 | ) -> Result<Response<Streaming<UpdateForWorker>>, Status> { |
71 | 0 | self.inner.connect_worker(request).await |
72 | 0 | } |
73 | | |
74 | 0 | async fn keep_alive(&mut self, request: KeepAliveRequest) -> Result<Response<()>, Status> { |
75 | 0 | self.inner.keep_alive(request).await |
76 | 0 | } |
77 | | |
78 | 0 | async fn going_away(&mut self, request: GoingAwayRequest) -> Result<Response<()>, Status> { |
79 | 0 | self.inner.going_away(request).await |
80 | 0 | } |
81 | | |
82 | 0 | async fn execution_response(&mut self, request: ExecuteResult) -> Result<Response<()>, Status> { |
83 | 0 | self.inner.execution_response(request).await |
84 | 0 | } |
85 | | |
86 | 0 | async fn execution_complete( |
87 | 0 | &mut self, |
88 | 0 | request: ExecuteComplete, |
89 | 0 | ) -> Result<Response<()>, Status> { |
90 | 0 | self.inner.execution_complete(request).await |
91 | 0 | } |
92 | | } |