awaken_stores/memory/

pending.rs

use std::collections::HashMap;
use std::collections::HashSet;

use async_trait::async_trait;
use awaken_server_contract::contract::message::{
    DeliveryBoundary, DeliveryMode, Message, MessageRecord, PendingMessageRecord,
    pending_queue_revision, select_pending_for_freeze, select_pending_for_freeze_for_run,
};
use awaken_server_contract::contract::storage::{
    RunRecord, StorageError, checkpoint_parent_thread_id, message_append,
};
use awaken_server_contract::thread::Thread;

use crate::PendingMessageStore;
use crate::pending_message_store::{
    validate_pending_message_record, validate_pending_message_records,
};

use super::validate_thread_hierarchy_map;
use super::{InMemoryStore, current_millis};

fn normalize_pending_positions(pending: &mut [PendingMessageRecord]) {
    for (index, record) in pending.iter_mut().enumerate() {
        record.position = index as u64 + 1;
    }
}

fn pending_not_found(thread_id: &str, pending_id: &str) -> StorageError {
    StorageError::NotFound(format!(
        "pending message '{pending_id}' in thread '{thread_id}'"
    ))
}

fn already_consumed(pending_id: &str) -> StorageError {
    StorageError::Validation(format!(
        "pending message '{pending_id}' is already consumed"
    ))
}

fn duplicate_pending_id(pending_id: &str) -> StorageError {
    StorageError::Validation(format!("pending message '{pending_id}' already exists"))
}

fn selected_pending_ids(
    pending: &[PendingMessageRecord],
    selected_indexes: &[usize],
) -> Vec<String> {
    selected_indexes
        .iter()
        .map(|index| pending[*index].pending_id.clone())
        .collect()
}

impl InMemoryStore {
    async fn committed_message_exists(
        &self,
        thread_id: &str,
        message_id: &str,
    ) -> Result<bool, StorageError> {
        let guard = self.messages.read().await;
        Ok(guard.get(thread_id).is_some_and(|messages| {
            messages
                .iter()
                .any(|message| message.id.as_deref() == Some(message_id))
        }))
    }
}

#[async_trait]
impl PendingMessageStore for InMemoryStore {
    async fn load_pending_message_records(
        &self,
        thread_id: &str,
    ) -> Result<Vec<PendingMessageRecord>, StorageError> {
        let guard = self.pending_messages.read().await;
        let records = guard.get(thread_id).cloned().unwrap_or_default();
        validate_pending_message_records(&records)?;
        Ok(records)
    }

    async fn list_threads_with_pending_messages(
        &self,
        limit: usize,
        after: Option<&str>,
    ) -> Result<Vec<String>, StorageError> {
        let guard = self.pending_messages.read().await;
        let mut ids: Vec<String> = guard
            .iter()
            .filter(|(_, records)| !records.is_empty())
            .map(|(thread_id, _)| thread_id.clone())
            .filter(|thread_id| after.is_none_or(|cursor| thread_id.as_str() > cursor))
            .collect();
        ids.sort();
        if limit > 0 {
            ids.truncate(limit);
        }
        Ok(ids)
    }

    async fn append_pending_message_records(
        &self,
        thread_id: &str,
        messages: &[Message],
        delivery_mode: DeliveryMode,
    ) -> Result<Vec<PendingMessageRecord>, StorageError> {
        let now = current_millis() / 1000;
        let committed_guard = self.messages.read().await;
        let mut guard = self.pending_messages.write().await;
        let pending = guard.entry(thread_id.to_owned()).or_default();
        let start_position = pending.len() as u64 + 1;
        let records = messages
            .iter()
            .cloned()
            .enumerate()
            .map(|(index, message)| {
                let mut record = PendingMessageRecord::from_message(
                    thread_id.to_owned(),
                    start_position + index as u64,
                    message,
                    delivery_mode.clone(),
                );
                record.created_at = Some(now);
                record.updated_at = Some(now);
                record
            })
            .collect::<Vec<_>>();
        let mut seen = pending
            .iter()
            .map(|record| record.pending_id.as_str())
            .collect::<HashSet<_>>();
        for record in &records {
            validate_pending_message_record(record)?;
            if !seen.insert(record.pending_id.as_str()) {
                return Err(duplicate_pending_id(&record.pending_id));
            }
            if committed_guard.get(thread_id).is_some_and(|committed| {
                committed
                    .iter()
                    .any(|message| message.id.as_deref() == Some(record.pending_id.as_str()))
            }) {
                return Err(already_consumed(&record.pending_id));
            }
        }
        pending.extend(records.iter().cloned());
        Ok(records)
    }

