path: root/luz/src/connection/mod.rs

// TODO: consider if this needs to be handled by a supervisor or could be handled by luz directly

use std::{
    collections::HashMap,
    ops::{Deref, DerefMut},
    sync::Arc,
    time::Duration,
};

use jabber::{connection::Tls, jabber_stream::bound_stream::BoundJabberStream};
use jid::JID;
use read::{ReadControl, ReadControlHandle};
use stanza::client::Stanza;
use tokio::{
    sync::{mpsc, oneshot, Mutex},
    task::{JoinHandle, JoinSet},
};
use write::{WriteControl, WriteControlHandle, WriteHandle, WriteMessage};

use crate::{
    db::Db,
    error::{Error, Reason},
    UpdateMessage,
};

mod read;
pub(crate) mod write;

pub struct Supervisor {
    connection_commands: mpsc::Receiver<SupervisorCommand>,
    writer_crash: oneshot::Receiver<(WriteMessage, mpsc::Receiver<WriteMessage>)>,
    reader_crash: oneshot::Receiver<(
        Db,
        mpsc::Sender<UpdateMessage>,
        tokio::task::JoinSet<()>,
        mpsc::Sender<SupervisorCommand>,
        WriteHandle,
        Arc<Mutex<HashMap<String, oneshot::Sender<Result<Stanza, Reason>>>>>,
    )>,
    sender: mpsc::Sender<UpdateMessage>,
    writer_handle: WriteControlHandle,
    reader_handle: ReadControlHandle,
    on_shutdown: oneshot::Sender<()>,
    jid: Arc<Mutex<JID>>,
    password: Arc<String>,
}

pub enum SupervisorCommand {
    Disconnect,
    // for if there was a stream error, require to reconnect
    // couldn't stream errors just cause a crash? lol
    Reconnect(State),
}

pub enum State {
    Write(mpsc::Receiver<WriteMessage>),
    Read(
        (
            Db,
            mpsc::Sender<UpdateMessage>,
            tokio::task::JoinSet<()>,
            mpsc::Sender<SupervisorCommand>,
            WriteHandle,
            Arc<Mutex<HashMap<String, oneshot::Sender<Result<Stanza, Reason>>>>>,
        ),
    ),
}

impl Supervisor {
    fn new(
        connection_commands: mpsc::Receiver<SupervisorCommand>,
        writer_crash: oneshot::Receiver<(WriteMessage, mpsc::Receiver<WriteMessage>)>,
        reader_crash: oneshot::Receiver<(
            Db,
            mpsc::Sender<UpdateMessage>,
            JoinSet<()>,
            mpsc::Sender<SupervisorCommand>,
            WriteHandle,
            Arc<Mutex<HashMap<String, oneshot::Sender<Result<Stanza, Reason>>>>>,
        )>,
        sender: mpsc::Sender<UpdateMessage>,
        writer_handle: WriteControlHandle,
        reader_handle: ReadControlHandle,
        on_shutdown: oneshot::Sender<()>,
        jid: Arc<Mutex<JID>>,
        password: Arc<String>,
    ) -> Self {
        Self {
            connection_commands,
            writer_crash,
            sender,
            writer_handle,
            reader_handle,
            reader_crash,
            on_shutdown,
            jid,
            password,
        }
    }

