Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

with_shutdown for Tokio app #113

Merged
merged 2 commits into from
Jan 5, 2024
Merged

with_shutdown for Tokio app #113

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a266ff0
with_shutdown for Tokio app
Grinkers Jan 4, 2024
c9df8ed
tokio app's shutdown now access by `ref`
Grinkers Jan 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion .github/workflows/examples.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ jobs:
with:
command: check
args: --manifest-path examples/host/Cargo.toml

- name: Check monitor example
if: always()
uses: actions-rs/cargo@v1
Expand Down Expand Up @@ -168,3 +168,10 @@ jobs:
with:
command: check
args: --manifest-path examples/shutdown/Cargo.toml

- name: Check shutdown-tokio example
if: always()
uses: actions-rs/cargo@v1
with:
command: check
args: --manifest-path examples/shutdown-tokio/Cargo.toml
11 changes: 11 additions & 0 deletions examples/shutdown-tokio/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
[package]
name = "shutdown-tokio"
version = "0.1.0"
edition = "2021"

[dependencies]
humphrey = { path = "../../humphrey", features = ["tokio"] }
tokio = { version = "1", features = ["full"] }
tokio-util = "0.7"

[workspace]
30 changes: 30 additions & 0 deletions examples/shutdown-tokio/src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
use humphrey::http::{Request, Response, StatusCode};
use humphrey::App;
use tokio_util::sync::CancellationToken;

use std::error::Error;
use std::thread::{sleep, spawn};
use std::time::Duration;

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
let cancel = CancellationToken::new();
let app: App<()> = App::new()
.with_shutdown(cancel.clone())
.with_stateless_route("/", hello);

// Shutdown the main app after 5 seconds
spawn(move || {
sleep(Duration::from_secs(5));
cancel.cancel();
});

// Returns after shutdown signal
app.run("0.0.0.0:8080").await?;

Ok(())
}

async fn hello(_: Request) -> Response {
Response::new(StatusCode::OK, "Hello, world! - tokio")
}
2 changes: 1 addition & 1 deletion examples/shutdown/src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ fn main() -> Result<(), Box<dyn Error>> {
.with_shutdown(app_rx)
.with_stateless_route("/hello", |_| Response::new(StatusCode::OK, "Hello world!"));

// Shutdown both the main app after 5 seconds
// Shutdown the main app after 5 seconds
spawn(move || {
sleep(Duration::from_secs(5));
let _ = shutdown_app.send(());
Expand Down
6 changes: 5 additions & 1 deletion humphrey/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,9 +37,13 @@ optional = true
version = "^0.24.1"
optional = true

[dependencies.tokio-util]
version = "0.7"
optional = true

[features]
tls = ["rustls", "rustls-native-certs", "rustls-pemfile"]
tokio = ["dep:tokio", "futures", "tokio-rustls"]
tokio = ["dep:tokio", "futures", "tokio-rustls", "tokio-util"]

[lib]
doctest = false
221 changes: 128 additions & 93 deletions humphrey/src/tokio/app.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ use std::sync::Arc;

use tokio::io::AsyncWriteExt;
use tokio::net::{TcpListener, TcpStream, ToSocketAddrs};
use tokio_util::sync::CancellationToken;

#[cfg(feature = "tls")]
use rustls::ServerConfig;
Expand All @@ -42,6 +43,7 @@ where
tls_config: Option<Arc<ServerConfig>>,
#[cfg(feature = "tls")]
force_https: bool,
shutdown: Option<CancellationToken>,
}

/// Represents a function able to calculate whether a connection will be accepted.
Expand Down Expand Up @@ -98,6 +100,7 @@ where
tls_config: None,
#[cfg(feature = "tls")]
force_https: false,
shutdown: None,
}
}

Expand All @@ -114,6 +117,7 @@ where
tls_config: None,
#[cfg(feature = "tls")]
force_https: false,
shutdown: None,
}
}

Expand All @@ -129,51 +133,63 @@ where
let error_handler = Arc::new(self.error_handler);

