Move tokio_kcp to local files
This commit is contained in:
KernelErr
256
src/plugins/kcp/session.rs
Normal file
256
src/plugins/kcp/session.rs
Normal file
@@ -0,0 +1,256 @@
|
||||
use std::{
|
||||
collections::{hash_map::Entry, HashMap},
|
||||
net::SocketAddr,
|
||||
sync::{
|
||||
atomic::{AtomicBool, Ordering},
|
||||
Arc,
|
||||
},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use byte_string::ByteStr;
|
||||
use kcp::KcpResult;
|
||||
use log::{error, trace};
|
||||
use tokio::{
|
||||
net::UdpSocket,
|
||||
sync::{mpsc, Mutex},
|
||||
time::{self, Instant},
|
||||
};
|
||||
|
||||
use crate::plugins::kcp::{skcp::KcpSocket, KcpConfig};
|
||||
|
||||
pub struct KcpSession {
|
||||
socket: Mutex<KcpSocket>,
|
||||
closed: AtomicBool,
|
||||
session_expire: Duration,
|
||||
session_close_notifier: Option<mpsc::Sender<u32>>,
|
||||
input_tx: mpsc::Sender<Vec<u8>>,
|
||||
}
|
||||
|
||||
impl KcpSession {
|
||||
fn new(
|
||||
socket: KcpSocket,
|
||||
session_expire: Duration,
|
||||
session_close_notifier: Option<mpsc::Sender<u32>>,
|
||||
input_tx: mpsc::Sender<Vec<u8>>,
|
||||
) -> KcpSession {
|
||||
KcpSession {
|
||||
socket: Mutex::new(socket),
|
||||
closed: AtomicBool::new(false),
|
||||
session_expire,
|
||||
session_close_notifier,
|
||||
input_tx,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn new_shared(
|
||||
socket: KcpSocket,
|
||||
session_expire: Duration,
|
||||
session_close_notifier: Option<mpsc::Sender<u32>>,
|
||||
) -> Arc<KcpSession> {
|
||||
let is_client = session_close_notifier.is_none();
|
||||
|
||||
let (input_tx, mut input_rx) = mpsc::channel(64);
|
||||
|
||||
let udp_socket = socket.udp_socket().clone();
|
||||
|
||||
let session = Arc::new(KcpSession::new(
|
||||
socket,
|
||||
session_expire,
|
||||
session_close_notifier,
|
||||
input_tx,
|
||||
));
|
||||
|
||||
{
|
||||
let session = session.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut input_buffer = [0u8; 65536];
|
||||
let update_timer = time::sleep(Duration::from_millis(10));
|
||||
tokio::pin!(update_timer);
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
// recv() then input()
|
||||
// Drives the KCP machine forward
|
||||
recv_result = udp_socket.recv(&mut input_buffer), if is_client => {
|
||||
match recv_result {
|
||||
Err(err) => {
|
||||
error!("[SESSION] UDP recv failed, error: {}", err);
|
||||
}
|
||||
Ok(n) => {
|
||||
let input_buffer = &input_buffer[..n];
|
||||
trace!("[SESSION] UDP recv {} bytes, going to input {:?}", n, ByteStr::new(input_buffer));
|
||||
|
||||
let mut socket = session.socket.lock().await;
|
||||
|
||||
match socket.input(input_buffer) {
|
||||
Ok(true) => {
|
||||
trace!("[SESSION] UDP input {} bytes and waked sender/receiver", n);
|
||||
}
|
||||
Ok(false) => {}
|
||||
Err(err) => {
|
||||
error!("[SESSION] UDP input {} bytes error: {}, input buffer {:?}", n, err, ByteStr::new(input_buffer));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// bytes received from listener socket
|
||||
input_opt = input_rx.recv() => {
|
||||
if let Some(input_buffer) = input_opt {
|
||||
let mut socket = session.socket.lock().await;
|
||||
match socket.input(&input_buffer) {
|
||||
Ok(..) => {
|
||||
trace!("[SESSION] UDP input {} bytes from channel {:?}", input_buffer.len(), ByteStr::new(&input_buffer));
|
||||
}
|
||||
Err(err) => {
|
||||
error!("[SESSION] UDP input {} bytes from channel failed, error: {}, input buffer {:?}",
|
||||
input_buffer.len(), err, ByteStr::new(&input_buffer));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Call update() in period
|
||||
_ = &mut update_timer => {
|
||||
let mut socket = session.socket.lock().await;
|
||||
|
||||
let is_closed = session.closed.load(Ordering::Acquire);
|
||||
if is_closed && socket.can_close() {
|
||||
trace!("[SESSION] KCP session closed");
|
||||
break;
|
||||
}
|
||||
|
||||
// server socket expires
|
||||
if !is_client {
|
||||
// If this is a server stream, close it automatically after a period of time
|
||||
let last_update_time = socket.last_update_time();
|
||||
let elapsed = last_update_time.elapsed();
|
||||
|
||||
if elapsed > session.session_expire {
|
||||
if elapsed > session.session_expire * 2 {
|
||||
// Force close. Client may have already gone.
|
||||
trace!(
|
||||
"[SESSION] force close inactive session, conv: {}, last_update: {}s ago",
|
||||
socket.conv(),
|
||||
elapsed.as_secs()
|
||||
);
|
||||
break;
|
||||
}
|
||||
|
||||
if !is_closed {
|
||||
trace!(
|
||||
"[SESSION] closing inactive session, conv: {}, last_update: {}s ago",
|
||||
socket.conv(),
|
||||
elapsed.as_secs()
|
||||
);
|
||||
session.closed.store(true, Ordering::Release);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
match socket.update() {
|
||||
Ok(next_next) => {
|
||||
update_timer.as_mut().reset(Instant::from_std(next_next));
|
||||
}
|
||||
Err(err) => {
|
||||
error!("[SESSION] KCP update failed, error: {}", err);
|
||||
update_timer.as_mut().reset(Instant::now() + Duration::from_millis(10));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
// Close the socket.
|
||||
// Wake all pending tasks and let all send/recv return EOF
|
||||
|
||||
let mut socket = session.socket.lock().await;
|
||||
socket.close();
|
||||
}
|
||||
|
||||
if let Some(ref notifier) = session.session_close_notifier {
|
||||
let socket = session.socket.lock().await;
|
||||
let _ = notifier.send(socket.conv()).await;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
session
|
||||
}
|
||||
|
||||
pub fn kcp_socket(&self) -> &Mutex<KcpSocket> {
|
||||
&self.socket
|
||||
}
|
||||
|
||||
pub fn close(&self) {
|
||||
self.closed.store(true, Ordering::Release);
|
||||
}
|
||||
|
||||
pub async fn input(&self, buf: &[u8]) {
|
||||
self.input_tx
|
||||
.send(buf.to_owned())
|
||||
.await
|
||||
.expect("input channel closed")
|
||||
}
|
||||
}
|
||||
|
||||
pub struct KcpSessionManager {
|
||||
sessions: HashMap<u32, Arc<KcpSession>>,
|
||||
next_free_conv: u32,
|
||||
}
|
||||
|
||||
impl KcpSessionManager {
|
||||
pub fn new() -> KcpSessionManager {
|
||||
KcpSessionManager {
|
||||
sessions: HashMap::new(),
|
||||
next_free_conv: 0,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn close_conv(&mut self, conv: u32) {
|
||||
self.sessions.remove(&conv);
|
||||
}
|
||||
|
||||
pub fn alloc_conv(&mut self) -> u32 {
|
||||
loop {
|
||||
let (mut c, _) = self.next_free_conv.overflowing_add(1);
|
||||
if c == 0 {
|
||||
let (nc, _) = c.overflowing_add(1);
|
||||
c = nc;
|
||||
}
|
||||
self.next_free_conv = c;
|
||||
|
||||
if self.sessions.get(&self.next_free_conv).is_none() {
|
||||
let conv = self.next_free_conv;
|
||||
return conv;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_or_create(
|
||||
&mut self,
|
||||
config: &KcpConfig,
|
||||
conv: u32,
|
||||
udp: &Arc<UdpSocket>,
|
||||
peer_addr: SocketAddr,
|
||||
session_close_notifier: &mpsc::Sender<u32>,
|
||||
) -> KcpResult<(Arc<KcpSession>, bool)> {
|
||||
match self.sessions.entry(conv) {
|
||||
Entry::Occupied(occ) => Ok((occ.get().clone(), false)),
|
||||
Entry::Vacant(vac) => {
|
||||
let socket = KcpSocket::new(config, conv, udp.clone(), peer_addr, config.stream)?;
|
||||
let session = KcpSession::new_shared(
|
||||
socket,
|
||||
config.session_expire,
|
||||
Some(session_close_notifier.clone()),
|
||||
);
|
||||
trace!("created session for conv: {}, peer: {}", conv, peer_addr);
|
||||
vac.insert(session.clone());
|
||||
Ok((session, true))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user