    async fn run(mut self) {
        loop {
            tokio::select! {
                Some(msg) = self.connection_commands.recv() => {
                    match msg {
                        SupervisorCommand::Disconnect => {
                            let _ = self.writer_handle.send(WriteControl::Disconnect).await;
                            let _ = self.reader_handle.send(ReadControl::Disconnect).await;
                            tokio::select! {
                                _ = async { tokio::join!(
                                    async { let _ = (&mut self.writer_handle.handle).await; },
                                    async { let _ = (&mut self.reader_handle.handle).await; }
                                ) } => {},
                                _ = async { tokio::time::sleep(Duration::from_secs(5)) } => {
                                    (&mut self.reader_handle.handle).abort();
                                    (&mut self.writer_handle.handle).abort();
                                }
                            }
                            break;
                        },
                        SupervisorCommand::Reconnect(state) => {
                            // TODO: please omfg
                            // send abort to read stream, as already done, consider
                            let (read_state, mut write_state);
                            match state {
                                // TODO: proper state things for read and write thread
                                State::Write(receiver) => {
                                    write_state = receiver;
                                    let (send, recv) = oneshot::channel();
                                    let _ = self.reader_handle.send(ReadControl::Abort(send)).await;
                                    if let Ok(state) = recv.await {
                                        read_state = state;
                                    } else {
                                        break
                                    }
                                },
                                State::Read(read) => {
                                    read_state = read;
                                    let (send, recv) = oneshot::channel();
                                    let _ = self.writer_handle.send(WriteControl::Abort(send)).await;
                                    // TODO: need a tokio select, in case the state arrives from somewhere else
                                    if let Ok(state) = recv.await {
                                        write_state = state;
                                    } else {
                                        break
                                    }
                                },
                            }

                            let mut jid = self.jid.lock().await;
                            let mut domain = jid.domainpart.clone();
                            let connection = jabber::connect_and_login(&mut jid, &*self.password, &mut domain).await;
                            match connection {
                                Ok(c) => {

                                    let (read, write) = c.split();
                                    let (send, recv) = oneshot::channel();
                                    self.writer_crash = recv;
                                    self.writer_handle =
                                        WriteControlHandle::reconnect(write, send, write_state);
                                    let (send, recv) = oneshot::channel();
                                    self.reader_crash = recv;
                                    let (db, update_sender, tasks, supervisor_command, write_sender, pending_iqs) = read_state;
                                    self.reader_handle = ReadControlHandle::reconnect(
                                        read,
                                        send,
                                        db,
                                        update_sender,
                                        supervisor_command,
                                        write_sender,
                                        tasks,
                                        pending_iqs,
                                    );
                                },
                                Err(e) => {
                                    // if reconnection failure, respond to all current write messages with lost connection error. the received processes should complete themselves.
                                    write_state.close();
                                    while let Some(msg) = write_state.recv().await {
                                        let _ = msg.respond_to.send(Err(Reason::LostConnection));
                                    }
                                    // TODO: is this the correct error?
                                    let _ = self.sender.send(UpdateMessage::Error(Error::LostConnection)).await;
                                    break;
                                },
                            }
                        },
                    }
                },
                Ok((write_msg, mut write_recv)) = &mut self.writer_crash => {
                    // consider awaiting/aborting the read and write threads
                    let (send, recv) = oneshot::channel();
                    let _ = self.reader_handle.send(ReadControl::Abort(send)).await;
                    let (db, update_sender, tasks, supervisor_command, write_sender, pending_iqs) = tokio::select! {
                        Ok(s) = recv => s,
                        Ok(s) = &mut self.reader_crash => s,
                        // in case, just break as irrecoverable
                        else => break,
                    };

                    let mut jid = self.jid.lock().await;
                    let mut domain = jid.domainpart.clone();
                    let connection = jabber::connect_and_login(&mut jid, &*self.password, &mut domain).await;
                    match connection {
                        Ok(c) => {
                            let (read, write) = c.split();
                            let (send, recv) = oneshot::channel();
                            self.writer_crash = recv;
                            self.writer_handle =
                                WriteControlHandle::reconnect_retry(write, send, write_msg, write_recv);
                            let (send, recv) = oneshot::channel();
                            self.reader_crash = recv;
                            self.reader_handle = ReadControlHandle::reconnect(
                                read,
                                send,
                                db,
                                update_sender,
                                supervisor_command,
                                write_sender,
                                tasks,
                                pending_iqs,
                            );
                        },
                        Err(e) => {
                            // if reconnection failure, respond to all current write messages with lost connection error. the received processes should complete themselves.
                            write_recv.close();
                            let _ = write_msg.respond_to.send(Err(Reason::LostConnection));
                            while let Some(msg) = write_recv.recv().await {
                                let _ = msg.respond_to.send(Err(Reason::LostConnection));
                            }
                            // TODO: is this the correct error to send?
                            let _ = self.sender.send(UpdateMessage::Error(Error::LostConnection)).await;
                            break;
                        },
                    }
                },
                Ok((db, update_sender, tasks, supervisor_control, write_handle, pending_iqs)) = &mut self.reader_crash => {
                    let (send, recv) = oneshot::channel();
                    let _ = self.writer_handle.send(WriteControl::Abort(send)).await;
                    let (retry_msg, mut write_receiver) = tokio::select! {
                        Ok(s) = recv => (None, s),
                        Ok(s) = &mut self.writer_crash => (Some(s.0), s.1),
                        // in case, just break as irrecoverable
                        else => break,
                    };

                    let mut jid = self.jid.lock().await;
                    let mut domain = jid.domainpart.clone();
                    let connection = jabber::connect_and_login(&mut jid, &*self.password, &mut domain).await;
                    match connection {
                        Ok(c) => {
                            let (read, write) = c.split();
                            let (send, recv) = oneshot::channel();
                            self.writer_crash = recv;
                            if let Some(msg) = retry_msg {
                                self.writer_handle =
                                    WriteControlHandle::reconnect_retry(write, send, msg, write_receiver);
                            } else {
                                self.writer_handle = WriteControlHandle::reconnect(write, send, write_receiver)
                            }
                            let (send, recv) = oneshot::channel();
                            self.reader_crash = recv;
                            self.reader_handle = ReadControlHandle::reconnect(
                                read,
                                send,
                                db,
                                update_sender,
                                supervisor_control,
                                write_handle,
                                tasks,
                                pending_iqs,
                            );
                        },
                        Err(e) => {
                            // if reconnection failure, respond to all current messages with lost connection error.
                            write_receiver.close();
                            if let Some(msg) = retry_msg {
                                msg.respond_to.send(Err(Reason::LostConnection));
                            }
                            while let Some(msg) = write_receiver.recv().await {
                                msg.respond_to.send(Err(Reason::LostConnection));
                            }
                            // TODO: is this the correct error?
                            let _ = self.sender.send(UpdateMessage::Error(Error::LostConnection)).await;
                            break;
                        },
                    }
                },
                else => break,
            }
        }
        let _ = self.on_shutdown.send(());
    }
}