    async fn update_pending_message_record_checked(
        &self,
        thread_id: &str,
        pending_id: &str,
        expected_revision: Option<u64>,
        mut message: Message,
    ) -> Result<PendingMessageRecord, StorageError> {
        let mut guard = self.pending_messages.write().await;
        if let Some(pending) = guard.get_mut(thread_id)
            && let Some(record) = pending
                .iter_mut()
                .find(|record| record.pending_id == pending_id)
        {
            if let Some(expected) = expected_revision
                && record.revision != expected
            {
                return Err(StorageError::VersionConflict {
                    expected,
                    actual: record.revision,
                });
            }
            match message.id.as_deref() {
                Some(message_id) if message_id != pending_id => {
                    return Err(StorageError::Validation(format!(
                        "pending message '{pending_id}' cannot change message id to '{message_id}'"
                    )));
                }
                Some(_) => {}
                None => message.id = Some(pending_id.to_owned()),
            }
            let mut updated = record.clone();
            updated.message = message;
            validate_pending_message_record(&updated)?;
            record.message = updated.message;
            record.revision += 1;
            record.updated_at = Some(current_millis() / 1000);
            return Ok(record.clone());
        }
        drop(guard);
        if self.committed_message_exists(thread_id, pending_id).await? {
            return Err(already_consumed(pending_id));
        }
        Err(pending_not_found(thread_id, pending_id))
    }

    async fn retract_pending_message_record_checked(
        &self,
        thread_id: &str,
        pending_id: &str,
        expected_revision: Option<u64>,
    ) -> Result<PendingMessageRecord, StorageError> {
        let mut guard = self.pending_messages.write().await;
        if let Some(pending) = guard.get_mut(thread_id)
            && let Some(index) = pending
                .iter()
                .position(|record| record.pending_id == pending_id)
        {
            if let Some(expected) = expected_revision
                && pending[index].revision != expected
            {
                return Err(StorageError::VersionConflict {
                    expected,
                    actual: pending[index].revision,
                });
            }
            let removed = pending.remove(index);
            normalize_pending_positions(pending);
            return Ok(removed);
        }
        drop(guard);
        if self.committed_message_exists(thread_id, pending_id).await? {
            return Err(already_consumed(pending_id));
        }
        Err(pending_not_found(thread_id, pending_id))
    }

    async fn reorder_pending_message_records_checked(
        &self,
        thread_id: &str,
        expected_queue_revision: Option<u64>,
        ordered_pending_ids: &[String],
    ) -> Result<Vec<PendingMessageRecord>, StorageError> {
        let mut guard = self.pending_messages.write().await;
        let Some(pending) = guard.get_mut(thread_id) else {
            drop(guard);
            for pending_id in ordered_pending_ids {
                if self.committed_message_exists(thread_id, pending_id).await? {
                    return Err(already_consumed(pending_id));
                }
            }
            return Err(StorageError::NotFound(thread_id.to_owned()));
        };
        let actual_queue_revision = pending_queue_revision(pending);
        if let Some(expected) = expected_queue_revision
            && expected != actual_queue_revision
        {
            return Err(StorageError::VersionConflict {
                expected,
                actual: actual_queue_revision,
            });
        }
        let pending_ids = pending
            .iter()
            .map(|record| record.pending_id.as_str())
            .collect::<HashSet<_>>();
        let consumed_id = ordered_pending_ids
            .iter()
            .find(|pending_id| !pending_ids.contains(pending_id.as_str()))
            .cloned();
        if let Some(pending_id) = consumed_id {
            drop(guard);
            if self
                .committed_message_exists(thread_id, &pending_id)
                .await?
            {
                return Err(already_consumed(&pending_id));
            }
            return Err(StorageError::NotFound(pending_id));
        }
        if pending.len() != ordered_pending_ids.len() {
            return Err(StorageError::VersionConflict {
                expected: ordered_pending_ids.len() as u64,
                actual: pending.len() as u64,
            });
        }
        let mut by_id = pending
            .iter()
            .cloned()
            .map(|record| (record.pending_id.clone(), record))
            .collect::<HashMap<_, _>>();
        let mut reordered = Vec::with_capacity(ordered_pending_ids.len());
        for pending_id in ordered_pending_ids {
            let record = by_id
                .remove(pending_id)
                .ok_or_else(|| StorageError::NotFound(pending_id.clone()))?;
            reordered.push(record);
        }
        if !by_id.is_empty() {
            return Err(StorageError::Validation(format!(
                "reorder for thread '{thread_id}' omitted pending ids"
            )));
        }
        let now = current_millis() / 1000;
        normalize_pending_positions(&mut reordered);
        for record in &mut reordered {
            record.revision += 1;
            record.updated_at = Some(now);
        }
        *pending = reordered.clone();
        Ok(reordered)
    }

