later 0.0.48

use super::{memory::MemoryStorage, RangeOrder, Storage, StorageOperation};
use std::{sync::Arc, time::Duration};

fn key(name: &str) -> String {
    format!("storage-test-{}-{}", name, crate::generate_id())
}

async fn run_storage_behavior(storage: Arc<dyn Storage>) -> anyhow::Result<()> {
    let value_key = key("value");
    storage.set(&value_key, b"value").await?;
    assert_eq!(storage.get(&value_key).await?, Some(b"value".to_vec()));
    assert!(storage.exist(&value_key).await?);
    storage.del(&value_key).await?;
    assert_eq!(storage.get(&value_key).await?, None);

    let ttl_key = key("ttl");
    storage.set(&ttl_key, b"expires").await?;
    storage.expire(&ttl_key, 1).await?;
    assert_eq!(storage.get(&ttl_key).await?, Some(b"expires".to_vec()));
    tokio::time::sleep(Duration::from_millis(1_100)).await;
    assert_eq!(storage.get(&ttl_key).await?, None);

    let range_key = key("range");
    storage
        .apply(vec![
            StorageOperation::RangeAdd {
                key: range_key.clone(),
                value: b"first".to_vec(),
            },
            StorageOperation::RangeAdd {
                key: range_key.clone(),
                value: b"second".to_vec(),
            },
            StorageOperation::RangeAdd {
                key: range_key.clone(),
                value: b"second".to_vec(),
            },
            StorageOperation::RangeAdd {
                key: range_key.clone(),
                value: b"third".to_vec(),
            },
        ])
        .await?;
    assert_eq!(storage.range_count(&range_key).await?, 3);

    let expiring_range_key = key("expiring-range");
    storage
        .apply(vec![
            StorageOperation::RangeAdd {
                key: expiring_range_key.clone(),
                value: b"expires".to_vec(),
            },
            StorageOperation::RangeExpire {
                key: expiring_range_key.clone(),
                value: b"expires".to_vec(),
                ttl_seconds: 1,
            },
            StorageOperation::RangeAdd {
                key: expiring_range_key.clone(),
                value: b"refreshed".to_vec(),
            },
            StorageOperation::RangeExpire {
                key: expiring_range_key.clone(),
                value: b"refreshed".to_vec(),
                ttl_seconds: 1,
            },
            StorageOperation::RangeAdd {
                key: expiring_range_key.clone(),
                value: b"refreshed".to_vec(),
            },
        ])
        .await?;
    assert_eq!(storage.range_count(&expiring_range_key).await?, 2);
    tokio::time::sleep(Duration::from_millis(1_100)).await;
    assert_eq!(storage.range_count(&expiring_range_key).await?, 1);
    assert_eq!(
        storage
            .range_page(&expiring_range_key, None, 10, RangeOrder::OldestFirst)
            .await?
            .items
            .into_iter()
            .map(|item| item.value)
            .collect::<Vec<_>>(),
        vec![b"refreshed".to_vec()]
    );

    let first_page = storage
        .range_page(&range_key, None, 2, RangeOrder::OldestFirst)
        .await?;
    let second_page = storage
        .range_page(
            &range_key,
            first_page.next_cursor,
            2,
            RangeOrder::OldestFirst,
        )
        .await?;
    let newest = storage
        .range_page(&range_key, None, 3, RangeOrder::NewestFirst)
        .await?;
    assert_eq!(
        vec![
            first_page
                .items
                .iter()
                .map(|item| item.value.clone())
                .collect::<Vec<_>>(),
            second_page
                .items
                .iter()
                .map(|item| item.value.clone())
                .collect::<Vec<_>>(),
            newest
                .items
                .iter()
                .map(|item| item.value.clone())
                .collect::<Vec<_>>(),
        ],
        vec![
            vec![b"first".to_vec(), b"second".to_vec()],
            vec![b"third".to_vec()],
            vec![b"third".to_vec(), b"second".to_vec(), b"first".to_vec()],
        ]
    );

    let changing_range_key = key("changing-range");
    storage
        .apply(
            ["a", "b", "c", "d"]
                .into_iter()
                .map(|value| StorageOperation::RangeAdd {
                    key: changing_range_key.clone(),
                    value: value.as_bytes().to_vec(),
                })
                .collect(),
        )
        .await?;
    let changing_first_page = storage
        .range_page(&changing_range_key, None, 2, RangeOrder::OldestFirst)
        .await?;
    storage
        .apply(vec![
            StorageOperation::RangeRemove {
                key: changing_range_key.clone(),
                value: b"a".to_vec(),
            },
            StorageOperation::RangeRemove {
                key: changing_range_key.clone(),
                value: b"c".to_vec(),
            },
        ])
        .await?;
    let changing_second_page = storage
        .range_page(
            &changing_range_key,
            changing_first_page.next_cursor,
            2,
            RangeOrder::OldestFirst,
        )
        .await?;
    assert_eq!(
        (
            changing_first_page
                .items
                .into_iter()
                .map(|item| item.value)
                .collect::<Vec<_>>(),
            changing_second_page
                .items
                .into_iter()
                .map(|item| item.value)
                .collect::<Vec<_>>(),
        ),
        (vec![b"a".to_vec(), b"b".to_vec()], vec![b"d".to_vec()],)
    );

    let metadata_key = key("metadata");
    let next_range_key = key("next-range");
    storage
        .apply(vec![
            StorageOperation::Set {
                key: metadata_key.clone(),
                value: b"moved".to_vec(),
            },
            StorageOperation::RangeRemove {
                key: range_key.clone(),
                value: b"second".to_vec(),
            },
            StorageOperation::RangeAdd {
                key: next_range_key.clone(),
                value: b"second".to_vec(),
            },
        ])
        .await?;
    assert_eq!(
        (
            storage.get(&metadata_key).await?,
            storage.range_count(&range_key).await?,
            storage.range_count(&next_range_key).await?,
        ),
        (Some(b"moved".to_vec()), 2, 1)
    );

    let concurrent_range_key = key("concurrent-range");
    let mut tasks = Vec::new();
    for index in 0..100 {
        let storage = storage.clone();
        let range_key = concurrent_range_key.clone();
        tasks.push(tokio::spawn(async move {
            storage
                .apply(vec![StorageOperation::RangeAdd {
                    key: range_key,
                    value: format!("item-{}", index % 20).into_bytes(),
                }])
                .await
        }));
    }
    for task in tasks {
        task.await??;
    }
    let concurrent_page = storage
        .range_page(&concurrent_range_key, None, 100, RangeOrder::OldestFirst)
        .await?;
    assert_eq!(
        (
            storage.range_count(&concurrent_range_key).await?,
            concurrent_page.items.len(),
        ),
        (20, 20)
    );

    let rollback_key = key("rollback");
    let rollback_result = storage
        .apply(vec![
            StorageOperation::Set {
                key: rollback_key.clone(),
                value: b"must-not-be-visible".to_vec(),
            },
            StorageOperation::Expire {
                key: rollback_key.clone(),
                ttl_seconds: usize::MAX,
            },
        ])
        .await;
    assert!(rollback_result.is_err());
    assert_eq!(storage.get(&rollback_key).await?, None);

    storage
        .apply(vec![
            StorageOperation::RangeClear { key: range_key },
            StorageOperation::RangeClear {
                key: next_range_key,
            },
            StorageOperation::Delete { key: metadata_key },
            StorageOperation::RangeClear {
                key: concurrent_range_key,
            },
            StorageOperation::RangeClear {
                key: expiring_range_key,
            },
            StorageOperation::RangeClear {
                key: changing_range_key,
            },
        ])
        .await?;
    Ok(())
}

