Log Server->FE alloc assignment deltas (#1611)

Author: earhart

server/src/executor_api.rs

@@ -19,12 +19,13 @@ use executor_api_pb::{
     HostResources,
     ReportExecutorStateRequest,
     ReportExecutorStateResponse,
+    TaskAllocation,
     TaskResult,
 };
-use tokio::sync::watch;
+use tokio::sync::watch::{self, Receiver, Sender};
 use tokio_stream::{wrappers::WatchStream, Stream};
 use tonic::{Request, Response, Status};
-use tracing::{debug, error, info, trace, warn};
+use tracing::{debug, error, info, instrument, trace, warn};
 
 use crate::{
     blob_store::{self, BlobStorage},
@@ -711,6 +712,105 @@ impl ExecutorAPIService {
     }
 }
 
+fn log_desired_executor_state_delta(
+    last_sent_state: &DesiredExecutorState,
+    desired_state: &DesiredExecutorState,
+) {
+    debug!(?desired_state, "got desired state");
+
+    let mut last_assignments: HashMap<String, String> = HashMap::default();
+    for ta in &last_sent_state.task_allocations {
+        if let (Some(fe_id), Some(alloc_id)) = (&ta.function_executor_id, &ta.allocation_id) {
+            last_assignments.insert(fe_id.clone(), alloc_id.clone());
+        }
+    }
+
+    for TaskAllocation {
+        function_executor_id: fn_executor_id_option,
+        allocation_id: allocation_id_option,
+        ..
+    } in &desired_state.task_allocations
+    {
+        if let (Some(fn_executor_id), Some(allocation_id)) =
+            (fn_executor_id_option, allocation_id_option)
+        {
+            match last_assignments.get(fn_executor_id) {
+                Some(last_allocation_id) => {
+                    if allocation_id != last_allocation_id {
+                        info!(
+                            %fn_executor_id,
+                            %allocation_id, %last_allocation_id, "re-assigning FE"
+                        )
+                    }
+                    last_assignments.remove(fn_executor_id);
+                }
+                None => {
+                    info!(%fn_executor_id, %allocation_id, "assigning FE")
+                }
+            }
+        }
+    }
+
+    for (fn_executor_id, last_allocation_id) in last_assignments {
+        info!(%fn_executor_id, %last_allocation_id, "idling FE")
+    }
+}
+
+#[instrument(skip_all, fields(executor_id = %executor_id))]
+async fn executor_update_loop(
+    executor_id: ExecutorId,
+    executor_manager: Arc<ExecutorManager>,
+    mut executor_state_rx: Receiver<()>,
+    grpc_tx: Sender<Result<DesiredExecutorState, Status>>,
+) {
+    // Mark the state as changed to trigger the first change
+    // notification to the executor. This is important because between
+    // the report_executor_state and the get_desired_executor_states
+    // requests, the executor state may received a new desired state.
+    executor_state_rx.mark_changed();
+
+    // Use the default/empty value for the last-sent desired state, so
+    // that the first change logged will be a delta from nothing (=>
+    // the complete executor state).
+    let mut last_sent_state = DesiredExecutorState::default();
+
+    loop {
+        // Wait for a state change for this executor or grpc stream closing.
+        tokio::select! {
+            _ = grpc_tx.closed() => {
+                info!(
+                    "get_desired_executor_states: grpc stream closed"
+                );
+                break;
+            }
+            result = executor_state_rx.changed() => {
+                if let Err(err) = result {
+                    info!(
+                        ?err,
+                        "state machine watcher closing"
+                    );
+                    break;
+                }
+            }
+        }
+
+        // Get the latest state
+        let desired_state = executor_manager.get_executor_state(&executor_id).await;
+
+        // Log the state delta
+        log_desired_executor_state_delta(&last_sent_state, &desired_state);
+
+        // Send the state to the executor
+        if let Err(err) = grpc_tx.send(Ok(desired_state.clone())) {
+            info!(?err, "grpc stream closing");
+            break;
+        }
+
+        // Store the sent state for next comparison
+        last_sent_state = desired_state;
+    }
+}
+
 #[tonic::async_trait]
 impl ExecutorApi for ExecutorAPIService {
     #[allow(non_camel_case_types)] // The autogenerated code in the trait uses snake_case types in some cases
@@ -810,7 +910,7 @@ impl ExecutorApi for ExecutorAPIService {
             "Got get_desired_executor_states request",
         );
 
-        let mut executor_state_rx = match self.executor_manager.subscribe(&executor_id).await {
+        let executor_state_rx = match self.executor_manager.subscribe(&executor_id).await {
             Some(state_rx) => state_rx,
             None => {
                 let msg = "executor not found, or not yet registered";
@@ -827,65 +927,12 @@ impl ExecutorApi for ExecutorAPIService {
             task_allocations: vec![],
             clock: Some(0),
         }));
-        let executor_manager = self.executor_manager.clone();
-        tokio::spawn(async move {
-            // 1. Mark the state as changed to trigger the first change notification to the
-            //    executor. This is important because between the report_executor_state and
-            //    the get_desired_executor_states requests, the executor state may received
-            //    a new desired state.
-            executor_state_rx.mark_changed();
-            loop {
-                // 2. Wait for a state change for this executor or grpc stream closing.
-                tokio::select! {
-                    _ = grpc_tx.closed() => {
-                        info!(
-                            executor_id = executor_id.get(),
-                            "get_desired_executor_states: grpc stream closed"
-                        );
-                        break;
-                    }
-                    result = executor_state_rx.changed() => {
-                        if let Err(err) = result {
-                            info!(
-                                executor_id = executor_id.get(),
-                                "get_desired_executor_states: state machine watcher closing: {}", err
-                            );
-                            break;
-                        }
-                    }
-                }
-
-                // 3. Get the latest state
-                let desired_state: DesiredExecutorState = match executor_manager
-                    .get_executor_state(&executor_id)
-                    .await
-                    .try_into()
-                {
-                    Ok(desired_state) => desired_state,
-                    Err(err) => {
-                        info!(
-                            executor_id = executor_id.get(),
-                            "get_desired_executor_states: failed to convert desired state, ignoring: {:?}", err
-                        );
-                        continue;
-                    }
-                };
-
-                debug!(
-                    executor_id = executor_id.get(),
-                    "get_desired_executor_states: got desired state: {:#?}", desired_state
-                );
-
-                // 4. Send the state to the executor
-                if let Err(err) = grpc_tx.send(Ok(desired_state)) {
-                    info!(
-                        executor_id = executor_id.get(),
-                        "get_desired_executor_states: grpc stream closing: {}", err
-                    );
-                    break;
-                }
-            }
-        });
+        tokio::spawn(executor_update_loop(
+            executor_id,
+            self.executor_manager.clone(),
+            executor_state_rx,
+            grpc_tx,
+        ));
 
         let grpc_stream = WatchStream::from_changes(grpc_rx);
         Ok(Response::new(