millipede_storage_fs/

queue.rs

//! Durable file-system request queue.

use crate::layout::{is_temporary_file, temporary_suffix};
use millipede_core::{
    request::{Method, Request, RequestBody, RequestId, RequestState, UserData},
    storage::{
        AddOptions, BatchAddHandle, Lease, LeaseId, ProcessedRequest, QueueOpInfo, ReclaimOptions,
        RequestQueue, RequestSource, StorageError, StorageResult,
    },
};
use serde::{Deserialize, Serialize};
use std::{
    collections::{BTreeMap, HashMap},
    fmt,
    path::{Path, PathBuf},
    sync::Arc,
    time::{Duration, Instant},
};
use tokio::sync::{Mutex, RwLock};

const LEASE_TTL: Duration = Duration::from_secs(180);
const STATE_VERSION: u8 = 1;

/// A durable FIFO request queue backed by atomic file replacements.
///
/// Request files, rather than `state.json`, are the source of truth. Leases are
/// deliberately process-local, so requests leased by a process that exits are
/// pending again when a new client rescans the queue. Writes use a temporary
/// file and same-directory rename, but do not call `fsync`; a machine-level
/// power loss can therefore lose writes that the operating system had not yet
/// flushed.
pub struct FsRequestQueue {
    name: String,
    path: PathBuf,
    requests_path: PathBuf,
    operations: Arc<RwLock<()>>,
    state: Mutex<QueueState>,
}

struct QueueState {
    pending: BTreeMap<i64, RequestId>,
    dedup: HashMap<String, (RequestId, bool)>,
    leases: HashMap<LeaseId, (RequestId, Instant)>,
    requests: HashMap<RequestId, StoredRequest>,
    handled_count: u64,
    next_order: i64,
    next_forefront_order: i64,
    next_lease_id: u64,
}

struct StoredRequest {
    request: Request,
    order_no: Option<i64>,
}

#[derive(Serialize)]
struct RequestEnvelope {
    id: String,
    url: String,
    #[serde(rename = "uniqueKey")]
    unique_key: String,
    method: String,
    #[serde(rename = "retryCount")]
    retry_count: u32,
    #[serde(rename = "orderNo")]
    order_no: Option<i64>,
    json: Request,
}

#[derive(Deserialize)]
struct ReadEnvelope {
    #[serde(rename = "orderNo")]
    order_no: Option<i64>,
    json: Request,
}

#[derive(Deserialize)]
struct CrawleeRequestEnvelope {
    id: RequestId,
    url: String,
    #[serde(rename = "uniqueKey")]
    unique_key: String,
    #[serde(default = "default_method")]
    method: String,
    #[serde(default)]
    headers: HashMap<String, String>,
    #[serde(default)]
    payload: Option<String>,
    #[serde(default, rename = "userData")]
    user_data: UserData,
    #[serde(default, rename = "retryCount")]
    retry_count: u32,
    #[serde(default, rename = "sessionRotationCount")]
    session_rotation_count: u32,
    #[serde(default, rename = "maxRetries")]
    max_retries: Option<u32>,
    #[serde(default, rename = "noRetry")]
    no_retry: bool,
    #[serde(default, rename = "errorMessages")]
    error_messages: Vec<String>,
    #[serde(default, rename = "loadedUrl")]
    loaded_url: Option<String>,
    #[serde(default, rename = "handledAt", with = "time::serde::rfc3339::option")]
    handled_at: Option<time::OffsetDateTime>,
    #[serde(default, rename = "crawlDepth")]
    crawl_depth: u32,
    #[serde(default, rename = "skipNavigation")]
    skip_navigation: bool,
    #[serde(rename = "orderNo")]
    order_no: Option<i64>,
}

fn default_method() -> String {
    "GET".to_owned()
}

