{/if}
-{#if $my_rollout_query.engine_state === "paused"}
+{#if $view.engine_state === "paused"}
{/if}
-{#each $my_rollout_query.rollouts as rollout}
+{#each $view.rollouts as rollout}
{#if (($url.hash === "" || $url.hash === "#active") && rollout.state !== "complete" && rollout.state !== "failed") || ($url.hash === "#complete" && rollout.state === "complete") || ($url.hash === "#failed" && rollout.state === "failed") || $url.hash === "#all"}
{/if}
diff --git a/rollout-dashboard/frontend/src/lib/stores.ts b/rollout-dashboard/frontend/src/lib/stores.ts
index df29864..ed74bf3 100644
--- a/rollout-dashboard/frontend/src/lib/stores.ts
+++ b/rollout-dashboard/frontend/src/lib/stores.ts
@@ -1,28 +1,17 @@
import { get, writable, type Writable } from 'svelte/store'
-import { type Rollout } from './types'
+import { type Rollout, type RolloutsViewDelta } from './types'
-const BACKEND_TIMEOUT = 15000
-
-type SseMessage = {
- error: [number, string] | null;
- rollouts: Rollout[];
- updated: Rollout[] | undefined;
- deleted: String[] | undefined;
- engine_state?: string;
-}
-
-export type RolloutResult = {
+export type RolloutsView = {
error: [number, string] | string | null;
rollouts: Rollout[];
engine_state?: string;
}
-
const API_URL = import.meta.env.BACKEND_API_PATH || "/api/v1";
-const url = API_URL + "/rollouts/sse?incremental"
+const url = API_URL + "/sse/rollouts_view?incremental"
var evtSource: null | EventSource = null;
-const rollout_store = writable
({ rollouts: [], error: "loading" })
+const airflow_state = writable({ rollouts: [], error: "loading" })
function resetupEventSource() {
if (null !== evtSource) {
@@ -39,11 +28,11 @@ function resetupEventSource() {
// var evtSourceGenerated = new Date();
evtSource = new EventSource(url);
evtSource.onmessage = async function (event) {
- var sse_message: SseMessage = JSON.parse(event.data);
+ var sse_message: RolloutsViewDelta = JSON.parse(event.data);
if (sse_message.error !== null) {
let status = sse_message.error[0];
if (status == 204) {
- rollout_store.set({ rollouts: [], error: "loading" })
+ airflow_state.set({ rollouts: [], error: "loading" })
} else {
let responseText = sse_message.error[1];
let errorText = status + " " + responseText;
@@ -52,21 +41,21 @@ function resetupEventSource() {
errorText = errorText + ": " + responseText
}
console.log('Request for rollout data failed: ' + errorText)
- rollout_store.set({
- rollouts: get(rollout_store).rollouts,
- engine_state: get(rollout_store).engine_state,
+ airflow_state.set({
+ rollouts: get(airflow_state).rollouts,
+ engine_state: get(airflow_state).engine_state,
error: errorText
})
}
} else if (sse_message.rollouts !== undefined) {
console.log("Full sync with " + sse_message.rollouts.length + " rollouts and engine state " + sse_message.engine_state);
- rollout_store.set({
+ airflow_state.set({
rollouts: sse_message.rollouts,
error: null,
engine_state: sse_message.engine_state,
})
} else {
- var rollouts: Rollout[] = get(rollout_store).rollouts;
+ var rollouts: Rollout[] = get(airflow_state).rollouts;
var updated: Rollout[] | undefined = sse_message["updated"];
var deleted: String[] | undefined = sse_message["deleted"];
if (updated !== undefined) {
@@ -97,7 +86,7 @@ function resetupEventSource() {
}
}
}
- rollout_store.set({
+ airflow_state.set({
rollouts: rollouts,
error: null,
engine_state: sse_message.engine_state,
@@ -108,15 +97,15 @@ function resetupEventSource() {
console.log({ message: "Disconnected from event source. Reconnecting in 5 seconds.", event: e })
if (evtSource !== null) { evtSource.close(); evtSource = null; }
var errorText = 'Rollout dashboard is down — reconnecting in 5 seconds'
- rollout_store.set({
- rollouts: get(rollout_store).rollouts,
+ airflow_state.set({
+ rollouts: get(airflow_state).rollouts,
error: errorText
})
setTimeout(resetupEventSource, 5000)
}
}
-export const rollout_query = ((): Writable => {
+export const rollouts_view = ((): Writable => {
resetupEventSource()
- return rollout_store
+ return airflow_state
});
diff --git a/rollout-dashboard/frontend/src/lib/types.ts b/rollout-dashboard/frontend/src/lib/types.ts
index bbbaf68..f03f4f2 100644
--- a/rollout-dashboard/frontend/src/lib/types.ts
+++ b/rollout-dashboard/frontend/src/lib/types.ts
@@ -77,3 +77,10 @@ export type Rollout = {
last_scheduling_decision?: Date;
batches: Batch[];
};
+export type RolloutsViewDelta = {
+ error: [number, string] | null;
+ rollouts: Rollout[];
+ updated: Rollout[] | undefined;
+ deleted: String[] | undefined;
+ engine_state?: string;
+}
diff --git a/rollout-dashboard/server/src/api_server.rs b/rollout-dashboard/server/src/api_server.rs
new file mode 100644
index 0000000..137e96b
--- /dev/null
+++ b/rollout-dashboard/server/src/api_server.rs
@@ -0,0 +1,213 @@
+use crate::live_state::{AirflowStateSyncer, Live, RolloutInfoCacheResponse, RolloutsView};
+use async_stream::try_stream;
+use axum::extract::Query;
+use axum::http::StatusCode;
+use axum::response::sse;
+use axum::response::Sse;
+use axum::routing::get;
+use axum::{Json, Router};
+use futures::stream::Stream;
+use log::debug;
+use rollout_dashboard::types::Rollout;
+use rollout_dashboard::types::RolloutEngineState;
+use rollout_dashboard::types::RolloutsViewDelta;
+use serde::{de, Deserialize, Deserializer};
+use std::collections::{HashMap, HashSet, VecDeque};
+use std::convert::Infallible;
+use std::fmt;
+use std::str::FromStr;
+use std::sync::Arc;
+use tokio::time::Duration;
+
+/// Serde deserialization decorator to map empty Strings to None,
+fn empty_value_as_true<'de, D, T>(de: D) -> Result, D::Error>
+where
+ D: Deserializer<'de>,
+ T: FromStr,
+ T::Err: fmt::Display,
+{
+ let opt = Option::::deserialize(de)?;
+ match opt.as_deref() {
+ Some("") | None => FromStr::from_str("true")
+ .map_err(de::Error::custom)
+ .map(Some),
+ Some(s) => FromStr::from_str(s).map_err(de::Error::custom).map(Some),
+ }
+}
+
+#[derive(Deserialize)]
+struct SseHandlerParameters {
+ #[serde(default, deserialize_with = "empty_value_as_true")]
+ incremental: Option,
+}
+
+pub(crate) struct ApiServer {
+ state_syncer: Arc>,
+}
+
+impl ApiServer {
+ pub fn new(state_syncer: Arc>) -> Self {
+ Self { state_syncer }
+ }
+
+ // #[debug_handler]
+ async fn get_rollout_data(&self) -> Result>, (StatusCode, String)> {
+ match self.state_syncer.get_current_rollout_status().await {
+ Ok((_, rollouts)) => Ok(Json(rollouts)),
+ Err(e) => Err(e),
+ }
+ }
+
+ // #[debug_handler]
+ async fn get_engine_state(&self) -> Result, (StatusCode, String)> {
+ match self.state_syncer.get_current_rollout_status().await {
+ Ok((state, _)) => Ok(Json(state)),
+ Err(e) => Err(e),
+ }
+ }
+
+ // #[debug_handler]
+ async fn get_cache(&self) -> Result>, (StatusCode, String)> {
+ Ok(Json(self.state_syncer.get_cache().await))
+ }
+
+ /// Produce an SSE stream structured as a dictionary:
+ /// * rollouts: appears and contains a list of Rollout when rollouts have been updated and
+ /// the caller did not indicate delta_support. If the caller indicated
+ /// delta_support, then this appears only on initial connection, or after a
+ /// backend error has been reported in a prior message.
+ /// * error: always appears, but only contains an error (is non-null) when there was an
+ /// error polling Airflow.
+ /// * updated: when the caller indicates delta_support, this appears and updated rollouts
+ /// are listed here.
+ /// * deleted: when the caller indicates delta_support, this appears and lists the names
+ /// of the rollouts that have disappeared.
+ /// * engine_state: explains the state of the rollout engine. Only present if no error
+ /// is present.
+ pub fn produce_rollouts_sse_stream(
+ &self,
+ delta_support: bool,
+ ) -> Sse>> {
+ struct DisconnectionGuard {}
+
+ impl Default for DisconnectionGuard {
+ fn default() -> Self {
+ debug!(target: "server::sse", "New client connected.");
+ Self {}
+ }
+ }
+
+ impl Drop for DisconnectionGuard {
+ fn drop(&mut self) {
+ debug!(target: "server::sse", "Client disconnected.");
+ }
+ }
+
+ let mut stream_rx = self.state_syncer.subscribe_to_state_updates();
+
+ let stream = try_stream! {
+ // Set up something that will be dropped (thus log) when SSE is disconnected.
+ let disconnection_guard = DisconnectionGuard::default();
+
+ // Set an initial message to diff the first broadcast message against.
+ let mut last_rollout_data: RolloutsView = Err((StatusCode::OK, "initial, ignored, nothing matches this".to_string()));
+
+ loop {
+ let current_rollout_data = stream_rx.borrow_and_update().clone();
+
+ let message = match (¤t_rollout_data, &last_rollout_data) {
+ // Error before. Send full sync.
+ (Ok((engine_state, new_rollouts)), Err(_)) => Some(RolloutsViewDelta::full(engine_state, new_rollouts)),
+ // Last time was a good update. Send differential sync.
+ (Ok((engine_state, new_rollouts)), Ok((_, old_rollouts))) => {
+ let new_names = new_rollouts.iter().map(|r| r.name.clone()).collect::>();
+ let old_rollouts_map = old_rollouts.iter().map(|r| (r.name.clone(), r)).collect::>();
+ let updated = new_rollouts.iter().filter_map(|r| match old_rollouts_map.get(&r.name) { None => Some(r.clone()), Some(old_rollout) => match r.update_count != old_rollout.update_count {true => Some(r.clone()), false => None}}).collect::>();
+ let deleted = old_rollouts.iter().filter_map(|r| match new_names.contains(&r.name) { true => None, false => Some(r.name.clone())}).collect::>();
+ match updated.is_empty() && deleted.is_empty() {
+ true => None,
+ false => match delta_support {
+ false => Some(RolloutsViewDelta::full(engine_state, new_rollouts)),
+ true => Some(RolloutsViewDelta::partial(
+ engine_state,
+ &updated,
+ &deleted,
+ )),
+ }
+ }
+ }
+ // Error after a good update. Send error.
+ (Err(e), Ok(_)) => Some(RolloutsViewDelta::error(e)),
+ // Error after an error. Only send update if errors differ.
+ (Err(e), Err(olde)) => match e == olde {
+ true => None,
+ false => Some(RolloutsViewDelta::error(e)),
+ },
+ };
+
+ match &message {
+ Some(m) => yield sse::Event::default().data(serde_json::to_string(m).unwrap()),
+ None => ()
+ }
+
+ last_rollout_data = current_rollout_data;
+ if stream_rx.changed().await.is_err() {
+ break;
+ }
+ }
+
+ // Drop the disconnection guard to log the message that the client disconnected.
+ drop(disconnection_guard);
+ };
+
+ Sse::new(stream).keep_alive(
+ axum::response::sse::KeepAlive::new()
+ .interval(Duration::from_secs(5))
+ .text("keepalive"),
+ )
+ }
+
+ fn v1_api(self: Arc) -> Router {
+ let rollouts_handler_ref = self.clone();
+ let engine_state_handler_ref = self.clone();
+ let compat_sse_handler_ref = self.clone();
+ let sse_handler_ref = self.clone();
+ Router::new()
+ .route(
+ "/rollouts",
+ get(move || async move { rollouts_handler_ref.get_rollout_data().await }),
+ )
+ .route(
+ "/engine_state",
+ get(move || async move { engine_state_handler_ref.get_engine_state().await }),
+ )
+ .route(
+ "/rollouts/sse",
+ get(move |options: Query| {
+ let options: SseHandlerParameters = options.0;
+ async move {
+ compat_sse_handler_ref
+ .produce_rollouts_sse_stream(options.incremental.unwrap_or_default())
+ }
+ }),
+ )
+ .route(
+ "/sse/rollouts_view",
+ get(move || async move { sse_handler_ref.produce_rollouts_sse_stream(true) }),
+ )
+ }
+
+ fn unstable_api(self: Arc) -> Router {
+ let cached_data_handler_ref = self.clone();
+ Router::new().route(
+ "/cache",
+ get(move || async move { cached_data_handler_ref.get_cache().await }),
+ )
+ }
+
+ pub fn routes(self: Arc) -> Router {
+ Router::new()
+ .nest("/api/v1", self.clone().v1_api())
+ .nest("/api/unstable", self.unstable_api())
+ }
+}
diff --git a/rollout-dashboard/server/src/lib.rs b/rollout-dashboard/server/src/lib.rs
index e70c44d..096f8d4 100644
--- a/rollout-dashboard/server/src/lib.rs
+++ b/rollout-dashboard/server/src/lib.rs
@@ -21,4 +21,5 @@
//! available by appending `/sse` to the rollouts API call. The server
//! will stream the initial state right away, then stream updates to
//! the rollout data as they happen.
+pub mod airflow_client;
pub mod types;
diff --git a/rollout-dashboard/server/src/frontend_api.rs b/rollout-dashboard/server/src/live_state/mod.rs
similarity index 85%
rename from rollout-dashboard/server/src/frontend_api.rs
rename to rollout-dashboard/server/src/live_state/mod.rs
index 03c9a32..e009259 100644
--- a/rollout-dashboard/server/src/frontend_api.rs
+++ b/rollout-dashboard/server/src/live_state/mod.rs
@@ -1,22 +1,23 @@
-use crate::airflow_client::{
- AirflowClient, AirflowError, DagRunState, DagRunsResponseItem, DagsQueryFilter, DagsResponse,
- DagsResponseItem, EventLogsResponseFilters, TaskInstanceRequestFilters, TaskInstanceState,
- TaskInstancesResponseItem, TasksResponse, TasksResponseItem,
-};
-use crate::python;
use chrono::{DateTime, Utc};
use futures::future::join_all;
use indexmap::IndexMap;
use lazy_static::lazy_static;
-use log::{debug, info, trace, warn};
+use log::{debug, error, info, trace, warn};
use regex::Regex;
+use reqwest::StatusCode;
+use rollout_dashboard::airflow_client::{
+ AirflowClient, AirflowError, DagRunState, DagRunsResponseItem, DagsQueryFilter,
+ EventLogsResponseFilters, TaskInstanceRequestFilters, TaskInstanceState,
+ TaskInstancesResponseItem, TasksResponse, TasksResponseItem,
+};
use rollout_dashboard::types::{
- Batch, Rollout, RolloutState, Rollouts, Subnet, SubnetRolloutState,
+ Batch, Rollout, RolloutEngineState, RolloutState, Rollouts, Subnet, SubnetRolloutState,
};
use serde::{Serialize, Serializer};
use std::cmp::{max, min};
use std::collections::hash_map::Entry::{Occupied, Vacant};
-use std::collections::{HashMap, HashSet};
+use std::collections::{HashMap, HashSet, VecDeque};
+use std::error::Error;
use std::fmt::{self, Display};
use std::future::Future;
use std::num::ParseIntError;
@@ -25,9 +26,15 @@ use std::rc::Rc;
use std::str::FromStr;
use std::sync::Arc;
use std::{vec, vec::Vec};
+use tokio::sync::watch::{self, Sender};
use tokio::sync::Mutex;
+use tokio::task::JoinHandle;
+use tokio::time::{sleep, Duration};
+use tokio::{select, spawn};
use topological_sort::TopologicalSort;
+mod python;
+
lazy_static! {
// unwrap() is legitimate here because we know these cannot fail to compile.
static ref SubnetGitRevisionRe: Regex = Regex::new("dfinity.ic_types.SubnetRolloutInstance.*@version=0[(]start_at=.*,subnet_id=([0-9-a-z-]+),git_revision=([0-9a-f]+)[)]").unwrap();
@@ -49,7 +56,7 @@ where
}
#[derive(Debug)]
-pub struct CyclicDependencyError {
+struct CyclicDependencyError {
message: String,
}
@@ -242,7 +249,7 @@ impl RolloutPlan {
}
#[derive(Debug)]
-pub enum RolloutDataGatherError {
+enum RolloutDataGatherError {
AirflowError(AirflowError),
CyclicDependency(CyclicDependencyError),
RolloutPlanParseError(RolloutPlanParseError),
@@ -474,7 +481,7 @@ impl AirflowIncrementalLogInspector {
|| (event.event == "failed" && event.extra.is_some())
|| (event.event == "clear" && event.extra.is_some());
- trace!(target: "frontend_api::log_inspector", "Processing event:\n{:#?}\n", event);
+ trace!(target: "rollout_data::log_inspector", "Processing event:\n{:#?}\n", event);
match task_instances_to_update_per_dag
.dag_runs
@@ -484,13 +491,13 @@ impl AirflowIncrementalLogInspector {
Vacant(ventry) => {
ventry.insert(match (&event.task_id, force_refresh_all_tasks) {
(Some(t), false) => {
- trace!(target: "frontend_api::log_inspector", "{}: initializing plan with a request to update task {}", event_run_id, t);
+ trace!(target: "rollout_data::log_inspector", "{}: initializing plan with a request to update task {}", event_run_id, t);
let mut init = HashSet::new();
init.insert(t.clone());
DagRunUpdateType::SomeTaskInstances(init)
},
_ => {
- trace!(target: "frontend_api::log_inspector", "{}: initializing plan with a request to update all tasks", event_run_id);
+ trace!(target: "rollout_data::log_inspector", "{}: initializing plan with a request to update all tasks", event_run_id);
DagRunUpdateType::AllTaskInstances
},
});
@@ -501,14 +508,14 @@ impl AirflowIncrementalLogInspector {
if let DagRunUpdateType::SomeTaskInstances(thevec) = entry.get_mut() {
let ts = t.to_string();
if !thevec.contains(&ts) {
- trace!(target: "frontend_api::log_inspector", "{}: adding task {} to plan", event_run_id, ts);
+ trace!(target: "rollout_data::log_inspector", "{}: adding task {} to plan", event_run_id, ts);
thevec.insert(ts);
}
}
}
_ => {
if let DagRunUpdateType::SomeTaskInstances(_) = entry.get() {
- trace!(target: "frontend_api::log_inspector", "{}: switching plan to request to update all tasks", event_run_id);
+ trace!(target: "rollout_data::log_inspector", "{}: switching plan to request to update all tasks", event_run_id);
entry.insert(DagRunUpdateType::AllTaskInstances);
}
}
@@ -518,7 +525,7 @@ impl AirflowIncrementalLogInspector {
// Now that we have a plan, we know what data to fetch from Airflow, minimizing the load on the server.
for (k, v) in task_instances_to_update_per_dag.dag_runs.iter() {
- debug!(target: "frontend_api::log_inspector", "{}: tasks that will be updated: {}", k, match v {
+ debug!(target: "rollout_data::log_inspector", "{}: tasks that will be updated: {}", k, match v {
DagRunUpdateType::AllTaskInstances => "all tasks".to_string(),
DagRunUpdateType::SomeTaskInstances(set_of_tasks) => set_of_tasks.iter().cloned().collect::>().join(", "),
});
@@ -527,7 +534,7 @@ impl AirflowIncrementalLogInspector {
&& !task_instances_to_update_per_dag.dag_runs.is_empty()
{
debug!(
- target: "frontend_api::log_inspector", "Setting incremental refresh date to {:?}",
+ target: "rollout_data::log_inspector", "Setting incremental refresh date to {:?}",
last_event_log_update
)
};
@@ -548,7 +555,7 @@ impl AirflowIncrementalLogInspector {
last_event_log_update = Some(event.when);
}
if !event_logs.event_logs.is_empty() {
- debug!(target: "frontend_api::log_inspector", "Setting initial refresh date to {:?}", last_event_log_update);
+ debug!(target: "rollout_data::log_inspector", "Setting initial refresh date to {:?}", last_event_log_update);
}
}
@@ -562,7 +569,7 @@ impl AirflowIncrementalLogInspector {
}
#[derive(Clone)]
-struct RolloutDataCache {
+struct RolloutInfoCache {
task_instances: HashMap, TaskInstancesResponseItem>>,
dispatch_time: DateTime,
note: Option,
@@ -590,7 +597,7 @@ where
}
#[derive(Serialize)]
-pub struct RolloutDataCacheResponse {
+pub struct RolloutInfoCacheResponse {
rollout_id: String,
dispatch_time: DateTime,
#[serde(serialize_with = "serialize_cache_response")]
@@ -600,46 +607,10 @@ pub struct RolloutDataCacheResponse {
linearized_task_instances: Vec,
}
-#[derive(Serialize, Debug, Clone)]
-#[serde(rename_all = "snake_case")]
-pub enum RolloutEngineState {
- Missing,
- Broken,
- Paused,
- Inactive,
- Active,
-}
-
-impl From for RolloutEngineState {
- fn from(resp: DagsResponse) -> Self {
- match resp.dags.len() {
- 0 => RolloutEngineState::Missing,
- _ => match resp.dags[0] {
- DagsResponseItem {
- has_import_errors: true,
- ..
- } => Self::Broken,
- DagsResponseItem {
- is_paused: true, ..
- } => Self::Paused,
- DagsResponseItem {
- is_active: false, ..
- } => Self::Inactive,
- DagsResponseItem {
- is_active: true,
- is_paused: false,
- has_import_errors: false,
- ..
- } => Self::Active,
- },
- }
- }
-}
-
-struct RolloutApiCache {
+struct RolloutStateCache {
/// Map from DAG run ID to task instance ID (with / without index)
/// to task instance.
- by_dag_run: HashMap,
+ by_dag_run: HashMap,
log_inspector: AirflowIncrementalLogInspector,
}
@@ -668,10 +639,10 @@ fn annotate_subnet_state(
if (only_decrease && new_state < subnet.state)
|| (!only_decrease && new_state != subnet.state)
{
- trace!(target: "frontend_api::annotate_subnet_state", "{}: {} {:?} transition {} => {} note: {}", task_instance.dag_run_id, task_instance.task_id, task_instance.map_index, subnet.state, new_state, subnet.comment);
+ trace!(target: "rollout_data::annotate_subnet_state", "{}: {} {:?} transition {} => {} note: {}", task_instance.dag_run_id, task_instance.task_id, task_instance.map_index, subnet.state, new_state, subnet.comment);
subnet.state = new_state.clone();
} else {
- trace!(target: "frontend_api::annotate_subnet_state", "{}: {} {:?} NO transition {} => {}", task_instance.dag_run_id, task_instance.task_id, task_instance.map_index, subnet.state, new_state);
+ trace!(target: "rollout_data::annotate_subnet_state", "{}: {} {:?} NO transition {} => {}", task_instance.dag_run_id, task_instance.task_id, task_instance.map_index, subnet.state, new_state);
}
if new_state == subnet.state {
subnet.comment = format!(
@@ -706,7 +677,7 @@ fn annotate_subnet_state(
struct RolloutUpdater<'a> {
dag_run: &'a DagRunsResponseItem,
- cache_entry: RolloutDataCache,
+ cache_entry: RolloutInfoCache,
update_type: DagRunUpdateType,
}
@@ -716,7 +687,7 @@ impl<'a> RolloutUpdater<'a> {
airflow_api: &AirflowClient,
sorter: TaskInstanceTopologicalSorter,
last_event_log_update: Option>,
- ) -> Result<(Rollout, RolloutDataCache), RolloutDataGatherError> {
+ ) -> Result<(Rollout, RolloutInfoCache), RolloutDataGatherError> {
// If the note of the rollout has changed,
// note that this has been updated.
let cache_entry = &mut self.cache_entry;
@@ -749,7 +720,7 @@ impl<'a> RolloutUpdater<'a> {
cache_entry.task_instances.is_empty(),
) {
(_, true) | (DagRunUpdateType::AllTaskInstances, _) => {
- debug!(target:"frontend_api::get_rollout_data", "{}: collecting data about all task instances", dag_run.dag_run_id);
+ debug!(target:"rollout_data::get_rollout_data", "{}: collecting data about all task instances", dag_run.dag_run_id);
vec![Box::pin(async move {
match airflow_api
.task_instances(
@@ -769,7 +740,7 @@ impl<'a> RolloutUpdater<'a> {
(DagRunUpdateType::SomeTaskInstances(updated_task_instances), false) => {
let updated_task_instances =
updated_task_instances.iter().cloned().collect::>();
- debug!(target:"frontend_api::get_rollout_data", "{}: collecting data about task instances updated since {:?} and a specific set of tasks too: {:?}", dag_run_id, last_update_time, updated_task_instances);
+ debug!(target:"rollout_data::get_rollout_data", "{}: collecting data about task instances updated since {:?} and a specific set of tasks too: {:?}", dag_run_id, last_update_time, updated_task_instances);
vec![
Box::pin(async move {
match airflow_api
@@ -852,7 +823,7 @@ impl<'a> RolloutUpdater<'a> {
}
debug!(
- target: "frontend_api::get_rollout_data", "{}: retrieved {} tasks",
+ target: "rollout_data::get_rollout_data", "{}: retrieved {} tasks",
dag_run_id, retrieved_task_instances.len()
);
@@ -911,7 +882,7 @@ impl<'a> RolloutUpdater<'a> {
if tasks.len() > 1 {
if let Occupied(_) = tasks.entry(None) {
debug!(
- target: "frontend_api::get_rollout_data", "formerly unmapped task {} is now mapped",
+ target: "rollout_data::get_rollout_data", "formerly unmapped task {} is now mapped",
task_instance_id
);
tasks.remove(&None);
@@ -927,7 +898,7 @@ impl<'a> RolloutUpdater<'a> {
.collect();
debug!(
- target: "frontend_api::get_rollout_data", "{}: total disambiguated tasks including locally cached ones: {}",
+ target: "rollout_data::get_rollout_data", "{}: total disambiguated tasks including locally cached ones: {}",
dag_run.dag_run_id, linearized_tasks.len(),
);
@@ -943,7 +914,7 @@ impl<'a> RolloutUpdater<'a> {
// * handle tasks corresponding to a batch/subnet in a special way
// (commented below in its pertinent section).
trace!(
- target: "frontend_api::get_rollout_data", "Processing task {}.{:?} in state {:?}",
+ target: "rollout_data::get_rollout_data", "Processing task {}.{:?} in state {:?}",
task_instance.task_id, task_instance.map_index, task_instance.state,
);
if task_instance.task_id == "schedule" {
@@ -970,7 +941,7 @@ impl<'a> RolloutUpdater<'a> {
}
ScheduleCacheValidity::Stale => {
// Same schedule task has been updated. Data may not be missing anymore.
- info!(target: "frontend_api::get_rollout_data", "{}: schedule task is outdated; requerying", dag_run.dag_run_id);
+ info!(target: "rollout_data::get_rollout_data", "{}: schedule task is outdated; requerying", dag_run.dag_run_id);
match airflow_api
.xcom_entry(
dag_id,
@@ -984,14 +955,14 @@ impl<'a> RolloutUpdater<'a> {
Ok(schedule) => {
match &RolloutPlan::from_python_string(&schedule.value) {
Ok(schedule) => {
- info!(target: "frontend_api::get_rollout_data", "{}: saving schedule cache", dag_run.dag_run_id);
+ info!(target: "rollout_data::get_rollout_data", "{}: saving schedule cache", dag_run.dag_run_id);
cache_entry
.schedule
.update(&task_instance, &schedule.batches);
schedule.batches.clone()
}
Err(e) => {
- warn!(target: "frontend_api::get_rollout_data", "{}: could not parse schedule data: {}", dag_run.dag_run_id, e);
+ warn!(target: "rollout_data::get_rollout_data", "{}: could not parse schedule data: {}", dag_run.dag_run_id, e);
cache_entry.schedule.invalidate(
&task_instance,
Some(schedule.value),
@@ -1006,7 +977,7 @@ impl<'a> RolloutUpdater<'a> {
// There is no schedule to be found.
// Or there was no schedule to be found last time
// it was queried.
- warn!(target: "frontend_api::get_rollout_data", "{}: no schedule despite schedule task finished", dag_run.dag_run_id);
+ warn!(target: "rollout_data::get_rollout_data", "{}: no schedule despite schedule task finished", dag_run.dag_run_id);
cache_entry.schedule.invalidate(&task_instance, None);
continue;
}
@@ -1050,7 +1021,7 @@ impl<'a> RolloutUpdater<'a> {
// * update the subnet link to the corresponding Airflow task if the
// state of the task (after update) corresponds to the expected state,
// * update rollout state to problem / error depending on the task state.
- trace!(target: "frontend_api::get_rollout_data::subnet_state", "{}: processing {} {:?} in state {:?}", task_instance.dag_run_id, task_instance.task_id, task_instance.map_index, task_instance.state);
+ trace!(target: "rollout_data::get_rollout_data::subnet_state", "{}: processing {} {:?} in state {:?}", task_instance.dag_run_id, task_instance.task_id, task_instance.map_index, task_instance.state);
let (batch, task_name) = (
// We get away with unwrap() here because we know we captured an integer.
match rollout
@@ -1059,7 +1030,7 @@ impl<'a> RolloutUpdater<'a> {
{
Some(batch) => batch,
None => {
- trace!(target: "frontend_api::get_rollout_data::subnet_state", "{}: no corresponding batch, continuing", task_instance.dag_run_id);
+ trace!(target: "rollout_data::get_rollout_data::subnet_state", "{}: no corresponding batch, continuing", task_instance.dag_run_id);
continue;
}
},
@@ -1090,7 +1061,7 @@ impl<'a> RolloutUpdater<'a> {
if task_name == "collect_batch_subnets" {
trans_exact!(SubnetRolloutState::Pending);
} else {
- trace!(target: "frontend_api::get_rollout_data::subnet_state", "{}: ignoring task instance {} {:?} with no state", task_instance.dag_run_id, task_instance.task_id, task_instance.map_index);
+ trace!(target: "rollout_data::get_rollout_data::subnet_state", "{}: ignoring task instance {} {:?} with no state", task_instance.dag_run_id, task_instance.task_id, task_instance.map_index);
}
}
Some(state) => match state {
@@ -1100,7 +1071,7 @@ impl<'a> RolloutUpdater<'a> {
// If a task is skipped, the next task (in state Running / Deferred)
// will pick up the slack for changing subnet state.
TaskInstanceState::Removed | TaskInstanceState::Skipped => {
- trace!(target: "frontend_api::get_rollout_data::subnet_state", "{}: ignoring task instance {} {:?} in state {:?}", task_instance.dag_run_id, task_instance.task_id, task_instance.map_index, task_instance.state);
+ trace!(target: "rollout_data::get_rollout_data::subnet_state", "{}: ignoring task instance {} {:?} in state {:?}", task_instance.dag_run_id, task_instance.task_id, task_instance.map_index, task_instance.state);
}
TaskInstanceState::UpForRetry | TaskInstanceState::Restarting => {
trans_min!(SubnetRolloutState::Error);
@@ -1148,7 +1119,7 @@ impl<'a> RolloutUpdater<'a> {
trans_min!(SubnetRolloutState::Complete);
}
&_ => {
- warn!(target: "frontend_api::get_rollout_data::subnet_state", "{}: no info on to handle task instance {} {:?} in state {:?}", task_instance.dag_run_id, task_instance.task_id, task_instance.map_index, task_instance.state);
+ warn!(target: "rollout_data::get_rollout_data::subnet_state", "{}: no info on to handle task instance {} {:?} in state {:?}", task_instance.dag_run_id, task_instance.task_id, task_instance.map_index, task_instance.state);
}
}
rollout.state = min(rollout.state, RolloutState::UpgradingSubnets)
@@ -1222,7 +1193,7 @@ impl<'a> RolloutUpdater<'a> {
batch.end_time = task_instance.end_date;
}
&_ => {
- warn!(target: "frontend_api::get_rollout_data::subnet_state", "{}: no info on how to handle task instance {} {:?} in state {:?}", task_instance.dag_run_id, task_instance.task_id, task_instance.map_index, task_instance.state);
+ warn!(target: "rollout_data::get_rollout_data::subnet_state", "{}: no info on how to handle task instance {} {:?} in state {:?}", task_instance.dag_run_id, task_instance.task_id, task_instance.map_index, task_instance.state);
}
},
},
@@ -1247,7 +1218,7 @@ impl<'a> RolloutUpdater<'a> {
}
}
} else {
- warn!(target: "frontend_api::get_rollout_data::subnet_state", "{}: unknown task {}", task_instance.dag_run_id, task_instance.task_id)
+ warn!(target: "rollout_data::get_rollout_data::subnet_state", "{}: unknown task {}", task_instance.dag_run_id, task_instance.task_id)
}
}
@@ -1271,23 +1242,23 @@ impl<'a> RolloutUpdater<'a> {
}
}
#[derive(Clone)]
-pub struct RolloutApi {
+pub(crate) struct AirflowStateUpdater {
airflow_api: Arc,
- cache: Arc>,
+ cache: Arc>,
}
-impl RolloutApi {
+impl AirflowStateUpdater {
pub fn new(client: AirflowClient) -> Self {
Self {
airflow_api: Arc::new(client),
- cache: Arc::new(Mutex::new(RolloutApiCache {
+ cache: Arc::new(Mutex::new(RolloutStateCache {
by_dag_run: HashMap::new(),
log_inspector: AirflowIncrementalLogInspector::default(),
})),
}
}
- pub async fn get_cache(&self) -> Vec {
+ pub async fn get_cache(&self) -> Vec {
let cache = self.cache.lock().await;
let mut result: Vec<_> = cache
.by_dag_run
@@ -1298,7 +1269,7 @@ impl RolloutApi {
.iter()
.flat_map(|(_, tasks)| tasks.iter().map(|(_, task)| task.clone()))
.collect();
- RolloutDataCacheResponse {
+ RolloutInfoCacheResponse {
rollout_id: k.clone(),
linearized_task_instances: linearized_tasks,
dispatch_time: v.dispatch_time,
@@ -1329,7 +1300,7 @@ impl RolloutApi {
/// true return in the update flag). Currently, only a change
/// in the rollout note, the state of any of its tasks, or
/// the rollout dispatch time are considered meaningful changes.
- pub async fn get_rollout_data(
+ async fn update(
&self,
max_rollouts: usize,
) -> Result<(RolloutEngineState, Rollouts), RolloutDataGatherError> {
@@ -1380,7 +1351,7 @@ impl RolloutApi {
.by_dag_run
.get(&dag_run.dag_run_id)
.map(ToOwned::to_owned)
- .unwrap_or(RolloutDataCache {
+ .unwrap_or(RolloutInfoCache {
task_instances: HashMap::new(),
dispatch_time: dag_run.logical_date,
note: dag_run.note.clone(),
@@ -1416,3 +1387,169 @@ impl RolloutApi {
Ok((engine_state, res))
}
}
+
+/// Contains either a list of rollouts ordered from newest to oldest,
+/// dated from the last time it was successfully updated, or an HTTP
+/// status code corresponding to -- and with -- a message for the last error.
+pub type RolloutsView = Result<(RolloutEngineState, VecDeque), (StatusCode, String)>;
+
+pub(crate) struct Initial;
+pub(crate) struct Live;
+
+pub(crate) struct AirflowStateSyncer {
+ state_updater: AirflowStateUpdater,
+ last_rollout_data: Arc>,
+ stream_tx: Sender,
+ refresh_interval: u64,
+ max_rollouts: usize,
+ #[allow(dead_code)]
+ state: S,
+}
+
+impl AirflowStateSyncer {
+ pub fn new(
+ rollout_api: AirflowStateUpdater,
+ max_rollouts: usize,
+ refresh_interval: u64,
+ ) -> Self {
+ let init: RolloutsView = Err((StatusCode::NO_CONTENT, "".to_string()));
+ let (stream_tx, _stream_rx) = watch::channel::(init.clone());
+ Self {
+ state_updater: rollout_api,
+ last_rollout_data: Arc::new(Mutex::new(init)),
+ stream_tx,
+ refresh_interval,
+ max_rollouts,
+ state: Initial,
+ }
+ }
+
+ pub fn start_syncing(
+ self,
+ mut cancel_receiver: watch::Receiver<()>,
+ ) -> (Arc>, JoinHandle<()>) {
+ let ret: Arc> = Arc::new(AirflowStateSyncer {
+ state: Live,
+ state_updater: self.state_updater,
+ last_rollout_data: self.last_rollout_data,
+ stream_tx: self.stream_tx,
+ refresh_interval: self.refresh_interval,
+ max_rollouts: self.max_rollouts,
+ });
+ let looper = ret.clone();
+ let background_loop_fut = spawn(async move {
+ looper
+ .periodically_sync_state(async move {
+ let _ = cancel_receiver.changed().await;
+ })
+ .await
+ });
+ (ret, background_loop_fut)
+ }
+}
+
+impl AirflowStateSyncer {
+ pub async fn get_current_rollout_status(&self) -> RolloutsView {
+ self.last_rollout_data.lock().await.clone()
+ }
+
+ pub async fn get_cache(&self) -> Vec {
+ self.state_updater.get_cache().await
+ }
+
+ /// Create a channel that will get state updates as soon as they are available.
+ /// This needs `periodically_refresh_state` running in a coroutine.
+ pub fn subscribe_to_state_updates(&self) -> watch::Receiver {
+ self.stream_tx.subscribe()
+ }
+
+ async fn sync_state(
+ &self,
+ max_rollouts: usize,
+ ) -> Result<(RolloutEngineState, VecDeque), (StatusCode, String)> {
+ match self.state_updater.update(max_rollouts).await {
+ Ok((engine_state, rollouts)) => Ok((engine_state, rollouts.into())),
+ Err(e) => {
+ let res = match e {
+ RolloutDataGatherError::AirflowError(AirflowError::StatusCode(c)) => {
+ (c, "Internal server error".to_string())
+ }
+ RolloutDataGatherError::AirflowError(AirflowError::ReqwestError(err)) => {
+ let mut explanation = format!("Cannot contact Airflow: {}", err);
+ let mut err = err.source();
+ loop {
+ match err {
+ None => break,
+ Some(e) => {
+ explanation = format!("{} -> {}", explanation.as_str(), e);
+ err = e.source();
+ }
+ }
+ }
+ (StatusCode::BAD_GATEWAY, explanation)
+ }
+ RolloutDataGatherError::AirflowError(AirflowError::Other(msg)) => {
+ (StatusCode::INTERNAL_SERVER_ERROR, msg)
+ }
+ RolloutDataGatherError::RolloutPlanParseError(parse_error) => (
+ StatusCode::INTERNAL_SERVER_ERROR,
+ format!("{}", parse_error),
+ ),
+ RolloutDataGatherError::CyclicDependency(dep) => (
+ StatusCode::INTERNAL_SERVER_ERROR,
+ format!("A cyclic dependency was found in the task graph: {:?}", dep),
+ ),
+ };
+ Err(res)
+ }
+ }
+ }
+
+ async fn periodically_sync_state(self: Arc, cancel: F)
+ where
+ F: Future