cloud_tasks_emulator/
lib.rs

1mod api;
2mod filter;
3mod names;
4mod oidc;
5mod rest;
6mod scheduler;
7mod state;
8
9use std::sync::{Arc, LazyLock, Mutex};
10use std::time::Duration;
11
12use axum::extract::{Path, State};
13use axum::http::StatusCode;
14use axum::response::{IntoResponse, Response};
15use axum::routing::{get, post};
16use prost_reflect::DescriptorPool;
17use tokio::net::TcpListener;
18use tokio::sync::Notify;
19use tonic::service::Routes;
20use tonic::transport::server::TcpIncoming;
21
22use crate::proto::google::cloud::tasks::v2::cloud_tasks_server::{CloudTasks, CloudTasksServer};
23use crate::proto::google::cloud::tasks::v2::{CreateQueueRequest, Queue};
24use crate::proto::google::longrunning::operations_server::OperationsServer;
25use crate::state::{Emulator, QueueState};
26
27pub mod proto {
28    pub mod google {
29        pub mod api {
30            tonic::include_proto!("google.api");
31        }
32        pub mod cloud {
33            pub mod tasks {
34                pub mod v2 {
35                    tonic::include_proto!("google.cloud.tasks.v2");
36                }
37            }
38        }
39        pub mod iam {
40            pub mod v1 {
41                tonic::include_proto!("google.iam.v1");
42            }
43        }
44        pub mod longrunning {
45            tonic::include_proto!("google.longrunning");
46        }
47        pub mod rpc {
48            tonic::include_proto!("google.rpc");
49        }
50        pub mod r#type {
51            tonic::include_proto!("google.r#type");
52        }
53    }
54}
55
56static DESCRIPTORS: LazyLock<DescriptorPool> = LazyLock::new(|| {
57    DescriptorPool::decode(include_bytes!(concat!(env!("OUT_DIR"), "/descriptors.bin")).as_slice())
58        .expect("build.rs encodes the descriptors")
59});
60
61/// Serves gRPC, REST and OIDC discovery for `issuer`, creating each of `queues` with default config
62pub async fn serve(
63    listener: TcpListener,
64    issuer: String,
65    queues: Vec<String>,
66) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
67    let emulator = Arc::new(Emulator {
68        state: Mutex::new(startup_state(&queues).await?),
69        wake: Notify::new(),
70    });
71    let issuer = Arc::new(oidc::Issuer::new(issuer));
72    tokio::spawn(scheduler::run(emulator.clone(), issuer.clone()));
73    let discovery = axum::Router::new()
74        .route("/.well-known/openid-configuration", get(oidc::discovery))
75        .route("/jwks", get(oidc::jwks))
76        .with_state(issuer);
77    let idle = axum::Router::new()
78        .route("/emulator/{*path}", get(wait_idle))
79        .with_state(emulator.clone());
80    let reset = axum::Router::new()
81        .route("/emulator:reset", post(reset))
82        .with_state((emulator.clone(), Arc::new(queues)));
83    let grpc = Routes::new(CloudTasksServer::from_arc(emulator.clone()))
84        .add_service(OperationsServer::from_arc(emulator));
85    let routes = grpc
86        .clone()
87        .into_axum_router()
88        .merge(discovery)
89        .merge(idle)
90        .merge(reset)
91        .merge(rest::router(grpc));
92    tonic::transport::Server::builder()
93        .accept_http1(true)
94        .add_routes(routes.into())
95        .serve_with_incoming(TcpIncoming::from(listener))
96        .await?;
97    Ok(())
98}
99
100/// State holding only `queues`, each created with default config
101async fn startup_state(queues: &[String]) -> Result<state::State, String> {
102    let emulator = Emulator::default();
103    for name in queues {
104        let request = CreateQueueRequest {
105            parent: name
106                .rsplit_once("/queues/")
107                .map_or(name.as_str(), |(parent, _)| parent)
108                .to_owned(),
109            queue: Some(Queue {
110                name: name.clone(),
111                ..Default::default()
112            }),
113        };
114        emulator
115            .create_queue(tonic::Request::new(request))
116            .await
117            .map_err(|status| format!("Cannot create {name}: {}", status.message()))?;
118    }
119    Ok(emulator.state.into_inner().unwrap())
120}
121
122/// Returns emulator to state it started with, keeping OIDC key
123async fn reset(State((emulator, queues)): State<(Arc<Emulator>, Arc<Vec<String>>)>) -> StatusCode {
124    let fresh = startup_state(&queues)
125        .await
126        .expect("startup created these queues already");
127    emulator.state.lock().unwrap().reset(fresh);
128    StatusCode::NO_CONTENT
129}
130
131/// Answers once queue in path such as `projects/p/locations/l/queues/q:waitIdle` holds no tasks
132async fn wait_idle(State(emulator): State<Arc<Emulator>>, Path(path): Path<String>) -> Response {
133    let Some(name) = path.strip_suffix(":waitIdle") else {
134        return StatusCode::NOT_FOUND.into_response();
135    };
136    loop {
137        {
138            let state = emulator.state.lock().unwrap();
139            let Some(queue) = state.queues.get(name) else {
140                return (
141                    StatusCode::NOT_FOUND,
142                    format!("Queue {name} does not exist"),
143                )
144                    .into_response();
145            };
146            if queue.tasks.is_empty() {
147                return StatusCode::NO_CONTENT.into_response();
148            }
149            if queue.state == QueueState::Paused {
150                return (
151                    StatusCode::CONFLICT,
152                    format!("Queue {name} is paused with tasks left"),
153                )
154                    .into_response();
155            }
156        }
157        tokio::time::sleep(Duration::from_millis(10)).await;
158    }
159}
cloud_tasks_emulator/lib.rs

cloud_tasks_emulator/
lib.rs
