ptth/crates/ptth_relay/src/relay_state.rs

122 lines
2.6 KiB
Rust

use std::{
collections::HashMap,
convert::TryFrom,
};
use chrono::{DateTime, Utc};
use dashmap::DashMap;
use tokio::sync::{
Mutex,
RwLock,
oneshot,
watch,
};
use crate::{
Body,
Config,
RelayError,
ShuttingDownError,
};
use ptth_core::http_serde;
/*
Here's what we need to handle:
When a request comes in:
- Park the client in response_rendezvous
- Look up the server ID in request_rendezvous
- If a server is parked, unpark it and send the request
- Otherwise, queue the request
When a server comes to listen:
- Look up the server ID in request_rendezvous
- Either return all pending requests, or park the server
When a server comes to respond:
- Look up the parked client in response_rendezvous
- Unpark the client and begin streaming
So we need these lookups to be fast:
- Server IDs, where (1 server) or (0 or many clients)
can be parked
- Request IDs, where 1 client is parked
*/
pub enum RequestRendezvous {
ParkedClients (Vec <http_serde::WrappedRequest>),
ParkedServer (oneshot::Sender <Result <http_serde::WrappedRequest, ShuttingDownError>>),
}
type ResponseRendezvous = oneshot::Sender <Result <(http_serde::ResponseParts, Body), ShuttingDownError>>;
#[derive (Clone)]
pub struct ServerStatus {
pub last_seen: DateTime <Utc>,
}
impl Default for ServerStatus {
fn default () -> Self {
Self {
last_seen: Utc::now (),
}
}
}
#[derive (Clone)]
pub struct RejectedServer {
pub name: String,
pub tripcode: blake3::Hash,
pub seen: DateTime <Utc>,
}
pub struct RelayState {
pub config: RwLock <Config>,
// Key: Server ID
pub request_rendezvous: Mutex <HashMap <String, RequestRendezvous>>,
pub server_status: Mutex <HashMap <String, ServerStatus>>,
// Key: Request ID
pub response_rendezvous: RwLock <DashMap <String, ResponseRendezvous>>,
pub shutdown_watch_tx: watch::Sender <bool>,
pub shutdown_watch_rx: watch::Receiver <bool>,
// List of recently rejected server names (Used to approve servers)
pub unregistered_servers: RwLock <Vec <RejectedServer>>,
}
impl TryFrom <Config> for RelayState {
type Error = RelayError;
fn try_from (config: Config) -> Result <Self, Self::Error> {
let (shutdown_watch_tx, shutdown_watch_rx) = watch::channel (false);
Ok (Self {
config: config.into (),
request_rendezvous: Default::default (),
server_status: Default::default (),
response_rendezvous: Default::default (),
shutdown_watch_tx,
shutdown_watch_rx,
unregistered_servers: Default::default (),
})
}
}
impl RelayState {
pub async fn list_servers (&self) -> Vec <String> {
self.request_rendezvous.lock ().await.iter ()
.map (|(k, _)| (*k).clone ())
.collect ()
}
}