From 9ec7c9e26732cfb1ed0f5f7fabfab1ef89d76d2e Mon Sep 17 00:00:00 2001 From: tristan-morris <44310937+tristan-morris@users.noreply.github.com> Date: Mon, 29 Jul 2024 15:31:49 +1000 Subject: [PATCH] Remove unwraps --- airmail/Cargo.toml | 1 - airmail/src/index.rs | 131 +++++++++++++++++------------------- airmail/src/lib.rs | 3 + airmail_service/Cargo.toml | 1 - airmail_service/src/api.rs | 8 ++- airmail_service/src/main.rs | 13 +++- 6 files changed, 84 insertions(+), 73 deletions(-) diff --git a/airmail/Cargo.toml b/airmail/Cargo.toml index 2c7875a..a2404b1 100644 --- a/airmail/Cargo.toml +++ b/airmail/Cargo.toml @@ -37,5 +37,4 @@ anyhow = "1.0.86" thiserror = "1.0.63" [features] -invasive_logging = [] remote_index = ["tantivy/quickwit"] diff --git a/airmail/src/index.rs b/airmail/src/index.rs index 3f3a8b1..df2abe8 100644 --- a/airmail/src/index.rs +++ b/airmail/src/index.rs @@ -2,10 +2,9 @@ use std::path::Path; use std::sync::Arc; use anyhow::Result; -use futures_util::future::join_all; use geo::Rect; use itertools::Itertools; -use log::debug; +use log::{trace, warn}; use s2::region::RegionCoverer; use std::collections::BTreeMap; use tantivy::schema::Value; @@ -346,70 +345,66 @@ impl AirmailIndex { let query_string = query.trim().replace("'s", "s"); let start = std::time::Instant::now(); - let (top_docs, searcher) = { - let query = self - .construct_query( - &searcher, - &query_string, - tags, - bbox, - boost_regions, - request_leniency, - ) - .await; - - #[cfg(feature = "invasive_logging")] - { - dbg!(&query); + + let query = self + .construct_query( + &searcher, + &query_string, + tags, + bbox, + boost_regions, + request_leniency, + ) + .await; + + trace!("Search query: {:?}", &query); + + // Perform the search and then resolve the returned documents + let top_docs: Result> = spawn_blocking(move || { + let doc_addresses = searcher.search(&query, &TopDocs::with_limit(10))?; + let mut docs = vec![]; + for (score, doc_address) in doc_addresses { + if let Ok(doc) = searcher.doc::(doc_address) { + docs.push((score, doc)); + } } - let (top_docs, searcher) = spawn_blocking(move || { - (searcher.search(&query, &TopDocs::with_limit(10)), searcher) - }) - .await?; - let top_docs = top_docs?; - debug!( - "Search took {:?} and yielded {} results", - start.elapsed(), - top_docs.len() - ); - (top_docs, searcher) - }; + Ok(docs) + }) + .await?; - let mut scores = Vec::new(); - let mut futures = Vec::new(); - for (score, doc_id) in top_docs { - let searcher = searcher.clone(); - let doc = spawn_blocking(move || searcher.doc::(doc_id)); - scores.push(score); - futures.push(doc); - } - let mut results = Vec::new(); - let top_docs = join_all(futures).await; - for (score, doc_future) in scores.iter().zip(top_docs) { - let doc = doc_future??; - let source = doc - .get_first(self.field_source()) - .map(|value| value.as_str().unwrap().to_string()) - .unwrap_or_default(); - let s2cell = doc - .get_first(self.field_s2cell()) - .unwrap() - .as_u64() - .unwrap(); - let cellid = s2::cellid::CellID(s2cell); - let latlng = s2::latlng::LatLng::from(cellid); - let tags: Vec<(String, String)> = doc - .get_first(self.field_tags()) - .unwrap() - .as_object() - .unwrap() - .map(|(k, v)| (k.to_string(), v.as_str().unwrap().to_string())) - .collect(); - - let poi = AirmailPoi::new(source, latlng.lat.deg(), latlng.lng.deg(), tags)?; - results.push((poi, *score)); - } + let top_docs = top_docs.map_err(|e| { + warn!("Search failed: {:?}", e); + e + })?; + + trace!( + "Search took {:?} and yielded {} results", + start.elapsed(), + top_docs.len() + ); + + let results = top_docs + .into_iter() + .flat_map(|(score, doc)| { + let source = doc + .get_first(self.field_source()) + .map(|value| value.as_str().unwrap_or_default().to_string()) + .unwrap_or_default(); + let s2cell = doc.get_first(self.field_s2cell())?.as_u64()?; + let cellid = s2::cellid::CellID(s2cell); + let latlng = s2::latlng::LatLng::from(cellid); + let tags: Vec<(String, String)> = doc + .get_first(self.field_tags())? + .as_object()? + .map(|(k, v)| (k.to_string(), v.as_str().unwrap_or_default().to_string())) + .collect(); + + AirmailPoi::new(source, latlng.lat.deg(), latlng.lng.deg(), tags) + .ok() + .map(|poi| (poi, score)) + }) + .collect::>(); Ok(results) } @@ -430,7 +425,7 @@ impl AirmailIndexWriter { for content in poi.content { self.process_field(&mut doc, &content); } - doc.add_text(self.schema.get_field(FIELD_SOURCE).unwrap(), source); + doc.add_text(self.schema.get_field(FIELD_SOURCE)?, source); let indexed_keys = [ "natural", "amenity", "shop", "leisure", "tourism", "historic", "cuisine", @@ -443,22 +438,22 @@ impl AirmailIndexWriter { .any(|prefix| key.starts_with(prefix)) { doc.add_text( - self.schema.get_field(FIELD_INDEXED_TAG).unwrap(), + self.schema.get_field(FIELD_INDEXED_TAG)?, format!("{}={}", key, value).as_str(), ); } } doc.add_object( - self.schema.get_field(FIELD_TAGS).unwrap(), + self.schema.get_field(FIELD_TAGS)?, poi.tags .iter() .map(|(k, v)| (k.to_string(), OwnedValue::Str(v.to_string()))) .collect::>(), ); - doc.add_u64(self.schema.get_field(FIELD_S2CELL).unwrap(), poi.s2cell); + doc.add_u64(self.schema.get_field(FIELD_S2CELL)?, poi.s2cell); for parent in poi.s2cell_parents { - doc.add_u64(self.schema.get_field(FIELD_S2CELL_PARENTS).unwrap(), parent); + doc.add_u64(self.schema.get_field(FIELD_S2CELL_PARENTS)?, parent); } self.tantivy_writer.add_document(doc)?; diff --git a/airmail/src/lib.rs b/airmail/src/lib.rs index 3970723..75fbf74 100644 --- a/airmail/src/lib.rs +++ b/airmail/src/lib.rs @@ -1,3 +1,6 @@ +#![forbid(unsafe_code)] +#![warn(clippy::missing_panics_doc)] + #[macro_use] extern crate lazy_static; diff --git a/airmail_service/Cargo.toml b/airmail_service/Cargo.toml index 246ca19..243f0c6 100644 --- a/airmail_service/Cargo.toml +++ b/airmail_service/Cargo.toml @@ -22,6 +22,5 @@ anyhow = "1.0.86" thiserror = "1.0.63" [features] -invasive_logging = ["airmail/invasive_logging"] remote_index = ["airmail/remote_index"] default = ["remote_index"] diff --git a/airmail_service/src/api.rs b/airmail_service/src/api.rs index 3996b6a..e92b57b 100644 --- a/airmail_service/src/api.rs +++ b/airmail_service/src/api.rs @@ -9,6 +9,7 @@ use axum::{ }; use deunicode::deunicode; use geo::{Coord, Rect}; +use log::debug; use serde::{Deserialize, Serialize}; use crate::error::AirmailServiceError; @@ -73,7 +74,12 @@ pub async fn search( let results = index.search(&query, leniency, tags, bbox, &[]).await?; - println!("{} results found in {:?}", results.len(), start.elapsed()); + debug!( + "Query: {:?} produced: {} results found in {:?}", + params, + results.len(), + start.elapsed() + ); let response = Response { metadata: MetadataResponse { query: params }, diff --git a/airmail_service/src/main.rs b/airmail_service/src/main.rs index cec98a0..4428157 100644 --- a/airmail_service/src/main.rs +++ b/airmail_service/src/main.rs @@ -1,6 +1,7 @@ #![forbid(unsafe_code)] #![warn(clippy::pedantic)] +use std::future::IntoFuture; use std::sync::Arc; use airmail::index::AirmailIndex; @@ -9,8 +10,9 @@ use api::search; use axum::{http::HeaderValue, routing::get, Router}; use clap::Parser; use env_logger::Env; -use log::{debug, info}; +use log::{debug, info, warn}; use tokio::net::TcpListener; +use tokio::select; use tower_http::cors::CorsLayer; mod api; @@ -60,7 +62,14 @@ async fn main() -> Result<()> { info!("Listening at: {}/search?q=query", args.bind); let listener = TcpListener::bind(args.bind).await?; - axum::serve(listener, app).await?; + let server = axum::serve(listener, app.into_make_service()).into_future(); + + select! { + _ = server => {} + _ = tokio::signal::ctrl_c() => { + warn!("Received ctrl-c, shutting down"); + } + } Ok(()) }