plexor_core/

thalamus.rs

// Copyright 2025 Alecks Gates
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use crate::codec::{Codec, CodecName};
use crate::erasure::payload::PayloadErased;
use crate::erasure::reactant::{ErrorReactantErased, ReactantErased};
use crate::ganglion::{Ganglion, GanglionError, GanglionInternal};
use crate::neuron::Neuron;
use crate::utils::struct_name_of_type;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::Mutex;
use uuid::Uuid;

/// Thalamus: routes each transmit to exactly one internal peer (round-robin),
/// avoiding cross-ganglia reactions/broadcasts. Reactants are attached to all peers so any
/// selected peer can process inbound messages.
pub struct Thalamus<G>
where
    G: GanglionInternal + Ganglion + Send + Sync + 'static,
{
    id: Uuid,
    peers: Vec<Arc<Mutex<G>>>,
    rr_idx: usize,
}

impl<G> Thalamus<G>
where
    G: GanglionInternal + Ganglion + Send + Sync + 'static,
{
    pub fn new(peers: Vec<Arc<Mutex<G>>>) -> Self {
        Self {
            id: Uuid::now_v7(),
            peers,
            rr_idx: 0,
        }
    }

    fn next_index(&mut self) -> usize {
        if self.peers.is_empty() {
            return 0;
        }
        let idx = self.rr_idx % self.peers.len();
        self.rr_idx = (idx + 1) % self.peers.len();
        idx
    }
}

impl<G> Ganglion for Thalamus<G>
where
    G: GanglionInternal + Ganglion + Send + Sync + 'static,
{
    fn capable<T, C>(&mut self, neuron: Arc<dyn Neuron<T, C> + Send + Sync>) -> bool
    where
        C: Codec<T> + CodecName + Send + Sync + 'static,
        T: Send + Sync + 'static,
    {
        for peer in &self.peers {
            // Try lock and ask capability. Errors are treated as not capable.
            if let Ok(mut guard) = peer.try_lock() {
                if guard.capable::<T, C>(neuron.clone()) {
                    return true;
                }
            }
        }
        false
    }

    fn adapt<T, C>(
        &mut self,
        neuron: Arc<dyn Neuron<T, C> + Send + Sync>,
    ) -> Pin<Box<dyn Future<Output = Result<(), GanglionError>> + Send + 'static>>
    where
        C: Codec<T> + CodecName + Send + Sync + 'static,
        T: Send + Sync + 'static,
    {
        // If no peer is capable, succeed without action (matches GanglionInprocess semantics)
        if !self.capable::<T, C>(neuron.clone()) {
            return Box::pin(async move { Ok(()) });
        }

        let peers = self.peers.clone();
        Box::pin(async move {
            // Forward adaptation to all capable peers
            for peer in peers.iter() {
                let mut p = peer.lock().await;
                if p.capable::<T, C>(neuron.clone()) {
                    p.adapt::<T, C>(neuron.clone()).await?;
                }
            }
            Ok(())
        })
    }
}