impl CrawleeRequestEnvelope {
    fn into_request(self) -> Result<(Request, Option<i64>), serde_json::Error> {
        let method = Method::from_bytes(self.method.as_bytes()).map_err(json_error)?;
        let mut builder = Request::builder()
            .url(self.url)
            .method(method)
            .unique_key(self.unique_key.clone())
            .user_data(self.user_data.clone())
            .crawl_depth(self.crawl_depth)
            .no_retry(self.no_retry)
            .skip_navigation(self.skip_navigation);
        if let Some(payload) = self.payload {
            builder = builder.body(RequestBody::Bytes(payload.into_bytes()));
        }
        if let Some(max_retries) = self.max_retries {
            builder = builder.max_retries(max_retries);
        }
        for (name, value) in self.headers {
            builder = builder.header(&name, &value);
        }
        if let Some(label) = self.user_data.get("label").and_then(|value| value.as_str()) {
            builder = builder.label(label.to_owned());
        }
        let mut request = builder.build().map_err(json_error)?;
        request.id = self.id;
        request.unique_key = self.unique_key;
        request.retry_count = self.retry_count;
        request.session_rotation_count = self.session_rotation_count;
        request.error_messages = self.error_messages;
        request.handled_at = self.handled_at;
        request.state = if self.order_no.is_some() {
            RequestState::Unprocessed
        } else {
            RequestState::Done
        };
        request.loaded_url = self
            .loaded_url
            .map(|url| url.parse().map_err(json_error))
            .transpose()?;
        Ok((request, self.order_no))
    }
}

fn json_error(error: impl fmt::Display) -> serde_json::Error {
    <serde_json::Error as serde::de::Error>::custom(error)
}

#[derive(Serialize, Deserialize)]
struct QueueCache {
    version: u8,
    #[serde(rename = "handledRequestCount")]
    handled_request_count: u64,
    #[serde(rename = "pendingRequestCount")]
    pending_request_count: u64,
    next_order: i64,
    next_forefront_order: i64,
}

impl fmt::Debug for FsRequestQueue {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("FsRequestQueue")
            .field("name", &self.name)
            .field("path", &self.path)
            .finish_non_exhaustive()
    }
}

impl FsRequestQueue {
    pub(crate) async fn open(
        name: String,
        path: PathBuf,
        operations: Arc<RwLock<()>>,
    ) -> StorageResult<Self> {
        let requests_path = path.join("requests");
        tokio::fs::create_dir_all(&requests_path).await?;
        let state = scan_queue(&name, &path, &requests_path).await?;
        let queue = Self {
            name,
            path,
            requests_path,
            operations,
            state: Mutex::new(state),
        };
        {
            let state = queue.state.lock().await;
            queue.write_cache_best_effort(&state).await;
        }
        Ok(queue)
    }

    pub(crate) async fn ensure_layout(&self) -> StorageResult<()> {
        tokio::fs::create_dir_all(&self.requests_path).await?;
        Ok(())
    }

    pub(crate) async fn reset(&self) {
        let mut state = self.state.lock().await;
        *state = QueueState::empty();
    }

    async fn write_request(&self, request: &Request, order_no: Option<i64>) -> StorageResult<()> {
        tokio::fs::create_dir_all(&self.requests_path).await?;
        let envelope = RequestEnvelope {
            id: request.id.to_string(),
            url: request.url.to_string(),
            unique_key: request.unique_key.clone(),
            method: request.method.as_str().to_owned(),
            retry_count: request.retry_count,
            order_no,
            json: request.clone(),
        };
        let destination = self
            .requests_path
            .join(format!("{}.json", request.id.as_str()));
        atomic_write(&destination, &serde_json::to_vec_pretty(&envelope)?).await
    }

    async fn write_cache_best_effort(&self, state: &QueueState) {
        let cache = QueueCache {
            version: STATE_VERSION,
            handled_request_count: state.handled_count,
            pending_request_count: state
                .requests
                .values()
                .filter(|request| request.order_no.is_some())
                .count() as u64,
            next_order: state.next_order,
            next_forefront_order: state.next_forefront_order,
        };
        let result = match serde_json::to_vec_pretty(&cache) {
            Ok(bytes) => atomic_write(&self.path.join("state.json"), &bytes).await,
            Err(error) => Err(error.into()),
        };
        if let Err(error) = result {
            tracing::warn!(
                queue = %self.name,
                %error,
                "failed to refresh request queue state cache"
            );
        }
    }

    fn lease_request_id(state: &QueueState, lease_id: &LeaseId) -> StorageResult<RequestId> {
        state
            .leases
            .get(lease_id)
            .map(|(request_id, _)| request_id.clone())
            .ok_or_else(|| StorageError::LeaseNotFound {
                lease_id: lease_id.clone(),
            })
    }

    fn validate_unique_key(
        state: &QueueState,
        request_id: &RequestId,
        request: &Request,
    ) -> StorageResult<()> {
        if let Some((known_id, _)) = state.dedup.get(&request.unique_key) {
            if known_id != request_id {
                return Err(StorageError::Backend(anyhow::anyhow!(
                    "request lease changed unique_key to an existing queue key"
                )));
            }
        }
        Ok(())
    }

