gtether 0.2.1

//! Network messaging utilities.
//!
//! This module contains some utilities for handling network messages, including processing and
//! dispatching of said messages.
//!
//! # Message composition
//!
//! A [Message] is composed of a [header](MessageHeader) and a [body](MessageBody). The header is
//! an internally managed type, and contains information common to all messages, while the body is
//! a user-implemented trait that describes unique message types. The body can be further described
//! with additional traits as follows:
//!  * [MessageSend] - These messages can be sent from a networking stack, and the trait describes
//!    how the message can be encoded into raw byte data.
//!  * [MessageRecv] - These messages can be received from a networking stack, and the trait
//!    describes how the message can be decoded from raw byte data.
//!  * [MessageRepliable] - These messages have a [reply type](MessageRepliable::Reply) associated
//!    with them.
//!
//! Generally, if a message body type implements [MessageSend], it usually also implements
//! [MessageRecv] and vice versa, but if the type differs between sending and receiving, only one
//! can be implemented, as long as the encoded/decoded byte data is compatible.
//!
//! # Handling messages
//!
//! Messages are handled on receipt by user-implemented [MessageHandlers](MessageHandler). These
//! handlers are registered with the [Networking](super::Networking) stack via
//! [Networking::insert_msg_handler()](super::Networking::insert_msg_handler).
//!
//! When handling a message, the handler may execute on the networking stack's own thread, so care
//! should be taken to avoid expensive computation directly in the handler. Instead, such
//! computation should be delegated to another thread, for example by using a
//! [MessageQueue](queue::MessageQueue).

use bitcode::{Decode, DecodeOwned, Encode};
use flagset::{flags, FlagSet};
use parking_lot::{Condvar, Mutex};
use std::convert::Infallible;
use std::error::Error;
use std::fmt::Debug;
use std::future::Future;
use std::num::NonZeroU64;
use std::pin::Pin;
use std::sync::{Arc, Weak};
use std::task::{Context, Poll, Waker};

use crate::net::{Connection, NetworkingError};

pub(super) mod dispatch;
pub mod queue;

#[cfg(feature = "derive")]
pub use gtether_derive::MessageBody;

/// Error that can occur when encoding messages.
#[derive(Debug, thiserror::Error)]
#[error("Could not serialize message data; {details}")]
pub struct MessageEncodeError {
    details: String,
    #[source]
    source: Box<dyn Error + Send + Sync>,
}

impl From<MessageEncodeError> for NetworkingError {
    #[inline]
    fn from(value: MessageEncodeError) -> Self {
        Self::MalformedMessage(Box::new(value))
    }
}

/// Error that can occur when decoding messages.
#[derive(Debug, thiserror::Error)]
#[error("Could not deserialize message data; {details}")]
pub struct MessageDecodeError {
    details: String,
    #[source]
    source: Option<Box<dyn Error + Send + Sync>>,
}

impl From<MessageDecodeError> for NetworkingError {
    #[inline]
    fn from(value: MessageDecodeError) -> Self {
        Self::MalformedMessage(Box::new(value))
    }
}

flags! {
    /// Flags that can change how [messages](Message) behave.
    ///
    /// Flags exist on a per-type basis, and are defined by [MessageBody::flags()].
    pub enum MessageFlags: u16 {
        /// The message should be sent/received reliably.
        Reliable,
    }
}

/// Header data for a [Message].
///
/// This type is managed internally, but can be accessed from a message in order to retrieve some
/// common data about the message.
#[derive(Encode, Decode, Debug, Clone, PartialEq, Eq)]
pub struct MessageHeader {
    key: String,
    msg_num: Option<NonZeroU64>,
    reply_num: Option<NonZeroU64>,
}

impl MessageHeader {
    /// The message key.
    ///
    /// This key is determined on a per-type basis, and is defined by [MessageBody::KEY].
    #[inline]
    pub fn key(&self) -> &str {
        &self.key
    }

    /// The message number.
    ///
    /// The number is generated by the networking stack, but is generally unique for each stack and
    /// message combination.
    ///
    /// If the networking stack has not assigned a number to the message yet (i.e. it has been
    /// created but not sent), this may yield `None`.
    #[inline]
    pub fn msg_num(&self) -> Option<NonZeroU64> {
        self.msg_num
    }

    /// The number of the message that this message is replying to.
    ///
    /// This number is optional, and may be `None`, but if it exists it is used to tie a message to
    /// the message it is responding to.
    #[inline]
    pub fn reply_num(&self) -> Option<NonZeroU64> {
        self.reply_num
    }

