dandri/mycelium
1
0
Fork 0
mirror of https://github.com/threefoldtech/mycelium.git synced 2026-07-28 14:29:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

Allow long poll when reading new messages

Browse Source 
Also unify the peek and pop handlers

Signed-off-by: Lee Smet <lee.smet@hotmail.com>
This commit is contained in:
Lee Smet
2023-10-27 15:45:51 +02:00
parent def11c5a6e
commit d0c8a813d5
5 changed files with 133 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+22
View File
@@ -168,6 +168,7 @@ dependencies = [
"serde",
"serde_json",
"serde_path_to_error",
"serde_urlencoded",
"sync_wrapper",
"tokio",
"tower",
@@ -494,6 +495,15 @@ version = "1.0.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
[[package]]
name = "form_urlencoded"
version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a62bc1cf6f830c2ec14a513a9fb124d0a213a629668a4186f329db21fe045652"
dependencies = [
"percent-encoding",
]
[[package]]
name = "futures"
version = "0.3.28"
@@ -1444,6 +1454,18 @@ dependencies = [
"serde",
]
[[package]]
name = "serde_urlencoded"
version = "0.7.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd"
dependencies = [
"form_urlencoded",
"itoa",
"ryu",
"serde",
]
[[package]]
name = "sharded-slab"
version = "0.1.7"
Cargo.toml
+1
View File
@@ -37,6 +37,7 @@ axum = { version = "0.6.20", default-features = false, features = [
"http1",
"http2",
"json",
"query",
"tokio",
] }
base64 = "0.21.5"
docs/api.yaml
+25 -21
View File
@@ -29,10 +29,33 @@ paths:
get:
tags:
- Message
summary: Pop a message from the inbound message queue
summary: Get a message from the inbound message queue
description: |
Pop a message from the inbound message queue. The message is removed from the queue and won't be shown again
Get a message from the inbound message queue. By default, the message is removed from the queue and won't be shown again.
If the peek query parameter is set to true, the message will be peeked, and the next call to this endpoint will show the same message.
This method returns immediately by default: a message is returned if one is ready, and if there isn't nothing is returned. If the timeout
query parameter is set, this call won't return for the given amount of seconds, unless a message is received
operationId: popMessage
parameters:
- in: query
name: peek
required: false
schema:
type: boolean
description: Whether to peek the message or not. If this is true, the message won't be removed from the inbound queue when it is read
example: true
- in: query
name: timeout
required: false
schema:
type: integer
format: int64
minimum: 0
description: |
Amount of seconds to wait for a message to arrive if one is not available. Setting this to 0 is valid and will return
a message if present, or return immediately if there isn't
example: 60
responses:
'200':
description: Message retrieved
@@ -63,25 +86,6 @@ paths:
schema:
$ref: '#/components/schemas/PushMessageResponse'
'/api/v1/messages/peek':
get:
tags:
- Message
summary: Peek a message from the inbound queue
description: |
Peek a message from the inbound message queue. The message is removed from the queue and will be shown
again on the next call again to peek or pop
operationId: peekMessage
responses:
'200':
description: Message peeked successfully
content:
application/json:
schema:
$ref: '#/components/schemas/InboundMessage'
'204':
description: No message ready
'/api/v1/messages/status/:id':
get:
tags:
src/api.rs
+43 -36
View File
@@ -4,7 +4,7 @@ use std::{
};
use axum::{
extract::{Path, State},
extract::{Path, Query, State},
http::StatusCode,
routing::get,
Json, Router,
@@ -74,8 +74,7 @@ impl Http {
pub fn spawn(message_stack: MessageStack, listen_addr: &SocketAddr) -> Self {
let server_state = HttpServerState { message_stack };
let msg_routes = Router::new()
.route("/messages", get(pop_message).post(push_message))
.route("/messages/peek", get(peek_message))
.route("/messages", get(get_message).post(push_message))
.route("/messages/status/:id", get(message_status))
.with_state(server_state);
let app = Router::new().nest("/api/v1", msg_routes);
@@ -95,44 +94,52 @@ impl Http {
}
}
async fn peek_message(
State(state): State<HttpServerState>,
) -> Result<Json<MessageReceiveInfo>, StatusCode> {
debug!("Attempt to peek message");
state
.message_stack
.peek_message()
.ok_or(StatusCode::NO_CONTENT)
.map(|m| {
Json(MessageReceiveInfo {
id: m.id,
src_ip: m.src_ip,
src_pk: m.src_pk,
dst_ip: m.dst_ip,
dst_pk: m.dst_pk,
payload: m.data,
})
})
#[derive(Deserialize)]
struct GetMessageQuery {
peek: Option<bool>,
timeout: Option<u64>,
}
async fn pop_message(
impl GetMessageQuery {
/// Did the query indicate we should peek the message instead of pop?
fn peek(&self) -> bool {
matches!(self.peek, Some(true))
}
/// Amount of seconds to hold and try and get values.
fn timeout_secs(&self) -> u64 {
self.timeout.unwrap_or(0)
}
}
async fn get_message(
State(state): State<HttpServerState>,
Query(query): Query<GetMessageQuery>,
) -> Result<Json<MessageReceiveInfo>, StatusCode> {
debug!("Attempt to pop message");
state
.message_stack
.pop_message()
.ok_or(StatusCode::NO_CONTENT)
.map(|m| {
Json(MessageReceiveInfo {
id: m.id,
src_ip: m.src_ip,
src_pk: m.src_pk,
dst_ip: m.dst_ip,
dst_pk: m.dst_pk,
payload: m.data,
})
debug!(
"Attempt to get message, peek {}, timeout {} seconds",
query.peek(),
query.timeout_secs()
);
// A timeout of 0 seconds essentially means get a message if there is one, and return
// immediatly if there isn't. This is the result of the implementation of Timeout, which does a
// poll of the internal future first, before polling the delay.
tokio::time::timeout(
Duration::from_secs(query.timeout_secs()),
state.message_stack.message(!query.peek()),
)
.await
.or(Err(StatusCode::NO_CONTENT))
.map(|m| {
Json(MessageReceiveInfo {
id: m.id,
src_ip: m.src_ip,
src_pk: m.src_pk,
dst_ip: m.dst_ip,
dst_pk: m.dst_pk,
payload: m.data,
})
})
}
#[derive(Serialize)]
src/message.rs
+42 -32
View File
@@ -19,6 +19,7 @@ use futures::{Stream, StreamExt};
use log::{debug, error, trace, warn};
use rand::Fill;
use serde::{de::Visitor, Deserialize, Deserializer, Serialize};
use tokio::sync::watch;
use crate::{
crypto::{PacketBuffer, PublicKey},
@@ -79,6 +80,8 @@ pub struct MessageStack {
data_plane: Arc<Mutex<DataPlane>>,
inbox: Arc<Mutex<MessageInbox>>,
outbox: Arc<Mutex<MessageOutbox>>,
/// Receiver handle for inbox listeners (basically a condvar).
subscriber: watch::Receiver<()>,
}
struct MessageOutbox {
@@ -91,6 +94,8 @@ struct MessageInbox {
pending_msges: HashMap<MessageId, ReceivedMessageInfo>,
/// Messages which have been completed.
complete_msges: VecDeque<ReceivedMessage>,
/// Notification sender used to allert subscribed listeners.
notify: watch::Sender<()>,
}
struct ReceivedMessageInfo {
@@ -163,10 +168,11 @@ enum TransmissionState {
}
impl MessageInbox {
fn new() -> Self {
fn new(notify: watch::Sender<()>) -> Self {
Self {
pending_msges: HashMap::new(),
complete_msges: VecDeque::new(),
notify,
}
}
}
@@ -193,10 +199,12 @@ impl MessageStack {
where
S: Stream<Item = (PacketBuffer, IpAddr, IpAddr)> + Send + Unpin + 'static,
{
let (notify, subscriber) = watch::channel(());
let ms = Self {
data_plane: Arc::new(Mutex::new(data_plane)),
inbox: Arc::new(Mutex::new(MessageInbox::new())),
inbox: Arc::new(Mutex::new(MessageInbox::new(notify))),
outbox: Arc::new(Mutex::new(MessageOutbox::new())),
subscriber,
};
tokio::task::spawn(
@@ -341,7 +349,7 @@ impl MessageStack {
// case we consider this to be a lingering chunk which was already accepted in the
// meantime (as the message is complete).
//
// SAFETY: a malcious node could send a lot of empty chunks, which trigger allocations
// SAFETY: a malicious node could send a lot of empty chunks, which trigger allocations
// to hold the chunk array, effectively exhausting memory. As such, we first need to
// determine if the chunk is feasible.
let mut inbox = self.inbox.lock().unwrap();
@@ -460,6 +468,8 @@ impl MessageStack {
// Move message to be read.
inbox.complete_msges.push_back(message);
inbox.pending_msges.remove(&message_id);
// Notify subscribers we have a new message.
inbox.notify.send_replace(());
Some(md.into_reply().into_inner())
} else {
@@ -774,35 +784,6 @@ impl MessageStack {
id
}
/// Peek a [`Message`] from the inbound queue.
///
/// The next call to [`MessageStack::peek_message`] or [`MessageStack::pop_message`] will
/// return the same message.
pub fn peek_message(&self) -> Option<ReceivedMessage> {
let inbox = self.inbox.lock().unwrap();
let maybe_msg = inbox.complete_msges.front().map(Clone::clone);
if let Some(ref msg) = maybe_msg {
// If a message is popped, notify the sender
self.notify_read(msg);
}
maybe_msg
}
/// Read a [`Message`] from the inbound queue and delete it.
pub fn pop_message(&self) -> Option<ReceivedMessage> {
let mut inbox = self.inbox.lock().unwrap();
let maybe_msg = inbox.complete_msges.pop_front();
if let Some(ref msg) = maybe_msg {
// If a message is popped, notify the sender
self.notify_read(msg);
}
maybe_msg
}
/// Get information about the status of an outbound message.
pub fn message_info(&self, id: MessageId) -> Option<MessageInfo> {
let outbox = self.outbox.lock().unwrap();
@@ -846,6 +827,35 @@ impl MessageStack {
})
}
/// A future which eventually resolves to a new (inbound message)[`ReceivedMessage`], if new messages come in.
///
/// If pop is false, the message is not removed and the next call of this method will return
/// the same message.
pub async fn message(&self, pop: bool) -> ReceivedMessage {
// Copy the subscriber since we need mutable access to it.
let mut subscriber = self.subscriber.clone();
loop {
// Scope to ensure we drop the lock after we checked for a message and don't hold
// it while waiting for a new notification.
{
let mut inbox = self.inbox.lock().unwrap();
if let Some(msg) = if pop {
inbox.complete_msges.pop_front()
} else {
inbox.complete_msges.front().map(Clone::clone)
} {
self.notify_read(&msg);
return msg;
};
}
// Sender can never be dropped since we hold a reference to self which contains the
// inbox.
let _ = subscriber.changed().await;
}
}
/// Notify the sender of a message that it has been read.
fn notify_read(&self, msg: &ReceivedMessage) {
let mut mp = MessagePacket::new(PacketBuffer::new());