    fn replace_stored_request(
        state: &mut QueueState,
        request_id: &RequestId,
        request: Request,
        order_no: Option<i64>,
        handled: bool,
    ) {
        if let Some(previous) = state.requests.get(request_id) {
            if previous.request.unique_key != request.unique_key {
                state.dedup.remove(&previous.request.unique_key);
            }
        }
        state
            .dedup
            .insert(request.unique_key.clone(), (request_id.clone(), handled));
        state
            .requests
            .insert(request_id.clone(), StoredRequest { request, order_no });
    }

    fn reclaim_expired(state: &mut QueueState) {
        let now = Instant::now();
        let expired: Vec<_> = state
            .leases
            .iter()
            .filter(|(_, (_, expires_at))| *expires_at <= now)
            .map(|(lease_id, _)| lease_id.clone())
            .collect();
        for lease_id in expired {
            if let Some((request_id, _)) = state.leases.remove(&lease_id) {
                if let Some(order_no) = state
                    .requests
                    .get(&request_id)
                    .and_then(|request| request.order_no)
                {
                    state.pending.insert(order_no, request_id);
                }
            }
        }
    }

    async fn add_locked(
        &self,
        state: &mut QueueState,
        request: Request,
        opts: &AddOptions,
    ) -> StorageResult<QueueOpInfo> {
        let unique_key = request.unique_key.clone();
        if let Some((request_id, handled)) = state.dedup.get(&unique_key) {
            return Ok(ProcessedRequest {
                request_id: request_id.clone(),
                unique_key,
                was_already_present: true,
                was_already_handled: *handled,
            });
        }

        let request_id = request.id.clone();
        let order_no = state.take_order(opts.forefront)?;
        self.write_request(&request, Some(order_no)).await?;
        state.pending.insert(order_no, request_id.clone());
        state
            .dedup
            .insert(unique_key.clone(), (request_id.clone(), false));
        state.requests.insert(
            request_id.clone(),
            StoredRequest {
                request,
                order_no: Some(order_no),
            },
        );
        self.write_cache_best_effort(state).await;
        Ok(ProcessedRequest {
            request_id,
            unique_key,
            was_already_present: false,
            was_already_handled: false,
        })
    }

    async fn requeue(
        &self,
        lease: Lease,
        forefront: bool,
        increment_retry: bool,
    ) -> StorageResult<()> {
        let _operation = self.operations.read().await;
        let mut state = self.state.lock().await;
        let request_id = Self::lease_request_id(&state, &lease.lease_id)?;
        let mut request = lease.request;
        Self::validate_unique_key(&state, &request_id, &request)?;
        if increment_retry {
            request.retry_count += 1;
        }
        let order_no = state.take_order(forefront)?;
        self.write_request(&request, Some(order_no)).await?;
        state.leases.remove(&lease.lease_id);
        state.pending.insert(order_no, request_id.clone());
        Self::replace_stored_request(&mut state, &request_id, request, Some(order_no), false);
        self.write_cache_best_effort(&state).await;
        Ok(())
    }
}

impl QueueState {
    fn empty() -> Self {
        Self {
            pending: BTreeMap::new(),
            dedup: HashMap::new(),
            leases: HashMap::new(),
            requests: HashMap::new(),
            handled_count: 0,
            next_order: 1,
            next_forefront_order: -1,
            next_lease_id: 0,
        }
    }

    fn take_order(&mut self, forefront: bool) -> StorageResult<i64> {
        if forefront {
            let order_no = self.next_forefront_order;
            self.next_forefront_order = order_no.checked_sub(1).ok_or_else(|| {
                StorageError::Backend(anyhow::anyhow!("request queue forefront order exhausted"))
            })?;
            Ok(order_no)
        } else {
            let order_no = self.next_order;
            self.next_order = order_no.checked_add(1).ok_or_else(|| {
                StorageError::Backend(anyhow::anyhow!("request queue order exhausted"))
            })?;
            Ok(order_no)
        }
    }
}

#[async_trait::async_trait]
impl RequestQueue for FsRequestQueue {
    async fn add(&self, request: Request, opts: AddOptions) -> StorageResult<QueueOpInfo> {
        let _operation = self.operations.read().await;
        let mut state = self.state.lock().await;
        self.add_locked(&mut state, request, &opts).await
    }

