ursula_runtime/engine/

in_memory.rs

use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;

use bytes::Bytes;
use ursula_shard::BucketStreamId;
use ursula_shard::ShardPlacement;
use ursula_stream::AppendStreamInput;
use ursula_stream::ProducerRequest;
use ursula_stream::StreamCommand;
use ursula_stream::StreamErrorCode;
use ursula_stream::StreamMessageRecord;
use ursula_stream::StreamReadPlan;
use ursula_stream::StreamReadSegment;
use ursula_stream::StreamResponse;
use ursula_stream::StreamSnapshot;
use ursula_stream::StreamStateMachine;

use super::GroupAckColdGcFuture;
use super::GroupAppendBatchFuture;
use super::GroupAppendBatchResponse;
use super::GroupAppendFuture;
use super::GroupBootstrapStreamFuture;
use super::GroupCloseStreamFuture;
use super::GroupColdHotBacklogFuture;
use super::GroupCreateStreamFuture;
use super::GroupDeleteSnapshotFuture;
use super::GroupDeleteStreamFuture;
use super::GroupEngine;
use super::GroupEngineCreateFuture;
use super::GroupEngineError;
use super::GroupEngineFactory;
use super::GroupEngineMetrics;
use super::GroupFlushColdFuture;
use super::GroupGetStreamAttrsFuture;
use super::GroupHeadStreamFuture;
use super::GroupInstallSnapshotFuture;
use super::GroupPlanColdFlushFuture;
use super::GroupPlanColdGcFuture;
use super::GroupPlanNextColdFlushBatchFuture;
use super::GroupPublishSnapshotFuture;
use super::GroupReadSnapshotFuture;
use super::GroupReadStreamFuture;
use super::GroupReadStreamPartsFuture;
use super::GroupSnapshotFuture;
use super::GroupTouchStreamAccessFuture;
use super::GroupUpdateStreamAttrsFuture;
use super::GroupWriteResponse;
use crate::cold_index::ColdIndexPageCache;
use crate::cold_index::ColdStoreColdIndexPageStore;
use crate::cold_index::rollback_cold_index_pages;
use crate::cold_index::write_cold_chunk_index_pages_with_rollback;
use crate::cold_index::write_external_segment_index_pages;
use crate::cold_store::ColdStoreHandle;
use crate::cold_store::DEFAULT_CONTENT_TYPE;
use crate::command::GroupSnapshot;
use crate::command::GroupWriteCommand;
use crate::request::AckColdGcResponse;
use crate::request::AppendBatchRequest;
use crate::request::AppendExternalRequest;
use crate::request::AppendRequest;
use crate::request::AppendResponse;
use crate::request::BootstrapStreamRequest;
use crate::request::BootstrapStreamResponse;
use crate::request::BootstrapUpdate;
use crate::request::CloseStreamRequest;
use crate::request::CloseStreamResponse;
use crate::request::ColdHotBacklog;
use crate::request::ColdWriteAdmission;
use crate::request::CreateStreamExternalRequest;
use crate::request::CreateStreamRequest;
use crate::request::CreateStreamResponse;
use crate::request::DeleteSnapshotRequest;
use crate::request::DeleteStreamRequest;
use crate::request::DeleteStreamResponse;
use crate::request::FlushColdRequest;
use crate::request::FlushColdResponse;
use crate::request::GetStreamAttrsRequest;
use crate::request::GetStreamAttrsResponse;
use crate::request::GroupReadStreamParts;
use crate::request::HeadStreamRequest;
use crate::request::HeadStreamResponse;
use crate::request::PlanColdFlushRequest;
use crate::request::PlanGroupColdFlushRequest;
use crate::request::PublishSnapshotRequest;
use crate::request::PublishSnapshotResponse;
use crate::request::ReadSnapshotRequest;
use crate::request::ReadSnapshotResponse;
use crate::request::ReadStreamRequest;
use crate::request::StreamAppendCount;
use crate::request::TouchStreamAccessResponse;
use crate::request::UpdateStreamAttrsRequest;
use crate::request::UpdateStreamAttrsResponse;

pub(crate) struct AppendPayloadInput<'a> {
    stream_id: BucketStreamId,
    content_type: Option<&'a str>,
    payload: &'a [u8],
    close_after: bool,
    stream_seq: Option<String>,
    producer: Option<ProducerRequest>,
    now_ms: u64,
    record_match: Option<u64>,
}

#[derive(Debug, Clone, Default)]
pub struct InMemoryGroupEngine {
    pub(crate) commit_index: u64,
    pub(crate) state_machine: StreamStateMachine,
    pub(crate) stream_append_counts: HashMap<BucketStreamId, u64>,
    pub(crate) cold_store: Option<ColdStoreHandle>,
    pub(crate) cold_index_cache: Option<Arc<ColdIndexPageCache<ColdStoreColdIndexPageStore>>>,
}

impl InMemoryGroupEngine {
    pub fn with_cold_store(cold_store: ColdStoreHandle) -> Self {
        let mut engine = Self::default();
        engine.set_cold_store(Some(cold_store));
        engine
    }

    pub fn cold_store(&self) -> Option<ColdStoreHandle> {
        self.cold_store.clone()
    }

    pub(crate) fn set_cold_store(&mut self, cold_store: Option<ColdStoreHandle>) {
        self.cold_index_cache = cold_store.as_ref().map(|cold_store| {
            Arc::new(ColdIndexPageCache::new(
                Arc::new(ColdStoreColdIndexPageStore::new(cold_store.clone())),
                1024,
            ))
        });
        self.cold_store = cold_store;
    }

    pub fn apply_committed_write(
        &mut self,
        command: GroupWriteCommand,
        placement: ShardPlacement,
    ) -> Result<GroupWriteResponse, GroupEngineError> {
        match command {
            GroupWriteCommand::Stream(command) => self.apply_stream_command(command, placement),
            GroupWriteCommand::Batch { commands } => Ok(GroupWriteResponse::Batch(
                commands
                    .into_iter()
                    .map(|command| self.apply_stream_command(command, placement))
                    .collect(),
            )),
        }
    }

    /// Applies one canonical [`StreamCommand`] to the deterministic state
    /// machine and lifts its [`StreamResponse`] into the group-level response,
    /// maintaining the group commit index and per-stream append counts.
    pub fn apply_stream_command(
        &mut self,
        command: StreamCommand,
        placement: ShardPlacement,
    ) -> Result<GroupWriteResponse, GroupEngineError> {
        match command {
            // Appends skip `StreamStateMachine::apply` to keep the exact
            // borrowed fast path (no TTL sweep on the append hot path).
            StreamCommand::Append {
                stream_id,
                content_type,
                payload,
                close_after,
                stream_seq,
                producer,
                now_ms,
                record_match,
            } => self
                .append_payload(
                    AppendPayloadInput {
                        stream_id,
                        content_type: content_type.as_deref(),
                        payload: &payload,
                        close_after,
                        stream_seq,
                        producer,
                        now_ms,
                        record_match,
                    },
                    placement,
                )
                .map(GroupWriteResponse::Append),
            StreamCommand::AppendBatch {
                stream_id,
                content_type,
                payloads,
                producer,
                now_ms,
            } => self.apply_append_batch(
                stream_id,
                content_type,
                payloads,
                producer,
                now_ms,
                placement,
            ),
            command => {
                let stream_id = command_stream_id(&command);
                let command_producer = command_producer(&command);
                if let StreamCommand::CreateStream { stream_id, .. }
                | StreamCommand::CreateExternal { stream_id, .. } = &command
                {
                    ensure_bucket_exists(&mut self.state_machine, stream_id)?;
                }
                let response = self.state_machine.apply(command);
                self.group_response_from_stream(response, stream_id, command_producer, placement)
            }
        }
    }