    pub(in crate::net) fn encode(&self) -> impl Iterator<Item=u8> {
        let bytes = bitcode::encode(self);
        let len = bytes.len() as u16;
        len.to_be_bytes().into_iter()
            .chain(bytes)
    }

    pub(in crate::net) fn decode(bytes: &[u8]) -> Result<(Self, &[u8]), MessageDecodeError> {
        let header_len = if bytes.len() >= 2 {
            u16::from_be_bytes(bytes[0..2].try_into()
                .map_err(|err| MessageDecodeError {
                    details: "Could not deserialize message header length".to_owned(),
                    source: Some(Box::new(err)),
                })?)
        } else {
            return Err(MessageDecodeError {
                details: format!("Not enough bytes to parse header length ({} < 2)", bytes.len()),
                source: None,
            });
        };
        // Include 2 for the original 2 bytes
        let full_len = header_len as usize + 2;
        if bytes.len() >= full_len {
            let header = bitcode::decode(&bytes[2..full_len])
                .map_err(|err| MessageDecodeError {
                    details: "Could not deserialize message header".to_owned(),
                    source: Some(Box::new(err)),
                })?;
            Ok((header, &bytes[full_len..]))
        } else {
            Err(MessageDecodeError {
                details: format!("Not enough bytes to parse header ({} < {})", bytes.len(), full_len),
                source: None,
            })
        }
    }
}

/// Body data for a [Message].
///
/// This type is user-defined.
///
/// See [module-level documentation](super::message#message-composition) for more.
///
/// # Implementation
///
/// The easiest way to implement MessageBody is with the derive macro:
/// ```
/// use gtether::net::message::MessageBody;
///
/// #[derive(MessageBody)]
/// struct MyMessage {}
///
/// #[derive(MessageBody)]
/// #[message_flag(Reliable)]
/// struct MyReliableMessage {}
/// ```
///
/// You can however implement MessageBody manually if needed:
/// ```
/// use gtether::net::message::{MessageBody, MessageFlags};
/// use gtether::util::FlagSet;
///
/// struct MyMessage {}
/// impl MessageBody for MyMessage {
///     const KEY: &'static str = "MyMessage";
/// }
///
/// struct MyReliableMessage {}
/// impl MessageBody for MyReliableMessage {
///     const KEY: &'static str = "MyReliableMessage";
///     fn flags() -> FlagSet<MessageFlags> {
///         MessageFlags::Reliable.into()
///     }
/// }
/// ```
pub trait MessageBody {
    /// The key identifying this message type.
    ///
    /// This is used to match to a [handler](super::message#handling-messages) when dispatching
    /// messages.
    const KEY: &'static str;

    /// Any additional flags describing how a message should be treated.
    ///
    /// By default, yields the empty set of flags.
    fn flags() -> FlagSet<MessageFlags> {
        FlagSet::default()
    }
}

/// Trait describing a [Message] that can be sent.
///
/// This trait is used both as a marker and to define how a [Message] can be
/// [encoded](MessageSend::encode).
///
/// Note that for convenience this trait is automatically implemented for any type that implements
/// both [MessageBody] and bitcode's `Encode` trait.
pub trait MessageSend: MessageBody {
    /// Error type that may occur while encoding.
    type EncodeError: Error + Send + Sync + 'static;

    /// Encode this message body in raw byte data.
    ///
    /// The encode method is expected to allocate its own buffer to fill and return.
    fn encode(&self) -> Result<Vec<u8>, Self::EncodeError>;
}

impl<M: MessageBody + Encode> MessageSend for M {
    type EncodeError = Infallible;

    #[inline]
    fn encode(&self) -> Result<Vec<u8>, Self::EncodeError> {
        Ok(bitcode::encode(self))
    }
}

/// Trait describing a [Message] that can be received.
///
/// This trait is used both as a marker and to define how a [Message] can be
/// [decoded](MessageRecv::decode).
///
/// Note that for convenience this trait is automatically implemented for any type that implements
/// both [MessageBody] and bitcode's `DecodeOwned` trait.
pub trait MessageRecv: MessageBody {
    /// Error tyep that may occur while decoding.
    type DecodeError: Error + Send + Sync + 'static;

    /// Decode raw byte data into a new message body.
    fn decode(bytes: &[u8]) -> Result<Self, Self::DecodeError>
    where
        Self: Sized;
}

impl<M: MessageBody + DecodeOwned> MessageRecv for M {
    type DecodeError = bitcode::Error;

    #[inline]
    fn decode(bytes: &[u8]) -> Result<Self, Self::DecodeError>
    where
        Self: Sized,
    {
        Ok(bitcode::decode(bytes)?)
    }
}

