andygrove commented on code in PR #1925:
URL:
https://github.com/apache/datafusion-ballista/pull/1925#discussion_r3721036551
##########
ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs:
##########
@@ -82,6 +110,88 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan>
tx_event: &mpsc::Sender<QueryStageSchedulerEvent>,
_rx_event: &mpsc::Receiver<QueryStageSchedulerEvent>,
) -> Result<()> {
+ #[cfg(feature = "rest-api")]
+ if let Some(log) = &self.event_log {
+ match &event {
+ QueryStageSchedulerEvent::JobSubmitted {
+ job_id,
+ queued_at,
+ submitted_at,
+ } => {
+ if let Ok(Some(graph)) = self
+ .state
+ .task_manager
+ .get_job_execution_graph(job_id)
+ .await
+ {
+ log.append(
+ job_id.as_str(),
+ event_log::job_start_event(&graph, *queued_at,
*submitted_at),
+ );
+ }
+ }
+ QueryStageSchedulerEvent::TaskUpdating(executor_id, statuses)
=> {
+ for (job_id, group) in group_by_job(statuses) {
+ for ev in event_log::task_end_events(executor_id,
&group) {
+ log.append(&job_id, ev);
+ }
+ }
+ }
+ QueryStageSchedulerEvent::JobFinished {
+ job_id,
+ queued_at,
+ completed_at,
+ } => {
+ if let Ok(Some(graph)) = self
+ .state
+ .task_manager
+ .get_job_execution_graph(job_id)
Review Comment:
The `finish_job` call is intentional there, and I have added a comment
explaining it. It closes the per-job file handle, so a job whose `JobEnd` could
not be built still releases its fd instead of leaking one. The resulting log
has no `JobEnd` line, and the reader treats a log without one as incomplete and
skips it, so the history server never serves a half-written job.
What was wrong is that the `Err` was swallowed silently. Pulled the lookup
into an `event_log_graph` helper that logs both the `Err` case and the
`Ok(None)` case at warn level, and reused it in all three arms. Event logging
still never fails scheduling, it just costs the job its event.
##########
ballista/history/src/writer.rs:
##########
@@ -0,0 +1,303 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Async, buffered event-log writer. Each job's events append to
+//! `<log_dir>/<job_id>.eventlog` as JSONL. Appends are non-blocking; a
background
+//! task performs the file I/O so the scheduler hot path never waits on disk.
+
+use crate::event::HistoryEvent;
+use std::collections::HashMap;
+use std::path::{Path, PathBuf};
+use tokio::io::AsyncWriteExt;
+use tokio::sync::{mpsc, oneshot};
+
+enum WriterMsg {
+ Event {
+ job_id: String,
+ event: Box<HistoryEvent>,
+ },
+ Flush {
+ job_id: String,
+ done: oneshot::Sender<()>,
+ },
+ Finish {
+ job_id: String,
+ done: oneshot::Sender<()>,
+ },
+}
+
+#[derive(Clone)]
+pub struct EventLogWriter {
+ tx: mpsc::Sender<WriterMsg>,
+}
+
+impl EventLogWriter {
+ pub fn new(log_dir: PathBuf, buffer: usize) -> EventLogWriter {
+ let (tx, rx) = mpsc::channel(buffer.max(1));
+ tokio::spawn(run(log_dir, rx));
+ EventLogWriter { tx }
+ }
+
+ /// Enqueue an event for `job_id`. Never blocks; drops (with a warning) if
the
+ /// channel is full, so logging cannot stall scheduling.
+ pub fn append(&self, job_id: &str, event: HistoryEvent) {
+ let msg = WriterMsg::Event {
+ job_id: job_id.to_string(),
+ event: Box::new(event),
+ };
+ if self.tx.try_send(msg).is_err() {
+ eprintln!(
+ "event-log writer: dropping event for {job_id} (channel full
or closed)"
+ );
+ }
+ }
+
+ /// Await all currently-enqueued writes for `job_id` (best effort).
+ pub async fn flush_job(&self, job_id: &str) {
+ let (done, wait) = oneshot::channel();
+ if self
+ .tx
+ .send(WriterMsg::Flush {
+ job_id: job_id.to_string(),
+ done,
+ })
+ .await
+ .is_ok()
+ {
+ let _ = wait.await;
+ }
+ }
+
+ /// Enqueue a terminal event (e.g. `JobEnd`) for `job_id`. Unlike
`append`, this
+ /// awaits channel capacity instead of dropping the event when the channel
is
+ /// full, so the terminal record is never silently lost. Still best-effort
at
+ /// the process boundary: if the channel is closed (background task gone)
this
+ /// logs and returns rather than panicking.
+ pub async fn append_final(&self, job_id: &str, event: HistoryEvent) {
+ let msg = WriterMsg::Event {
+ job_id: job_id.to_string(),
+ event: Box::new(event),
+ };
+ if self.tx.send(msg).await.is_err() {
+ eprintln!(
+ "event-log writer: failed to enqueue terminal event for
{job_id} (channel closed)"
+ );
+ }
+ }
+
+ /// Flush and close the per-job file handle for `job_id`. Must be called
after
+ /// the terminal event has been enqueued (e.g. via `append_final`) so it is
+ /// ordered after it on the single-consumer FIFO channel. Best effort: if
the
+ /// channel is closed this logs and returns.
+ pub async fn finish_job(&self, job_id: &str) {
+ let (done, wait) = oneshot::channel();
+ if self
+ .tx
+ .send(WriterMsg::Finish {
+ job_id: job_id.to_string(),
+ done,
+ })
+ .await
+ .is_ok()
+ {
+ let _ = wait.await;
+ } else {
+ eprintln!(
+ "event-log writer: failed to enqueue finish for {job_id}
(channel closed)"
+ );
+ }
+ }
+}
+
+async fn run(log_dir: PathBuf, mut rx: mpsc::Receiver<WriterMsg>) {
+ if let Err(e) = tokio::fs::create_dir_all(&log_dir).await {
+ eprintln!("event-log writer: cannot create {}: {e}",
log_dir.display());
+ return;
+ }
+ // One open append handle per job for the life of the process.
+ let mut handles: HashMap<String, tokio::fs::File> = HashMap::new();
+
+ while let Some(msg) = rx.recv().await {
+ match msg {
+ WriterMsg::Event { job_id, event } => {
+ let file = match open_for(&log_dir, &mut handles,
&job_id).await {
+ Some(f) => f,
+ None => continue,
+ };
+ match serde_json::to_string(&*event) {
+ Ok(mut line) => {
+ line.push('\n');
+ if let Err(e) = file.write_all(line.as_bytes()).await {
+ eprintln!("event-log writer: write failed for
{job_id}: {e}");
+ }
+ }
+ Err(e) => eprintln!("event-log writer: serialize failed:
{e}"),
+ }
+ }
+ WriterMsg::Flush { job_id, done } => {
+ if let Some(file) = handles.get_mut(&job_id) {
+ let _ = file.flush().await;
+ }
+ let _ = done.send(());
+ }
+ WriterMsg::Finish { job_id, done } => {
+ if let Some(mut file) = handles.remove(&job_id) {
+ let _ = file.flush().await;
+ // Dropping `file` here closes the fd.
+ }
+ let _ = done.send(());
+ }
+ }
+ }
+}
+
+async fn open_for<'a>(
+ log_dir: &Path,
+ handles: &'a mut HashMap<String, tokio::fs::File>,
+ job_id: &str,
+) -> Option<&'a mut tokio::fs::File> {
+ if !handles.contains_key(job_id) {
+ let path = log_dir.join(format!("{job_id}.eventlog"));
+ match tokio::fs::OpenOptions::new()
+ .create(true)
+ .append(true)
+ .open(&path)
+ .await
+ {
+ Ok(f) => {
+ handles.insert(job_id.to_string(), f);
+ }
+ Err(e) => {
+ eprintln!("event-log writer: cannot open {}: {e}",
path.display());
+ return None;
+ }
+ }
+ }
+ handles.get_mut(job_id)
Review Comment:
I would rather keep this one as is. `entry()` takes an owned key, so it
allocates a `String` on every event, including the already-open case which is
the overwhelming majority of calls. The `contains_key` plus `get_mut` shape is
the usual workaround for the borrow checker rejecting the `get_mut`-then-insert
form, and the extra lookup only happens on the first event for a job.
Happy to switch if you feel strongly, the allocation is small, it just
seemed like paying it on every event to save one lookup per job.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