    async fn freeze_pending_message_records(
        &self,
        thread_id: &str,
        boundary: DeliveryBoundary,
        expected_message_version: Option<u64>,
    ) -> Result<Vec<MessageRecord>, StorageError> {
        let mut messages_guard = self.messages.write().await;
        let mut pending_guard = self.pending_messages.write().await;
        let committed = messages_guard.entry(thread_id.to_owned()).or_default();
        let actual = committed.len() as u64;
        if let Some(expected) = expected_message_version
            && expected != actual
        {
            return Err(StorageError::VersionConflict { expected, actual });
        }
        let Some(pending) = pending_guard.get_mut(thread_id) else {
            return Ok(Vec::new());
        };
        let selected_indexes = select_pending_for_freeze(pending, boundary);
        if selected_indexes.is_empty() {
            return Ok(Vec::new());
        }
        let selected_messages = selected_indexes
            .iter()
            .map(|index| pending[*index].message.clone())
            .collect::<Vec<_>>();
        message_append::validate_append_only_delta(committed, &selected_messages)?;

        let mut selected = Vec::with_capacity(selected_indexes.len());
        for index in selected_indexes.iter().rev() {
            selected.push(pending.remove(*index));
        }
        selected.reverse();
        normalize_pending_positions(pending);
        let start_seq = committed.len() as u64 + 1;
        let appended = selected
            .into_iter()
            .enumerate()
            .map(|(index, record)| {
                let message = record.message;
                committed.push(message.clone());
                MessageRecord::from_message(thread_id.to_owned(), start_seq + index as u64, message)
            })
            .collect();
        Ok(appended)
    }

    async fn freeze_pending_message_records_with_run(
        &self,
        thread_id: &str,
        boundary: DeliveryBoundary,
        expected_message_version: Option<u64>,
        expected_pending_ids: &[String],
        run: &RunRecord,
    ) -> Result<Vec<MessageRecord>, StorageError> {
        run.validate_for_persist()?;
        let now = current_millis();
        let mut thread_guard = self.threads.write().await;
        let existing_thread = thread_guard.get(thread_id).cloned();
        validate_thread_hierarchy_map(
            &thread_guard,
            thread_id,
            checkpoint_parent_thread_id(existing_thread.as_ref(), run),
        )?;
        let mut messages_guard = self.messages.write().await;
        let mut pending_guard = self.pending_messages.write().await;
        let mut run_guard = self.runs.write().await;
        let actual = messages_guard
            .get(thread_id)
            .map(|messages| messages.len() as u64)
            .unwrap_or(0);
        if let Some(expected) = expected_message_version
            && expected != actual
        {
            return Err(StorageError::VersionConflict { expected, actual });
        }
        let pending = pending_guard.entry(thread_id.to_owned()).or_default();
        let selected_indexes =
            select_pending_for_freeze_for_run(pending, boundary, Some(&run.run_id));
        let selected_ids = selected_pending_ids(pending, &selected_indexes);
        if selected_ids != expected_pending_ids {
            return Err(StorageError::PendingSelectionConflict {
                expected_ids: expected_pending_ids.to_vec(),
                actual_ids: selected_ids,
            });
        }
        let committed = messages_guard.entry(thread_id.to_owned()).or_default();
        let selected_messages = selected_indexes
            .iter()
            .map(|index| pending[*index].message.clone())
            .collect::<Vec<_>>();
        message_append::validate_append_only_delta(committed, &selected_messages)?;

        let mut selected = Vec::with_capacity(selected_indexes.len());
        for index in selected_indexes.iter().rev() {
            selected.push(pending.remove(*index));
        }
        selected.reverse();
        normalize_pending_positions(pending);
        let start_seq = committed.len() as u64 + 1;
        let appended = selected
            .into_iter()
            .enumerate()
            .map(|(index, record)| {
                let message = record.message;
                committed.push(message.clone());
                MessageRecord::from_message(thread_id.to_owned(), start_seq + index as u64, message)
            })
            .collect::<Vec<_>>();
        let mut thread = existing_thread.unwrap_or_else(|| Thread::with_id(thread_id));
        thread.touch(now);
        thread.apply_run_projection(run);
        thread.normalize_lineage();
        thread_guard.insert(thread_id.to_owned(), thread);
        run_guard.insert(run.run_id.clone(), run.clone());
        self.run_insertion
            .write()
            .await
            .insert(run.run_id.clone(), self.next_run_seq());
        Ok(appended)
    }

