Add a handler to kill a query

ipa-core/src/app.rs

@@ -203,6 +203,10 @@
  let query_id = ext_query_id(&req)?;
  HelperResponse::from(qp.complete(query_id).await?)
  }
+ RouteId::KillQuery => {
+ let query_id = ext_query_id(&req)?;
+ HelperResponse::from(qp.kill(query_id)?)
+ }
  })
  }
 }

ipa-core/src/helpers/transport/handler.rs

@@ -12,7 +12,7 @@ use crate::{
  },
  query::{
  NewQueryError, PrepareQueryError, ProtocolResult, QueryCompletionError, QueryInputError,
- QueryStatus, QueryStatusError,
+ QueryKillStatus, QueryKilled, QueryStatus, QueryStatusError,
  },
  sync::{Arc, Mutex, Weak},
 };
@@ -135,6 +135,13 @@ impl From<QueryStatus> for HelperResponse {
  }
 }
 
+impl From<QueryKilled> for HelperResponse {
+ fn from(value: QueryKilled) -> Self {
+ let v = serde_json::to_vec(&json!({"query_id": value.0, "status": "killed"})).unwrap();
+ Self { body: v }
+ }
+}
+
 impl<R: AsRef<dyn ProtocolResult>> From<R> for HelperResponse {
  fn from(value: R) -> Self {
  let v = value.as_ref().to_bytes();
@@ -156,6 +163,8 @@ pub enum Error {
  #[error(transparent)]
  QueryStatus(#[from] QueryStatusError),
  #[error(transparent)]
+ QueryKill(#[from] QueryKillStatus),
+ #[error(transparent)]
  DeserializationFailure(#[from] serde_json::Error),
  #[error("MalformedRequest: {0}")]
  BadRequest(BoxError),

ipa-core/src/helpers/transport/in_memory/transport.rs

@@ -119,7 +119,8 @@ impl<I: TransportIdentity> InMemoryTransport<I> {
  | RouteId::PrepareQuery
  | RouteId::QueryInput
  | RouteId::QueryStatus
- | RouteId::CompleteQuery => {
+ | RouteId::CompleteQuery
+ | RouteId::KillQuery => {
  handler
  .as_ref()
  .expect("Handler is set")

ipa-core/src/helpers/transport/routing.rs

@@ -16,6 +16,7 @@ pub enum RouteId {
  QueryInput,
  QueryStatus,
  CompleteQuery,
+ KillQuery,
 }
 
 /// The header/metadata of the incoming request.

ipa-core/src/helpers/transport/stream/collection.rs

@@ -114,6 +114,26 @@ impl<I: TransportIdentity, S: Stream> StreamCollection<I, S> {
  let mut streams = self.inner.lock().unwrap();
  streams.clear();
  }
+
+ /// Returns the number of streams inside this collection.
+ ///
+ /// ## Panics
+ /// if mutex is poisoned.
+ #[cfg(test)]
+ #[must_use]
+ pub fn len(&self) -> usize {
+ self.inner.lock().unwrap().len()
+ }
+
+ /// Returns `true` if this collection is empty.
+ ///
+ /// ## Panics
+ /// if mutex is poisoned.
+ #[must_use]
+ #[cfg(test)]
+ pub fn is_empty(&self) -> bool {
+ self.len() == 0
+ }
 }
 
 /// Describes the lifecycle of records stream inside [`StreamCollection`]

ipa-core/src/net/http_serde.rs

@@ -533,4 +533,82 @@ pub mod query {
 
  pub const AXUM_PATH: &str = "/:query_id/complete";
  }
+
+ pub mod kill {
+ use serde::{Deserialize, Serialize};
+
+ use crate::{
+ helpers::{routing::RouteId, HelperResponse, NoStep, RouteParams},
+ protocol::QueryId,
+ };
+
+ pub struct Request {
+ pub query_id: QueryId,
+ }
+
+ impl RouteParams<RouteId, QueryId, NoStep> for Request {
+ type Params = String;
+
+ fn resource_identifier(&self) -> RouteId {
+ RouteId::KillQuery
+ }
+
+ fn query_id(&self) -> QueryId {
+ self.query_id
+ }
+
+ fn gate(&self) -> NoStep {
+ NoStep
+ }
+
+ fn extra(&self) -> Self::Params {
+ String::new()
+ }
+ }
+
+ impl Request {
+ /// Currently, it is only possible to kill
+ /// a query by issuing an HTTP request manually.
+ /// Maybe report collector can support this API,
+ /// but for now, only tests exercise this path
+ /// hence methods here are hidden behind feature
+ /// flags
+ #[cfg(all(test, unit_test))]
+ pub fn new(query_id: QueryId) -> Self {
+ Self { query_id }
+ }
+
+ #[cfg(all(test, unit_test))]
+ pub fn try_into_http_request(
+ self,
+ scheme: axum::http::uri::Scheme,
+ authority: axum::http::uri::Authority,
+ ) -> crate::net::http_serde::OutgoingRequest {
+ let uri = axum::http::uri::Uri::builder()
+ .scheme(scheme)
+ .authority(authority)
+ .path_and_query(format!(
+ "{}/{}/kill",
+ crate::net::http_serde::query::BASE_AXUM_PATH,
+ self.query_id.as_ref()
+ ))
+ .build()?;
+ Ok(hyper::Request::post(uri).body(axum::body::Body::empty())?)
+ }
+ }
+
+ #[derive(Clone, Debug, Serialize, Deserialize)]
+ pub struct ResponseBody {
+ pub query_id: QueryId,
+ pub status: String,
+ }
+
+ impl From<HelperResponse> for ResponseBody {
+ fn from(value: HelperResponse) -> Self {
+ serde_json::from_slice(value.into_body().as_slice()).unwrap()
+ }
+ }
+
+ pub const AXUM_PATH: &str = "/:query_id/kill";
+ }
 }

ipa-core/src/net/server/handlers/query/kill.rs

@@ -0,0 +1,136 @@
+use axum::{extract::Path, routing::post, Extension, Json, Router};
+use hyper::StatusCode;
+
+use crate::{
+ helpers::{ApiError, BodyStream, Transport},
+ net::{
+ http_serde::query::{kill, kill::Request},
+ server::Error,
+ Error::QueryIdNotFound,
+ HttpTransport,
+ },
+ protocol::QueryId,
+ query::QueryKillStatus,
+ sync::Arc,
+};
+
+async fn handler(
+ transport: Extension<Arc<HttpTransport>>,
+ Path(query_id): Path<QueryId>,
+) -> Result<Json<kill::ResponseBody>, Error> {
+ let req = Request { query_id };
+ let transport = Transport::clone_ref(&*transport);
+ match transport.dispatch(req, BodyStream::empty()).await {
+ Ok(state) => Ok(Json(kill::ResponseBody::from(state))),
+ Err(ApiError::QueryKill(QueryKillStatus::NoSuchQuery(query_id))) => Err(
+ Error::application(StatusCode::NOT_FOUND, QueryIdNotFound(query_id)),
+ ),
+ Err(e) => Err(Error::application(StatusCode::INTERNAL_SERVER_ERROR, e)),
+ }
+}
+
+pub fn router(transport: Arc<HttpTransport>) -> Router {
+ Router::new()
+ .route(kill::AXUM_PATH, post(handler))
+ .layer(Extension(transport))
+}
+
+#[cfg(all(test, unit_test))]
+mod tests {
+ use axum::{
+ body::Body,
+ http::uri::{Authority, Scheme},
+ };
+ use hyper::StatusCode;
+
+ use crate::{
+ helpers::{
+ make_owned_handler,
+ routing::{Addr, RouteId},
+ ApiError, BodyStream, HelperIdentity, HelperResponse,
+ },
+ net::{
+ http_serde,
+ server::handlers::query::test_helpers::{
+ assert_fails_with, assert_fails_with_handler, assert_success_with,
+ },
+ },
+ protocol::QueryId,
+ query::{QueryKillStatus, QueryKilled},
+ };
+
+ #[tokio::test]
+ async fn calls_kill() {
+ let expected_query_id = QueryId;
+
+ let handler = make_owned_handler(
+ move |addr: Addr<HelperIdentity>, _data: BodyStream| async move {
+ let RouteId::KillQuery = addr.route else {
+ panic!("unexpected call: {addr:?}");
+ };
+ assert_eq!(addr.query_id, Some(expected_query_id));
+ Ok(HelperResponse::from(QueryKilled(expected_query_id)))
+ },
+ );
+
+ let req = http_serde::query::kill::Request::new(QueryId);
+ let req = req
+ .try_into_http_request(Scheme::HTTP, Authority::from_static("localhost"))
+ .unwrap();
+ assert_success_with(req, handler).await;
+ }
+
+ #[tokio::test]
+ async fn no_such_query() {
+ let handler = make_owned_handler(
+ move |_addr: Addr<HelperIdentity>, _data: BodyStream| async move {
+ Err(QueryKillStatus::NoSuchQuery(QueryId).into())
+ },
+ );
+
+ let req = http_serde::query::kill::Request::new(QueryId)
+ .try_into_http_request(Scheme::HTTP, Authority::from_static("localhost"))
+ .unwrap();
+ assert_fails_with_handler(req, handler, StatusCode::NOT_FOUND).await;
+ }
+
+ #[tokio::test]
+ async fn unknown_error() {
+ let handler = make_owned_handler(
+ move |_addr: Addr<HelperIdentity>, _data: BodyStream| async move {
+ Err(ApiError::DeserializationFailure(
+ serde_json::from_str::<()>("not-a-json").unwrap_err(),
+ ))
+ },
+ );
+
+ let req = http_serde::query::kill::Request::new(QueryId)
+ .try_into_http_request(Scheme::HTTP, Authority::from_static("localhost"))
+ .unwrap();
+ assert_fails_with_handler(req, handler, StatusCode::INTERNAL_SERVER_ERROR).await;
+ }
+
+ struct OverrideReq {
+ query_id: String,
+ }
+
+ impl From<OverrideReq> for hyper::Request<Body> {
+ fn from(val: OverrideReq) -> Self {
+ let uri = format!(
+ "http://localhost{}/{}/kill",
+ http_serde::query::BASE_AXUM_PATH,
+ val.query_id
+ );
+ hyper::Request::post(uri).body(Body::empty()).unwrap()
+ }
+ }
+
+ #[tokio::test]
+ async fn malformed_query_id() {
+ let req = OverrideReq {
+ query_id: "not-a-query-id".into(),
+ };
+
+ assert_fails_with(req.into(), StatusCode::BAD_REQUEST).await;
+ }
+}

ipa-core/src/net/server/handlers/query/mod.rs

@@ -1,5 +1,6 @@
 mod create;
 mod input;
+mod kill;
 mod prepare;
 mod results;
 mod status;
@@ -31,6 +32,7 @@ pub fn query_router(transport: Arc<HttpTransport>) -> Router {
  .merge(create::router(Arc::clone(&transport)))
  .merge(input::router(Arc::clone(&transport)))
  .merge(status::router(Arc::clone(&transport)))
+ .merge(kill::router(Arc::clone(&transport)))
  .merge(results::router(transport))
 }
 
@@ -139,6 +141,19 @@ pub mod test_helpers {
  assert_eq!(resp.status(), expected_status);
  }
 
+ pub async fn assert_fails_with_handler(
+ req: hyper::Request<Body>,
+ handler: Arc<dyn RequestHandler<Identity = HelperIdentity>>,
+ expected_status: StatusCode,
+ ) {
+ let test_server = TestServer::builder()
+ .with_request_handler(handler)
+ .build()
+ .await;
+ let resp = test_server.server.handle_req(req).await;
+ assert_eq!(resp.status(), expected_status);
+ }
+
  pub async fn assert_success_with(
  req: hyper::Request<Body>,
  handler: Arc<dyn RequestHandler<Identity = HelperIdentity>>,