loop {
match socket.accept().await {
Ok((mut stream, _)) => {
let cloned_state = self.state.clone();

// Check that the client is allowed to connect
if (self.connection_condition)(&mut stream, cloned_state) {
let cloned_state = self.state.clone();
let cloned_monitor = self.monitor.clone();
let cloned_subapps = subapps.clone();
let cloned_default_subapp = default_subapp.clone();
let cloned_error_handler = error_handler.clone();

cloned_monitor.send(
Event::new(EventType::ConnectionSuccess)
.with_peer_result(stream.peer_addr()),
);

// Spawn a new thread to handle the connection
tokio::spawn(async move {
cloned_monitor.send(
Event::new(EventType::ThreadPoolProcessStarted)
.with_peer_result(stream.peer_addr()),
);
let shutdown = async {
if let Some(sd) = self.shutdown.as_ref() {
sd.cancelled().await
} else {
futures::future::pending().await
}
};
tokio::select! {
() = shutdown => { break Ok(()); }
s = socket.accept() => {
match s {
Ok((mut stream, _)) => {
let cloned_state = self.state.clone();

// Check that the client is allowed to connect
if (self.connection_condition)(&mut stream, cloned_state) {
let cloned_state = self.state.clone();
let cloned_monitor = self.monitor.clone();
let cloned_subapps = subapps.clone();
let cloned_default_subapp = default_subapp.clone();
let cloned_error_handler = error_handler.clone();

cloned_monitor.send(
Event::new(EventType::ConnectionSuccess)
.with_peer_result(stream.peer_addr()),
);

// Spawn a new thread to handle the connection
tokio::spawn(async move {
cloned_monitor.send(
Event::new(EventType::ThreadPoolProcessStarted)
.with_peer_result(stream.peer_addr()),
);

client_handler(
Stream::Tcp(stream),
cloned_subapps,
cloned_default_subapp,
cloned_error_handler,
cloned_state,
cloned_monitor,
)
.await
});
} else {
self.monitor.send(
Event::new(EventType::ConnectionDenied)
.with_peer_result(stream.peer_addr()),
);
client_handler(
Stream::Tcp(stream),
cloned_subapps,
cloned_default_subapp,
cloned_error_handler,
cloned_state,
cloned_monitor,
)
.await
});
} else {
self.monitor.send(
Event::new(EventType::ConnectionDenied)
.with_peer_result(stream.peer_addr()),
);
}
}
Err(e) => self
.monitor
.send(Event::new(EventType::ConnectionError).with_info(e.to_string())),
}
}
Err(e) => self
.monitor
.send(Event::new(EventType::ConnectionError).with_info(e.to_string())),
}
};
}
}

Expand Down Expand Up @@ -203,60 +219,73 @@ where
let acceptor = TlsAcceptor::from(tls_config);

loop {
match socket.accept().await {
Ok((mut sock, _)) => {
let cloned_state = self.state.clone();

// Check that the client is allowed to connect
if (self.connection_condition)(&mut sock, cloned_state) {
let cloned_state = self.state.clone();
let cloned_subapps = subapps.clone();
let cloned_default_subapp = default_subapp.clone();
let cloned_error_handler = error_handler.clone();
let cloned_monitor = self.monitor.clone();
let cloned_acceptor = acceptor.clone();

cloned_monitor.send(
Event::new(EventType::ConnectionSuccess)
.with_peer_result(sock.peer_addr()),
);

// Spawn a new thread to handle the connection
tokio::spawn(async move {
cloned_monitor.send(
Event::new(EventType::ThreadPoolProcessStarted)
.with_peer_result(sock.peer_addr()),
);

match cloned_acceptor.accept(sock).await {
Ok(tls_stream) => {
let stream = Stream::Tls(tls_stream);

client_handler(
stream,
cloned_subapps,
cloned_default_subapp,
cloned_error_handler,
cloned_state,
cloned_monitor,
)
.await
}
Err(e) => cloned_monitor.send(
Event::new(EventType::ConnectionError).with_info(e.to_string()),
),
let shutdown = async {
if let Some(sd) = self.shutdown.as_ref() {
sd.cancelled().await
} else {
futures::future::pending().await
}
};

tokio::select! {
() = shutdown => { break Ok(()); }
s = socket.accept() => {
match s {
Ok((mut sock, _)) => {
let cloned_state = self.state.clone();

// Check that the client is allowed to connect
if (self.connection_condition)(&mut sock, cloned_state) {
let cloned_state = self.state.clone();
let cloned_subapps = subapps.clone();
let cloned_default_subapp = default_subapp.clone();
let cloned_error_handler = error_handler.clone();
let cloned_monitor = self.monitor.clone();
let cloned_acceptor = acceptor.clone();

cloned_monitor.send(
Event::new(EventType::ConnectionSuccess)
.with_peer_result(sock.peer_addr()),
);

// Spawn a new thread to handle the connection
tokio::spawn(async move {
cloned_monitor.send(
Event::new(EventType::ThreadPoolProcessStarted)
.with_peer_result(sock.peer_addr()),
);

match cloned_acceptor.accept(sock).await {
Ok(tls_stream) => {
let stream = Stream::Tls(tls_stream);

client_handler(
stream,
cloned_subapps,
cloned_default_subapp,
cloned_error_handler,
cloned_state,
cloned_monitor,
)
.await
}
Err(e) => cloned_monitor.send(
Event::new(EventType::ConnectionError).with_info(e.to_string()),
),
}
});
} else {
self.monitor.send(
Event::new(EventType::ConnectionDenied)
.with_peer_result(sock.peer_addr()),
);
}
});
} else {
self.monitor.send(
Event::new(EventType::ConnectionDenied)
.with_peer_result(sock.peer_addr()),
);
}
Err(e) => self
.monitor
.send(Event::new(EventType::ConnectionError).with_info(e.to_string())),
}
}
Err(e) => self
.monitor
.send(Event::new(EventType::ConnectionError).with_info(e.to_string())),
}
}
}
Expand Down Expand Up @@ -417,6 +446,12 @@ where
self
}

/// Registers a shutdown signal to gracefully shutdown the app, ending the run/run_tls loop.
pub fn with_shutdown(mut self, cancel_token: CancellationToken) -> Self {
self.shutdown = Some(cancel_token);
self
}

/// Gets a reference to the app's state.
/// This should only be used in the main thread, as the state is passed to request handlers otherwise.
pub fn get_state(&self) -> Arc<State> {
Expand Down
Loading