|
1 | | -use crate::api::v2ray_rust_api::{GetStatsRequest, GetStatsResponse}; |
| 1 | +use crate::api::v2ray_rust_api::{ |
| 2 | + GetLatencyRequest, GetLatencyResponse, GetStatsRequest, GetStatsResponse, |
| 3 | +}; |
2 | 4 | use crate::config::COUNTER_MAP; |
| 5 | +use crate::proxy::{Address, ChainStreamBuilder}; |
3 | 6 |
|
| 7 | +use std::collections::HashMap; |
| 8 | +use std::str::FromStr; |
4 | 9 | use std::sync::atomic::Ordering::Relaxed; |
| 10 | +use std::sync::Arc; |
| 11 | +use std::time::{Duration, Instant}; |
5 | 12 |
|
| 13 | +use futures_util::future::join_all; |
| 14 | +use tokio::io::AsyncWriteExt; |
| 15 | +use tokio::time::timeout; |
6 | 16 | use tonic::{Request, Response, Status}; |
7 | 17 |
|
8 | 18 | pub mod v2ray_rust_api { |
9 | 19 | tonic::include_proto!("v2ray.core.app.stats.command"); |
10 | 20 | } |
| 21 | + |
| 22 | +use v2ray_rust_api::latency_service_server::{LatencyService, LatencyServiceServer}; |
11 | 23 | use v2ray_rust_api::stats_service_server::{StatsService, StatsServiceServer}; |
12 | 24 |
|
13 | 25 | #[derive(Default)] |
@@ -45,3 +57,79 @@ impl StatsService for ApiServer { |
45 | 57 | })) |
46 | 58 | } |
47 | 59 | } |
| 60 | + |
| 61 | +pub struct ApiLatencyServer { |
| 62 | + inner_map: Arc<HashMap<String, ChainStreamBuilder>>, |
| 63 | +} |
| 64 | +impl ApiLatencyServer { |
| 65 | + pub(crate) fn new_server( |
| 66 | + inner_map: Arc<HashMap<String, ChainStreamBuilder>>, |
| 67 | + ) -> LatencyServiceServer<Self> { |
| 68 | + LatencyServiceServer::new(Self { inner_map }) |
| 69 | + } |
| 70 | +} |
| 71 | +#[tonic::async_trait] |
| 72 | +impl LatencyService for ApiLatencyServer { |
| 73 | + async fn get_latency( |
| 74 | + &self, |
| 75 | + request: Request<GetLatencyRequest>, |
| 76 | + ) -> Result<Response<GetLatencyResponse>, Status> { |
| 77 | + if let Some(req) = &request.get_ref().outbound_name { |
| 78 | + if let Some(b) = self.inner_map.get(req) { |
| 79 | + let start = Instant::now(); |
| 80 | + let stream = b |
| 81 | + .build_tcp( |
| 82 | + Address::from_str(request.get_ref().test_url.as_str()).map_err(|_| { |
| 83 | + Status::new(tonic::Code::InvalidArgument, "test_url is invalid") |
| 84 | + })?, |
| 85 | + ) |
| 86 | + .await; |
| 87 | + let duration = start.elapsed(); |
| 88 | + let mut latency_res = HashMap::new(); |
| 89 | + if stream.is_err() { |
| 90 | + latency_res.insert(req.clone(), -1i64); |
| 91 | + return Ok(Response::new(GetLatencyResponse { latency_res })); |
| 92 | + } else { |
| 93 | + latency_res.insert(req.clone(), duration.as_millis() as i64); |
| 94 | + return Ok(Response::new(GetLatencyResponse { latency_res })); |
| 95 | + } |
| 96 | + } else { |
| 97 | + return Err(Status::new(tonic::Code::InvalidArgument, "name is invalid")); |
| 98 | + } |
| 99 | + } else { |
| 100 | + let mut vec_fut = Vec::new(); |
| 101 | + let test_url = request.get_ref().test_url.as_str(); |
| 102 | + let addr = Address::from_str(test_url) |
| 103 | + .map_err(|_| std::io::Error::new(std::io::ErrorKind::Other, "invalid test url"))?; |
| 104 | + for (name, _) in self.inner_map.iter() { |
| 105 | + let name = name.clone(); |
| 106 | + let addr = addr.clone(); |
| 107 | + vec_fut.push(async move { |
| 108 | + let b = self.inner_map.get(&name).unwrap(); |
| 109 | + let start = Instant::now(); |
| 110 | + let timeout_stream = timeout(Duration::from_secs(5), async move { |
| 111 | + b.build_tcp(addr).await?.write_u128(u128::MAX).await |
| 112 | + }) |
| 113 | + .await; |
| 114 | + let duration = start.elapsed(); |
| 115 | + if timeout_stream.is_err() { |
| 116 | + return Ok::<(String, i64), std::io::Error>((name, -1i64)); |
| 117 | + } else { |
| 118 | + let stream = timeout_stream?; |
| 119 | + if stream.is_err() { |
| 120 | + return Ok((name, -1i64)); |
| 121 | + } |
| 122 | + return Ok((name, duration.as_millis() as i64)); |
| 123 | + } |
| 124 | + }); |
| 125 | + } |
| 126 | + let vec_res = join_all(vec_fut).await; |
| 127 | + let mut latency_res = HashMap::new(); |
| 128 | + for v in vec_res.into_iter() { |
| 129 | + let (k, v) = v?; |
| 130 | + latency_res.insert(k, v); |
| 131 | + } |
| 132 | + return Ok(Response::new(GetLatencyResponse { latency_res })); |
| 133 | + } |
| 134 | + } |
| 135 | +} |
0 commit comments