    fn apply_append_batch(
        &mut self,
        stream_id: BucketStreamId,
        content_type: Option<String>,
        payloads: Vec<Bytes>,
        producer: Option<ProducerRequest>,
        now_ms: u64,
        placement: ShardPlacement,
    ) -> Result<GroupWriteResponse, GroupEngineError> {
        if let Some(producer) = producer {
            let payload_refs = payloads.iter().map(Bytes::as_ref).collect::<Vec<_>>();
            let batch = self
                .state_machine
                .append_batch_borrowed(
                    stream_id.clone(),
                    content_type.as_deref(),
                    &payload_refs,
                    Some(producer.clone()),
                    now_ms,
                )
                .map_err(stream_response_error)?;
            let old_commit_index = self.commit_index;
            let old_append_count = *self.stream_append_counts.get(&stream_id).unwrap_or(&0);
            if !batch.deduplicated {
                let count = u64::try_from(batch.items.len()).expect("item count fits u64");
                self.commit_index += count;
                *self
                    .stream_append_counts
                    .entry(stream_id.clone())
                    .or_insert(0) += count;
            }
            let items = batch
                .items
                .into_iter()
                .enumerate()
                .map(|(index, item)| {
                    let item_index = u64::try_from(index + 1).expect("item index fits u64");
                    Ok(AppendResponse {
                        placement,
                        start_offset: item.offset,
                        next_offset: item.next_offset,
                        stream_append_count: if item.deduplicated {
                            old_append_count
                        } else {
                            old_append_count + item_index
                        },
                        group_commit_index: if item.deduplicated {
                            old_commit_index
                        } else {
                            old_commit_index + item_index
                        },
                        closed: item.closed,
                        deduplicated: item.deduplicated,
                        producer: None,
                        record_range: self
                            .state_machine
                            .record_range_for_append(
                                &stream_id,
                                item.offset,
                                item.next_offset,
                                Some(&producer),
                            )
                            .map_err(|err| {
                                GroupEngineError::new(format!("record range: {err:?}"))
                            })?,
                    })
                })
                .collect();
            return Ok(GroupWriteResponse::AppendBatch(GroupAppendBatchResponse {
                placement,
                items,
            }));
        }

        let mut items = Vec::with_capacity(payloads.len());
        for payload in payloads {
            if payload.is_empty() {
                items.push(Err(GroupEngineError::stream(
                    StreamErrorCode::EmptyAppend,
                    "append payload must be non-empty",
                )));
                continue;
            }
            items.push(self.append_payload(
                AppendPayloadInput {
                    stream_id: stream_id.clone(),
                    content_type: content_type.as_deref(),
                    payload: &payload,
                    close_after: false,
                    stream_seq: None,
                    producer: None,
                    now_ms,
                    record_match: None,
                },
                placement,
            ));
        }
        Ok(GroupWriteResponse::AppendBatch(GroupAppendBatchResponse {
            placement,
            items,
        }))
    }

    /// Lifts a [`StreamResponse`] into the matching [`GroupWriteResponse`],
    /// advancing the group commit index for every mutating outcome.
    fn group_response_from_stream(
        &mut self,
        response: StreamResponse,
        stream_id: Option<BucketStreamId>,
        command_producer: Option<ProducerRequest>,
        placement: ShardPlacement,
    ) -> Result<GroupWriteResponse, GroupEngineError> {
        match response {
            StreamResponse::Created {
                next_offset,
                closed,
                ..
            } => {
                let stream_id = require_response_stream_id(stream_id, "created")?;
                self.commit_index += 1;
                Ok(GroupWriteResponse::CreateStream(CreateStreamResponse {
                    placement,
                    next_offset,
                    closed,
                    already_exists: false,
                    group_commit_index: self.commit_index,
                    record_range: self
                        .state_machine
                        .record_range(&stream_id)
                        .map_err(|err| GroupEngineError::new(format!("record range: {err:?}")))?,
                }))
            }
            StreamResponse::AlreadyExists {
                next_offset,
                closed,
                ..
            } => Ok(GroupWriteResponse::CreateStream(CreateStreamResponse {
                placement,
                next_offset,
                closed,
                already_exists: true,
                group_commit_index: self.commit_index,
                record_range: None,
            })),
            StreamResponse::Appended {
                offset,
                next_offset,
                closed,
                deduplicated,
                producer,
            } => {
                let stream_id = require_response_stream_id(stream_id, "appended")?;
                let record_range = self
                    .state_machine
                    .record_range_for_append(&stream_id, offset, next_offset, producer.as_ref())
                    .map_err(|err| GroupEngineError::new(format!("record range: {err:?}")))?;
                let stream_append_count = self.stream_append_counts.entry(stream_id).or_insert(0);
                if !deduplicated {
                    self.commit_index += 1;
                    *stream_append_count += 1;
                }
                Ok(GroupWriteResponse::Append(AppendResponse {
                    placement,
                    start_offset: offset,
                    next_offset,
                    stream_append_count: *stream_append_count,
                    group_commit_index: self.commit_index,
                    closed,
                    deduplicated,
                    producer,
                    record_range,
                }))
            }
            StreamResponse::SnapshotPublished {
                snapshot_offset,
                record_range,
            } => {
                self.commit_index += 1;
                Ok(GroupWriteResponse::PublishSnapshot(
                    PublishSnapshotResponse {
                        placement,
                        snapshot_offset,
                        group_commit_index: self.commit_index,
                        record_range,
                    },
                ))
            }
            StreamResponse::Accessed { changed, expired } => {
                if changed || expired {
                    self.commit_index += 1;
                }
                Ok(GroupWriteResponse::TouchStreamAccess(
                    TouchStreamAccessResponse {
                        placement,
                        changed,
                        expired,
                        group_commit_index: self.commit_index,
                    },
                ))
            }
            StreamResponse::AttrsUpdated { changed } => {
                if changed {
                    self.commit_index += 1;
                }
                Ok(GroupWriteResponse::UpdateStreamAttrs(
                    UpdateStreamAttrsResponse {
                        placement,
                        changed,
                        group_commit_index: self.commit_index,
                    },
                ))
            }
            StreamResponse::ColdFlushed { hot_start_offset } => {
                self.commit_index += 1;
                Ok(GroupWriteResponse::FlushCold(FlushColdResponse {
                    placement,
                    hot_start_offset,
                    group_commit_index: self.commit_index,
                }))
            }
            StreamResponse::Closed {
                next_offset,
                deduplicated,
                ..
            } => {
                let stream_id = require_response_stream_id(stream_id, "closed")?;
                let record_range = self
                    .state_machine
                    .record_range_for_append(
                        &stream_id,
                        next_offset,
                        next_offset,
                        command_producer.as_ref(),
                    )
                    .map_err(|err| GroupEngineError::new(format!("record range: {err:?}")))?;
                if !deduplicated {
                    self.commit_index += 1;
                }
                Ok(GroupWriteResponse::CloseStream(CloseStreamResponse {
                    placement,
                    next_offset,
                    group_commit_index: self.commit_index,
                    deduplicated,
                    record_range,
                }))
            }
            StreamResponse::Deleted => {
                let stream_id = require_response_stream_id(stream_id, "deleted")?;
                self.commit_index += 1;
                // Stream is gone: drop its runtime append count so the map
                // stays bounded under delete churn.
                self.stream_append_counts.remove(&stream_id);
                Ok(GroupWriteResponse::DeleteStream(DeleteStreamResponse {
                    placement,
                    group_commit_index: self.commit_index,
                }))
            }
            StreamResponse::ColdGcAcked { removed } => {
                self.commit_index += 1;
                Ok(GroupWriteResponse::AckColdGc(AckColdGcResponse {
                    placement,
                    removed,
                    group_commit_index: self.commit_index,
                }))
            }
            StreamResponse::Error {
                code,
                message,
                next_offset,
                context,
            } => Err(GroupEngineError::stream_with_context(
                code,
                message,
                next_offset,
                context,
            )),
            other @ (StreamResponse::BucketCreated { .. }
            | StreamResponse::BucketAlreadyExists { .. }
            | StreamResponse::BucketDeleted { .. }) => Err(GroupEngineError::new(format!(
                "unexpected group write response: {other:?}"
            ))),
        }
    }

