Revert "Run message processing tasks in parallel (#878)"

This reverts commit bb8a64f.

Author: rolyatmax

common/src/typed_socket/mod.rs

@@ -26,35 +26,12 @@ pub struct TypedSocket<T: ChannelMessage> {
 }
 
 #[derive(Clone)]
-pub struct TypedSocketSender<T: ChannelMessage> {
-    sender: Sender<SocketAction<T>>,
+pub struct TypedSocketSender<A> {
+    inner_send:
+        Arc<dyn Fn(SocketAction<A>) -> Result<(), TypedSocketError> + 'static + Send + Sync>,
 }
 
-#[derive(Clone)]
-pub struct WrappedTypedSocketSender<K> {
-    send: Arc<dyn Fn(K) -> Result<(), TypedSocketError> + 'static + Send + Sync>,
-}
-
-impl<K> WrappedTypedSocketSender<K> {
-    pub fn new<T: ChannelMessage, F>(sender: Sender<SocketAction<T>>, transform: F) -> Self
-    where
-        F: (Fn(K) -> T) + 'static + Send + Sync,
-    {
-        Self {
-            send: Arc::new(move |message| {
-                sender
-                    .try_send(SocketAction::Send(transform(message)))
-                    .map_err(TypedSocketError::from)
-            }),
-        }
-    }
-
-    pub fn send(&self, message: K) -> Result<(), TypedSocketError> {
-        (self.send)(message)
-    }
-}
-
-impl<T: ChannelMessage> Debug for TypedSocketSender<T> {
+impl<T> Debug for TypedSocketSender<T> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.write_str("typed socket sender")
     }
@@ -77,24 +54,16 @@ impl<A> From<TrySendError<A>> for TypedSocketError {
     }
 }
 
