1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

//! Async request/reply

use super::*;
use crate::ctx::NngCtx;

#[derive(Debug, PartialEq)]
enum ReplyState {
    Receiving,
    Wait,
    Sending,
}

#[derive(Debug)]
struct ReplyContextAioArg {
    aio: NngAio,
    ctx: NngCtx,
    state: ReplyState,
    request_sender: mpsc::Sender<Result<NngMsg>>,
    reply_sender: Option<oneshot::Sender<Result<()>>>,
}

impl ReplyContextAioArg {
    pub fn new(
        socket: NngSocket,
        request_sender: mpsc::Sender<Result<NngMsg>>,
    ) -> Result<AioArg<Self>> {
        let ctx = NngCtx::new(socket)?;
        NngAio::create(
            |aio| Self {
                aio,
                ctx,
                state: ReplyState::Receiving,
                request_sender,
                reply_sender: None,
            },
            reply_callback,
        )
    }

    fn start_receive(&mut self) {
        if self.state != ReplyState::Receiving && self.state != ReplyState::Sending {
            panic!();
        }
        self.state = ReplyState::Receiving;
        unsafe {
            nng_ctx_recv(self.ctx.ctx(), self.aio.nng_aio());
        }
    }

    pub fn reply(&mut self, msg: NngMsg, sender: oneshot::Sender<Result<()>>) {
        if self.state != ReplyState::Wait {
            panic!();
        }

        self.reply_sender = Some(sender);
        unsafe {
            let aio = self.aio.nng_aio();

            self.state = ReplyState::Sending;
            // Nng assumes ownership of the message
            nng_aio_set_msg(aio, msg.take());
            nng_ctx_send(self.ctx.ctx(), aio);
        }
    }
}

impl Aio for ReplyContextAioArg {
    fn aio(&self) -> &NngAio {
        &self.aio
    }
    fn aio_mut(&mut self) -> &mut NngAio {
        &mut self.aio
    }
}

/// Asynchronous context for reply socket.
#[derive(Debug)]
pub struct ReplyStreamHandle {
    aio_arg: AioArg<ReplyContextAioArg>,
    receiver: Option<mpsc::Receiver<Result<NngMsg>>>,
}

impl AsyncStreamContext for ReplyStreamHandle {
    fn new(socket: NngSocket, buffer: usize) -> Result<Self> {
        let (sender, receiver) = mpsc::channel(buffer);
        let aio_arg = ReplyContextAioArg::new(socket, sender)?;
        let receiver = Some(receiver);
        let ctx = Self { aio_arg, receiver };
        Ok(ctx)
    }
}

/// Trait for asynchronous contexts that can receive a request and then send a reply.
pub trait AsyncReply {
    /// Asynchronously receive a request.
    fn receive(&mut self) -> Option<mpsc::Receiver<Result<NngMsg>>>;
    /// Asynchronously reply to previously received request.
    fn reply(&mut self, msg: NngMsg) -> oneshot::Receiver<Result<()>>;
}

impl AsyncReply for ReplyStreamHandle {
    fn receive(&mut self) -> Option<mpsc::Receiver<Result<NngMsg>>> {
        let receiver = self.receiver.take();
        if receiver.is_some() {
            self.aio_arg.start_receive();
        }
        receiver
    }

    fn reply(&mut self, msg: NngMsg) -> oneshot::Receiver<Result<()>> {
        let (sender, receiver) = oneshot::channel();
        self.aio_arg.reply(msg, sender);
        receiver
    }
}

unsafe extern "C" fn reply_callback(arg: AioArgPtr) {
    let ctx = &mut *(arg as *mut ReplyContextAioArg);
    let aio_nng = ctx.aio.nng_aio();
    trace!("reply_callback::{:?}", ctx.state);
    match ctx.state {
        ReplyState::Receiving => {
            let res = nng_int_to_result(nng_aio_result(aio_nng));
            match res {
                Err(res) => {
                    match res {
                        Error::Errno(NngErrno::ECLOSED) | Error::Errno(NngErrno::ECANCELED) => {
                            debug!("reply_callback {:?}", res);
                        }
                        _ => {
                            trace!("reply_callback::Err({:?})", res);
                            ctx.start_receive();
                        }
                    }

                    try_signal_complete(&mut ctx.request_sender, Err(res));
                }
                Ok(()) => {
                    let msg = NngMsg::from_raw(nng_aio_get_msg(aio_nng));
                    // Reset state before signaling completion
                    ctx.state = ReplyState::Wait;
                    try_signal_complete(&mut ctx.request_sender, Ok(msg));
                }
            }
        }
        ReplyState::Wait => panic!(),
        ReplyState::Sending => {
            let res = nng_int_to_result(nng_aio_result(aio_nng));
            if res.is_err() {
                // Nng requires we resume ownership of the message
                let _ = NngMsg::from_raw(nng_aio_get_msg(aio_nng));
            }

            let sender = ctx.reply_sender.take().unwrap();
            // Reset state and start receiving again before
            // signaling completion to avoid race condition where we say we're done, but
            // not yet ready for receive() to be called.
            ctx.start_receive();
            sender.send(res).unwrap();
        }
    }
}