    pub(crate) fn cold_hot_backlog_for(
        &self,
        stream_id: BucketStreamId,
    ) -> Result<ColdHotBacklog, GroupEngineError> {
        let stream_hot_bytes = self.state_machine.hot_payload_len(&stream_id).unwrap_or(0);
        Ok(ColdHotBacklog {
            stream_id,
            stream_hot_bytes,
            group_hot_bytes: self.state_machine.total_hot_payload_bytes(),
        })
    }

    pub fn check_cold_write_admission_bytes(
        &self,
        stream_id: &BucketStreamId,
        admission: ColdWriteAdmission,
        incoming_bytes: u64,
    ) -> Result<(), GroupEngineError> {
        let Some(limit) = admission.max_hot_bytes_per_group else {
            return Ok(());
        };
        if incoming_bytes == 0 {
            return Ok(());
        }
        let before = self.state_machine.total_hot_payload_bytes();
        let after = before.saturating_add(incoming_bytes);
        if after <= limit {
            return Ok(());
        }
        Err(GroupEngineError::cold_backpressure(
            stream_id.clone(),
            before,
            after,
            limit,
        ))
    }

    pub(crate) fn create_stream_with_admission_inner(
        &mut self,
        request: CreateStreamRequest,
        placement: ShardPlacement,
        admission: ColdWriteAdmission,
    ) -> Result<CreateStreamResponse, GroupEngineError> {
        let stream_id = request.stream_id.clone();
        if admission.is_enabled() {
            let mut preview = self.clone();
            let preview_response = match preview
                .apply_committed_write(GroupWriteCommand::from(request.clone()), placement)?
            {
                GroupWriteResponse::CreateStream(response) => response,
                other => {
                    return Err(GroupEngineError::new(format!(
                        "unexpected create stream preview response: {other:?}"
                    )));
                }
            };
            if !preview_response.already_exists {
                self.check_cold_write_admission_bytes(
                    &stream_id,
                    admission,
                    u64::try_from(request.initial_payload.len()).expect("payload len fits u64"),
                )?;
            }
        }
        let response =
            match self.apply_committed_write(GroupWriteCommand::from(request), placement)? {
                GroupWriteResponse::CreateStream(response) => response,
                other => {
                    return Err(GroupEngineError::new(format!(
                        "unexpected create stream write response: {other:?}"
                    )));
                }
            };
        Ok(response)
    }

    pub(crate) fn append_with_admission_inner(
        &mut self,
        request: AppendRequest,
        placement: ShardPlacement,
        admission: ColdWriteAdmission,
    ) -> Result<AppendResponse, GroupEngineError> {
        let stream_id = request.stream_id.clone();
        if admission.is_enabled() {
            let mut preview = self.clone();
            let preview_response = match preview
                .apply_committed_write(GroupWriteCommand::from(request.clone()), placement)?
            {
                GroupWriteResponse::Append(response) => response,
                other => {
                    return Err(GroupEngineError::new(format!(
                        "unexpected append preview response: {other:?}"
                    )));
                }
            };
            if !preview_response.deduplicated {
                self.check_cold_write_admission_bytes(
                    &stream_id,
                    admission,
                    u64::try_from(request.payload.len()).expect("payload len fits u64"),
                )?;
            }
        }
        let response =
            match self.apply_committed_write(GroupWriteCommand::from(request), placement)? {
                GroupWriteResponse::Append(response) => response,
                other => {
                    return Err(GroupEngineError::new(format!(
                        "unexpected append write response: {other:?}"
                    )));
                }
            };
        Ok(response)
    }

    pub(crate) fn append_batch_with_admission_inner(
        &mut self,
        request: AppendBatchRequest,
        placement: ShardPlacement,
        admission: ColdWriteAdmission,
    ) -> Result<GroupAppendBatchResponse, GroupEngineError> {
        let stream_id = request.stream_id.clone();
        let incoming_bytes = request
            .payloads
            .iter()
            .map(|payload| u64::try_from(payload.len()).expect("payload len fits u64"))
            .sum();
        if admission.is_enabled() {
            let mut preview = self.clone();
            let preview_response = match preview
                .apply_committed_write(GroupWriteCommand::from(request.clone()), placement)?
            {
                GroupWriteResponse::AppendBatch(response) => response,
                other => {
                    return Err(GroupEngineError::new(format!(
                        "unexpected append batch preview response: {other:?}"
                    )));
                }
            };
            let mutates = preview_response
                .items
                .iter()
                .any(|item| matches!(item, Ok(response) if !response.deduplicated));
            if mutates {
                self.check_cold_write_admission_bytes(&stream_id, admission, incoming_bytes)?;
            }
        }
        let response =
            match self.apply_committed_write(GroupWriteCommand::from(request), placement)? {
                GroupWriteResponse::AppendBatch(response) => response,
                other => {
                    return Err(GroupEngineError::new(format!(
                        "unexpected append batch write response: {other:?}"
                    )));
                }
            };
        Ok(response)
    }

    pub fn access_requires_write(
        &self,
        stream_id: &BucketStreamId,
        now_ms: u64,
        renew_ttl: bool,
    ) -> Result<bool, GroupEngineError> {
        self.state_machine
            .access_requires_write(stream_id, now_ms, renew_ttl)
            .map_err(stream_response_error)
    }

    pub(crate) fn apply_access_command(
        &mut self,
        stream_id: BucketStreamId,
        now_ms: u64,
        renew_ttl: bool,
        placement: ShardPlacement,
    ) -> Result<TouchStreamAccessResponse, GroupEngineError> {
        match self.apply_committed_write(
            GroupWriteCommand::Stream(StreamCommand::TouchStreamAccess {
                stream_id,
                now_ms,
                renew_ttl,
            }),
            placement,
        )? {
            GroupWriteResponse::TouchStreamAccess(response) => Ok(response),
            other => Err(GroupEngineError::new(format!(
                "unexpected touch stream access write response: {other:?}"
            ))),
        }
    }