impl<G> GanglionInternal for Thalamus<G>
where
    G: GanglionInternal + Ganglion + Send + Sync + 'static,
{
    fn transmit(
        &mut self,
        payload: Arc<dyn PayloadErased + Send + Sync + 'static>,
    ) -> Pin<Box<dyn Future<Output = Result<Vec<()>, GanglionError>> + Send + 'static>> {
        let peers = self.peers.clone();
        let start_idx = self.next_index();
        let neuron_name = payload.get_neuron_name();
        let ganglion_id = self.id;
        let ganglion_name = struct_name_of_type::<Self>().to_string();

        Box::pin(async move {
            if peers.is_empty() {
                return Err(GanglionError::SynapseNotFound {
                    neuron_name,
                    ganglion_name,
                    ganglion_id,
                });
            }

            let mut last_err: Option<GanglionError> = None;
            // Try each peer at most once, starting from start_idx
            for off in 0..peers.len() {
                let idx = (start_idx + off) % peers.len();
                let peer = &peers[idx];
                let future = {
                    let mut p = peer.lock().await;
                    p.transmit(payload.clone())
                };
                match future.await {
                    Ok(acks) => {
                        // Success. Return immediately.
                        return Ok(acks);
                    }
                    Err(e) => {
                        // Save last error and try next peer.
                        last_err = Some(e);
                        continue;
                    }
                }
            }
            Err(last_err.unwrap_or(GanglionError::Transmit {
                neuron_name,
                ganglion_name,
                ganglion_id,
                message: "No peers available or all transmissions failed".to_string(),
            }))
        })
    }

    fn react(
        &mut self,
        neuron_name: String,
        reactants: Vec<Arc<dyn ReactantErased + Send + Sync + 'static>>,
        error_reactants: Vec<Arc<dyn ErrorReactantErased + Send + Sync>>,
    ) -> Pin<Box<dyn Future<Output = Result<(), GanglionError>> + Send + 'static>> {
        let peers = self.peers.clone();
        Box::pin(async move {
            if peers.is_empty() {
                return Err(GanglionError::SynapseNotFound {
                    neuron_name,
                    ganglion_name: struct_name_of_type::<Self>().to_string(),
                    ganglion_id: Uuid::nil(),
                });
            }

            let mut at_least_one_ok = false;
            let mut last_err: Option<GanglionError> = None;
            for peer in peers.iter() {
                let future = {
                    let mut p = peer.lock().await;
                    p.react(
                        neuron_name.clone(),
                        reactants.clone(),
                        error_reactants.clone(),
                    )
                };
                match future.await {
                    Ok(()) => at_least_one_ok = true,
                    Err(e) => last_err = Some(e),
                }
            }
            if at_least_one_ok {
                Ok(())
            } else {
                Err(last_err.unwrap_or(GanglionError::SynapseNotFound {
                    neuron_name,
                    ganglion_name: struct_name_of_type::<Self>().to_string(),
                    ganglion_id: Uuid::nil(),
                }))
            }
        })
    }

    fn unique_id(&self) -> Uuid {
        self.id
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::erasure::payload::erase_payload;
    use crate::erasure::reactant::erase_reactant;
    use crate::ganglion::GanglionInprocess;
    use crate::logging::TraceContext;
    use crate::neuron::NeuronImpl;
    use crate::payload::Payload;
    use crate::reactant::Reactant;
    use crate::test_utils::{
        DebugCodec, DebugStruct, ResponseCodec, ResponseStruct, TokioMpscReactant, test_namespace,
    };
    use std::sync::Arc;
    use tokio::sync::Mutex;
    use tokio::sync::mpsc::channel;
    use uuid::Uuid;

    #[tokio::test]
    async fn test_thalamus_round_robin_basic() {
        let ns = test_namespace();
        let neuron: NeuronImpl<DebugStruct, DebugCodec> = NeuronImpl::new(ns.clone());
        let neuron_name = neuron.name();
        let neuron_arc: Arc<dyn Neuron<DebugStruct, DebugCodec> + Send + Sync> = Arc::new(neuron);

        let g1 = Arc::new(Mutex::new(GanglionInprocess::new()));
        let g2 = Arc::new(Mutex::new(GanglionInprocess::new()));

        let mut thalamus = Thalamus::new(vec![g1.clone(), g2.clone()]);
        // Adapt neuron across peers
        thalamus
            .adapt::<DebugStruct, DebugCodec>(neuron_arc.clone())
            .await
            .unwrap();

        let (tx, mut rx) = channel::<Arc<Payload<DebugStruct, DebugCodec>>>(10);
        let reactants = vec![erase_reactant::<DebugStruct, DebugCodec, _>(Box::new(
            TokioMpscReactant::new(tx),
        ))];

        // Attach reactants across all peers via thalamus
        thalamus
            .react(neuron_name.clone(), reactants, vec![])
            .await
            .unwrap();

        // Create payloads and send twice; expect two receives (not broadcast per transmit)
        let correlation_id = Uuid::now_v7();
        let span_id = correlation_id.as_u128() as u64;
        let payload1 = Arc::new(Payload::from_parts(
            Arc::new(DebugStruct {
                foo: 1,
                bar: "a".to_string(),
            }),
            neuron_arc.clone(),
            TraceContext::from_parts(correlation_id, span_id, None),
        ));
        let payload2 = Payload::new(
            DebugStruct {
                foo: 2,
                bar: "b".to_string(),
            },
            neuron_arc.clone(),
        );

        thalamus.transmit(erase_payload(payload1)).await.unwrap();
        thalamus.transmit(erase_payload(payload2)).await.unwrap();

        // Should receive exactly two messages
        let _m1 = rx.recv().await.expect("expected first message");
        let _m2 = rx.recv().await.expect("expected second message");
    }

    #[tokio::test]
    async fn test_thalamus_even_work_distribution() {
        let ns = test_namespace();
        let neuron: NeuronImpl<DebugStruct, DebugCodec> = NeuronImpl::new(ns.clone());
        let neuron_name = neuron.name();
        let neuron_arc: Arc<dyn Neuron<DebugStruct, DebugCodec> + Send + Sync> = Arc::new(neuron);

        let g1 = Arc::new(Mutex::new(GanglionInprocess::new()));
        let g2 = Arc::new(Mutex::new(GanglionInprocess::new()));
        let g3 = Arc::new(Mutex::new(GanglionInprocess::new()));

        let mut thalamus = Thalamus::new(vec![g1.clone(), g2.clone(), g3.clone()]);
        // Adapt neuron across peers
        thalamus
            .adapt::<DebugStruct, DebugCodec>(neuron_arc.clone())
            .await
            .unwrap();

        // Create separate channels for each ganglion to track work distribution
        let (tx1, mut rx1) = channel::<Arc<Payload<DebugStruct, DebugCodec>>>(10);
        let (tx2, mut rx2) = channel::<Arc<Payload<DebugStruct, DebugCodec>>>(10);
        let (tx3, mut rx3) = channel::<Arc<Payload<DebugStruct, DebugCodec>>>(10);

        // Attach reactants to each peer individually to track which peer handles what
        {
            let mut g1_guard = g1.lock().await;
            let reactants1 = vec![erase_reactant::<DebugStruct, DebugCodec, _>(Box::new(
                TokioMpscReactant::new(tx1),
            ))];
            g1_guard
                .react(neuron_name.clone(), reactants1, vec![])
                .await
                .unwrap();
        }
        {
            let mut g2_guard = g2.lock().await;
            let reactants2 = vec![erase_reactant::<DebugStruct, DebugCodec, _>(Box::new(
                TokioMpscReactant::new(tx2),
            ))];
            g2_guard
                .react(neuron_name.clone(), reactants2, vec![])
                .await
                .unwrap();
        }
        {
            let mut g3_guard = g3.lock().await;
            let reactants3 = vec![erase_reactant::<DebugStruct, DebugCodec, _>(Box::new(
                TokioMpscReactant::new(tx3),
            ))];
            g3_guard
                .react(neuron_name.clone(), reactants3, vec![])
                .await
                .unwrap();
        }

        // Send 6 messages through thalamus (2 per peer expected with round-robin)
        for i in 0..6 {
            let test_data = DebugStruct {
                foo: i,
                bar: format!("msg{i}"),
            };

            thalamus
                .transmit(erase_payload(Payload::new(test_data, neuron_arc.clone())))
                .await
                .expect("Failed to transmit");
        }

        // Count messages received by each ganglion
        let mut count1 = 0;
        let mut count2 = 0;
        let mut count3 = 0;

        // Use try_recv to avoid blocking, and give some time for async processing
        tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;

        while rx1.try_recv().is_ok() {
            count1 += 1;
        }
        while rx2.try_recv().is_ok() {
            count2 += 1;
        }
        while rx3.try_recv().is_ok() {
            count3 += 1;
        }

        // With round-robin, each ganglion should get exactly 2 messages
        assert_eq!(count1, 2, "Ganglion 1 should receive 2 messages");
        assert_eq!(count2, 2, "Ganglion 2 should receive 2 messages");
        assert_eq!(count3, 2, "Ganglion 3 should receive 2 messages");
        assert_eq!(count1 + count2 + count3, 6, "Total messages should be 6");
    }

    #[tokio::test]
    async fn test_thalamus_work_distribution_with_responses() {
        let ns = test_namespace();

        // Create request neuron for sending work
        let request_neuron: NeuronImpl<DebugStruct, DebugCodec> = NeuronImpl::new(ns.clone());
        let request_neuron_name = request_neuron.name();
        let request_neuron_arc: Arc<dyn Neuron<DebugStruct, DebugCodec> + Send + Sync> =
            Arc::new(request_neuron);

        let g1 = Arc::new(Mutex::new(GanglionInprocess::new()));
        let g2 = Arc::new(Mutex::new(GanglionInprocess::new()));
        let g3 = Arc::new(Mutex::new(GanglionInprocess::new()));

        let mut thalamus = Thalamus::new(vec![g1.clone(), g2.clone(), g3.clone()]);

        // Create response neuron for receiving responses using ResponseStruct/ResponseCodec
        let response_neuron: NeuronImpl<ResponseStruct, ResponseCodec> =
            NeuronImpl::new(ns.clone());
        let response_neuron_name = response_neuron.name();
        let response_neuron_arc: Arc<dyn Neuron<ResponseStruct, ResponseCodec> + Send + Sync> =
            Arc::new(response_neuron);

        // Adapt both request and response neurons across all peers
        thalamus
            .adapt::<DebugStruct, DebugCodec>(request_neuron_arc.clone())
            .await
            .unwrap();
        thalamus
            .adapt::<ResponseStruct, ResponseCodec>(response_neuron_arc.clone())
            .await
            .unwrap();

        // Channel to receive response payloads from thalamus
        let (response_tx, mut response_rx) =
            channel::<Arc<Payload<ResponseStruct, ResponseCodec>>>(20);

        // Create a custom reactant to capture responses using ResponseStruct
        #[derive(Clone)]
        struct ResponseCaptureReactant {
            sender: tokio::sync::mpsc::Sender<Arc<Payload<ResponseStruct, ResponseCodec>>>,
        }

        impl ResponseCaptureReactant {
            fn new(
                sender: tokio::sync::mpsc::Sender<Arc<Payload<ResponseStruct, ResponseCodec>>>,
            ) -> Self {
                Self { sender }
            }
        }

        impl Reactant<ResponseStruct, ResponseCodec> for ResponseCaptureReactant {
            fn react(
                &self,
                payload: Arc<Payload<ResponseStruct, ResponseCodec>>,
            ) -> Pin<
                Box<
                    dyn Future<Output = Result<(), crate::reactant::ReactantError>>
                        + Send
                        + 'static,
                >,
            > {
                let sender = self.sender.clone();
                let payload_clone = payload.clone();

                Box::pin(async move {
                    // No filtering needed - all ResponseStruct payloads are responses by definition
                    let _ = sender.try_send(payload_clone);
                    Ok(())
                })
            }

            fn erase(self: Box<Self>) -> Arc<dyn ReactantErased + Send + Sync + 'static> {
                erase_reactant(self)
            }
        }

        let response_capture_reactant = ResponseCaptureReactant::new(response_tx.clone());

        // Create a custom reactant that queues responses instead of transmitting directly
        #[derive(Clone)]
        struct ResponseGeneratingReactant {
            ganglion_id: u32,
            response_neuron: Arc<dyn Neuron<ResponseStruct, ResponseCodec> + Send + Sync>,
            queue_sender: tokio::sync::mpsc::Sender<Arc<Payload<ResponseStruct, ResponseCodec>>>,
        }

        impl ResponseGeneratingReactant {
            fn new(
                ganglion_id: u32,
                response_neuron: Arc<dyn Neuron<ResponseStruct, ResponseCodec> + Send + Sync>,
                queue_sender: tokio::sync::mpsc::Sender<
                    Arc<Payload<ResponseStruct, ResponseCodec>>,
                >,
            ) -> Self {
                Self {
                    ganglion_id,
                    response_neuron,
                    queue_sender,
                }
            }
        }

        impl Reactant<DebugStruct, DebugCodec> for ResponseGeneratingReactant {
            fn react(
                &self,
                payload: Arc<Payload<DebugStruct, DebugCodec>>,
            ) -> Pin<
                Box<
                    dyn Future<Output = Result<(), crate::reactant::ReactantError>>
                        + Send
                        + 'static,
                >,
            > {
                let ganglion_id = self.ganglion_id;
                let response_neuron = self.response_neuron.clone();
                let queue_sender = self.queue_sender.clone();
                let original_value = payload.value.clone();

                Box::pin(async move {
                    // Process all DebugStruct payloads as requests (no filtering needed)
                    // Create ResponseStruct payload indicating which ganglion processed this request
                    let response_payload = Payload::new(
                        ResponseStruct {
                            ganglion_id,
                            response_message: format!(
                                "response_from_ganglion_{}_for_{}",
                                ganglion_id, original_value.bar
                            ),
                        },
                        response_neuron,
                    );

                    // Add response payload to queue instead of transmitting directly
                    let _ = queue_sender.try_send(response_payload);
                    Ok(())
                })
            }

            fn erase(self: Box<Self>) -> Arc<dyn ReactantErased + Send + Sync + 'static> {
                erase_reactant(self)
            }
        }

        // Create queues for each ganglion's transmission loop
        let (queue1_tx, mut queue1_rx) = channel::<Arc<Payload<ResponseStruct, ResponseCodec>>>(10);
        let (queue2_tx, mut queue2_rx) = channel::<Arc<Payload<ResponseStruct, ResponseCodec>>>(10);
        let (queue3_tx, mut queue3_rx) = channel::<Arc<Payload<ResponseStruct, ResponseCodec>>>(10);

        let thalamus_arc = Arc::new(Mutex::new(thalamus));

        // Add the response capture reactant to the thalamus for the response neuron
        {
            let mut thalamus_guard = thalamus_arc.lock().await;
            let response_reactants = vec![erase_reactant::<ResponseStruct, ResponseCodec, _>(
                Box::new(response_capture_reactant),
            )];
            let future =
                thalamus_guard.react(response_neuron_name.clone(), response_reactants, vec![]);
            drop(thalamus_guard);
            future.await.unwrap();
        }

        // Start transmission loops for each ganglion
        let g1_clone = g1.clone();
        tokio::spawn(async move {
            while let Some(payload) = queue1_rx.recv().await {
                let future = {
                    let mut ganglion_guard = g1_clone.lock().await;
                    ganglion_guard.transmit(erase_payload(payload))
                };
                let _ = future.await;
            }
        });

        let g2_clone = g2.clone();
        tokio::spawn(async move {
            while let Some(payload) = queue2_rx.recv().await {
                let future = {
                    let mut ganglion_guard = g2_clone.lock().await;
                    ganglion_guard.transmit(erase_payload(payload))
                };
                let _ = future.await;
            }
        });

        let g3_clone = g3.clone();
        tokio::spawn(async move {
            while let Some(payload) = queue3_rx.recv().await {
                let future = {
                    let mut ganglion_guard = g3_clone.lock().await;
                    ganglion_guard.transmit(erase_payload(payload))
                };
                let _ = future.await;
            }
        });

        // Attach response-generating reactants to each child ganglion
        // Each reactant sends to its corresponding queue
        {
            let mut g1_guard = g1.lock().await;
            let reactants1 = vec![erase_reactant::<DebugStruct, DebugCodec, _>(Box::new(
                ResponseGeneratingReactant::new(1, response_neuron_arc.clone(), queue1_tx),
            ))];
            let future = g1_guard.react(request_neuron_name.clone(), reactants1, vec![]);
            drop(g1_guard);
            future.await.unwrap();
        }
        {
            let mut g2_guard = g2.lock().await;
            let reactants2 = vec![erase_reactant::<DebugStruct, DebugCodec, _>(Box::new(
                ResponseGeneratingReactant::new(2, response_neuron_arc.clone(), queue2_tx),
            ))];
            let future = g2_guard.react(request_neuron_name.clone(), reactants2, vec![]);
            drop(g2_guard);
            future.await.unwrap();
        }
        {
            let mut g3_guard = g3.lock().await;
            let reactants3 = vec![erase_reactant::<DebugStruct, DebugCodec, _>(Box::new(
                ResponseGeneratingReactant::new(3, response_neuron_arc.clone(), queue3_tx),
            ))];
            let future = g3_guard.react(request_neuron_name.clone(), reactants3, vec![]);
            drop(g3_guard);
            future.await.unwrap();
        }

        // Send 6 request messages through thalamus (2 per peer expected with round-robin)
        {
            for i in 0..6 {
                let payload = Payload::new(
                    DebugStruct {
                        foo: i,
                        bar: format!("request_{i}"),
                    },
                    request_neuron_arc.clone(),
                );

                let future = {
                    let mut thalamus_guard = thalamus_arc.lock().await;
                    thalamus_guard.transmit(erase_payload(payload))
                };
                future.await.unwrap();
            }
        }

        // Count response messages received by ganglion ID
        let mut count_g1 = 0;
        let mut count_g2 = 0;
        let mut count_g3 = 0;

        // Give some time for async processing and responses
        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;

        // Collect all response messages using proper async coordination like plexus tests
        let mut total_received = 0;
        while total_received < 6 && !response_rx.is_empty() {
            if let Ok(payload) = response_rx.try_recv() {
                match payload.value.ganglion_id {
                    1 => count_g1 += 1,
                    2 => count_g2 += 1,
                    3 => count_g3 += 1,
                    _ => panic!(
                        "Unexpected ganglion ID in response: {}",
                        payload.value.ganglion_id
                    ),
                }
                total_received += 1;
            } else {
                break;
            }
        }

        // With round-robin, each ganglion should send exactly 2 responses
        assert_eq!(count_g1, 2, "Should receive 2 responses from ganglion 1");
        assert_eq!(count_g2, 2, "Should receive 2 responses from ganglion 2");
        assert_eq!(count_g3, 2, "Should receive 2 responses from ganglion 3");
        assert_eq!(
            count_g1 + count_g2 + count_g3,
            6,
            "Total responses should be 6"
        );
    }
}