neutron 0.0.2

A Rust client library for Pulsar
Documentation

use futures::FutureExt;

use crate::{
    broker_address::BrokerAddress,
    command_resolver::CommandResolver,
    engine::EngineConnection,
    message::{Command, Inbound, Outbound},
    NeutronError,
};

#[derive(Clone)]
pub struct ClientData {
    pub id: u64,
    pub broker_address: BrokerAddress,
    pub topic: String,
}

#[derive(Debug)]
pub(crate) struct ClientConnection {
    pub(crate) id: u64,
    pub(crate) connection: EngineConnection<Inbound, Command<Outbound, Inbound>>,
    pub(crate) topic: String,
    pub(crate) broker_address: BrokerAddress,
}

impl From<&ClientConnection> for ClientData {
    fn from(val: &ClientConnection) -> Self {
        ClientData {
            id: val.id,
            broker_address: val.broker_address.clone(),
            topic: val.topic.clone(),
        }
    }
}

impl ClientConnection {
    fn get_connection(&self) -> &EngineConnection<Inbound, Command<Outbound, Inbound>> {
        &self.connection
    }

    fn broker_address(&self) -> &BrokerAddress {
        &self.broker_address
    }
}

pub(crate) struct ClientManager {
    clients: Vec<ClientConnection>,
    pub(crate) command_resolver: CommandResolver<Outbound, Inbound>,
}

impl ClientManager {
    pub(crate) fn new() -> Self {
        ClientManager {
            clients: Vec::new(),
            command_resolver: CommandResolver::new(),
        }
    }

    pub(crate) fn add_client(&mut self, client: ClientConnection) {
        self.clients.push(client);
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.clients.is_empty()
    }

    pub(crate) async fn next(&self) -> (ClientData, Result<Outbound, NeutronError>) {
        let (next, _, _) = futures::future::select_all(self.clients.iter().map(|client| {
            async {
                let connection = client.get_connection();
                let message = connection.recv().await;
                (client.into(), message)
            }
            .boxed()
        }))
        .await;

        let outbound = match next.1 {
            Ok(cmd) => match cmd {
                Command::RequestResponse(outbound, sender) => {
                    self.command_resolver.put(outbound.clone(), sender).await;
                    Ok(outbound)
                }
                _ => Ok(cmd.get_outbound()),
            },
            Err(err) => Err(err),
        };

        (next.0, outbound)
    }

    pub(crate) async fn send(
        &self,
        inbound: &Inbound,
        broker_address: &BrokerAddress,
    ) -> Result<(), NeutronError> {
        if self.command_resolver.try_resolve(inbound.clone()).await {
            return Ok(());
        }

        let consumer_or_producer_id = inbound.try_consumer_or_producer_id();

        let clients = self
            .clients
            .iter()
            .filter(|client| client.broker_address() == broker_address)
            .filter(|client| {
                if let Some(id) = consumer_or_producer_id {
                    client.id == id
                } else {
                    true
                }
            });

        futures::future::join_all(clients.map(|client| async {
            let connection = client.get_connection();
            connection.send(Ok(inbound.clone())).await
        }))
        .await;

        Ok(())
    }

    #[allow(dead_code)]
    pub(crate) async fn send_all(
        &self,
        inbound: &Result<Inbound, NeutronError>,
    ) -> Result<(), NeutronError> {
        let clients = self.clients.iter().map(|client| {
            let connection = client.get_connection();
            connection.send(inbound.clone())
        });

        futures::future::join_all(clients).await;

        Ok(())
    }

    pub(crate) fn move_client_to_broker(&mut self, id: u64, broker_address: &BrokerAddress) {
        for client in self.clients.iter_mut() {
            if client.id == id {
                client.broker_address = broker_address.clone();
            }
        }
    }

    pub(crate) fn get_client(&self, id: u64) -> Option<&ClientConnection> {
        self.clients.iter().find(|client| client.id == id)
    }
}