    pub(crate) fn ensure_stream_access(
        &mut self,
        stream_id: &BucketStreamId,
        now_ms: u64,
        renew_ttl: bool,
        placement: ShardPlacement,
    ) -> Result<Option<TouchStreamAccessResponse>, GroupEngineError> {
        if !self.access_requires_write(stream_id, now_ms, renew_ttl)? {
            return Ok(None);
        }
        let response =
            self.apply_access_command(stream_id.clone(), now_ms, renew_ttl, placement)?;
        if response.expired {
            return Err(GroupEngineError::stream(
                StreamErrorCode::StreamNotFound,
                format!("stream '{stream_id}' does not exist"),
            ));
        }
        Ok(Some(response))
    }

    pub(crate) fn append_payload(
        &mut self,
        input: AppendPayloadInput<'_>,
        placement: ShardPlacement,
    ) -> Result<AppendResponse, GroupEngineError> {
        let AppendPayloadInput {
            stream_id,
            content_type,
            payload,
            close_after,
            stream_seq,
            producer,
            now_ms,
            record_match,
        } = input;
        let stream_count_key = stream_id.clone();
        let response = self.state_machine.append_borrowed(AppendStreamInput {
            stream_id,
            content_type,
            payload,
            close_after,
            stream_seq,
            producer,
            now_ms,
            record_match,
        });
        match response {
            StreamResponse::Appended {
                offset,
                next_offset,
                closed,
                deduplicated,
                producer,
                ..
            } => {
                let stream_append_count = self
                    .stream_append_counts
                    .entry(stream_count_key.clone())
                    .or_insert(0);
                let record_range = self
                    .state_machine
                    .record_range_for_append(
                        &stream_count_key,
                        offset,
                        next_offset,
                        producer.as_ref(),
                    )
                    .map_err(|err| GroupEngineError::new(format!("record range: {err:?}")))?;
                if !deduplicated {
                    self.commit_index += 1;
                    *stream_append_count += 1;
                }
                Ok(AppendResponse {
                    placement,
                    start_offset: offset,
                    next_offset,
                    stream_append_count: *stream_append_count,
                    group_commit_index: self.commit_index,
                    closed,
                    deduplicated,
                    producer,
                    record_range,
                })
            }
            StreamResponse::Error {
                code,
                message,
                next_offset,
                context,
            } => Err(GroupEngineError::stream_with_context(
                code,
                message,
                next_offset,
                context,
            )),
            other => Err(GroupEngineError::new(format!(
                "unexpected append response: {other:?}"
            ))),
        }
    }

    pub fn read_stream_plan(
        &mut self,
        request: &ReadStreamRequest,
        placement: ShardPlacement,
    ) -> Result<StreamReadPlan, GroupEngineError> {
        self.ensure_stream_access(&request.stream_id, request.now_ms, true, placement)?;
        self.read_stream_plan_after_access(request)
    }

    pub fn read_stream_plan_after_access(
        &self,
        request: &ReadStreamRequest,
    ) -> Result<StreamReadPlan, GroupEngineError> {
        let Some(record) = request.record else {
            let mut plan = self
                .state_machine
                .read_plan_at(
                    &request.stream_id,
                    request.offset,
                    request.max_len,
                    request.now_ms,
                )
                .map_err(stream_response_error)?;
            plan.retained_record_range = self
                .state_machine
                .record_range(&request.stream_id)
                .map_err(|err| GroupEngineError::new(format!("record range: {err:?}")))?;
            return Ok(plan);
        };
        let retained_record_range = self
            .state_machine
            .record_range(&request.stream_id)
            .map_err(|err| GroupEngineError::new(format!("record range: {err:?}")))?
            .ok_or_else(|| {
                GroupEngineError::stream(
                    StreamErrorCode::InvalidRecordBoundaries,
                    "record coordinates are inactive for this stream",
                )
            })?;
        if record < retained_record_range.first_record {
            return Err(GroupEngineError::stream(
                StreamErrorCode::StreamGone,
                format!(
                    "record {record} is older than first retained record {}",
                    retained_record_range.first_record
                ),
            ));
        }
        if record > retained_record_range.next_record {
            return Err(GroupEngineError::stream(
                StreamErrorCode::InvalidRecordBoundaries,
                format!(
                    "record {record} is beyond record tail {}",
                    retained_record_range.next_record
                ),
            ));
        }
        let next_record = request
            .max_records
            .map(|limit| record.saturating_add(limit))
            .unwrap_or(retained_record_range.next_record)
            .min(retained_record_range.next_record);
        let offset = self
            .state_machine
            .offset_for_record(&request.stream_id, record)
            .map_err(|err| GroupEngineError::new(format!("record offset: {err:?}")))?
            .ok_or_else(|| GroupEngineError::new("record stream disappeared"))?;
        let next_offset = self
            .state_machine
            .offset_for_record(&request.stream_id, next_record)
            .map_err(|err| GroupEngineError::new(format!("record offset: {err:?}")))?
            .ok_or_else(|| GroupEngineError::new("record stream disappeared"))?;
        let max_len = usize::try_from(next_offset.saturating_sub(offset))
            .map_err(|_| GroupEngineError::new("record read window exceeds usize"))?;
        let mut plan = self
            .state_machine
            .read_plan_at(&request.stream_id, offset, max_len, request.now_ms)
            .map_err(stream_response_error)?;
        plan.retained_record_range = Some(retained_record_range);
        plan.record_range = Some(ursula_stream::StreamRecordRange {
            first_record: record,
            next_record,
        });
        Ok(plan)
    }

    pub fn head_stream_after_access(
        &mut self,
        request: &HeadStreamRequest,
        placement: ShardPlacement,
    ) -> Result<HeadStreamResponse, GroupEngineError> {
        let Some(metadata) = self
            .state_machine
            .head_at(&request.stream_id, request.now_ms)
        else {
            return Err(GroupEngineError::stream(
                StreamErrorCode::StreamNotFound,
                format!("stream '{}' does not exist", request.stream_id),
            ));
        };
        let content_type = metadata.content_type.clone();
        let tail_offset = metadata.tail_offset;
        let closed = metadata.status == ursula_stream::StreamStatus::Closed;
        let stream_ttl_seconds = metadata.stream_ttl_seconds;
        let stream_expires_at_ms = metadata.stream_expires_at_ms;
        let _ = metadata;
        Ok(HeadStreamResponse {
            placement,
            content_type,
            tail_offset,
            cold_hot_start_offset: self.state_machine.hot_start_offset(&request.stream_id),
            closed,
            stream_ttl_seconds,
            stream_expires_at_ms,
            snapshot_offset: self
                .state_machine
                .latest_snapshot(&request.stream_id)
                .map_err(stream_response_error)?
                .map(|snapshot| snapshot.offset),
            integrity: self
                .state_machine
                .integrity_snapshot(&request.stream_id)
                .map_err(stream_response_error)?,
            record_range: self
                .state_machine
                .record_range(&request.stream_id)
                .map_err(|err| GroupEngineError::new(format!("record range: {err:?}")))?,
        })
    }