#[tokio::test]
async fn memory_storage_behavior() -> anyhow::Result<()> {
    run_storage_behavior(Arc::new(MemoryStorage::new())).await
}

#[tokio::test]
async fn memory_storage_sweep_expired_is_a_no_op() -> anyhow::Result<()> {
    // In-memory has no separate physically-present-but-expired rows to
    // reclaim (see `Storage::sweep_expired`'s default) - worth pinning down
    // explicitly since a caller looping on the count until it hits zero
    // must not spin forever against a backend that never removes anything.
    assert_eq!(MemoryStorage::new().sweep_expired(100).await?, 0);
    Ok(())
}

#[cfg(feature = "sqlite")]
#[tokio::test]
async fn sqlite_storage_behavior() -> anyhow::Result<()> {
    run_storage_behavior(Arc::new(super::Sqlite::new("sqlite::memory:").await?)).await
}

#[cfg(feature = "postgres")]
#[tokio::test]
async fn postgres_storage_behavior() -> anyhow::Result<()> {
    let url = std::env::var("LATER_POSTGRES_TEST_URL")
        .unwrap_or_else(|_| "postgres://test:test@127.0.0.1:55432/later_test".to_string());
    run_storage_behavior(Arc::new(super::Postgres::new(&url).await?)).await
}