    async fn append_and_freeze_pending_message_records_with_run(
        &self,
        thread_id: &str,
        new_messages: &[Message],
        append_delivery_mode: DeliveryMode,
        boundary: DeliveryBoundary,
        expected_message_version: Option<u64>,
        expected_pending_ids: &[String],
        run: &RunRecord,
    ) -> Result<Vec<MessageRecord>, StorageError> {
        run.validate_for_persist()?;
        let now = current_millis();
        let mut thread_guard = self.threads.write().await;
        let existing_thread = thread_guard.get(thread_id).cloned();
        validate_thread_hierarchy_map(
            &thread_guard,
            thread_id,
            checkpoint_parent_thread_id(existing_thread.as_ref(), run),
        )?;
        let mut messages_guard = self.messages.write().await;
        let mut pending_guard = self.pending_messages.write().await;
        let mut run_guard = self.runs.write().await;
        let actual = messages_guard
            .get(thread_id)
            .map(|messages| messages.len() as u64)
            .unwrap_or(0);
        if let Some(expected) = expected_message_version
            && expected != actual
        {
            return Err(StorageError::VersionConflict { expected, actual });
        }
        let pending = pending_guard.entry(thread_id.to_owned()).or_default();
        // Append the new messages to pending *under the same held locks* as the
        // freeze below, so the two are one atomic boundary (ADR-0042 D7).
        let append_now = now / 1000;
        let start_position = pending.len() as u64 + 1;
        let mut seen = pending
            .iter()
            .map(|record| record.pending_id.clone())
            .collect::<HashSet<_>>();
        let mut appended_pending = Vec::with_capacity(new_messages.len());
        for (index, message) in new_messages.iter().cloned().enumerate() {
            let mut record = PendingMessageRecord::from_message(
                thread_id.to_owned(),
                start_position + index as u64,
                message,
                append_delivery_mode.clone(),
            );
            record.created_at = Some(append_now);
            record.updated_at = Some(append_now);
            validate_pending_message_record(&record)?;
            if !seen.insert(record.pending_id.clone()) {
                return Err(duplicate_pending_id(&record.pending_id));
            }
            if messages_guard.get(thread_id).is_some_and(|committed| {
                committed
                    .iter()
                    .any(|message| message.id.as_deref() == Some(record.pending_id.as_str()))
            }) {
                return Err(already_consumed(&record.pending_id));
            }
            appended_pending.push(record);
        }
        pending.extend(appended_pending);

        // Freeze the caller's selection over the now-appended pending. Mirrors
        // `freeze_pending_message_records_with_run`, but with append folded into
        // the same locked region.
        let selected_indexes =
            select_pending_for_freeze_for_run(pending, boundary, Some(&run.run_id));
        let selected_ids = selected_pending_ids(pending, &selected_indexes);
        if selected_ids != expected_pending_ids {
            return Err(StorageError::PendingSelectionConflict {
                expected_ids: expected_pending_ids.to_vec(),
                actual_ids: selected_ids,
            });
        }
        let committed = messages_guard.entry(thread_id.to_owned()).or_default();
        let selected_messages = selected_indexes
            .iter()
            .map(|index| pending[*index].message.clone())
            .collect::<Vec<_>>();
        message_append::validate_append_only_delta(committed, &selected_messages)?;
        let mut selected = Vec::with_capacity(selected_indexes.len());
        for index in selected_indexes.iter().rev() {
            selected.push(pending.remove(*index));
        }
        selected.reverse();
        normalize_pending_positions(pending);
        let start_seq = committed.len() as u64 + 1;
        let appended = selected
            .into_iter()
            .enumerate()
            .map(|(index, record)| {
                let message = record.message;
                committed.push(message.clone());
                MessageRecord::from_message(thread_id.to_owned(), start_seq + index as u64, message)
            })
            .collect::<Vec<_>>();
        let mut thread = existing_thread.unwrap_or_else(|| Thread::with_id(thread_id));
        thread.touch(now);
        thread.apply_run_projection(run);
        thread.normalize_lineage();
        thread_guard.insert(thread_id.to_owned(), thread);
        run_guard.insert(run.run_id.clone(), run.clone());
        self.run_insertion
            .write()
            .await
            .insert(run.run_id.clone(), self.next_run_seq());
        Ok(appended)
    }
}