impl UDP sward
This commit is contained in:
2
Cargo.lock
generated
2
Cargo.lock
generated
@@ -1568,9 +1568,11 @@ version = "0.1.0"
|
|||||||
dependencies = [
|
dependencies = [
|
||||||
"bytes",
|
"bytes",
|
||||||
"compact_str",
|
"compact_str",
|
||||||
|
"libc",
|
||||||
"rand",
|
"rand",
|
||||||
"regex",
|
"regex",
|
||||||
"reqwest",
|
"reqwest",
|
||||||
|
"socket2",
|
||||||
"thiserror",
|
"thiserror",
|
||||||
"tokio",
|
"tokio",
|
||||||
"tokio-util",
|
"tokio-util",
|
||||||
|
|||||||
@@ -12,6 +12,8 @@ rand = {workspace = true}
|
|||||||
thiserror = {workspace = true}
|
thiserror = {workspace = true}
|
||||||
compact_str = {workspace = true}
|
compact_str = {workspace = true}
|
||||||
bytes = {workspace = true}
|
bytes = {workspace = true}
|
||||||
|
libc = "0.2"
|
||||||
|
socket2 = "0.6"
|
||||||
wyrand = "0.3"
|
wyrand = "0.3"
|
||||||
regex = "1.11"
|
regex = "1.11"
|
||||||
tokio-util = "0.7"
|
tokio-util = "0.7"
|
||||||
@@ -0,0 +1,69 @@
|
|||||||
|
use std::pin::Pin;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::task::{Context, Poll};
|
||||||
|
use tokio::net::UdpSocket;
|
||||||
|
use tower::Service;
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct UdpSward {
|
||||||
|
socket: Arc<UdpSocket>,
|
||||||
|
sent_count: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl UdpSward {
|
||||||
|
pub fn new(socket: Arc<UdpSocket>) -> Self {
|
||||||
|
Self {
|
||||||
|
socket,
|
||||||
|
sent_count: 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Copy)]
|
||||||
|
pub struct SizedUdpRequest<const N: usize> {
|
||||||
|
pub bytes: [u8; N],
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<const N: usize> Service<SizedUdpRequest<N>> for UdpSward {
|
||||||
|
type Response = usize;
|
||||||
|
type Error = std::io::Error;
|
||||||
|
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
|
||||||
|
|
||||||
|
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||||
|
Poll::Ready(Ok(()))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn call(&mut self, req: SizedUdpRequest<N>) -> Self::Future {
|
||||||
|
self.sent_count += 1;
|
||||||
|
let socket = self.socket.clone();
|
||||||
|
Box::pin(async move {
|
||||||
|
socket.send(&req.bytes).await
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct BoxedUdpRequest(pub Box<[u8]>);
|
||||||
|
|
||||||
|
impl Service<BoxedUdpRequest> for UdpSward {
|
||||||
|
type Response = usize;
|
||||||
|
type Error = std::io::Error;
|
||||||
|
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
|
||||||
|
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||||
|
Poll::Ready(Ok(()))
|
||||||
|
}
|
||||||
|
fn call(&mut self, req: BoxedUdpRequest) -> Self::Future {
|
||||||
|
self.sent_count += 1;
|
||||||
|
let socket = self.socket.clone();
|
||||||
|
Box::pin(async move {
|
||||||
|
socket.send(&req.0).await
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl tower::load::Load for UdpSward {
|
||||||
|
type Metric = usize;
|
||||||
|
fn load(&self) -> Self::Metric {
|
||||||
|
self.sent_count
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user