    async fn add_batch(
        &self,
        requests: Vec<RequestSource>,
        opts: AddOptions,
    ) -> StorageResult<BatchAddHandle> {
        let _operation = self.operations.read().await;
        let mut state = self.state.lock().await;
        let mut processed = Vec::with_capacity(requests.len());
        for source in requests {
            let request = match source {
                RequestSource::Request(request) => request,
                _ => {
                    return Err(StorageError::Backend(anyhow::anyhow!(
                        "unsupported request source"
                    )));
                }
            };
            processed.push(self.add_locked(&mut state, request, &opts).await?);
        }
        Ok(BatchAddHandle::ready(processed))
    }

    async fn fetch_next(&self) -> StorageResult<Option<Lease>> {
        let _operation = self.operations.read().await;
        let mut state = self.state.lock().await;
        Self::reclaim_expired(&mut state);
        let Some((_, request_id)) = state.pending.pop_first() else {
            return Ok(None);
        };
        let request = state
            .requests
            .get(&request_id)
            .expect("pending request has a stored request")
            .request
            .clone();
        let raw_lease_id = state.next_lease_id;
        state.next_lease_id = raw_lease_id.checked_add(1).ok_or_else(|| {
            StorageError::Backend(anyhow::anyhow!("request queue lease identifiers exhausted"))
        })?;
        let lease_id = LeaseId::new(raw_lease_id);
        let expires_at = Instant::now() + LEASE_TTL;
        state
            .leases
            .insert(lease_id.clone(), (request_id, expires_at));
        Ok(Some(Lease {
            request,
            lease_id,
            expires_at,
        }))
    }

    async fn mark_handled(&self, lease: Lease) -> StorageResult<()> {
        let _operation = self.operations.read().await;
        let mut state = self.state.lock().await;
        let request_id = Self::lease_request_id(&state, &lease.lease_id)?;
        Self::validate_unique_key(&state, &request_id, &lease.request)?;
        self.write_request(&lease.request, None).await?;
        state.leases.remove(&lease.lease_id);
        state.handled_count = state.handled_count.saturating_add(1);
        Self::replace_stored_request(&mut state, &request_id, lease.request, None, true);
        self.write_cache_best_effort(&state).await;
        Ok(())
    }

    async fn reclaim(&self, lease: Lease, opts: ReclaimOptions) -> StorageResult<()> {
        self.requeue(lease, opts.forefront, opts.increment_retry)
            .await
    }

    async fn renew(&self, lease_id: &LeaseId, extend_by: Duration) -> StorageResult<()> {
        let _operation = self.operations.read().await;
        let mut state = self.state.lock().await;
        let expired = state
            .leases
            .get(lease_id)
            .is_some_and(|(_, expires_at)| *expires_at <= Instant::now());
        if expired {
            if let Some((request_id, _)) = state.leases.remove(lease_id) {
                if let Some(order_no) = state
                    .requests
                    .get(&request_id)
                    .and_then(|request| request.order_no)
                {
                    state.pending.insert(order_no, request_id);
                }
            }
            return Err(StorageError::LeaseNotFound {
                lease_id: lease_id.clone(),
            });
        }
        let (_, expires_at) =
            state
                .leases
                .get_mut(lease_id)
                .ok_or_else(|| StorageError::LeaseNotFound {
                    lease_id: lease_id.clone(),
                })?;
        *expires_at += extend_by;
        Ok(())
    }

    async fn abandon(&self, lease: Lease) -> StorageResult<()> {
        self.requeue(lease, true, false).await
    }

    async fn is_empty(&self) -> StorageResult<bool> {
        let _operation = self.operations.read().await;
        Ok(self.state.lock().await.pending.is_empty())
    }

    async fn is_finished(&self) -> StorageResult<bool> {
        let _operation = self.operations.read().await;
        let state = self.state.lock().await;
        Ok(state.pending.is_empty() && state.leases.is_empty())
    }

    async fn handled_count(&self) -> StorageResult<u64> {
        let _operation = self.operations.read().await;
        Ok(self.state.lock().await.handled_count)
    }

    async fn pending_count(&self) -> StorageResult<u64> {
        let _operation = self.operations.read().await;
        Ok(self.state.lock().await.pending.len() as u64)
    }
}