/// Trait describing a [Message] that can be replied to.
///
/// # Implementation
///
/// The easiest way to implement MessageRepliable is with the [MessageBody] derive macro:
/// ```
/// use gtether::net::message::MessageBody;
///
/// #[derive(MessageBody)]
/// #[message_reply(MyMessageReply)]
/// struct MyMessage {}
///
/// #[derive(MessageBody)]
/// struct MyMessageReply {}
/// ```
///
/// You can however implement MessageBody manually if needed:
/// ```
/// use gtether::net::message::{MessageBody, MessageRepliable};
///
/// #[derive(MessageBody)]
/// struct MyMessage {}
/// impl MessageRepliable for MyMessage {
///     type Reply = MyMessageReply;
/// }
///
/// #[derive(MessageBody)]
/// struct MyMessageReply {}
/// ```
pub trait MessageRepliable: MessageBody {
    /// The message type of the reply.
    type Reply: MessageBody;
}

/// A full message, including both a [header](MessageHeader) and a [body](MessageBody).
#[derive(Debug, PartialEq, Eq)]
pub struct Message<M: MessageBody> {
    header: MessageHeader,
    body: M,
}

impl<M: MessageBody> Message<M> {
    /// Create a new [Message] from a given `body`.
    ///
    /// This new message will not have a number assigned to it until it is given to a networking
    /// stack.
    ///
    /// Note that you can also convert a [MessageBody] into a [Message] using `into()`/`from()`.
    #[inline]
    pub fn new(body: M) -> Self {
        let header = MessageHeader {
            key: M::KEY.to_owned(),
            msg_num: None,
            reply_num: None,
        };

        Self {
            header,
            body,
        }
    }

    #[inline]
    pub(super) fn with_header(header: MessageHeader, body: M) -> Self {
        Self {
            header,
            body,
        }
    }

    /// Get a reference to the [header](MessageHeader).
    #[inline]
    pub fn header(&self) -> &MessageHeader {
        &self.header
    }

    /// Get a reference to the [body](MessageBody).
    #[inline]
    pub fn body(&self) -> &M {
        &self.body
    }

    /// Unwrap into the [body](MessageBody).
    #[inline]
    pub fn into_body(self) -> M {
        self.body
    }

    pub(super) fn set_msg_num(&mut self, msg_num: NonZeroU64) {
        self.header.msg_num = Some(msg_num);
    }
}

impl<M: MessageSend> Message<M> {
    /// Encode the entirety of this message into raw bytes, including the header.
    ///
    /// Will allocate and yield a new buffer.
    pub fn encode(&self) -> Result<Vec<u8>, MessageEncodeError> {
        let body = self.body.encode()
            .map_err(|err| MessageEncodeError {
                details: "Could not serialize message body".to_owned(),
                source: Box::new(err)
            })?;

        Ok(self.header.encode()
            .chain(body)
            .collect::<Vec<_>>())
    }
}

impl<M: MessageRepliable> Message<M> {
    /// Create a new message as a reply to this message.
    ///
    /// The new message will have its reply number set to this message's number.
    ///
    /// # Examples
    /// ```
    /// use std::convert::Infallible;
    /// use gtether::net::message::{Message, MessageBody, MessageRepliable};
    ///
    /// struct MyMessage {}
    /// struct MyMessageReply {}
    ///
    /// impl MessageBody for MyMessage {
    ///     const KEY: &'static str = "my-message";
    /// }
    ///
    /// impl MessageBody for MyMessageReply {
    ///     const KEY: &'static str = "my-message-reply";
    /// }
    ///
    /// impl MessageRepliable for MyMessage {
    ///     type Reply = MyMessageReply;
    /// }
    ///
    /// fn msg_handler(msg: Message<MyMessage>) -> Result<(), Infallible> {
    ///     let reply = msg.reply(MyMessageReply {});
    ///     // Send the reply...
    ///     Ok(())
    /// }
    /// ```
    #[inline]
    pub fn reply(&self, body: M::Reply) -> Message<M::Reply> {
        let header = MessageHeader {
            key: M::Reply::KEY.to_owned(),
            msg_num: None,
            reply_num: self.header.msg_num,
        };

        Message {
            header,
            body,
        }
    }
}

impl<M: MessageBody> From<M> for Message<M> {
    #[inline]
    fn from(value: M) -> Self {
        Message::new(value)
    }
}

pub(in crate::net) struct MessageReplyContext<M: MessageRecv> {
    value: Mutex<Option<Message<M>>>,
    waker: Mutex<Option<Waker>>,
    cvar: Condvar,
}

impl<M: MessageRecv> MessageReplyContext<M> {
    fn new() -> Self {
        Self {
            value: Mutex::new(None),
            waker: Mutex::new(None),
            cvar: Condvar::new(),
        }
    }

