Improve fairing example.

This commit is contained in:
Sergio Benitez 2021-03-09 21:13:26 -08:00
parent cce9ea1e1e
commit ab3413826c
1 changed files with 13 additions and 17 deletions

View File

@ -2,8 +2,9 @@
use std::io::Cursor; use std::io::Cursor;
use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use rocket::{Request, State, Data, Response}; use rocket::{Rocket, Request, State, Data, Response};
use rocket::fairing::{AdHoc, Fairing, Info, Kind}; use rocket::fairing::{AdHoc, Fairing, Info, Kind};
use rocket::http::{Method, ContentType, Status}; use rocket::http::{Method, ContentType, Status};
@ -11,10 +12,10 @@ struct Token(i64);
#[cfg(test)] mod tests; #[cfg(test)] mod tests;
#[derive(Default)] #[derive(Default, Clone)]
struct Counter { struct Counter {
get: AtomicUsize, get: Arc<AtomicUsize>,
post: AtomicUsize, post: Arc<AtomicUsize>,
} }
#[rocket::async_trait] #[rocket::async_trait]
@ -22,7 +23,7 @@ impl Fairing for Counter {
fn info(&self) -> Info { fn info(&self) -> Info {
Info { Info {
name: "GET/POST Counter", name: "GET/POST Counter",
kind: Kind::Request | Kind::Response kind: Kind::Attach | Kind::Request
} }
} }
@ -34,20 +35,15 @@ impl Fairing for Counter {
} }
} }
async fn on_response<'r>(&self, req: &'r Request<'_>, res: &mut Response<'r>) { async fn on_attach(&self, rocket: Rocket) -> Result<Rocket, Rocket> {
if res.status() != Status::NotFound { #[get("/counts")]
return fn counts(counts: State<'_, Counter>) -> String {
let get_count = counts.get.load(Ordering::Relaxed);
let post_count = counts.post.load(Ordering::Relaxed);
format!("Get: {}\nPost: {}", get_count, post_count)
} }
if req.method() == Method::Get && req.uri().path() == "/counts" { Ok(rocket.manage(self.clone()).mount("/", routes![counts]))
let get_count = self.get.load(Ordering::Relaxed);
let post_count = self.post.load(Ordering::Relaxed);
let body = format!("Get: {}\nPost: {}", get_count, post_count);
res.set_status(Status::Ok);
res.set_header(ContentType::Plain);
res.set_sized_body(body.len(), Cursor::new(body));
}
} }
} }