387 lines
13 KiB
Rust
387 lines
13 KiB
Rust
use chrono::Utc;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
use std::time::Duration;
|
|
use uuid::Uuid;
|
|
use redis::AsyncCommands;
|
|
use thiserror::Error;
|
|
|
|
/// Redis namespace prefix for all Hero job-related keys
|
|
pub const NAMESPACE_PREFIX: &str = "hero:job:";
|
|
|
|
/// Script type enumeration for different worker types
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub enum ScriptType {
|
|
/// OSIS - A worker that executes Rhai/HeroScript
|
|
OSIS,
|
|
/// SAL - A worker that executes system abstraction layer functionalities in rhai
|
|
SAL,
|
|
/// V - A worker that executes heroscript in V
|
|
V,
|
|
/// Python - A worker that executes heroscript in python
|
|
Python,
|
|
}
|
|
|
|
impl ScriptType {
|
|
/// Get the worker queue suffix for this script type
|
|
pub fn worker_queue_suffix(&self) -> &'static str {
|
|
match self {
|
|
ScriptType::OSIS => "osis",
|
|
ScriptType::SAL => "sal",
|
|
ScriptType::V => "v",
|
|
ScriptType::Python => "python",
|
|
}
|
|
}
|
|
|
|
pub fn as_str(&self) -> &'static str {
|
|
match self {
|
|
ScriptType::OSIS => "osis",
|
|
ScriptType::SAL => "sal",
|
|
ScriptType::V => "v",
|
|
ScriptType::Python => "python",
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Job status enumeration
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub enum JobStatus {
|
|
Dispatched,
|
|
WaitingForPrerequisites,
|
|
Started,
|
|
Error,
|
|
Finished,
|
|
}
|
|
|
|
impl JobStatus {
|
|
pub fn as_str(&self) -> &'static str {
|
|
match self {
|
|
JobStatus::Dispatched => "dispatched",
|
|
JobStatus::WaitingForPrerequisites => "waiting_for_prerequisites",
|
|
JobStatus::Started => "started",
|
|
JobStatus::Error => "error",
|
|
JobStatus::Finished => "finished",
|
|
}
|
|
}
|
|
|
|
pub fn from_str(s: &str) -> Option<Self> {
|
|
match s {
|
|
"dispatched" => Some(JobStatus::Dispatched),
|
|
"waiting_for_prerequisites" => Some(JobStatus::WaitingForPrerequisites),
|
|
"started" => Some(JobStatus::Started),
|
|
"error" => Some(JobStatus::Error),
|
|
"finished" => Some(JobStatus::Finished),
|
|
_ => None,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Representation of a script execution request.
|
|
///
|
|
/// This structure contains all the information needed to execute a script
|
|
/// on a worker service, including the script content, dependencies, and metadata.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct Job {
|
|
pub id: String,
|
|
pub caller_id: String,
|
|
pub context_id: String,
|
|
pub script: String,
|
|
pub script_type: ScriptType,
|
|
pub timeout: Duration,
|
|
pub retries: u8, // retries on script execution
|
|
pub concurrent: bool, // whether to execute script in separate thread
|
|
pub log_path: Option<String>, // path to write logs of script execution to
|
|
pub env_vars: HashMap<String, String>, // environment variables for script execution
|
|
pub prerequisites: Vec<String>, // job IDs that must complete before this job can run
|
|
pub dependents: Vec<String>, // job IDs that depend on this job completing
|
|
pub created_at: chrono::DateTime<chrono::Utc>,
|
|
pub updated_at: chrono::DateTime<chrono::Utc>,
|
|
}
|
|
|
|
/// Error types for job operations
|
|
#[derive(Error, Debug)]
|
|
pub enum JobError {
|
|
#[error("Redis error: {0}")]
|
|
RedisError(#[from] redis::RedisError),
|
|
#[error("Serialization error: {0}")]
|
|
SerializationError(#[from] serde_json::Error),
|
|
#[error("Job not found: {0}")]
|
|
JobNotFound(String),
|
|
#[error("Invalid job data: {0}")]
|
|
InvalidJobData(String),
|
|
#[error("Missing required field: {0}")]
|
|
MissingField(String),
|
|
}
|
|
|
|
impl Job {
|
|
/// Create a new job with the given parameters
|
|
pub fn new(
|
|
caller_id: String,
|
|
context_id: String,
|
|
script: String,
|
|
script_type: ScriptType,
|
|
) -> Self {
|
|
let now = Utc::now();
|
|
Self {
|
|
id: Uuid::new_v4().to_string(),
|
|
caller_id,
|
|
context_id,
|
|
script,
|
|
script_type,
|
|
timeout: Duration::from_secs(30),
|
|
retries: 0,
|
|
concurrent: false,
|
|
log_path: None,
|
|
env_vars: HashMap::new(),
|
|
prerequisites: Vec::new(),
|
|
dependents: Vec::new(),
|
|
created_at: now,
|
|
updated_at: now,
|
|
}
|
|
}
|
|
|
|
/// Store this job in Redis
|
|
pub async fn store_in_redis(&self, conn: &mut redis::aio::MultiplexedConnection) -> Result<(), JobError> {
|
|
let job_key = format!("{}{}", NAMESPACE_PREFIX, self.id);
|
|
|
|
let mut hset_args: Vec<(String, String)> = vec![
|
|
("jobId".to_string(), self.id.clone()),
|
|
("script".to_string(), self.script.clone()),
|
|
("script_type".to_string(), format!("{:?}", self.script_type)),
|
|
("callerId".to_string(), self.caller_id.clone()),
|
|
("contextId".to_string(), self.context_id.clone()),
|
|
("status".to_string(), "pending".to_string()),
|
|
("timeout".to_string(), self.timeout.as_secs().to_string()),
|
|
("retries".to_string(), self.retries.to_string()),
|
|
("concurrent".to_string(), self.concurrent.to_string()),
|
|
("createdAt".to_string(), self.created_at.to_rfc3339()),
|
|
("updatedAt".to_string(), self.updated_at.to_rfc3339()),
|
|
];
|
|
|
|
// Add optional log path
|
|
if let Some(log_path) = &self.log_path {
|
|
hset_args.push(("log_path".to_string(), log_path.clone()));
|
|
}
|
|
|
|
// Add environment variables as JSON string if any are provided
|
|
if !self.env_vars.is_empty() {
|
|
let env_vars_json = serde_json::to_string(&self.env_vars)?;
|
|
hset_args.push(("env_vars".to_string(), env_vars_json));
|
|
}
|
|
|
|
// Add prerequisites as JSON string if any are provided
|
|
if !self.prerequisites.is_empty() {
|
|
let prerequisites_json = serde_json::to_string(&self.prerequisites)?;
|
|
hset_args.push(("prerequisites".to_string(), prerequisites_json));
|
|
}
|
|
|
|
// Add dependents as JSON string if any are provided
|
|
if !self.dependents.is_empty() {
|
|
let dependents_json = serde_json::to_string(&self.dependents)?;
|
|
hset_args.push(("dependents".to_string(), dependents_json));
|
|
}
|
|
|
|
conn.hset_multiple::<_, _, _, ()>(&job_key, &hset_args).await?;
|
|
Ok(())
|
|
}
|
|
|
|
/// Load a job from Redis by ID
|
|
pub async fn load_from_redis(
|
|
conn: &mut redis::aio::MultiplexedConnection,
|
|
job_id: &str,
|
|
) -> Result<Self, JobError> {
|
|
let job_key = format!("{}{}", NAMESPACE_PREFIX, job_id);
|
|
let job_data: HashMap<String, String> = conn.hgetall(&job_key).await?;
|
|
|
|
if job_data.is_empty() {
|
|
return Err(JobError::JobNotFound(job_id.to_string()));
|
|
}
|
|
|
|
// Parse required fields
|
|
let id = job_data.get("jobId")
|
|
.ok_or_else(|| JobError::MissingField("jobId".to_string()))?
|
|
.clone();
|
|
|
|
let script = job_data.get("script")
|
|
.ok_or_else(|| JobError::MissingField("script".to_string()))?
|
|
.clone();
|
|
|
|
let script_type_str = job_data.get("script_type")
|
|
.ok_or_else(|| JobError::MissingField("script_type".to_string()))?;
|
|
|
|
let script_type = match script_type_str.as_str() {
|
|
"OSIS" => ScriptType::OSIS,
|
|
"SAL" => ScriptType::SAL,
|
|
"V" => ScriptType::V,
|
|
"Python" => ScriptType::Python,
|
|
_ => return Err(JobError::InvalidJobData(format!("Unknown script type: {}", script_type_str))),
|
|
};
|
|
|
|
let caller_id = job_data.get("callerId")
|
|
.ok_or_else(|| JobError::MissingField("callerId".to_string()))?
|
|
.clone();
|
|
|
|
let context_id = job_data.get("contextId")
|
|
.ok_or_else(|| JobError::MissingField("contextId".to_string()))?
|
|
.clone();
|
|
|
|
let timeout_secs: u64 = job_data.get("timeout")
|
|
.ok_or_else(|| JobError::MissingField("timeout".to_string()))?
|
|
.parse()
|
|
.map_err(|_| JobError::InvalidJobData("Invalid timeout value".to_string()))?;
|
|
|
|
let retries: u8 = job_data.get("retries")
|
|
.unwrap_or(&"0".to_string())
|
|
.parse()
|
|
.map_err(|_| JobError::InvalidJobData("Invalid retries value".to_string()))?;
|
|
|
|
let concurrent: bool = job_data.get("concurrent")
|
|
.unwrap_or(&"false".to_string())
|
|
.parse()
|
|
.map_err(|_| JobError::InvalidJobData("Invalid concurrent value".to_string()))?;
|
|
|
|
let created_at = job_data.get("createdAt")
|
|
.ok_or_else(|| JobError::MissingField("createdAt".to_string()))?
|
|
.parse()
|
|
.map_err(|_| JobError::InvalidJobData("Invalid createdAt timestamp".to_string()))?;
|
|
|
|
let updated_at = job_data.get("updatedAt")
|
|
.ok_or_else(|| JobError::MissingField("updatedAt".to_string()))?
|
|
.parse()
|
|
.map_err(|_| JobError::InvalidJobData("Invalid updatedAt timestamp".to_string()))?;
|
|
|
|
// Parse optional fields
|
|
let log_path = job_data.get("log_path").cloned();
|
|
|
|
let env_vars = if let Some(env_vars_json) = job_data.get("env_vars") {
|
|
serde_json::from_str(env_vars_json)?
|
|
} else {
|
|
HashMap::new()
|
|
};
|
|
|
|
let prerequisites = if let Some(prerequisites_json) = job_data.get("prerequisites") {
|
|
serde_json::from_str(prerequisites_json)?
|
|
} else {
|
|
Vec::new()
|
|
};
|
|
|
|
let dependents = if let Some(dependents_json) = job_data.get("dependents") {
|
|
serde_json::from_str(dependents_json)?
|
|
} else {
|
|
Vec::new()
|
|
};
|
|
|
|
Ok(Self {
|
|
id,
|
|
caller_id,
|
|
context_id,
|
|
script,
|
|
script_type,
|
|
timeout: Duration::from_secs(timeout_secs),
|
|
retries,
|
|
concurrent,
|
|
log_path,
|
|
env_vars,
|
|
prerequisites,
|
|
dependents,
|
|
created_at,
|
|
updated_at,
|
|
})
|
|
}
|
|
|
|
/// Update job status in Redis
|
|
pub async fn update_status(
|
|
conn: &mut redis::aio::MultiplexedConnection,
|
|
job_id: &str,
|
|
status: JobStatus,
|
|
) -> Result<(), JobError> {
|
|
let job_key = format!("{}{}", NAMESPACE_PREFIX, job_id);
|
|
let now = Utc::now();
|
|
|
|
conn.hset::<_, _, _, ()>(&job_key, "status", status.as_str()).await?;
|
|
conn.hset::<_, _, _, ()>(&job_key, "updatedAt", now.to_rfc3339()).await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Get job status from Redis
|
|
pub async fn get_status(
|
|
conn: &mut redis::aio::MultiplexedConnection,
|
|
job_id: &str,
|
|
) -> Result<JobStatus, JobError> {
|
|
let job_key = format!("{}{}", NAMESPACE_PREFIX, job_id);
|
|
let status_str: String = conn.hget(&job_key, "status").await?;
|
|
|
|
JobStatus::from_str(&status_str)
|
|
.ok_or_else(|| JobError::InvalidJobData(format!("Unknown status: {}", status_str)))
|
|
}
|
|
|
|
/// Set job result in Redis
|
|
pub async fn set_result(
|
|
conn: &mut redis::aio::MultiplexedConnection,
|
|
job_id: &str,
|
|
result: &str,
|
|
) -> Result<(), JobError> {
|
|
let job_key = format!("{}{}", NAMESPACE_PREFIX, job_id);
|
|
let now = Utc::now();
|
|
|
|
conn.hset::<_, _, _, ()>(&job_key, "output", result).await?;
|
|
conn.hset::<_, _, _, ()>(&job_key, "status", JobStatus::Finished.as_str()).await?;
|
|
conn.hset::<_, _, _, ()>(&job_key, "updatedAt", now.to_rfc3339()).await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Set job error in Redis
|
|
pub async fn set_error(
|
|
conn: &mut redis::aio::MultiplexedConnection,
|
|
job_id: &str,
|
|
error: &str,
|
|
) -> Result<(), JobError> {
|
|
let job_key = format!("{}{}", NAMESPACE_PREFIX, job_id);
|
|
let now = Utc::now();
|
|
|
|
conn.hset::<_, _, _, ()>(&job_key, "error", error).await?;
|
|
conn.hset::<_, _, _, ()>(&job_key, "status", JobStatus::Error.as_str()).await?;
|
|
conn.hset::<_, _, _, ()>(&job_key, "updatedAt", now.to_rfc3339()).await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Delete job from Redis
|
|
pub async fn delete_from_redis(
|
|
conn: &mut redis::aio::MultiplexedConnection,
|
|
job_id: &str,
|
|
) -> Result<(), JobError> {
|
|
let job_key = format!("{}{}", NAMESPACE_PREFIX, job_id);
|
|
conn.del::<_, ()>(&job_key).await?;
|
|
Ok(())
|
|
}
|
|
|
|
/// List all job IDs from Redis
|
|
pub async fn list_all_job_ids(
|
|
conn: &mut redis::aio::MultiplexedConnection,
|
|
) -> Result<Vec<String>, JobError> {
|
|
// Search specifically for job keys with the exact job pattern
|
|
let job_keys: Vec<String> = conn.keys(format!("{}*", NAMESPACE_PREFIX)).await?;
|
|
let job_ids: Vec<String> = job_keys
|
|
.iter()
|
|
.filter_map(|key| {
|
|
// Only include keys that exactly match the job key pattern hero:job:*
|
|
if key.starts_with(NAMESPACE_PREFIX) {
|
|
let potential_id = key.strip_prefix(NAMESPACE_PREFIX)?;
|
|
// Validate that this looks like a UUID (job IDs are UUIDs)
|
|
if potential_id.len() == 36 && potential_id.chars().filter(|&c| c == '-').count() == 4 {
|
|
Some(potential_id.to_string())
|
|
} else {
|
|
None
|
|
}
|
|
} else {
|
|
None
|
|
}
|
|
})
|
|
.collect();
|
|
Ok(job_ids)
|
|
}
|
|
}
|