    pub fn get_stream_attrs_after_access(
        &mut self,
        request: &GetStreamAttrsRequest,
        placement: ShardPlacement,
    ) -> Result<GetStreamAttrsResponse, GroupEngineError> {
        if self
            .state_machine
            .head_at(&request.stream_id, request.now_ms)
            .is_none()
        {
            return Err(GroupEngineError::stream(
                StreamErrorCode::StreamNotFound,
                format!("stream '{}' does not exist", request.stream_id),
            ));
        }
        Ok(GetStreamAttrsResponse {
            placement,
            attrs: self.state_machine.stream_attrs(&request.stream_id).cloned(),
        })
    }

    pub async fn read_payload_from_plan(
        cold_store: Option<&ColdStoreHandle>,
        cold_index_cache: Option<&Arc<ColdIndexPageCache<ColdStoreColdIndexPageStore>>>,
        stream_id: &BucketStreamId,
        plan: &StreamReadPlan,
    ) -> Result<Vec<u8>, GroupEngineError> {
        let mut payload = Vec::new();
        for segment in &plan.segments {
            match segment {
                StreamReadSegment::Hot(bytes) => payload.extend_from_slice(bytes),
                StreamReadSegment::ColdIndex(segment) => {
                    let Some(cold_store) = cold_store else {
                        return Err(GroupEngineError::stream_with_next_offset(
                            StreamErrorCode::InvalidColdFlush,
                            format!("stream '{stream_id}' read requires object payload store"),
                            Some(plan.next_offset),
                        ));
                    };
                    let Some(cache) = cold_index_cache else {
                        return Err(GroupEngineError::stream_with_next_offset(
                            StreamErrorCode::InvalidColdFlush,
                            format!("stream '{stream_id}' read requires cold index page cache"),
                            Some(plan.next_offset),
                        ));
                    };
                    let objects = cache
                        .object_segments_for_read(stream_id, segment)
                        .await
                        .map_err(|err| GroupEngineError::new(err.to_string()))?;
                    let segment_end = segment
                        .read_start_offset
                        .saturating_add(u64::try_from(segment.len).expect("read len fits u64"));
                    for object in objects {
                        let start = object.start_offset.max(segment.read_start_offset);
                        let end = object.end_offset.min(segment_end);
                        if start >= end {
                            continue;
                        }
                        let bytes = cold_store
                            .read_object_range_for_stream(
                                stream_id,
                                &object,
                                start,
                                usize::try_from(end - start).expect("object read len fits usize"),
                            )
                            .await
                            .map_err(|err| GroupEngineError::new(err.to_string()))?;
                        payload.extend_from_slice(&bytes);
                    }
                }
                StreamReadSegment::Object(segment) => {
                    let Some(cold_store) = cold_store else {
                        return Err(GroupEngineError::stream_with_next_offset(
                            StreamErrorCode::InvalidColdFlush,
                            format!("stream '{stream_id}' read requires object payload store"),
                            Some(plan.next_offset),
                        ));
                    };
                    let bytes = cold_store
                        .read_object_range_for_stream(
                            stream_id,
                            &segment.object,
                            segment.read_start_offset,
                            segment.len,
                        )
                        .await
                        .map_err(|err| GroupEngineError::new(err.to_string()))?;
                    payload.extend_from_slice(&bytes);
                }
            }
        }
        Ok(payload)
    }

    pub(crate) async fn read_own_payload_from_plan(
        &self,
        stream_id: &BucketStreamId,
        plan: &StreamReadPlan,
    ) -> Result<Vec<u8>, GroupEngineError> {
        Self::read_payload_from_plan(
            self.cold_store.as_ref(),
            self.cold_index_cache.as_ref(),
            stream_id,
            plan,
        )
        .await
    }

    pub(crate) async fn bootstrap_updates(
        &self,
        stream_id: &BucketStreamId,
        records: &[StreamMessageRecord],
        content_type: &str,
        now_ms: u64,
    ) -> Result<Vec<BootstrapUpdate>, GroupEngineError> {
        let mut updates = Vec::with_capacity(records.len());
        for record in records {
            let len = usize::try_from(record.end_offset - record.start_offset).map_err(|_| {
                GroupEngineError::stream(
                    StreamErrorCode::InvalidSnapshot,
                    format!(
                        "bootstrap message [{}..{}) for stream '{stream_id}' is too large",
                        record.start_offset, record.end_offset
                    ),
                )
            })?;
            let plan = self
                .state_machine
                .read_plan_at(stream_id, record.start_offset, len, now_ms)
                .map_err(stream_response_error)?;
            let payload = self.read_own_payload_from_plan(stream_id, &plan).await?;
            updates.push(BootstrapUpdate {
                start_offset: record.start_offset,
                next_offset: record.end_offset,
                content_type: content_type.to_owned(),
                payload,
            });
        }
        Ok(updates)
    }

    pub(crate) fn build_snapshot(&self, placement: ShardPlacement) -> GroupSnapshot {
        let stream_snapshot = self.state_machine.snapshot();
        let stream_append_counts = self.stream_append_counts_snapshot(&stream_snapshot);
        GroupSnapshot {
            placement,
            group_commit_index: self.commit_index,
            stream_snapshot,
            stream_append_counts,
        }
    }

    pub(crate) fn stream_append_counts_snapshot(
        &self,
        stream_snapshot: &ursula_stream::StreamSnapshot,
    ) -> Vec<StreamAppendCount> {
        // Only emit append counts for streams actually present in the snapshot.
        // A deleted/expired stream can leave a stale entry in the runtime map;
        // emitting it would make every follower's `install_snapshot` fail the
        // `restore_stream_append_counts` consistency check, so a lagging node
        // could never catch up (and leadership transfer, which catches the
        // target up via a snapshot, could never complete).
        let live: HashSet<&BucketStreamId> = stream_snapshot
            .streams
            .iter()
            .map(|entry| &entry.metadata.stream_id)
            .collect();
        let mut counts = self
            .stream_append_counts
            .iter()
            .filter(|(stream_id, _)| live.contains(stream_id))
            .map(|(stream_id, append_count)| StreamAppendCount {
                stream_id: stream_id.clone(),
                append_count: *append_count,
            })
            .collect::<Vec<_>>();
        counts.sort_by(|left, right| compare_stream_ids(&left.stream_id, &right.stream_id));
        counts
    }

    pub fn stream_tail_offset(&self, stream_id: &BucketStreamId) -> Option<u64> {
        self.state_machine
            .head(stream_id)
            .map(|metadata| metadata.tail_offset)
    }

    pub(crate) fn install_snapshot_inner(
        &mut self,
        snapshot: GroupSnapshot,
    ) -> Result<(), GroupEngineError> {
        let GroupSnapshot {
            placement: _,
            group_commit_index,
            stream_snapshot,
            stream_append_counts,
        } = snapshot;
        self.install_snapshot_parts(group_commit_index, stream_snapshot, stream_append_counts)
    }