    fn accept(&self, msg: Message<M>) {
        let mut value = self.value.lock();
        *value = Some(msg);
        if let Some(waker) = self.waker.lock().take() {
            waker.wake();
        }
        self.cvar.notify_one();
    }
}

/// A future indicating a reply that is yet to be received.
///
/// Yielded when sending a message and expecting a reply back. In addition to implementing the std's
/// [Future], can also be waited synchronously via [MessageReplyFuture::wait()].
pub struct MessageReplyFuture<M: MessageRecv> {
    ctx: Arc<MessageReplyContext<M>>,
}

impl<M: MessageRecv> MessageReplyFuture<M> {
    /// Block and synchronously wait for a reply.
    pub fn wait(self) -> Message<M> {
        let mut msg = self.ctx.value.lock();
        if let Some(msg) = msg.take() {
            msg
        } else {
            self.ctx.cvar.wait(&mut msg);
            msg.take().expect("'msg' should be set after cvar wakes")
        }
    }
}

impl<M: MessageRecv> Future for MessageReplyFuture<M> {
    type Output = Message<M>;

    fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        if let Some(msg) = self.ctx.value.lock().take() {
            Poll::Ready(msg)
        } else {
            let mut waker = self.ctx.waker.lock();
            if let Some(waker) = waker.as_mut() {
                waker.clone_from(ctx.waker());
            } else {
                *waker = Some(ctx.waker().clone());
            }
            Poll::Pending
        }
    }
}


/// Errors that can occur when dispatching messages.
#[derive(Debug, thiserror::Error)]
pub enum MessageDispatchError {
    #[error(transparent)]
    DecodeError(#[from] MessageDecodeError),
    #[error("No handler was identified for key: '{key}'")]
    NoHandler {
        /// `key` that is missing a handler.
        key: String,
    },
    #[error("Could not dispatch message: {0}")]
    DispatchFailed(#[source] Box<dyn Error + Send + Sync>),
}

/// Interface for handling received messages.
///
/// Handles a single [Message], and optionally returns an error. The [Connection] that the message
/// was sent from is also given, to distinguish which connection it came from in networking stacks
/// that have multiple connections.
///
/// Message handlers may be executed from networking [driver](super::driver::NetDriver) internal
/// threads. Because of this, any work done in the handler should be light, as it has the potential
/// to block further networking processing. If heavy work needs to be done, it should be delegated
/// to a separate thread. For example, a [MessageQueue](queue::MessageQueue) can be used to put all
/// messages in a queue, where they can be polled and processed from another thread.
pub trait MessageHandler<M, E: Error>: Send + Sync + 'static
where
    M: MessageBody,
    E: Error,
{
    /// Process and handle a [Message].
    ///
    /// The [Connection] represents which connection sent the message.
    fn handle(
        &self,
        connection: Connection,
        msg: Message<M>,
    ) -> Result<(), E>;

    /// Whether this handler is still valid.
    ///
    /// If it is not valid, it will be removed and no longer used to handle messages. This is
    /// primarily used with e.g. handlers implemented for weak references. Note that MessageHandler
    /// is automatically implemented for any `Weak<H: MessageHandler>`, where this method will
    /// return `false` if the weak reference is no longer valid.
    #[inline]
    fn is_valid(&self) -> bool { true }
}

impl<M, E, F> MessageHandler<M, E> for F
where
    M: MessageBody,
    E: Error,
    F: (Fn(Connection, Message<M>) -> Result<(), E>) + Send + Sync + 'static,
{
    #[inline]
    fn handle(
        &self,
        connection: Connection,
        msg: Message<M>,
    ) -> Result<(), E> {
        self(connection, msg)
    }
}

impl<M, E, H> MessageHandler<M, E> for Arc<H>
where
    M: MessageBody,
    E: Error,
    H: MessageHandler<M, E>,
{
    #[inline]
    fn handle(&self, connection: Connection, msg: Message<M>) -> Result<(), E> {
        (**self).handle(connection, msg)
    }
}

impl<M, E, H> MessageHandler<M, E> for Weak<H>
where
    M: MessageBody,
    E: Error,
    H: MessageHandler<M, E>,
{
    #[inline]
    fn handle(&self, connection: Connection, msg: Message<M>) -> Result<(), E> {
        if let Some(handler) = self.upgrade() {
            handler.handle(connection, msg)
        } else {
            Ok(())
        }
    }

    #[inline]
    fn is_valid(&self) -> bool {
        if let Some(handler) = self.upgrade() {
            handler.is_valid()
        } else {
            false
        }
    }
}