126 lines
3.7 KiB
Rust
126 lines
3.7 KiB
Rust
/*
|
|
* Author: Ruben Fiszel
|
|
* Copyright: Windmill Labs, Inc 2022
|
|
* This file and its contents are licensed under the AGPLv3 License.
|
|
* Please see the included NOTICE for copyright information and
|
|
* LICENSE-AGPL for a copy of the license.
|
|
*/
|
|
|
|
use axum::{
|
|
extract::{Extension, Query},
|
|
routing::get,
|
|
Json, Router,
|
|
};
|
|
|
|
use itertools::Itertools;
|
|
use regex::Regex;
|
|
use serde::{Deserialize, Serialize};
|
|
use sqlx::FromRow;
|
|
use windmill_common::{
|
|
db::UserDB,
|
|
error::JsonResult,
|
|
utils::{paginate, Pagination},
|
|
};
|
|
|
|
use std::collections::HashMap;
|
|
#[cfg(feature = "benchmark")]
|
|
use std::sync::atomic::Ordering;
|
|
#[cfg(feature = "benchmark")]
|
|
use windmill_queue::IDLE_WORKERS;
|
|
|
|
use crate::db::ApiAuthed;
|
|
|
|
#[cfg(not(feature = "benchmark"))]
|
|
pub fn global_service() -> Router {
|
|
Router::new()
|
|
.route("/list", get(list_worker_pings))
|
|
.route("/custom_tags", get(get_custom_tags))
|
|
}
|
|
|
|
#[cfg(feature = "benchmark")]
|
|
pub fn global_service() -> Router {
|
|
Router::new()
|
|
.route("/toggle", get(toggle))
|
|
.route("/list", get(list_worker_pings))
|
|
.route("/custom_tags", get(get_custom_tags))
|
|
}
|
|
|
|
lazy_static::lazy_static! {
|
|
pub static ref CUSTOM_TAGS: Vec<String> = std::env::var("CUSTOM_TAGS")
|
|
.ok()
|
|
.map(|x| x.split(',').map(|x| x.to_string()).collect::<Vec<_>>()).unwrap_or_default();
|
|
|
|
pub static ref CUSTOM_TAGS_PER_WORKSPACE: (Vec<String>, HashMap<String, Vec<String>>) = process_custom_tags(std::env::var("CUSTOM_TAGS")
|
|
.ok());
|
|
|
|
pub static ref ALL_TAGS: Vec<String> = [CUSTOM_TAGS_PER_WORKSPACE.0.clone(), CUSTOM_TAGS_PER_WORKSPACE.1.keys().map(|x| x.to_string()).collect_vec()].concat();
|
|
|
|
}
|
|
|
|
fn process_custom_tags(o: Option<String>) -> (Vec<String>, HashMap<String, Vec<String>>) {
|
|
let regex = Regex::new(r"^(\w+)\(((?:\w+)\+?)+\)$").unwrap();
|
|
if let Some(s) = o {
|
|
let mut global = vec![];
|
|
let mut specific: HashMap<String, Vec<String>> = HashMap::new();
|
|
for e in s.split(",") {
|
|
if let Some(cap) = regex.captures(e) {
|
|
let tag = cap.get(1).unwrap().as_str().to_string();
|
|
let workspaces = cap.get(2).unwrap().as_str().split("+");
|
|
specific.insert(tag, workspaces.map(|x| x.to_string()).collect_vec());
|
|
} else {
|
|
global.push(e.to_string());
|
|
}
|
|
}
|
|
(global, specific)
|
|
} else {
|
|
(vec![], HashMap::new())
|
|
}
|
|
}
|
|
|
|
#[derive(FromRow, Serialize, Deserialize)]
|
|
struct WorkerPing {
|
|
worker: String,
|
|
worker_instance: String,
|
|
last_ping: Option<i32>,
|
|
started_at: chrono::DateTime<chrono::Utc>,
|
|
ip: String,
|
|
jobs_executed: i32,
|
|
custom_tags: Option<Vec<String>>,
|
|
}
|
|
|
|
#[derive(Serialize, Deserialize)]
|
|
struct EnableWorkerQuery {
|
|
disable: bool,
|
|
}
|
|
|
|
async fn list_worker_pings(
|
|
authed: ApiAuthed,
|
|
Extension(user_db): Extension<UserDB>,
|
|
Query(pagination): Query<Pagination>,
|
|
) -> JsonResult<Vec<WorkerPing>> {
|
|
let mut tx = user_db.begin(&authed).await?;
|
|
|
|
let (per_page, offset) = paginate(pagination);
|
|
|
|
let rows = sqlx::query_as!(
|
|
WorkerPing,
|
|
"SELECT worker, worker_instance, EXTRACT(EPOCH FROM (now() - ping_at))::integer as last_ping, started_at, ip, jobs_executed, custom_tags FROM worker_ping ORDER BY ping_at desc LIMIT $1 OFFSET $2",
|
|
per_page as i64,
|
|
offset as i64
|
|
)
|
|
.fetch_all(&mut *tx)
|
|
.await?;
|
|
tx.commit().await?;
|
|
Ok(Json(rows))
|
|
}
|
|
|
|
#[cfg(feature = "benchmark")]
|
|
async fn toggle(Query(query): Query<EnableWorkerQuery>) -> JsonResult<bool> {
|
|
IDLE_WORKERS.store(query.disable, Ordering::Relaxed);
|
|
Ok(Json(IDLE_WORKERS.load(Ordering::Relaxed)))
|
|
}
|
|
|
|
async fn get_custom_tags() -> Json<Vec<String>> {
|
|
Json(ALL_TAGS.clone())
|
|
}
|