    pub(crate) fn install_snapshot_parts(
        &mut self,
        group_commit_index: u64,
        stream_snapshot: StreamSnapshot,
        stream_append_counts: Vec<StreamAppendCount>,
    ) -> Result<(), GroupEngineError> {
        let stream_ids = stream_snapshot
            .streams
            .iter()
            .map(|entry| entry.metadata.stream_id.clone())
            .collect::<HashSet<_>>();
        let state_machine = StreamStateMachine::restore(stream_snapshot)
            .map_err(|err| GroupEngineError::new(format!("restore stream snapshot: {err}")))?;
        let stream_append_counts = restore_stream_append_counts(stream_append_counts, &stream_ids)?;

        self.commit_index = group_commit_index;
        self.state_machine = state_machine;
        self.stream_append_counts = stream_append_counts;
        Ok(())
    }
}

impl GroupEngine for InMemoryGroupEngine {
    fn create_stream<'a>(
        &'a mut self,
        request: CreateStreamRequest,
        placement: ShardPlacement,
        admission: ColdWriteAdmission,
    ) -> GroupCreateStreamFuture<'a> {
        if admission.is_enabled() {
            return Box::pin(async move {
                self.create_stream_with_admission_inner(request, placement, admission)
            });
        }
        let command = GroupWriteCommand::from(request);
        Box::pin(async move {
            match self.apply_committed_write(command, placement)? {
                GroupWriteResponse::CreateStream(response) => Ok(response),
                other => Err(GroupEngineError::new(format!(
                    "unexpected create stream write response: {other:?}"
                ))),
            }
        })
    }

    fn create_stream_external<'a>(
        &'a mut self,
        request: CreateStreamExternalRequest,
        placement: ShardPlacement,
    ) -> GroupCreateStreamFuture<'a> {
        Box::pin(async move {
            if let Some(cold_store) = self.cold_store.as_ref() {
                let store = ColdStoreColdIndexPageStore::new(cold_store.clone());
                write_external_segment_index_pages(
                    &store,
                    &request.stream_id,
                    0,
                    &request.initial_payload,
                )
                .await
                .map_err(|err| GroupEngineError::new(err.to_string()))?;
            }
            let command = GroupWriteCommand::from(request);
            match self.apply_committed_write(command, placement)? {
                GroupWriteResponse::CreateStream(response) => Ok(response),
                other => Err(GroupEngineError::new(format!(
                    "unexpected external create stream write response: {other:?}"
                ))),
            }
        })
    }

    fn read_stream<'a>(
        &'a mut self,
        request: ReadStreamRequest,
        placement: ShardPlacement,
    ) -> GroupReadStreamFuture<'a> {
        Box::pin(async move {
            self.read_stream_parts(request, placement)
                .await?
                .into_response()
                .await
        })
    }

    fn read_stream_parts<'a>(
        &'a mut self,
        request: ReadStreamRequest,
        placement: ShardPlacement,
    ) -> GroupReadStreamPartsFuture<'a> {
        Box::pin(async move {
            let stream_id = request.stream_id.clone();
            let plan = self.read_stream_plan(&request, placement)?;
            Ok(GroupReadStreamParts::from_plan(
                placement,
                stream_id,
                plan,
                self.cold_store(),
                self.cold_index_cache.clone(),
            ))
        })
    }

    fn publish_snapshot<'a>(
        &'a mut self,
        request: PublishSnapshotRequest,
        placement: ShardPlacement,
    ) -> GroupPublishSnapshotFuture<'a> {
        Box::pin(async move {
            self.ensure_stream_access(&request.stream_id, request.now_ms, false, placement)?;
            let command = GroupWriteCommand::from(request);
            match self.apply_committed_write(command, placement)? {
                GroupWriteResponse::PublishSnapshot(response) => Ok(response),
                other => Err(GroupEngineError::new(format!(
                    "unexpected publish snapshot write response: {other:?}"
                ))),
            }
        })
    }

    fn read_snapshot<'a>(
        &'a mut self,
        request: ReadSnapshotRequest,
        placement: ShardPlacement,
    ) -> GroupReadSnapshotFuture<'a> {
        Box::pin(async move {
            self.ensure_stream_access(&request.stream_id, request.now_ms, true, placement)?;
            let snapshot = match request.snapshot_offset {
                Some(offset) => self
                    .state_machine
                    .read_snapshot(&request.stream_id, offset)
                    .map_err(stream_response_error)?,
                None => self
                    .state_machine
                    .latest_snapshot(&request.stream_id)
                    .map_err(stream_response_error)?
                    .ok_or_else(|| {
                        GroupEngineError::stream(
                            StreamErrorCode::SnapshotNotFound,
                            format!("stream '{}' has no visible snapshot", request.stream_id),
                        )
                    })?,
            };
            let tail_offset = self
                .state_machine
                .head_at(&request.stream_id, request.now_ms)
                .map(|metadata| metadata.tail_offset)
                .unwrap_or(snapshot.offset);
            Ok(ReadSnapshotResponse {
                placement,
                snapshot_offset: snapshot.offset,
                next_offset: snapshot.offset,
                content_type: snapshot.content_type,
                payload: snapshot.payload,
                up_to_date: snapshot.offset == tail_offset,
                record_range: self
                    .state_machine
                    .record_range(&request.stream_id)
                    .map_err(|err| GroupEngineError::new(format!("record range: {err:?}")))?,
            })
        })
    }

    fn delete_snapshot<'a>(
        &'a mut self,
        request: DeleteSnapshotRequest,
        placement: ShardPlacement,
    ) -> GroupDeleteSnapshotFuture<'a> {
        Box::pin(async move {
            self.ensure_stream_access(&request.stream_id, request.now_ms, false, placement)?;
            match self
                .state_machine
                .delete_snapshot(&request.stream_id, request.snapshot_offset)
            {
                StreamResponse::Error {
                    code,
                    message,
                    next_offset,
                    context,
                } => Err(GroupEngineError::stream_with_context(
                    code,
                    message,
                    next_offset,
                    context,
                )),
                other => Err(GroupEngineError::new(format!(
                    "unexpected delete snapshot response: {other:?}"
                ))),
            }
        })
    }

    fn bootstrap_stream<'a>(
        &'a mut self,
        request: BootstrapStreamRequest,
        placement: ShardPlacement,
    ) -> GroupBootstrapStreamFuture<'a> {
        Box::pin(async move {
            self.ensure_stream_access(&request.stream_id, request.now_ms, true, placement)?;
            let plan = self
                .state_machine
                .bootstrap_plan(&request.stream_id)
                .map_err(stream_response_error)?;
            let snapshot_offset = plan.snapshot.as_ref().map(|snapshot| snapshot.offset);
            let snapshot_content_type = plan
                .snapshot
                .as_ref()
                .map(|snapshot| snapshot.content_type.clone())
                .unwrap_or_else(|| DEFAULT_CONTENT_TYPE.to_owned());
            let snapshot_payload = plan
                .snapshot
                .as_ref()
                .map(|snapshot| snapshot.payload.clone())
                .unwrap_or_default();
            let updates = self
                .bootstrap_updates(
                    &request.stream_id,
                    &plan.updates,
                    &plan.content_type,
                    request.now_ms,
                )
                .await?;
            Ok(BootstrapStreamResponse {
                placement,
                snapshot_offset,
                snapshot_content_type,
                snapshot_payload,
                updates,
                next_offset: plan.next_offset,
                up_to_date: plan.up_to_date,
                closed: plan.closed,
                record_range: self
                    .state_machine
                    .record_range(&request.stream_id)
                    .map_err(|err| GroupEngineError::new(format!("record range: {err:?}")))?,
            })
        })
    }

    fn touch_stream_access<'a>(
        &'a mut self,
        stream_id: BucketStreamId,
        now_ms: u64,
        renew_ttl: bool,
        placement: ShardPlacement,
    ) -> GroupTouchStreamAccessFuture<'a> {
        Box::pin(async move { self.apply_access_command(stream_id, now_ms, renew_ttl, placement) })
    }

    fn head_stream<'a>(
        &'a mut self,
        request: HeadStreamRequest,
        placement: ShardPlacement,
    ) -> GroupHeadStreamFuture<'a> {
        Box::pin(async move {
            self.ensure_stream_access(&request.stream_id, request.now_ms, false, placement)?;
            self.head_stream_after_access(&request, placement)
        })
    }

    fn get_stream_attrs<'a>(
        &'a mut self,
        request: GetStreamAttrsRequest,
        placement: ShardPlacement,
    ) -> GroupGetStreamAttrsFuture<'a> {
        Box::pin(async move {
            self.ensure_stream_access(&request.stream_id, request.now_ms, false, placement)?;
            self.get_stream_attrs_after_access(&request, placement)
        })
    }

    fn update_stream_attrs<'a>(
        &'a mut self,
        request: UpdateStreamAttrsRequest,
        placement: ShardPlacement,
    ) -> GroupUpdateStreamAttrsFuture<'a> {
        Box::pin(async move {
            match self.apply_committed_write(GroupWriteCommand::from(request), placement)? {
                GroupWriteResponse::UpdateStreamAttrs(response) => Ok(response),
                other => Err(GroupEngineError::new(format!(
                    "unexpected update stream attrs write response: {other:?}"
                ))),
            }
        })
    }

    fn close_stream<'a>(
        &'a mut self,
        request: CloseStreamRequest,
        placement: ShardPlacement,
    ) -> GroupCloseStreamFuture<'a> {
        Box::pin(async move {
            self.ensure_stream_access(&request.stream_id, request.now_ms, false, placement)?;
            let command = GroupWriteCommand::from(request);
            match self.apply_committed_write(command, placement)? {
                GroupWriteResponse::CloseStream(response) => Ok(response),
                other => Err(GroupEngineError::new(format!(
                    "unexpected close stream write response: {other:?}"
                ))),
            }
        })
    }

    fn delete_stream<'a>(
        &'a mut self,
        request: DeleteStreamRequest,
        placement: ShardPlacement,
    ) -> GroupDeleteStreamFuture<'a> {
        let command = GroupWriteCommand::from(request);
        Box::pin(async move {
            match self.apply_committed_write(command, placement)? {
                GroupWriteResponse::DeleteStream(response) => Ok(response),
                other => Err(GroupEngineError::new(format!(
                    "unexpected delete stream write response: {other:?}"
                ))),
            }
        })
    }

    fn ack_cold_gc<'a>(
        &'a mut self,
        up_to_seq: u64,
        placement: ShardPlacement,
    ) -> GroupAckColdGcFuture<'a> {
        Box::pin(async move {
            match self.apply_committed_write(
                GroupWriteCommand::Stream(StreamCommand::AckColdGc { up_to_seq }),
                placement,
            )? {
                GroupWriteResponse::AckColdGc(response) => Ok(response),
                other => Err(GroupEngineError::new(format!(
                    "unexpected ack cold gc write response: {other:?}"
                ))),
            }
        })
    }

    fn plan_cold_gc<'a>(
        &'a mut self,
        max: usize,
        _placement: ShardPlacement,
    ) -> GroupPlanColdGcFuture<'a> {
        let entries = self.state_machine.pending_cold_gc_batch(max);
        Box::pin(async move { Ok(entries) })
    }

    fn append<'a>(
        &'a mut self,
        request: AppendRequest,
        placement: ShardPlacement,
        admission: ColdWriteAdmission,
    ) -> GroupAppendFuture<'a> {
        if admission.is_enabled() {
            return Box::pin(async move {
                self.append_with_admission_inner(request, placement, admission)
            });
        }
        Box::pin(async move {
            self.ensure_stream_access(&request.stream_id, request.now_ms, false, placement)?;
            let command = GroupWriteCommand::from(request);
            match self.apply_committed_write(command, placement)? {
                GroupWriteResponse::Append(response) => Ok(response),
                other => Err(GroupEngineError::new(format!(
                    "unexpected append write response: {other:?}"
                ))),
            }
        })
    }

    fn append_external<'a>(
        &'a mut self,
        request: AppendExternalRequest,
        placement: ShardPlacement,
    ) -> GroupAppendFuture<'a> {
        Box::pin(async move {
            self.ensure_stream_access(&request.stream_id, request.now_ms, false, placement)?;
            if let Some(cold_store) = self.cold_store.as_ref() {
                let start_offset = self
                    .state_machine
                    .head(&request.stream_id)
                    .map(|metadata| metadata.tail_offset)
                    .ok_or_else(|| {
                        GroupEngineError::stream(
                            ursula_stream::StreamErrorCode::StreamNotFound,
                            format!("stream '{}' does not exist", request.stream_id),
                        )
                    })?;
                let store = ColdStoreColdIndexPageStore::new(cold_store.clone());
                write_external_segment_index_pages(
                    &store,
                    &request.stream_id,
                    start_offset,
                    &request.payload,
                )
                .await
                .map_err(|err| GroupEngineError::new(err.to_string()))?;
            }
            let command = GroupWriteCommand::from(request);
            match self.apply_committed_write(command, placement)? {
                GroupWriteResponse::Append(response) => Ok(response),
                other => Err(GroupEngineError::new(format!(
                    "unexpected external append write response: {other:?}"
                ))),
            }
        })
    }

    fn append_batch<'a>(
        &'a mut self,
        request: AppendBatchRequest,
        placement: ShardPlacement,
        admission: ColdWriteAdmission,
    ) -> GroupAppendBatchFuture<'a> {
        if admission.is_enabled() {
            return Box::pin(async move {
                self.append_batch_with_admission_inner(request, placement, admission)
            });
        }
        Box::pin(async move {
            self.ensure_stream_access(&request.stream_id, request.now_ms, false, placement)?;
            let command = GroupWriteCommand::from(request);
            match self.apply_committed_write(command, placement)? {
                GroupWriteResponse::AppendBatch(response) => Ok(response),
                other => Err(GroupEngineError::new(format!(
                    "unexpected append batch write response: {other:?}"
                ))),
            }
        })
    }

    fn flush_cold<'a>(
        &'a mut self,
        request: FlushColdRequest,
        placement: ShardPlacement,
    ) -> GroupFlushColdFuture<'a> {
        Box::pin(async move {
            let mut index_rollback = None;
            if let Some(cold_store) = self.cold_store.as_ref() {
                let store = ColdStoreColdIndexPageStore::new(cold_store.clone());
                let rollback = write_cold_chunk_index_pages_with_rollback(
                    &store,
                    &request.stream_id,
                    &request.chunk,
                )
                .await
                .map_err(|err| GroupEngineError::new(err.to_string()))?;
                index_rollback = Some((store, rollback));
            }
            let command = GroupWriteCommand::from(request);
            match self.apply_committed_write(command, placement) {
                Ok(GroupWriteResponse::FlushCold(response)) => Ok(response),
                Ok(other) => {
                    if let Some((store, rollback)) = index_rollback {
                        rollback_cold_index_pages(&store, rollback)
                            .await
                            .map_err(|err| GroupEngineError::new(err.to_string()))?;
                    }
                    Err(GroupEngineError::new(format!(
                        "unexpected flush cold write response: {other:?}"
                    )))
                }
                Err(err) => {
                    if let Some((store, rollback)) = index_rollback {
                        rollback_cold_index_pages(&store, rollback).await.map_err(
                            |rollback_err| {
                                GroupEngineError::new(format!(
                                    "rollback cold index after flush failure: {rollback_err}"
                                ))
                            },
                        )?;
                    }
                    Err(err)
                }
            }
        })
    }

    fn plan_cold_flush<'a>(
        &'a mut self,
        request: PlanColdFlushRequest,
        _placement: ShardPlacement,
    ) -> GroupPlanColdFlushFuture<'a> {
        Box::pin(async move {
            self.state_machine
                .plan_cold_flush(
                    &request.stream_id,
                    request.min_hot_bytes,
                    request.max_flush_bytes,
                )
                .map_err(stream_response_error)
        })
    }

    fn plan_next_cold_flush_batch<'a>(
        &'a mut self,
        request: PlanGroupColdFlushRequest,
        _placement: ShardPlacement,
        max_candidates: usize,
    ) -> GroupPlanNextColdFlushBatchFuture<'a> {
        Box::pin(async move {
            self.state_machine
                .plan_next_cold_flush_batch(
                    request.min_hot_bytes,
                    request.max_flush_bytes,
                    max_candidates,
                )
                .map_err(stream_response_error)
        })
    }

    fn cold_hot_backlog<'a>(
        &'a mut self,
        stream_id: BucketStreamId,
        _placement: ShardPlacement,
    ) -> GroupColdHotBacklogFuture<'a> {
        Box::pin(async move { self.cold_hot_backlog_for(stream_id) })
    }

    fn snapshot<'a>(&'a mut self, placement: ShardPlacement) -> GroupSnapshotFuture<'a> {
        Box::pin(async move { Ok(self.build_snapshot(placement)) })
    }

    fn install_snapshot<'a>(
        &'a mut self,
        snapshot: GroupSnapshot,
    ) -> GroupInstallSnapshotFuture<'a> {
        Box::pin(async move { self.install_snapshot_inner(snapshot) })
    }
}

