diff --git a/binaries/coordinator/src/lib.rs b/binaries/coordinator/src/lib.rs index fd2b45f9..d977126d 100644 --- a/binaries/coordinator/src/lib.rs +++ b/binaries/coordinator/src/lib.rs @@ -7,7 +7,8 @@ use dora_core::{ config::{NodeId, OperatorId}, coordinator_messages::RegisterResult, daemon_messages::{DaemonCoordinatorEvent, DaemonCoordinatorReply}, - descriptor::Descriptor, + descriptor::{self, Descriptor, ResolvedNode}, + message::Metadata, topics::{ control_socket_addr, ControlRequest, ControlRequestReply, DataflowId, DORA_COORDINATOR_PORT_DEFAULT, @@ -522,6 +523,7 @@ struct RunningDataflow { machines: BTreeSet, /// IDs of machines that are waiting until all nodes are started. pending_machines: BTreeSet, + nodes: Vec, } impl PartialEq for RunningDataflow { @@ -625,24 +627,33 @@ async fn retrieve_logs( })?; let mut reply_logs = Vec::new(); - for machine_id in &dataflow.machines { - let daemon_connection = daemon_connections - .get_mut(machine_id) - .wrap_err("no daemon connection")?; // TODO: take from dataflow spec - tcp_send(daemon_connection, &message) - .await - .wrap_err("failed to send logs message to daemon")?; + let nodes = &dataflow.nodes; + let machine_ids: Vec = nodes + .iter() + .filter(|node| node.id == node_id) + .map(|node| node.deploy.machine.clone()) + .collect(); - // wait for reply - let reply_raw = tcp_receive(daemon_connection) - .await - .wrap_err("failed to retrieve logs reply from daemon")?; - match serde_json::from_slice(&reply_raw) - .wrap_err("failed to deserialize logs reply from daemon")? - { - DaemonCoordinatorReply::Logs { logs } => reply_logs = logs, - other => bail!("unexpected reply after sending reload: {other:?}"), - } + let machine_id = machine_ids + .first() + .wrap_err("Did not find node in dataflow")?; + + let daemon_connection = daemon_connections + .get_mut(machine_id.as_str()) + .wrap_err("no daemon connection")?; // TODO: take from dataflow spec + tcp_send(daemon_connection, &message) + .await + .wrap_err("failed to send logs message to daemon")?; + + // wait for reply + let reply_raw = tcp_receive(daemon_connection) + .await + .wrap_err("failed to retrieve logs reply from daemon")?; + match serde_json::from_slice(&reply_raw) + .wrap_err("failed to deserialize logs reply from daemon")? + { + DaemonCoordinatorReply::Logs { logs } => reply_logs = logs, + other => bail!("unexpected reply after sending reload: {other:?}"), } tracing::info!("successfully retrieved logs for `{dataflow_id}/{node_id}`"); @@ -655,8 +666,11 @@ async fn start_dataflow( name: Option, daemon_connections: &mut HashMap, ) -> eyre::Result { - let SpawnedDataflow { uuid, machines } = - spawn_dataflow(dataflow, working_dir, daemon_connections).await?; + let SpawnedDataflow { + uuid, + machines, + nodes, + } = spawn_dataflow(path, working_dir, daemon_connections).await?; Ok(RunningDataflow { uuid, name, @@ -666,6 +680,7 @@ async fn start_dataflow( BTreeSet::new() }, machines, + nodes, }) } diff --git a/binaries/coordinator/src/run/mod.rs b/binaries/coordinator/src/run/mod.rs index 6647548a..4e2843a1 100644 --- a/binaries/coordinator/src/run/mod.rs +++ b/binaries/coordinator/src/run/mod.rs @@ -5,7 +5,7 @@ use crate::{ use dora_core::{ daemon_messages::{DaemonCoordinatorEvent, DaemonCoordinatorReply, SpawnDataflowNodes}, - descriptor::Descriptor, + descriptor::{Descriptor, ResolvedNode}, }; use eyre::{bail, eyre, ContextCompat, WrapErr}; use std::{ @@ -39,7 +39,7 @@ pub(super) async fn spawn_dataflow( let spawn_command = SpawnDataflowNodes { dataflow_id: uuid, working_dir, - nodes, + nodes: nodes.clone(), communication: dataflow.communication, machine_listen_ports, }; @@ -54,7 +54,11 @@ pub(super) async fn spawn_dataflow( tracing::info!("successfully spawned dataflow `{uuid}`"); - Ok(SpawnedDataflow { uuid, machines }) + Ok(SpawnedDataflow { + uuid, + machines, + nodes, + }) } async fn spawn_dataflow_on_machine( @@ -85,4 +89,5 @@ async fn spawn_dataflow_on_machine( pub struct SpawnedDataflow { pub uuid: Uuid, pub machines: BTreeSet, + pub nodes: Vec, }