/// Creates one already-expired plain key and one already-expired range
/// member, then confirms `sweep_expired` both respects its `limit` (never
/// removes more than asked in one call, regardless of how large the total
/// backlog is - including from other tests sharing the same database) and
/// eventually clears a backlog fully when called repeatedly with a
/// generous limit.
async fn run_sweep_expired_behavior(storage: Arc<dyn Storage>) -> anyhow::Result<()> {
    let expired_plain = key("sweep-plain-expired");
    storage.set(&expired_plain, b"x").await?;
    storage.expire(&expired_plain, 1).await?;

    let range_key = key("sweep-range");
    storage
        .apply(vec![
            StorageOperation::RangeAdd {
                key: range_key.clone(),
                value: b"expired-member".to_vec(),
            },
            StorageOperation::RangeExpire {
                key: range_key.clone(),
                value: b"expired-member".to_vec(),
                ttl_seconds: 1,
            },
        ])
        .await?;

    tokio::time::sleep(Duration::from_millis(1_100)).await;

    let bounded = storage.sweep_expired(1).await?;
    assert!(
        bounded <= 1,
        "sweep_expired(1) removed {bounded}, expected at most 1"
    );

    let mut total = bounded;
    for _ in 0..20 {
        if total >= 2 {
            break;
        }
        total += storage.sweep_expired(1_000).await?;
    }
    assert!(
        total >= 2,
        "expected our expired plain key and range member both swept eventually, only got {total}"
    );
    Ok(())
}

#[cfg(feature = "sqlite")]
#[tokio::test]
async fn sqlite_sweep_expired_is_bounded_and_eventually_catches_up() -> anyhow::Result<()> {
    run_sweep_expired_behavior(Arc::new(super::Sqlite::new("sqlite::memory:").await?)).await
}

#[cfg(feature = "postgres")]
#[tokio::test]
async fn postgres_sweep_expired_is_bounded_and_eventually_catches_up() -> anyhow::Result<()> {
    let url = std::env::var("LATER_POSTGRES_TEST_URL")
        .unwrap_or_else(|_| "postgres://test:test@127.0.0.1:55432/later_test".to_string());
    run_sweep_expired_behavior(Arc::new(super::Postgres::new(&url).await?)).await
}

#[cfg(feature = "postgres")]
#[tokio::test]
async fn postgres_coordinates_independent_pools() -> anyhow::Result<()> {
    let url = std::env::var("LATER_POSTGRES_TEST_URL")
        .unwrap_or_else(|_| "postgres://test:test@127.0.0.1:55432/later_test".to_string());
    let first = Arc::new(super::Postgres::new(&url).await?);
    let second = Arc::new(super::Postgres::new(&url).await?);
    let range_key = key("postgres-independent-pools");
    let mut tasks = Vec::new();

    for index in 0..100 {
        let storage = if index % 2 == 0 {
            first.clone()
        } else {
            second.clone()
        };
        let range_key = range_key.clone();
        tasks.push(tokio::spawn(async move {
            storage
                .apply(vec![StorageOperation::RangeAdd {
                    key: range_key,
                    value: format!("item-{}", index % 25).into_bytes(),
                }])
                .await
        }));
    }
    for task in tasks {
        task.await??;
    }

    assert_eq!(first.range_count(&range_key).await?, 25);
    assert_eq!(
        second
            .range_page(&range_key, None, 100, RangeOrder::OldestFirst)
            .await?
            .items
            .len(),
        25
    );
    first
        .apply(vec![StorageOperation::RangeClear { key: range_key }])
        .await
}