#[derive(Debug, Clone, Default)]
pub struct InMemoryGroupEngineFactory {
    cold_store: Option<ColdStoreHandle>,
}

impl InMemoryGroupEngineFactory {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_cold_store(cold_store: Option<ColdStoreHandle>) -> Self {
        Self { cold_store }
    }
}

impl GroupEngineFactory for InMemoryGroupEngineFactory {
    fn create<'a>(
        &'a self,
        _placement: ShardPlacement,
        _metrics: GroupEngineMetrics,
    ) -> GroupEngineCreateFuture<'a> {
        Box::pin(async move {
            let mut engine = InMemoryGroupEngine::default();
            engine.set_cold_store(self.cold_store.clone());
            let engine: Box<dyn GroupEngine> = Box::new(engine);
            Ok(engine)
        })
    }
}

pub(crate) fn compare_stream_ids(
    left: &BucketStreamId,
    right: &BucketStreamId,
) -> std::cmp::Ordering {
    left.bucket_id
        .cmp(&right.bucket_id)
        .then_with(|| left.stream_id.cmp(&right.stream_id))
}
pub(crate) fn ensure_bucket_exists(
    state_machine: &mut StreamStateMachine,
    stream_id: &BucketStreamId,
) -> Result<(), GroupEngineError> {
    if state_machine.bucket_exists(&stream_id.bucket_id) {
        return Ok(());
    }

    match state_machine.apply(StreamCommand::CreateBucket {
        bucket_id: stream_id.bucket_id.clone(),
    }) {
        StreamResponse::BucketCreated { .. } | StreamResponse::BucketAlreadyExists { .. } => Ok(()),
        StreamResponse::Error {
            code,
            message,
            next_offset,
            context,
        } => Err(GroupEngineError::stream_with_context(
            code,
            message,
            next_offset,
            context,
        )),
        other => Err(GroupEngineError::new(format!(
            "unexpected create bucket response: {other:?}"
        ))),
    }
}

