460 lines
No EOL
16 KiB
Rust
460 lines
No EOL
16 KiB
Rust
use crate::*;
|
|
use crate::event_log::{BuildEventLogError, Result, create_build_event, current_timestamp_nanos, generate_event_id, query_engine::BELQueryEngine};
|
|
use std::sync::Arc;
|
|
use log::debug;
|
|
|
|
/// Common interface for writing events to the build event log with validation
|
|
pub struct EventWriter {
|
|
query_engine: Arc<BELQueryEngine>,
|
|
}
|
|
|
|
impl EventWriter {
|
|
/// Create a new EventWriter with the specified query engine
|
|
pub fn new(query_engine: Arc<BELQueryEngine>) -> Self {
|
|
Self { query_engine }
|
|
}
|
|
|
|
/// Append an event directly to the event log
|
|
pub async fn append_event(&self, event: BuildEvent) -> Result<()> {
|
|
self.query_engine.append_event(event).await.map(|_| ())
|
|
}
|
|
|
|
/// Get access to the underlying query engine for direct operations
|
|
pub fn query_engine(&self) -> &BELQueryEngine {
|
|
self.query_engine.as_ref()
|
|
}
|
|
|
|
/// Request a new build for the specified partitions
|
|
pub async fn request_build(
|
|
&self,
|
|
build_request_id: String,
|
|
requested_partitions: Vec<PartitionRef>,
|
|
) -> Result<()> {
|
|
debug!("Writing build request event for build: {}", build_request_id);
|
|
|
|
let event = create_build_event(
|
|
build_request_id,
|
|
build_event::EventType::BuildRequestEvent(BuildRequestEvent {
|
|
status: Some(BuildRequestStatusCode::BuildRequestReceived.status()),
|
|
requested_partitions,
|
|
message: "Build request received".to_string(),
|
|
comment: None,
|
|
want_id: None,
|
|
}),
|
|
);
|
|
|
|
self.query_engine.append_event(event).await.map(|_| ())
|
|
}
|
|
|
|
/// Update build request status
|
|
pub async fn update_build_status(
|
|
&self,
|
|
build_request_id: String,
|
|
status: BuildRequestStatus,
|
|
message: String,
|
|
) -> Result<()> {
|
|
debug!("Updating build status for {}: {:?}", build_request_id, status);
|
|
|
|
let event = create_build_event(
|
|
build_request_id,
|
|
build_event::EventType::BuildRequestEvent(BuildRequestEvent {
|
|
status: Some(status),
|
|
requested_partitions: vec![],
|
|
message,
|
|
comment: None,
|
|
want_id: None,
|
|
}),
|
|
);
|
|
|
|
self.query_engine.append_event(event).await.map(|_| ())
|
|
}
|
|
|
|
/// Update build request status with partition list
|
|
pub async fn update_build_status_with_partitions(
|
|
&self,
|
|
build_request_id: String,
|
|
status: BuildRequestStatus,
|
|
requested_partitions: Vec<PartitionRef>,
|
|
message: String,
|
|
) -> Result<()> {
|
|
debug!("Updating build status for {}: {:?}", build_request_id, status);
|
|
|
|
let event = create_build_event(
|
|
build_request_id,
|
|
build_event::EventType::BuildRequestEvent(BuildRequestEvent {
|
|
status: Some(status),
|
|
requested_partitions,
|
|
message,
|
|
comment: None,
|
|
want_id: None,
|
|
}),
|
|
);
|
|
|
|
self.query_engine.append_event(event).await.map(|_| ())
|
|
}
|
|
|
|
/// Update partition status
|
|
pub async fn update_partition_status(
|
|
&self,
|
|
build_request_id: String,
|
|
partition_ref: PartitionRef,
|
|
status: PartitionStatus,
|
|
message: String,
|
|
job_run_id: Option<String>,
|
|
) -> Result<()> {
|
|
debug!("Updating partition status for {}: {:?}", partition_ref.str, status);
|
|
|
|
let event = BuildEvent {
|
|
event_id: generate_event_id(),
|
|
timestamp: current_timestamp_nanos(),
|
|
build_request_id: Some(build_request_id),
|
|
event_type: Some(build_event::EventType::PartitionEvent(PartitionEvent {
|
|
partition_ref: Some(partition_ref),
|
|
status_code: status as i32,
|
|
status_name: status.to_display_string(),
|
|
message,
|
|
job_run_id: job_run_id.unwrap_or_default(),
|
|
})),
|
|
};
|
|
|
|
self.query_engine.append_event(event).await.map(|_| ())
|
|
}
|
|
|
|
/// Invalidate a partition with a reason
|
|
pub async fn invalidate_partition(
|
|
&self,
|
|
build_request_id: String,
|
|
partition_ref: PartitionRef,
|
|
reason: String,
|
|
) -> Result<()> {
|
|
// First validate that the partition exists by checking its current status
|
|
let current_status = self.query_engine.get_latest_partition_status(&partition_ref.str).await?;
|
|
|
|
if current_status.is_none() {
|
|
return Err(BuildEventLogError::QueryError(
|
|
format!("Cannot invalidate non-existent partition: {}", partition_ref.str)
|
|
));
|
|
}
|
|
|
|
let event = BuildEvent {
|
|
event_id: generate_event_id(),
|
|
timestamp: current_timestamp_nanos(),
|
|
build_request_id: Some(build_request_id),
|
|
event_type: Some(build_event::EventType::PartitionInvalidationEvent(
|
|
PartitionInvalidationEvent {
|
|
partition_ref: Some(partition_ref),
|
|
reason,
|
|
}
|
|
)),
|
|
};
|
|
|
|
self.query_engine.append_event(event).await.map(|_| ())
|
|
}
|
|
|
|
/// Schedule a job for execution
|
|
pub async fn schedule_job(
|
|
&self,
|
|
build_request_id: String,
|
|
job_run_id: String,
|
|
job_label: JobLabel,
|
|
target_partitions: Vec<PartitionRef>,
|
|
config: JobConfig,
|
|
) -> Result<()> {
|
|
debug!("Scheduling job {} for partitions: {:?}", job_label.label, target_partitions);
|
|
|
|
let event = BuildEvent {
|
|
event_id: generate_event_id(),
|
|
timestamp: current_timestamp_nanos(),
|
|
build_request_id: Some(build_request_id),
|
|
event_type: Some(build_event::EventType::JobEvent(JobEvent {
|
|
job_run_id,
|
|
job_label: Some(job_label),
|
|
target_partitions,
|
|
status_code: JobStatus::JobScheduled as i32,
|
|
status_name: JobStatus::JobScheduled.to_display_string(),
|
|
message: "Job scheduled for execution".to_string(),
|
|
config: Some(config),
|
|
manifests: vec![],
|
|
})),
|
|
};
|
|
|
|
self.query_engine.append_event(event).await.map(|_| ())
|
|
}
|
|
|
|
/// Update job status
|
|
pub async fn update_job_status(
|
|
&self,
|
|
build_request_id: String,
|
|
job_run_id: String,
|
|
job_label: JobLabel,
|
|
target_partitions: Vec<PartitionRef>,
|
|
status: JobStatus,
|
|
message: String,
|
|
manifests: Vec<PartitionManifest>,
|
|
) -> Result<()> {
|
|
debug!("Updating job {} status to {:?}", job_run_id, status);
|
|
|
|
let event = BuildEvent {
|
|
event_id: generate_event_id(),
|
|
timestamp: current_timestamp_nanos(),
|
|
build_request_id: Some(build_request_id),
|
|
event_type: Some(build_event::EventType::JobEvent(JobEvent {
|
|
job_run_id,
|
|
job_label: Some(job_label),
|
|
target_partitions,
|
|
status_code: status as i32,
|
|
status_name: status.to_display_string(),
|
|
message,
|
|
config: None,
|
|
manifests,
|
|
})),
|
|
};
|
|
|
|
self.query_engine.append_event(event).await.map(|_| ())
|
|
}
|
|
|
|
/// Cancel a task (job run) with a reason
|
|
pub async fn cancel_task(
|
|
&self,
|
|
build_request_id: String,
|
|
job_run_id: String,
|
|
reason: String,
|
|
) -> Result<()> {
|
|
// Validate that the job run exists and is in a cancellable state
|
|
let job_events = self.query_engine.get_job_run_events(&job_run_id).await?;
|
|
|
|
if job_events.is_empty() {
|
|
return Err(BuildEventLogError::QueryError(
|
|
format!("Cannot cancel non-existent job run: {}", job_run_id)
|
|
));
|
|
}
|
|
|
|
// Find the latest job status
|
|
let latest_status = job_events.iter()
|
|
.rev()
|
|
.find_map(|e| match &e.event_type {
|
|
Some(build_event::EventType::JobEvent(job)) => Some(job.status_code),
|
|
_ => None,
|
|
});
|
|
|
|
match latest_status {
|
|
Some(status) if status == JobStatus::JobCompleted as i32 => {
|
|
return Err(BuildEventLogError::QueryError(
|
|
format!("Cannot cancel completed job run: {}", job_run_id)
|
|
));
|
|
}
|
|
Some(status) if status == JobStatus::JobFailed as i32 => {
|
|
return Err(BuildEventLogError::QueryError(
|
|
format!("Cannot cancel failed job run: {}", job_run_id)
|
|
));
|
|
}
|
|
Some(status) if status == JobStatus::JobCancelled as i32 => {
|
|
return Err(BuildEventLogError::QueryError(
|
|
format!("Job run already cancelled: {}", job_run_id)
|
|
));
|
|
}
|
|
_ => {}
|
|
}
|
|
|
|
let event = BuildEvent {
|
|
event_id: generate_event_id(),
|
|
timestamp: current_timestamp_nanos(),
|
|
build_request_id: Some(build_request_id),
|
|
event_type: Some(build_event::EventType::JobRunCancelEvent(JobRunCancelEvent {
|
|
job_run_id,
|
|
reason,
|
|
})),
|
|
};
|
|
|
|
self.query_engine.append_event(event).await.map(|_| ())
|
|
}
|
|
|
|
/// Cancel a build request with a reason
|
|
pub async fn cancel_build(
|
|
&self,
|
|
build_request_id: String,
|
|
reason: String,
|
|
) -> Result<()> {
|
|
// Validate that the build exists and is in a cancellable state
|
|
let build_events = self.query_engine.get_build_request_events(&build_request_id, None).await?;
|
|
|
|
if build_events.is_empty() {
|
|
return Err(BuildEventLogError::QueryError(
|
|
format!("Cannot cancel non-existent build: {}", build_request_id)
|
|
));
|
|
}
|
|
|
|
// Find the latest build status
|
|
let latest_status = build_events.iter()
|
|
.rev()
|
|
.find_map(|e| match &e.event_type {
|
|
Some(build_event::EventType::BuildRequestEvent(br)) => Some(br.clone().status.unwrap().code),
|
|
_ => None,
|
|
});
|
|
|
|
match latest_status {
|
|
Some(status) if status == BuildRequestStatusCode::BuildRequestCompleted as i32 => {
|
|
return Err(BuildEventLogError::QueryError(
|
|
format!("Cannot cancel completed build: {}", build_request_id)
|
|
));
|
|
}
|
|
Some(status) if status == BuildRequestStatusCode::BuildRequestFailed as i32 => {
|
|
return Err(BuildEventLogError::QueryError(
|
|
format!("Cannot cancel failed build: {}", build_request_id)
|
|
));
|
|
}
|
|
Some(status) if status == BuildRequestStatusCode::BuildRequestCancelled as i32 => {
|
|
return Err(BuildEventLogError::QueryError(
|
|
format!("Build already cancelled: {}", build_request_id)
|
|
));
|
|
}
|
|
_ => {}
|
|
}
|
|
|
|
let event = BuildEvent {
|
|
event_id: generate_event_id(),
|
|
timestamp: current_timestamp_nanos(),
|
|
build_request_id: Some(build_request_id.clone()),
|
|
event_type: Some(build_event::EventType::BuildCancelEvent(BuildCancelEvent {
|
|
reason,
|
|
})),
|
|
};
|
|
|
|
self.query_engine.append_event(event).await.map(|_| ())?;
|
|
|
|
// Also emit a build request status update
|
|
self.update_build_status(
|
|
build_request_id,
|
|
BuildRequestStatusCode::BuildRequestCancelled.status(),
|
|
"Build cancelled by user".to_string(),
|
|
).await
|
|
}
|
|
|
|
/// Record a delegation event when a partition build is delegated to another build
|
|
pub async fn record_delegation(
|
|
&self,
|
|
build_request_id: String,
|
|
partition_ref: PartitionRef,
|
|
delegated_to_build_request_id: String,
|
|
message: String,
|
|
) -> Result<()> {
|
|
debug!("Recording delegation of {} to build {}", partition_ref.str, delegated_to_build_request_id);
|
|
|
|
let event = create_build_event(
|
|
build_request_id,
|
|
build_event::EventType::DelegationEvent(DelegationEvent {
|
|
partition_ref: Some(partition_ref),
|
|
delegated_to_build_request_id,
|
|
message,
|
|
}),
|
|
);
|
|
|
|
self.query_engine.append_event(event).await.map(|_| ())
|
|
}
|
|
|
|
/// Record the analyzed job graph
|
|
pub async fn record_job_graph(
|
|
&self,
|
|
build_request_id: String,
|
|
job_graph: JobGraph,
|
|
message: String,
|
|
) -> Result<()> {
|
|
debug!("Recording job graph for build: {}", build_request_id);
|
|
|
|
let event = BuildEvent {
|
|
event_id: generate_event_id(),
|
|
timestamp: current_timestamp_nanos(),
|
|
build_request_id: Some(build_request_id),
|
|
event_type: Some(build_event::EventType::JobGraphEvent(JobGraphEvent {
|
|
job_graph: Some(job_graph),
|
|
message,
|
|
})),
|
|
};
|
|
|
|
self.query_engine.append_event(event).await.map(|_| ())
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::event_log::mock::create_mock_bel_query_engine;
|
|
|
|
#[tokio::test]
|
|
async fn test_event_writer_build_lifecycle() {
|
|
let query_engine = create_mock_bel_query_engine().await.unwrap();
|
|
let writer = EventWriter::new(query_engine);
|
|
|
|
let build_id = "test-build-123".to_string();
|
|
let partitions = vec![PartitionRef { str: "test/partition".to_string() }];
|
|
|
|
// Test build request
|
|
writer.request_build(build_id.clone(), partitions.clone()).await.unwrap();
|
|
|
|
// Test status updates
|
|
writer.update_build_status(
|
|
build_id.clone(),
|
|
BuildRequestStatusCode::BuildRequestPlanning.status(),
|
|
"Starting planning".to_string(),
|
|
).await.unwrap();
|
|
|
|
writer.update_build_status(
|
|
build_id.clone(),
|
|
BuildRequestStatusCode::BuildRequestExecuting.status(),
|
|
"Starting execution".to_string(),
|
|
).await.unwrap();
|
|
|
|
writer.update_build_status(
|
|
build_id.clone(),
|
|
BuildRequestStatusCode::BuildRequestCompleted.status(),
|
|
"Build completed successfully".to_string(),
|
|
).await.unwrap();
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_event_writer_partition_and_job() {
|
|
let query_engine = create_mock_bel_query_engine().await.unwrap();
|
|
let writer = EventWriter::new(query_engine);
|
|
|
|
let build_id = "test-build-456".to_string();
|
|
let partition = PartitionRef { str: "data/users".to_string() };
|
|
let job_run_id = "job-run-789".to_string();
|
|
let job_label = JobLabel { label: "//:test_job".to_string() };
|
|
|
|
// Test partition status update
|
|
writer.update_partition_status(
|
|
build_id.clone(),
|
|
partition.clone(),
|
|
PartitionStatus::PartitionBuilding,
|
|
"Building partition".to_string(),
|
|
Some(job_run_id.clone()),
|
|
).await.unwrap();
|
|
|
|
// Test job scheduling
|
|
let config = JobConfig {
|
|
outputs: vec![partition.clone()],
|
|
inputs: vec![],
|
|
args: vec!["test".to_string()],
|
|
env: std::collections::HashMap::new(),
|
|
};
|
|
|
|
writer.schedule_job(
|
|
build_id.clone(),
|
|
job_run_id.clone(),
|
|
job_label.clone(),
|
|
vec![partition.clone()],
|
|
config,
|
|
).await.unwrap();
|
|
|
|
// Test job status update
|
|
writer.update_job_status(
|
|
build_id.clone(),
|
|
job_run_id,
|
|
job_label,
|
|
vec![partition],
|
|
JobStatus::JobCompleted,
|
|
"Job completed successfully".to_string(),
|
|
vec![],
|
|
).await.unwrap();
|
|
}
|
|
} |