async fn scan_queue(name: &str, path: &Path, requests_path: &Path) -> StorageResult<QueueState> {
    let mut entries = tokio::fs::read_dir(requests_path).await?;
    let mut files = Vec::new();
    while let Some(entry) = entries.next_entry().await? {
        if !entry.file_type().await?.is_file() {
            continue;
        }
        let file_name = entry.file_name();
        let Some(file_name) = file_name.to_str() else {
            continue;
        };
        if is_temporary_file(file_name) || !file_name.ends_with(".json") {
            continue;
        }
        files.push(entry.path());
    }
    files.sort_unstable();

    let mut decoded = Vec::with_capacity(files.len());
    for file in files {
        let bytes = tokio::fs::read(&file).await?;
        match decode_request(&bytes) {
            Ok(request) => decoded.push(request),
            Err(error) => {
                tracing::warn!(
                    queue = name,
                    path = %file.display(),
                    %error,
                    "skipping unreadable request queue item"
                );
            }
        }
    }

    let max_normal = decoded
        .iter()
        .filter_map(|(_, order_no, _)| order_no.filter(|order_no| *order_no > 0))
        .max()
        .unwrap_or(0);
    let min_forefront = decoded
        .iter()
        .filter_map(|(_, order_no, _)| order_no.filter(|order_no| *order_no < 0))
        .min()
        .unwrap_or(0);
    let mut state = QueueState::empty();
    state.next_order = max_normal.saturating_add(1);
    state.next_forefront_order = if min_forefront < 0 {
        min_forefront.saturating_sub(1)
    } else {
        -1
    };

    for (request, disk_order_no, bare) in decoded {
        let request_id = request.id.clone();
        let order_no = if bare {
            Some(state.take_order(false)?)
        } else if let Some(order_no) = disk_order_no {
            if state.pending.contains_key(&order_no) {
                Some(state.take_order(false)?)
            } else {
                Some(order_no)
            }
        } else {
            None
        };
        let handled = order_no.is_none();
        if let Some(order_no) = order_no {
            state.pending.insert(order_no, request_id.clone());
        } else {
            state.handled_count = state.handled_count.saturating_add(1);
        }
        state
            .dedup
            .insert(request.unique_key.clone(), (request_id.clone(), handled));
        state
            .requests
            .insert(request_id, StoredRequest { request, order_no });
    }

    compare_cache(name, path, &state).await;
    Ok(state)
}

fn decode_request(bytes: &[u8]) -> Result<(Request, Option<i64>, bool), serde_json::Error> {
    let value: serde_json::Value = serde_json::from_slice(bytes)?;
    if value.get("json").is_some() {
        let envelope: ReadEnvelope = serde_json::from_value(value)?;
        Ok((envelope.json, envelope.order_no, false))
    } else if value.get("uniqueKey").is_some() || value.get("orderNo").is_some() {
        let envelope: CrawleeRequestEnvelope = serde_json::from_value(value)?;
        let (request, order_no) = envelope.into_request()?;
        Ok((request, order_no, false))
    } else {
        let request = serde_json::from_value(value)?;
        Ok((request, None, true))
    }
}

async fn compare_cache(name: &str, path: &Path, state: &QueueState) {
    let bytes = match tokio::fs::read(path.join("state.json")).await {
        Ok(bytes) => bytes,
        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return,
        Err(error) => {
            tracing::warn!(queue = name, %error, "failed to read request queue state cache");
            return;
        }
    };
    let cache: QueueCache = match serde_json::from_slice::<QueueCache>(&bytes) {
        Ok(cache) if cache.version == STATE_VERSION => cache,
        Ok(_) => {
            tracing::warn!(
                queue = name,
                "ignoring unsupported request queue state cache"
            );
            return;
        }
        Err(error) => {
            tracing::warn!(queue = name, %error, "ignoring unreadable request queue state cache");
            return;
        }
    };
    let pending_count = state
        .requests
        .values()
        .filter(|request| request.order_no.is_some())
        .count() as u64;
    if cache.handled_request_count != state.handled_count
        || cache.pending_request_count != pending_count
    {
        tracing::warn!(
            queue = name,
            cached_handled = cache.handled_request_count,
            scanned_handled = state.handled_count,
            cached_pending = cache.pending_request_count,
            scanned_pending = pending_count,
            "request queue state cache disagrees with request files"
        );
    }
}

async fn atomic_write(destination: &Path, bytes: &[u8]) -> StorageResult<()> {
    let file_name = destination
        .file_name()
        .and_then(|name| name.to_str())
        .ok_or_else(|| StorageError::Backend(anyhow::anyhow!("invalid queue file path")))?;
    let temporary = destination.with_file_name(format!("{file_name}.{}", temporary_suffix()));
    tokio::fs::write(&temporary, bytes).await?;
    if let Err(error) = tokio::fs::rename(&temporary, destination).await {
        let _ = tokio::fs::remove_file(&temporary).await;
        return Err(error.into());
    }
    Ok(())
}