/// Stream id a command targets, if any (bucket and GC commands have none).
fn command_stream_id(command: &StreamCommand) -> Option<BucketStreamId> {
    match command {
        StreamCommand::CreateBucket { .. }
        | StreamCommand::DeleteBucket { .. }
        | StreamCommand::AckColdGc { .. } => None,
        StreamCommand::CreateStream { stream_id, .. }
        | StreamCommand::CreateExternal { stream_id, .. }
        | StreamCommand::Append { stream_id, .. }
        | StreamCommand::AppendExternal { stream_id, .. }
        | StreamCommand::AppendBatch { stream_id, .. }
        | StreamCommand::PublishSnapshot { stream_id, .. }
        | StreamCommand::TouchStreamAccess { stream_id, .. }
        | StreamCommand::UpdateStreamAttrs { stream_id, .. }
        | StreamCommand::FlushCold { stream_id, .. }
        | StreamCommand::Close { stream_id, .. }
        | StreamCommand::DeleteStream { stream_id } => Some(stream_id.clone()),
    }
}

fn command_producer(command: &StreamCommand) -> Option<ProducerRequest> {
    match command {
        StreamCommand::Close { producer, .. } => producer.clone(),
        _ => None,
    }
}

fn require_response_stream_id(
    stream_id: Option<BucketStreamId>,
    response: &str,
) -> Result<BucketStreamId, GroupEngineError> {
    stream_id.ok_or_else(|| {
        GroupEngineError::new(format!(
            "{response} response for a command without a stream id"
        ))
    })
}

pub(crate) fn stream_response_error(response: StreamResponse) -> GroupEngineError {
    match response {
        StreamResponse::Error {
            code,
            message,
            next_offset,
            context,
        } => GroupEngineError::stream_with_context(code, message, next_offset, context),
        other => GroupEngineError::new(format!("unexpected stream response error: {other:?}")),
    }
}

pub(crate) fn restore_stream_append_counts(
    counts: Vec<StreamAppendCount>,
    snapshot_stream_ids: &HashSet<BucketStreamId>,
) -> Result<HashMap<BucketStreamId, u64>, GroupEngineError> {
    let mut restored = HashMap::with_capacity(counts.len());
    for count in counts {
        if !snapshot_stream_ids.contains(&count.stream_id) {
            return Err(GroupEngineError::new(format!(
                "append count references missing snapshot stream '{}'",
                count.stream_id
            )));
        }
        if restored
            .insert(count.stream_id.clone(), count.append_count)
            .is_some()
        {
            return Err(GroupEngineError::new(format!(
                "snapshot contains duplicate append count for stream '{}'",
                count.stream_id
            )));
        }
    }
    Ok(restored)
}