#[cfg(feature = "redis")]
#[tokio::test]
async fn redis_storage_behavior() -> anyhow::Result<()> {
    let url = std::env::var("LATER_REDIS_TEST_URL")
        .unwrap_or_else(|_| "redis://127.0.0.1:56379/".to_string());
    run_storage_behavior(Arc::new(super::Redis::new(&url).await?)).await
}

#[cfg(feature = "redis")]
#[tokio::test]
async fn redis_data_survives_a_new_connection() -> anyhow::Result<()> {
    let url = std::env::var("LATER_REDIS_TEST_URL")
        .unwrap_or_else(|_| "redis://127.0.0.1:56379/".to_string());
    let value_key = key("redis-reconnect");
    {
        let storage = super::Redis::new(&url).await?;
        storage.set(&value_key, b"persisted").await?;
    }

    let reconnected = super::Redis::new(&url).await?;
    assert_eq!(
        reconnected.get(&value_key).await?,
        Some(b"persisted".to_vec())
    );
    reconnected.del(&value_key).await
}

#[cfg(feature = "sqlite")]
#[tokio::test]
async fn sqlite_file_survives_restart() -> anyhow::Result<()> {
    let directory = tempfile::tempdir()?;
    let database_path = directory.path().join("later.sqlite3");
    let url = format!("sqlite://{}", database_path.display());
    let value_key = key("sqlite-restart-value");
    let range_key = key("sqlite-restart-range");

    {
        let storage = super::Sqlite::new(&url).await?;
        storage
            .apply(vec![
                StorageOperation::Set {
                    key: value_key.clone(),
                    value: b"persisted".to_vec(),
                },
                StorageOperation::RangeAdd {
                    key: range_key.clone(),
                    value: b"range-item".to_vec(),
                },
            ])
            .await?;
    }

    let reopened = super::Sqlite::new(&url).await?;
    assert_eq!(reopened.get(&value_key).await?, Some(b"persisted".to_vec()));
    assert_eq!(
        reopened
            .range_page(&range_key, None, 10, RangeOrder::OldestFirst)
            .await?
            .items
            .into_iter()
            .map(|item| item.value)
            .collect::<Vec<_>>(),
        vec![b"range-item".to_vec()]
    );
    Ok(())
}

/// `Sqlite::new`/`from_pool` must tolerate sharing a database file with an
/// application's own, entirely unrelated sqlx migration history - the
/// scenario `later::backend::SqliteBackend::from_pool`/`enqueue_in` exist
/// for. Later tracks its own migrations in `later_schema_migrations`
/// (see `run_migrations`), never touching `_sqlx_migrations` at all, so an
/// application's own migrator (which almost certainly *does* use the
/// sqlx-default `_sqlx_migrations` table) never sees a row it doesn't
/// recognize either - the earlier fix for this (`Migrator::set_ignore_missing`
/// on Later's own migrator) only avoided the symptom on Later's side; the
/// application's migrator would still choke on Later's rows on its own next
/// run, since sqlx's migrator assumes it owns the entire table.
#[cfg(feature = "sqlite")]
#[tokio::test]
async fn sqlite_tolerates_an_unrelated_migration_history_and_uses_its_own_table(
) -> anyhow::Result<()> {
    let directory = tempfile::tempdir()?;
    let database_path = directory.path().join("shared.sqlite3");
    let url = format!("sqlite://{}", database_path.display());

    let pool = sqlx::sqlite::SqlitePoolOptions::new()
        .max_connections(1)
        .connect_with(
            sqlx::sqlite::SqliteConnectOptions::new()
                .filename(&database_path)
                .create_if_missing(true),
        )
        .await?;

    // Simulates an application's own migrator having already run against
    // this file, recording a version Later's own migration set has never
    // heard of.
    sqlx::query(
        "CREATE TABLE _sqlx_migrations ( \
            version BIGINT PRIMARY KEY, \
            description TEXT NOT NULL, \
            installed_on TIMESTAMP NOT NULL DEFAULT (datetime('now')), \
            success BOOLEAN NOT NULL, \
            checksum BLOB NOT NULL, \
            execution_time BIGINT NOT NULL \
        )",
    )
    .execute(&pool)
    .await?;
    sqlx::query(
        "INSERT INTO _sqlx_migrations \
            (version, description, success, checksum, execution_time) \
            VALUES (1, 'an applications own migration', 1, x'00', 0)",
    )
    .execute(&pool)
    .await?;
    pool.close().await;

    let storage = super::Sqlite::new(&url).await?;
    storage
        .apply(vec![StorageOperation::Set {
            key: key("sqlite-shared-migrations-value"),
            value: b"ok".to_vec(),
        }])
        .await?;

    // Later's own history lives in its own table, and the application's
    // foreign row is exactly as it left it - Later never touched it.
    let later_migrations: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM later_schema_migrations")
        .fetch_one(storage.pool())
        .await?;
    assert!(later_migrations > 0);
    let app_migrations: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM _sqlx_migrations")
        .fetch_one(storage.pool())
        .await?;
    assert_eq!(app_migrations, 1);

    // Reopening must not re-apply (or fail re-applying) anything already
    // recorded in later_schema_migrations.
    super::Sqlite::new(&url).await?;

    Ok(())
}

