use super::{
config::Assistant,
session::{lock, session_id},
};
use crate::serve::storage::Failure;
use serde_json::{Value, json};
use std::{
collections::BTreeMap,
io::Write,
os::unix::net::UnixStream,
sync::{Arc, Mutex, mpsc},
time::Duration,
};
#[path = "assistant_replay.rs"]
mod replay;
use replay::Replay;
pub(super) type Pair = (String, String);
pub(super) enum Event {
Item(Value),
Closed,
}
pub(super) struct Seat {
pub agent: String,
pub generation: String,
pub socket: Arc<Mutex<UnixStream>>,
pub pending: BTreeMap<String, mpsc::SyncSender<Result<(), Failure>>>,
pub presence: Value,
}
struct Watch {
pair: Pair,
sender: Option<mpsc::SyncSender<Event>>,
replay: Option<Replay>,
}
#[derive(Default)]
struct State {
seats: BTreeMap<Pair, Seat>,
watches: BTreeMap<String, Watch>,
}
pub(super) struct Hub {
pub config: Assistant,
pub timeout: Duration,
state: Mutex<State>,
}
pub(super) fn bound(name: &str, maximum: usize, count: usize) -> Failure {
Failure::new(
"assistant_bound",
format!("{name}={maximum}; count={count}"),
)
}
impl Hub {
pub fn new(config: Assistant, timeout: Duration) -> Self {
Self {
config,
timeout,
state: Mutex::default(),
}
}
fn check(&self, pair: &Pair, pairs: impl Iterator<Item = Pair>) -> Result<(), Failure> {
let (mut person, mut store) = (0, 0);
for (who, place) in pairs {
person += usize::from(who == pair.0);
store += usize::from(place == pair.1);
}
for (name, maximum, count) in [
("max_per_person", self.config.max_per_person.get(), person),
("max_per_store", self.config.max_per_store.get(), store),
] {
if count >= maximum {
return Err(bound(name, maximum, count));
}
}
Ok(())
}
pub fn admit(&self, pair: Pair, seat: Seat) -> Result<(), Failure> {
let mut state = lock(&self.state)?;
if state.seats.contains_key(&pair) {
return Err(Failure::new(
"seat_taken",
"this person already has a live seat for this store",
));
}
self.check(&pair, state.seats.keys().cloned())?;
self.publish(&mut state, &pair, &seat.presence);
state.seats.insert(pair, seat);
drop(state);
Ok(())
}
pub fn depart(&self, pair: &Pair, generation: &str) -> Result<(), Failure> {
let mut state = lock(&self.state)?;
if state
.seats
.get(pair)
.is_some_and(|seat| seat.generation == generation)
{
if let Some(seat) = state.seats.remove(pair) {
for pending in seat.pending.into_values() {
match pending.try_send(Err(Failure::new(
"assistant_disconnected",
"the seat closed before acknowledging",
))) {
Ok(()) | Err(mpsc::TrySendError::Disconnected(_)) => {}
Err(mpsc::TrySendError::Full(_)) => {
return Err(Failure::new(
"assistant_ack_conflict",
"an acknowledgement was already waiting",
));
}
}
}
}
self.publish(
&mut state,
pair,
&json!({"kind":"presence","present":false}),
);
}
drop(state);
Ok(())
}
pub fn item(&self, pair: &Pair, item: &Value) -> Result<(), Failure> {
let mut state = lock(&self.state)?;
self.publish(&mut state, pair, item);
Ok(())
}
fn publish(&self, state: &mut State, pair: &Pair, item: &Value) {
for watch in state.watches.values_mut() {
if &watch.pair != pair {
continue;
}
if watch
.replay
.as_mut()
.is_some_and(|replay| replay.capture(item, self.config.events_after_bytes.get()))
{
continue;
}
if let Some(sender) = &watch.sender
&& sender.try_send(Event::Item(item.clone())).is_err()
{
watch.sender = None;
}
}
}
#[cfg(test)]
pub fn subscribe(self: &Arc<Self>, pair: Pair) -> Result<Subscription, Failure> {
self.subscribe_replay(pair, None)
}
pub fn subscribe_replay(
self: &Arc<Self>,
pair: Pair,
conversation: Option<String>,
) -> Result<Subscription, Failure> {
let mut state = lock(&self.state)?;
self.check(
&pair,
state.watches.values().map(|watch| watch.pair.clone()),
)?;
let id = session_id()?;
let (sender, input) = mpsc::sync_channel(self.config.max_per_store.get().saturating_add(1));
let presence = state.seats.get(&pair).map_or_else(
|| json!({"kind":"presence","present":false}),
|seat| seat.presence.clone(),
);
sender
.try_send(Event::Item(presence))
.map_err(|_| Failure::new("assistant_stream_closed", "stream could not attach"))?;
state.watches.insert(
id.clone(),
Watch {
pair,
sender: Some(sender),
replay: conversation.map(Replay::new),
},
);
drop(state);
Ok(Subscription {
hub: Arc::clone(self),
id,
input,
})
}
pub fn finish_replay(&self, id: &str) -> Result<Option<Vec<Value>>, Failure> {
let mut state = lock(&self.state)?;
let watch = state.watches.get_mut(id).ok_or_else(|| {
Failure::new("assistant_stream_closed", "the browser closed its stream")
})?;
if let Some(replay) = watch.replay.as_mut() {
if let Some(error) = replay.refused.take() {
return Err(error);
}
if replay.dirty {
replay.dirty = false;
return Ok(None);
}
}
let held = watch
.replay
.take()
.map_or_else(Vec::new, |replay| replay.held);
drop(state);
Ok(Some(held))
}
pub fn seat_identity(&self, pair: &Pair) -> Result<(String, String), Failure> {
lock(&self.state)?
.seats
.get(pair)
.map(|seat| (seat.agent.clone(), seat.generation.clone()))
.ok_or_else(|| Failure::new("assistant_not_started", "the seat is no longer connected"))
}
pub fn close(&self, id: &str) -> Result<(), Failure> {
if let Some(watch) = lock(&self.state)?.watches.remove(id)
&& let Some(sender) = watch.sender
{
match sender.try_send(Event::Closed) {
Ok(()) | Err(mpsc::TrySendError::Disconnected(_) | mpsc::TrySendError::Full(_)) => {
}
}
}
Ok(())
}
pub fn operation(&self, pair: &Pair, body: Value) -> Result<String, Failure> {
self.operation_for(pair, None, body)
}
pub fn operation_for(
&self,
pair: &Pair,
generation: Option<&str>,
mut body: Value,
) -> Result<String, Failure> {
let id = session_id()?;
body["id"] = json!(id);
body["person"] = json!(pair.0);
let mut bytes = serde_json::to_vec(&json!({"op":body}))
.map_err(|_| Failure::new("assistant_op_invalid", "operation could not be encoded"))?;
bytes.push(b'\n');
if bytes.len() > self.config.frame_bytes.get() {
return Err(bound(
"frame_bytes",
self.config.frame_bytes.get(),
bytes.len(),
));
}
let (send, receive) = mpsc::sync_channel(1);
let socket = {
let mut state = lock(&self.state)?;
self.check(
pair,
state
.seats
.iter()
.flat_map(|(pair, seat)| std::iter::repeat_n(pair.clone(), seat.pending.len())),
)?;
let seat = state.seats.get_mut(pair).ok_or_else(|| {
Failure::new("assistant_not_started", "start the assistant in Lys")
})?;
if generation.is_some_and(|expected| expected != seat.generation) {
return Err(Failure::new(
"assistant_disconnected",
"the seat changed after the draft was checked",
));
}
seat.pending.insert(id.clone(), send);
let socket = Arc::clone(&seat.socket);
drop(state);
socket
};
let result = (|| {
lock(&socket)?.write_all(&bytes).map_err(|_| {
Failure::new("assistant_disconnected", "the operation could not be sent")
})?;
receive.recv_timeout(self.timeout).map_err(|_| {
Failure::new(
"assistant_ack_unknown",
"the operation has no acknowledged outcome",
)
})?
})();
if let Some(seat) = lock(&self.state)?.seats.get_mut(pair) {
seat.pending.remove(&id);
}
result.map(|()| id)
}
pub fn acknowledge(
&self,
pair: &Pair,
id: &str,
result: Result<(), Failure>,
) -> Result<(), Failure> {
let send = lock(&self.state)?
.seats
.get_mut(pair)
.and_then(|seat| seat.pending.remove(id));
if let Some(send) = send {
match send.try_send(result) {
Ok(()) | Err(mpsc::TrySendError::Disconnected(_)) => {}
Err(mpsc::TrySendError::Full(_)) => {
return Err(Failure::new(
"assistant_ack_conflict",
"an acknowledgement was already waiting",
));
}
}
}
Ok(())
}
}
pub(super) struct Subscription {
pub hub: Arc<Hub>,
pub id: String,
pub input: mpsc::Receiver<Event>,
}
impl Drop for Subscription {
fn drop(&mut self) {
if let Err(error) = self.hub.close(&self.id) {
eprintln!("haematite assistant_stream_release_failed: {}", error.code);
}
}
}