-impl<T: ChannelMessage> TypedSocketSender<T> {
-    pub fn send(&self, message: T) -> Result<(), TypedSocketError> {
-        self.sender.try_send(SocketAction::Send(message))?;
+impl<A: Debug> TypedSocketSender<A> {
+    pub fn send(&self, message: A) -> Result<(), TypedSocketError> {
+        (self.inner_send)(SocketAction::Send(message))?;
         Ok(())
     }
 
     pub fn close(&mut self) -> Result<(), TypedSocketError> {
-        self.sender.try_send(SocketAction::Close)?;
+        (self.inner_send)(SocketAction::Close)?;
         Ok(())
     }
-
-    /// Wrap the sender with a transform function.
-    pub fn wrap<K, F>(&self, transform: F) -> WrappedTypedSocketSender<K>
-    where
-        F: (Fn(K) -> T) + 'static + Send + Sync,
-    {
-        WrappedTypedSocketSender::new(self.sender.clone(), transform)
-    }
 }
 
 impl<T: ChannelMessage> TypedSocket<T> {
@@ -109,10 +78,22 @@ impl<T: ChannelMessage> TypedSocket<T> {
         self.recv.recv().await
     }
 
-    pub fn sender(&self) -> TypedSocketSender<T> {
+    pub fn sender<A, F>(&self, transform: F) -> TypedSocketSender<A>
+    where
+        F: (Fn(A) -> T) + 'static + Send + Sync,
+    {
         let sender = self.send.clone();
-
-        TypedSocketSender { sender }
+        let inner_send = move |message: SocketAction<A>| {
+            let message = match message {
+                SocketAction::Close => SocketAction::Close,
+                SocketAction::Send(message) => SocketAction::Send(transform(message)),
+            };
+            sender.try_send(message).map_err(|e| e.into())
+        };
+
+        TypedSocketSender {
+            inner_send: Arc::new(inner_send),
+        }
     }
 
     pub async fn close(&mut self) {

plane/plane-tests/tests/cert_manager.rs

@@ -11,7 +11,6 @@ use std::sync::Arc;
 mod common;
 
 #[plane_test]
-#[ignore = "Doesn't work"]
 async fn cert_manager_does_refresh(env: TestEnvironment) {
     let controller = env.controller().await;
 
@@ -57,7 +56,6 @@ async fn cert_manager_does_refresh(env: TestEnvironment) {
 }
 
 #[plane_test(500)]
-#[ignore = "Doesn't work"]
 async fn cert_manager_does_refresh_eab(env: TestEnvironment) {
     let certs_dir = env.scratch_dir.join("certs");
 

plane/plane-tests/tests/proxy_cors.rs

@@ -126,6 +126,6 @@ async fn proxy_valid_request_has_cors_headers(env: TestEnvironment) {
             .unwrap()
             .to_str()
             .unwrap(),
-        "*, Authorization"
+        "authorization, accept, content-type"
     );
 }

plane/src/controller/drone.rs

@@ -9,7 +9,7 @@ use plane_common::{
         ApiErrorKind, BackendAction, BackendActionMessage, Heartbeat, KeyDeadlines,
         MessageFromDrone, MessageToDrone, RenewKeyResponse,
     },
-    typed_socket::{server::new_server, TypedSocketSender},
+    typed_socket::{server::new_server, TypedSocket},
     types::{
         backend_state::TerminationReason, ClusterName, DronePoolName, NodeId, TerminationKind,
     },
@@ -41,7 +41,7 @@ pub async fn handle_message_from_drone(
     msg: MessageFromDrone,
     drone_id: NodeId,
     controller: &Controller,
-    sender: TypedSocketSender<MessageToDrone>,
+    sender: &mut TypedSocket<MessageToDrone>,
 ) -> anyhow::Result<()> {
     match msg {
         MessageFromDrone::BackendMetrics(metrics_msg) => {
@@ -146,7 +146,7 @@ pub async fn sweep_loop(db: PlaneDatabase, drone_id: NodeId) {
 
 pub async fn process_pending_actions(
     db: &PlaneDatabase,
-    socket: TypedSocketSender<MessageToDrone>,
+    socket: &mut TypedSocket<MessageToDrone>,
     drone_id: &NodeId,
 ) -> Result<(), anyhow::Error> {
     let mut count = 0;
@@ -200,7 +200,7 @@ pub async fn drone_socket_inner(
     let mut backend_actions: Subscription<BackendActionMessage> =
         controller.db.subscribe_with_key(&drone_id.to_string());
 
-    process_pending_actions(&controller.db, socket.sender(), &drone_id).await?;
+    process_pending_actions(&controller.db, &mut socket, &drone_id).await?;
 
     let mut interval = tokio::time::interval(Duration::from_secs(5));
     interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
@@ -212,13 +212,7 @@ pub async fn drone_socket_inner(
     loop {
         tokio::select! {
             _ = interval.tick() => {
-                let sender = socket.sender();
-                let db = controller.db.clone();
-                tokio::spawn(async move {
-                    if let Err(err) = process_pending_actions(&db, sender, &drone_id).await {
-                        tracing::error!(?err, "Error processing pending actions");
-                    }
-                });
+                process_pending_actions(&controller.db, &mut socket, &drone_id).await?;
             }
             _ = log_interval.tick() => {
                 let (outgoing, incoming) = socket.channel_depths();
@@ -269,14 +263,9 @@ pub async fn drone_socket_inner(
                                 *message_counts.entry("backend_metrics").or_insert(0) += 1;
                             }
                         }
-
-                        let sender = socket.sender();
-                        let controller = controller.clone();
-                        tokio::spawn(async move {
-                            if let Err(err) = handle_message_from_drone(message_from_drone, drone_id, &controller, sender).await {
-                                tracing::error!(?err, "Error handling message from drone");
-                            }
-                        });
+                        if let Err(err) = handle_message_from_drone(message_from_drone, drone_id, &controller, &mut socket).await {
+                            tracing::error!(?err, "Error handling message from drone");
+                        }
                     }
                     None => {
                         tracing::info!("Drone socket closed");

plane/src/controller/proxy.rs

@@ -14,7 +14,7 @@ use plane_common::{
         ApiErrorKind, CertManagerRequest, CertManagerResponse, MessageFromProxy, MessageToProxy,
         RouteInfoRequest, RouteInfoResponse,
     },
-    typed_socket::{server::new_server, TypedSocketSender},
+    typed_socket::{server::new_server, TypedSocket},
     types::{BackendState, BearerToken, ClusterName, NodeId},
 };
 use std::{
@@ -28,7 +28,7 @@ use valuable::Valuable;
 pub async fn handle_route_info_request(
     token: BearerToken,
     controller: &Controller,
-    socket: TypedSocketSender<MessageToProxy>,
+    socket: &mut TypedSocket<MessageToProxy>,
 ) -> anyhow::Result<()> {
     match controller.db.backend().route_info_for_token(&token).await {
         // When a proxy requests a route, either:
@@ -79,58 +79,65 @@ pub async fn handle_route_info_request(
                 }
             }
 
-            let socket = socket.wrap(MessageToProxy::RouteInfoResponse);
-
-            loop {
-                // Note: this timeout is arbitrary to avoid a memory leak. Under normal system operation, the critical
-                // timeout will be that of the backend failing to start. We use a large timeout to avoid it becoming
-                // the critical timeout when the system is functioning.
-                let result = match tokio::time::timeout(
-                    std::time::Duration::from_secs(30 * 60 /* 30 minutes */),
-                    sub.next(),
-                )
-                .await
-                {
-                    Ok(Some(result)) => result,
-                    Ok(None) => {
-                        tracing::error!("Event subscription closed!");
-                        break;
-                    }
-                    Err(_) => {
-                        tracing::error!("Timeout waiting for backend state");
-                        break;
-                    }
-                };
+            let socket = socket.sender(MessageToProxy::RouteInfoResponse);
+            tokio::spawn(async move {
+                loop {
+                    // Note: this timeout is arbitrary to avoid a memory leak. Under normal system operation, the critical
+                    // timeout will be that of the backend failing to start. We use a large timeout to avoid it becoming
+                    // the critical timeout when the system is functioning.
+                    let result = match tokio::time::timeout(
+                        std::time::Duration::from_secs(30 * 60 /* 30 minutes */),
+                        sub.next(),
+                    )
+                    .await
+                    {
+                        Ok(Some(result)) => result,
+                        Ok(None) => {
+                            tracing::error!("Event subscription closed!");
+                            break;
+                        }
+                        Err(_) => {
+                            tracing::error!("Timeout waiting for backend state");
+                            break;
+                        }
+                    };
 
-                let Notification { payload, .. } = result;
+                    let Notification { payload, .. } = result;
 
-                match payload {
-                    BackendState::Ready { address } => {
-                        let route_info = partial_route_info.set_address(address);
-                        let response = RouteInfoResponse {
-                            token,
-                            route_info: Some(route_info),
-                        };
-                        if let Err(err) = socket.send(response) {
-                            tracing::error!(?err, "Error sending route info response to proxy.");
+                    match payload {
+                        BackendState::Ready { address } => {
+                            let route_info = partial_route_info.set_address(address);
+                            let response = RouteInfoResponse {
+                                token,
+                                route_info: Some(route_info),
+                            };
+                            if let Err(err) = socket.send(response) {
+                                tracing::error!(
+                                    ?err,
+                                    "Error sending route info response to proxy."
+                                );
+                            }
+                            break;
                         }
-                        break;
-                    }
-                    BackendState::Terminated { .. }
-                    | BackendState::Terminating { .. }
-                    | BackendState::HardTerminating { .. } => {
-                        let response = RouteInfoResponse {
-                            token,
-                            route_info: None,
-                        };
-                        if let Err(err) = socket.send(response) {
-                            tracing::error!(?err, "Error sending route info response to proxy.");
+                        BackendState::Terminated { .. }
+                        | BackendState::Terminating { .. }
+                        | BackendState::HardTerminating { .. } => {
+                            let response = RouteInfoResponse {
+                                token,
+                                route_info: None,
+                            };
+                            if let Err(err) = socket.send(response) {
+                                tracing::error!(
+                                    ?err,
+                                    "Error sending route info response to proxy."
+                                );
+                            }
+                            break;
                         }
-                        break;
+                        _ => {}
                     }
-                    _ => {}
                 }
-            }
+            });
         }
         Ok(RouteInfoResult::NotFound) => {
             let response = RouteInfoResponse {
@@ -152,7 +159,7 @@ pub async fn handle_route_info_request(
 pub async fn handle_message_from_proxy(
     message: MessageFromProxy,
     controller: &Controller,
-    socket: TypedSocketSender<MessageToProxy>,
+    socket: &mut TypedSocket<MessageToProxy>,
     cluster: &ClusterName,
     node_id: NodeId,
 ) -> anyhow::Result<()> {
@@ -294,15 +301,7 @@ pub async fn proxy_socket_inner(
                                 *message_counts.entry("cert_manager_request").or_insert(0) += 1;
                             }
                         }
-
-                        let sender = socket.sender();
-                        let controller = controller.clone();
-                        let cluster = cluster.clone();
-                        tokio::spawn(async move {
-                            if let Err(err) = handle_message_from_proxy(message, &controller, sender, &cluster, node_guard.id).await {
-                                tracing::error!(?err, "Error handling message from proxy");
-                            }
-                        });
+                        handle_message_from_proxy(message, &controller, &mut socket, &cluster, node_guard.id).await?
                     }
                     None => {
                         tracing::info!("Proxy socket closed");

plane/src/drone/heartbeat.rs

@@ -1,8 +1,6 @@
 use crate::heartbeat_consts::HEARTBEAT_INTERVAL;
 use chrono::Utc;
-use plane_common::{
-    log_types::LoggableTime, protocol::Heartbeat, typed_socket::WrappedTypedSocketSender,
-};
+use plane_common::{log_types::LoggableTime, protocol::Heartbeat, typed_socket::TypedSocketSender};
 use tokio::task::JoinHandle;
 
 /// A background task that sends heartbeats to the server.
@@ -11,7 +9,7 @@ pub struct HeartbeatLoop {
 }
 
 impl HeartbeatLoop {
-    pub fn start(sender: WrappedTypedSocketSender<Heartbeat>) -> Self {
+    pub fn start(sender: TypedSocketSender<Heartbeat>) -> Self {
         let handle = tokio::spawn(async move {
             loop {
                 let local_time = LoggableTime(Utc::now());

plane/src/drone/key_manager.rs

@@ -5,7 +5,7 @@ use plane_common::{
     log_types::LoggableTime,
     names::BackendName,
     protocol::{AcquiredKey, BackendAction, KeyDeadlines, RenewKeyRequest},
-    typed_socket::WrappedTypedSocketSender,
+    typed_socket::TypedSocketSender,
     types::{backend_state::TerminationReason, TerminationKind},
 };
 use std::{collections::HashMap, sync::Arc, time::Duration};
@@ -20,13 +20,13 @@ pub struct KeyManager {
     /// and terminating the backend if the key cannot be renewed.
     handles: HashMap<BackendName, (AcquiredKey, GuardHandle)>,
 
-    sender: Option<WrappedTypedSocketSender<RenewKeyRequest>>,
+    sender: Option<TypedSocketSender<RenewKeyRequest>>,
 }
 
 async fn renew_key_loop(
     key: AcquiredKey,
     backend: BackendName,
-    sender: Option<WrappedTypedSocketSender<RenewKeyRequest>>,
+    sender: Option<TypedSocketSender<RenewKeyRequest>>,
     executor: Arc<Executor>,
 ) {
     loop {
@@ -120,7 +120,7 @@ impl KeyManager {
         }
     }
 
-    pub fn set_sender(&mut self, sender: WrappedTypedSocketSender<RenewKeyRequest>) {
+    pub fn set_sender(&mut self, sender: TypedSocketSender<RenewKeyRequest>) {
         self.sender.replace(sender);
 
         for (backend, (acquired_key, handle)) in self.handles.iter_mut() {