#[cfg(feature = "postgres")]
#[tokio::test]
async fn postgres_stage_counters_follow_the_index_rows_exactly() -> anyhow::Result<()> {
    use crate::storage::{JobIndexRow, Storage};
    let url = std::env::var("LATER_POSTGRES_TEST_URL")
        .unwrap_or_else(|_| "postgres://test:test@127.0.0.1:55432/later_test".to_string());
    let storage = super::Postgres::new(&url).await?;
    let namespace = format!("counters-{}", crate::generate_id());
    let now = chrono::Utc::now();
    let row = |id: &str, stage: &str, revision: i64| JobIndexRow {
        job_id: id.to_string(),
        payload_type: "t".to_string(),
        stage: stage.to_string(),
        stage_date: now,
        revision,
        created_at: now,
        wait_ms: None,
        wait_mode: None,
        topic: None,
        partition: None,
        sequence: None,
        parent_job_id: None,
        date_expire: None,
    };
    let count = |stage: &'static str| {
        let (storage, namespace) = (&storage, namespace.clone());
        async move {
            anyhow::Ok(
                storage
                    .job_index_stage_counts(&namespace)
                    .await?
                    .get(stage)
                    .copied()
                    .unwrap_or(0),
            )
        }
    };

    storage
        .job_index_upsert(&namespace, row("a", "enqueued", 0))
        .await?;
    storage
        .job_index_upsert(&namespace, row("b", "enqueued", 0))
        .await?;
    // Re-saving a job in the same stage must not inflate its counter.
    for _ in 0..5 {
        storage
            .job_index_upsert(&namespace, row("a", "enqueued", 0))
            .await?;
    }
    assert_eq!(count("enqueued").await?, 2);

    storage
        .job_index_upsert(&namespace, row("a", "running", 1))
        .await?;
    assert_eq!((count("enqueued").await?, count("running").await?), (1, 1));

    // A write the revision guard rejects changes no row, so no counter.
    storage
        .job_index_upsert(&namespace, row("a", "enqueued", 0))
        .await?;
    assert_eq!((count("enqueued").await?, count("running").await?), (1, 1));

    sqlx::query("DELETE FROM later_jobs_index WHERE namespace = $1 AND job_id = 'b'")
        .bind(&namespace)
        .execute(storage.pool())
        .await?;
    assert_eq!(count("enqueued").await?, 0);

    // A recount replaces drifted counters with what the index holds.
    sqlx::query("UPDATE later_stage_counts SET count = 999 WHERE namespace = $1")
        .bind(&namespace)
        .execute(storage.pool())
        .await?;
    storage.job_index_reconcile_stage_counts(&namespace).await?;
    assert_eq!((count("enqueued").await?, count("running").await?), (0, 1));
    Ok(())
}