version = "0.1.0"
dependencies = [
"anyhow",
+ "pretty_env_logger",
"serde",
"shared",
"tokio",
"hyper",
"inner",
"log",
- "pretty_env_logger",
"prometheus",
"redis",
"serde",
dependencies = [
"anyhow",
"ed25519-dalek",
+ "futures-util",
"hex",
"hyper",
"lazy_static",
nats_crate::Client,
payloads::{CachePayload, DispatchEventTagged, Tracing},
};
+use tokio::sync::oneshot;
use std::{convert::TryFrom, pin::Pin};
use twilight_gateway::{Event, Shard};
mod config;
-use futures::{Future, StreamExt};
+use futures::{Future, StreamExt, select};
use twilight_model::gateway::event::DispatchEvent;
+use futures::FutureExt;
struct GatewayServer {}
impl Component for GatewayServer {
type Config = GatewayConfig;
const SERVICE_NAME: &'static str = "gateway";
- fn start(&self, settings: Settings<Self::Config>) -> AnyhowResultFuture<()> {
+ fn start(
+ &self,
+ settings: Settings<Self::Config>,
+ stop: oneshot::Receiver<()>,
+ ) -> AnyhowResultFuture<()> {
Box::pin(async move {
let (shard, mut events) = Shard::builder(settings.token.to_owned(), settings.intents)
.shard(settings.shard, settings.shard_total)?
shard.start().await?;
- while let Some(event) = events.next().await {
- match event {
- Event::Ready(ready) => {
- info!("Logged in as {}", ready.user.name);
- }
+ let mut stop = stop.fuse();
+ loop {
- _ => {
- let name = event.kind().name();
- if let Ok(dispatch_event) = DispatchEvent::try_from(event) {
- let data = CachePayload {
- tracing: Tracing {
- node_id: "".to_string(),
- span: None,
- },
- data: DispatchEventTagged {
- data: dispatch_event,
- },
- };
- let value = serde_json::to_string(&data)?;
- debug!("nats send: {}", value);
- let bytes = bytes::Bytes::from(value);
- nats.publish(format!("nova.cache.dispatch.{}", name.unwrap()), bytes)
- .await?;
+ select! {
+ event = events.next().fuse() => {
+ if let Some(event) = event {
+ match event {
+ Event::Ready(ready) => {
+ info!("Logged in as {}", ready.user.name);
+ }
+
+ _ => {
+ let name = event.kind().name();
+ if let Ok(dispatch_event) = DispatchEvent::try_from(event) {
+ let data = CachePayload {
+ tracing: Tracing {
+ node_id: "".to_string(),
+ span: None,
+ },
+ data: DispatchEventTagged {
+ data: dispatch_event,
+ },
+ };
+ let value = serde_json::to_string(&data)?;
+ debug!("nats send: {}", value);
+ let bytes = bytes::Bytes::from(value);
+ nats.publish(format!("nova.cache.dispatch.{}", name.unwrap()), bytes)
+ .await?;
+ }
+ }
+ }
+ } else {
+ break
}
- }
- }
+ },
+ _ = stop => break
+ };
}
+ info!("stopping shard...");
+ shard.shutdown();
+
Ok(())
})
}
use serde::Deserialize;
fn default_listening_address() -> SocketAddr {
- SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 8080))
+ SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 8080))
}
#[derive(Debug, Deserialize, Clone)]
convert::TryFrom,
hash::{Hash, Hasher},
str::FromStr,
+ time::Instant,
};
use anyhow::bail;
pub async fn handle_request(
client: Client<HttpsConnector<HttpConnector>, Body>,
ratelimiter: RemoteRatelimiter,
- token: String,
+ token: &str,
mut request: Request<Body>,
) -> Result<Response<Body>, anyhow::Error> {
let (hash, uri_string) = {
let request_path = request.uri().path();
let (api_path, trimmed_path) = normalize_path(&request_path);
- let mut uri_string = format!("http://192.168.0.27:8000{}{}", api_path, trimmed_path);
+ let mut uri_string = format!("https://discord.com{}{}", api_path, trimmed_path);
if let Some(query) = request.uri().query() {
uri_string.push('?');
uri_string.push_str(query);
(hash.finish().to_string(), uri_string)
};
+ let start_ticket_request = Instant::now();
let header_sender = match ratelimiter.ticket(hash).await {
Ok(sender) => sender,
Err(e) => {
bail!("failed to reteive ticket");
}
};
+ let time_took_ticket = Instant::now() - start_ticket_request;
request.headers_mut().insert(
AUTHORIZATION,
request.headers_mut().remove(AUTHORIZATION);
request.headers_mut().append(
AUTHORIZATION,
- HeaderValue::from_static(
- "Bot ODA3MTg4MzM1NzE3Mzg0MjEy.G3sXFM.8gY2sVYDAq2WuPWwDskAAEFLfTg8htooxME-LE",
- ),
+ HeaderValue::from_str(&format!("Bot {}", token))?,
);
let uri = match Uri::from_str(&uri_string) {
}
};
*request.uri_mut() = uri;
- let resp = match client.request(request).await {
+
+ let start_upstream_req = Instant::now();
+ let mut resp = match client.request(request).await {
Ok(response) => response,
Err(e) => {
error!("Error when requesting the Discord API: {:?}", e);
bail!("failed to request the discord api");
}
};
+ let upstream_time_took = Instant::now() - start_upstream_req;
+ resp.headers_mut().append(
+ "X-TicketRequest-Ms",
+ HeaderValue::from_str(&time_took_ticket.as_millis().to_string()).unwrap(),
+ );
+ resp.headers_mut().append(
+ "X-Upstream-Ms",
+ HeaderValue::from_str(&upstream_time_took.as_millis().to_string()).unwrap(),
+ );
+
let ratelimit_headers = resp
.headers()
.into_iter()
use hyper_tls::HttpsConnector;
use leash::{ignite, AnyhowResultFuture, Component};
use shared::config::Settings;
-use std::convert::Infallible;
+use std::{convert::Infallible, sync::Arc};
+use tokio::sync::oneshot;
mod config;
mod handler;
type Config = ReverseProxyConfig;
const SERVICE_NAME: &'static str = "rest";
- fn start(&self, settings: Settings<Self::Config>) -> AnyhowResultFuture<()> {
+ fn start(
+ &self,
+ settings: Settings<Self::Config>,
+ stop: oneshot::Receiver<()>,
+ ) -> AnyhowResultFuture<()> {
Box::pin(async move {
// Client to the remote ratelimiters
let ratelimiter = ratelimit_client::RemoteRatelimiter::new();
let client = Client::builder().build(HttpsConnector::new());
+ let token = Arc::new(settings.discord.token.clone());
let service_fn = make_service_fn(move |_: &AddrStream| {
let client = client.clone();
let ratelimiter = ratelimiter.clone();
+ let token = token.clone();
async move {
Ok::<_, Infallible>(service_fn(move |request: Request<Body>| {
let client = client.clone();
let ratelimiter = ratelimiter.clone();
+ let token = token.clone();
async move {
- handle_request(client, ratelimiter, "token".to_string(), request).await
+ let token = token.as_str();
+ handle_request(client, ratelimiter, token, request).await
}
}))
}
let server = Server::bind(&settings.config.server.listening_adress).serve(service_fn);
- server.await?;
+ server
+ .with_graceful_shutdown(async {
+ stop.await.expect("should not fail");
+ })
+ .await?;
Ok(())
})
obj_clone.get_ratelimiters().await.unwrap();
tokio::select! {
() = &mut sleep => {
- println!("timer elapsed");
+ debug!("timer elapsed");
},
_ = tx.recv() => {}
}
ed25519-dalek = "1"
twilight-model = { version = "0.14" }
anyhow = "1.0.68"
+futures-util = "0.3.25"
[[bin]]
name = "webhook"
use serde::{Deserialize, Deserializer};
fn default_listening_address() -> SocketAddr {
- SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 8080))
+ SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 8080))
}
#[derive(Debug, Deserialize, Clone, Copy)]
use hyper::Server;
use leash::{ignite, AnyhowResultFuture, Component};
use shared::{config::Settings, log::info, nats_crate::Client};
+use tokio::sync::oneshot;
#[derive(Clone, Copy)]
struct WebhookServer {}
type Config = WebhookConfig;
const SERVICE_NAME: &'static str = "webhook";
- fn start(&self, settings: Settings<Self::Config>) -> AnyhowResultFuture<()> {
+ fn start(
+ &self,
+ settings: Settings<Self::Config>,
+ stop: oneshot::Receiver<()>,
+ ) -> AnyhowResultFuture<()> {
Box::pin(async move {
info!("Starting server on {}", settings.server.listening_adress);
let server = Server::bind(&bind).serve(make_service);
- server.await?;
+ server.with_graceful_shutdown(async {
+ stop.await.expect("should not fail");
+ }).await?;
Ok(())
})
shared = { path = "../shared" }
anyhow = "1.0.68"
tokio = { version = "1.23.0", features = ["full"] }
-
+pretty_env_logger = "0.4"
serde = "1.0.152"
\ No newline at end of file
use anyhow::Result;
use serde::de::DeserializeOwned;
-use shared::config::Settings;
+use shared::{
+ config::Settings,
+ log::{error, info},
+};
use std::{future::Future, pin::Pin};
+use tokio::{signal::{unix::SignalKind}, sync::oneshot};
pub type AnyhowResultFuture<T> = Pin<Box<dyn Future<Output = Result<T>>>>;
pub trait Component: Send + Sync + 'static + Sized {
type Config: Default + Clone + DeserializeOwned;
const SERVICE_NAME: &'static str;
- fn start(&self, settings: Settings<Self::Config>) -> AnyhowResultFuture<()>;
+ fn start(
+ &self,
+ settings: Settings<Self::Config>,
+ stop: oneshot::Receiver<()>,
+ ) -> AnyhowResultFuture<()>;
fn new() -> Self;
fn _internal_start(self) -> AnyhowResultFuture<()> {
Box::pin(async move {
+ pretty_env_logger::init();
let settings = Settings::<Self::Config>::new(Self::SERVICE_NAME);
+ let (stop, stop_channel) = oneshot::channel();
// Start the grpc healthcheck
tokio::spawn(async move {});
// Start the prometheus monitoring job
tokio::spawn(async move {});
- self.start(settings?).await
+ tokio::spawn(async move {
+ match tokio::signal::unix::signal(SignalKind::terminate()).unwrap().recv().await {
+ Some(()) => {
+ info!("Stopping program.");
+
+ stop.send(()).unwrap();
+ }
+ None => {
+ error!("Unable to listen for shutdown signal");
+ // we also shut down in case of error
+ }
+ }
+ });
+
+ self.start(settings?, stop_channel).await
})
}
}
#[cfg(test)]
mod test {
use serde::Deserialize;
+ use tokio::sync::oneshot;
use crate::Component;
fn start(
&self,
_settings: shared::config::Settings<Self::Config>,
+ _stop: oneshot::Receiver<()>,
) -> crate::AnyhowResultFuture<()> {
Box::pin(async move { Ok(()) })
}
Self {}
}
}
-
+
ignite!(TestComponent);
}
edition = "2021"
[dependencies]
-pretty_env_logger = "0.4"
log = { version = "0.4", features = ["std"] }
serde = { version = "1.0.8", features = ["derive"] }
serde_repr = "0.1"
projects / matthieu / nova.git / commitdiff