63 lines
2.0 KiB
Rust
63 lines
2.0 KiB
Rust
use tonic::{Request, Response, Status};
|
|
use crate::coordinator::coordinator_client::CoordinatorClient as GeneratedCoordinatorClient;
|
|
use crate::coordinator::{RegisterRequest, RegisterResponse, HeartbeatRequest, HeartbeatResponse, PeerRequest, PeerResponse};
|
|
|
|
pub struct CoordinatorClient {
|
|
client: GeneratedCoordinatorClient<tonic::transport::Channel>,
|
|
}
|
|
|
|
impl CoordinatorClient {
|
|
pub async fn new(coordinator_addr: String) -> Result<Self, tonic::transport::Error> {
|
|
let client = GeneratedCoordinatorClient::connect(coordinator_addr).await?;
|
|
Ok(Self { client })
|
|
}
|
|
|
|
pub async fn register_node(
|
|
&mut self,
|
|
network_key: String,
|
|
node_id: String,
|
|
public_key: String,
|
|
local_ip: String,
|
|
) -> Result<RegisterResponse, Status> {
|
|
let request = Request::new(RegisterRequest {
|
|
network_key,
|
|
node_id,
|
|
public_key,
|
|
local_ip,
|
|
});
|
|
let response = self.client.register_node(request).await?;
|
|
Ok(response.into_inner())
|
|
}
|
|
|
|
pub async fn heartbeat(
|
|
&mut self,
|
|
node_id: String,
|
|
session_token: String,
|
|
public_endpoint: String,
|
|
) -> Result<HeartbeatResponse, Status> {
|
|
let request = Request::new(HeartbeatRequest {
|
|
node_id,
|
|
session_token,
|
|
public_endpoint,
|
|
});
|
|
let response = self.client.heartbeat(request).await?;
|
|
Ok(response.into_inner())
|
|
}
|
|
|
|
pub async fn get_peer_endpoint(
|
|
&mut self,
|
|
target_node_id: String,
|
|
session_token: String,
|
|
) -> Result<PeerResponse, Status> {
|
|
let request = Request::new(PeerRequest {
|
|
target_node_id,
|
|
session_token,
|
|
});
|
|
let response = self.client.get_peer_endpoint(request).await?;
|
|
Ok(response.into_inner())
|
|
}
|
|
}
|
|
|
|
// Optional: You can add a method to update node name if needed in the future.
|
|
// pub async fn update_node_name(&mut self, ...) -> ... { ... }
|