pub struct SupervisorHandle {
    sender: SupervisorSender,
    handle: JoinHandle<()>,
}

impl Deref for SupervisorHandle {
    type Target = SupervisorSender;

    fn deref(&self) -> &Self::Target {
        &self.sender
    }
}

impl DerefMut for SupervisorHandle {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.sender
    }
}

#[derive(Clone)]
pub struct SupervisorSender {
    sender: mpsc::Sender<SupervisorCommand>,
}

impl Deref for SupervisorSender {
    type Target = mpsc::Sender<SupervisorCommand>;

    fn deref(&self) -> &Self::Target {
        &self.sender
    }
}

impl DerefMut for SupervisorSender {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.sender
    }
}

impl SupervisorHandle {
    pub fn new(
        streams: BoundJabberStream<Tls>,
        update_sender: mpsc::Sender<UpdateMessage>,
        db: Db,
        on_shutdown: oneshot::Sender<()>,
        jid: Arc<Mutex<JID>>,
        password: Arc<String>,
        pending_iqs: Arc<Mutex<HashMap<String, oneshot::Sender<Result<Stanza, Reason>>>>>,
    ) -> (WriteHandle, Self) {
        let (command_sender, command_receiver) = mpsc::channel(20);
        let (writer_error_sender, writer_error_receiver) = oneshot::channel();
        let (reader_crash_sender, reader_crash_receiver) = oneshot::channel();

        let (reader, writer) = streams.split();
        let (write_handle, write_control_handle) =
            WriteControlHandle::new(writer, writer_error_sender);
        let jabber_reader_control_handle = ReadControlHandle::new(
            reader,
            reader_crash_sender,
            db,
            update_sender.clone(),
            command_sender.clone(),
            write_handle.clone(),
            pending_iqs,
        );

        let actor = Supervisor::new(
            command_receiver,
            writer_error_receiver,
            reader_crash_receiver,
            update_sender,
            write_control_handle,
            jabber_reader_control_handle,
            on_shutdown,
            jid,
            password,
        );

        let handle = tokio::spawn(async move { actor.run().await });

        (
            write_handle,
            Self {
                sender: SupervisorSender {
                    sender: command_sender,
                },
                handle,
            },
        )
    }

    pub fn sender(&self) -> SupervisorSender {
        self.sender.clone()
    }
}