From e1161cab0e71ec604e376d2d87f7d1226f3f0244 Mon Sep 17 00:00:00 2001 From: Quentin Dufour Date: Wed, 17 Jan 2024 16:56:05 +0100 Subject: [PATCH] idle sync --- src/bayou.rs | 36 ++++++++++++++++++--- src/imap/command/selected.rs | 2 +- src/imap/flow.rs | 4 +-- src/imap/mod.rs | 62 +++++++++++++++++++++++++++++------- src/imap/request.rs | 1 - src/imap/response.rs | 4 ++- src/imap/session.rs | 14 ++++++-- src/mail/mailbox.rs | 10 ++++++ 8 files changed, 110 insertions(+), 23 deletions(-) diff --git a/src/bayou.rs b/src/bayou.rs index d77e9dc..14f9728 100644 --- a/src/bayou.rs +++ b/src/bayou.rs @@ -238,6 +238,22 @@ impl Bayou { Ok(()) } + pub async fn idle_sync(&mut self) -> Result<()> { + tracing::debug!("start idle_sync"); + loop { + tracing::trace!("idle_sync loop"); + let fut_notif = self.watch.learnt_remote_update.notified(); + + if self.last_sync_watch_ct != *self.watch.rx.borrow() { + break + } + fut_notif.await; + } + tracing::trace!("idle_sync done"); + self.sync().await?; + Ok(()) + } + /// Applies a new operation on the state. Once this function returns, /// the operation has been safely persisted to storage backend. /// Make sure to call `.opportunistic_sync()` before doing this, @@ -257,7 +273,7 @@ impl Bayou { seal_serialize(&op, &self.key)?, ); self.storage.row_insert(vec![row_val]).await?; - self.watch.notify.notify_one(); + self.watch.propagate_local_update.notify_one(); let new_state = self.state().apply(&op); self.history.push((ts, op, Some(new_state))); @@ -423,7 +439,8 @@ impl Bayou { struct K2vWatch { target: storage::RowRef, rx: watch::Receiver, - notify: Notify, + propagate_local_update: Notify, + learnt_remote_update: Notify, } impl K2vWatch { @@ -434,9 +451,10 @@ impl K2vWatch { let storage = creds.storage.build().await?; let (tx, rx) = watch::channel::(target.clone()); - let notify = Notify::new(); + let propagate_local_update = Notify::new(); + let learnt_remote_update = Notify::new(); - let watch = Arc::new(K2vWatch { target, rx, notify }); + let watch = Arc::new(K2vWatch { target, rx, propagate_local_update, learnt_remote_update }); tokio::spawn(Self::background_task(Arc::downgrade(&watch), storage, tx)); @@ -459,7 +477,12 @@ impl K2vWatch { this.target.uid.shard, this.target.uid.sort ); tokio::select!( + // Needed to exit: will force a loop iteration every minutes, + // that will stop the loop if other Arc references have been dropped + // and free resources. Otherwise we would be blocked waiting forever... _ = tokio::time::sleep(Duration::from_secs(60)) => continue, + + // Watch if another instance has modified the log update = storage.row_poll(&row) => { match update { Err(e) => { @@ -471,10 +494,13 @@ impl K2vWatch { if tx.send(row.clone()).is_err() { break; } + this.learnt_remote_update.notify_waiters(); } } } - _ = this.notify.notified() => { + + // It appears we have modified the log, informing other people + _ = this.propagate_local_update.notified() => { let rand = u128::to_be_bytes(thread_rng().gen()).to_vec(); let row_val = storage::RowVal::new(row.clone(), rand); if let Err(e) = storage.row_insert(vec![row_val]).await diff --git a/src/imap/command/selected.rs b/src/imap/command/selected.rs index b62e2cb..4eb4e61 100644 --- a/src/imap/command/selected.rs +++ b/src/imap/command/selected.rs @@ -82,7 +82,7 @@ pub async fn dispatch<'a>( // IDLE extension (rfc2177) CommandBody::Idle => { Ok(( - Response::build().to_req(ctx.req).message("DUMMY response due to anti-pattern").ok()?, + Response::build().to_req(ctx.req).message("DUMMY command due to anti-pattern in the code").ok()?, flow::Transition::Idle(tokio::sync::Notify::new()), )) } diff --git a/src/imap/flow.rs b/src/imap/flow.rs index d1e27d4..37f225b 100644 --- a/src/imap/flow.rs +++ b/src/imap/flow.rs @@ -21,7 +21,7 @@ pub enum State { NotAuthenticated, Authenticated(Arc), Selected(Arc, MailboxView, MailboxPerm), - Idle(Arc, MailboxView, MailboxPerm, Notify), + Idle(Arc, MailboxView, MailboxPerm, Arc), Logout, } @@ -56,7 +56,7 @@ impl State { State::Authenticated(u.clone()) } (State::Selected(u, m, p), Transition::Idle(s)) => { - State::Idle(u, m, p, s) + State::Idle(u, m, p, Arc::new(s)) }, (State::Idle(u, m, p, _), Transition::UnIdle) => { State::Selected(u, m, p) diff --git a/src/imap/mod.rs b/src/imap/mod.rs index edfbbc4..c50c3fc 100644 --- a/src/imap/mod.rs +++ b/src/imap/mod.rs @@ -97,11 +97,14 @@ impl Server { } } -use tokio::sync::mpsc::*; +use tokio::sync::mpsc::*; +use tokio_util::bytes::BytesMut; +use tokio::sync::Notify; +use std::sync::Arc; enum LoopMode { Quit, Interactive, - Idle, + Idle(BytesMut, Arc), } // @FIXME a full refactor of this part of the code will be needed sooner or later @@ -190,7 +193,7 @@ impl NetLoop { loop { mode = match mode { LoopMode::Interactive => self.interactive_mode().await?, - LoopMode::Idle => self.idle_mode().await?, + LoopMode::Idle(buff, stop) => self.idle_mode(buff, stop).await?, LoopMode::Quit => break, } } @@ -238,11 +241,11 @@ impl NetLoop { } self.server.enqueue_status(response.completion); }, - Some(ResponseOrIdle::StartIdle) => { + Some(ResponseOrIdle::StartIdle(stop)) => { let cr = CommandContinuationRequest::basic(None, "Idling")?; self.server.enqueue_continuation(cr); self.cmd_tx.try_send(Request::Idle)?; - return Ok(LoopMode::Idle) + return Ok(LoopMode::Idle(BytesMut::new(), stop)) }, None => { self.server.enqueue_status(Status::bye(None, "Internal session exited").unwrap()); @@ -260,9 +263,19 @@ impl NetLoop { Ok(LoopMode::Interactive) } - async fn idle_mode(&mut self) -> Result { + async fn idle_mode(&mut self, mut buff: BytesMut, stop: Arc) -> Result { + // Flush send + loop { + match self.server.progress_send().await? { + Some(..) => continue, + None => break, + } + } + tokio::select! { + // Receiving IDLE event from background maybe_msg = self.resp_rx.recv() => match maybe_msg { + // Session decided idle is terminated Some(ResponseOrIdle::Response(response)) => { for body_elem in response.body.into_iter() { let _handle = match body_elem { @@ -273,6 +286,7 @@ impl NetLoop { self.server.enqueue_status(response.completion); return Ok(LoopMode::Interactive) }, + // Session has some information for user Some(ResponseOrIdle::IdleEvent(elems)) => { for body_elem in elems.into_iter() { let _handle = match body_elem { @@ -280,17 +294,43 @@ impl NetLoop { Body::Status(s) => self.server.enqueue_status(s), }; } - return Ok(LoopMode::Idle) + self.cmd_tx.try_send(Request::Idle)?; + return Ok(LoopMode::Idle(buff, stop)) }, + + // Session crashed None => { self.server.enqueue_status(Status::bye(None, "Internal session exited").unwrap()); tracing::error!("session task exited for {:?}, quitting", self.ctx.addr); return Ok(LoopMode::Interactive) }, - Some(ResponseOrIdle::StartIdle) => unreachable!(), - } + + // Session can't start idling while already idling, it's a logic error! + Some(ResponseOrIdle::StartIdle(..)) => bail!("can't start idling while already idling!"), + }, + + // User is trying to interact with us + _read_client_bytes = self.server.stream.read(&mut buff) => { + use imap_codec::decode::Decoder; + let codec = imap_codec::IdleDoneCodec::new(); + match codec.decode(&buff) { + Ok(([], imap_codec::imap_types::extensions::idle::IdleDone)) => { + // Session will be informed that it must stop idle + // It will generate the "done" message and change the loop mode + stop.notify_one() + }, + Err(_) => (), + _ => bail!("Client sent data after terminating the continuation without waiting for the server. This is an unsupported behavior and bug in Aerogramme, quitting."), + }; + + return Ok(LoopMode::Idle(buff, stop)) + }, + + // When receiving a CTRL+C + _ = self.ctx.must_exit.changed() => { + self.server.enqueue_status(Status::bye(None, "Server is being shutdown").unwrap()); + return Ok(LoopMode::Interactive) + }, }; - /*self.cmd_tx.try_send(Request::Idle).unwrap(); - Ok(LoopMode::Idle)*/ } } diff --git a/src/imap/request.rs b/src/imap/request.rs index 2382b09..49b4992 100644 --- a/src/imap/request.rs +++ b/src/imap/request.rs @@ -1,5 +1,4 @@ use imap_codec::imap_types::command::Command; -use tokio::sync::Notify; #[derive(Debug)] pub enum Request { diff --git a/src/imap/response.rs b/src/imap/response.rs index 7b7f92d..afcb29f 100644 --- a/src/imap/response.rs +++ b/src/imap/response.rs @@ -1,4 +1,6 @@ +use std::sync::Arc; use anyhow::Result; +use tokio::sync::Notify; use imap_codec::imap_types::command::Command; use imap_codec::imap_types::core::Tag; use imap_codec::imap_types::response::{Code, Data, Status}; @@ -116,6 +118,6 @@ impl<'a> Response<'a> { #[derive(Debug)] pub enum ResponseOrIdle { Response(Response<'static>), - StartIdle, + StartIdle(Arc), IdleEvent(Vec>), } diff --git a/src/imap/session.rs b/src/imap/session.rs index d15016f..1d473ed 100644 --- a/src/imap/session.rs +++ b/src/imap/session.rs @@ -38,6 +38,16 @@ impl Instance { _ => unreachable!(), }; + tokio::select! { + _ = stop.notified() => { + return Response::build() + .tag(imap_codec::imap_types::core::Tag::try_from("FIXME").unwrap()) + .message("IDLE completed") + .ok() + .unwrap() + } + } + unimplemented!(); } @@ -108,8 +118,8 @@ impl Instance { .unwrap()); } - match self.state { - flow::State::Idle(..) => ResponseOrIdle::StartIdle, + match &self.state { + flow::State::Idle(_, _, _, n) => ResponseOrIdle::StartIdle(n.clone()), _ => ResponseOrIdle::Response(resp), } } diff --git a/src/mail/mailbox.rs b/src/mail/mailbox.rs index 5e95f32..4310a73 100644 --- a/src/mail/mailbox.rs +++ b/src/mail/mailbox.rs @@ -67,6 +67,11 @@ impl Mailbox { self.mbox.write().await.opportunistic_sync().await } + /// Block until a sync has been done (due to changes in the event log) + pub async fn idle_sync(&self) -> Result<()> { + self.mbox.write().await.idle_sync().await + } + // ---- Functions for reading the mailbox ---- /// Get a clone of the current UID Index of this mailbox @@ -199,6 +204,11 @@ impl MailboxInternal { Ok(()) } + async fn idle_sync(&mut self) -> Result<()> { + self.uid_index.idle_sync().await?; + Ok(()) + } + // ---- Functions for reading the mailbox ---- async fn fetch_meta(&self, ids: &[UniqueIdent]) -> Result> {