cron ジョブ

use std::time::Duration;

use wakuwaku::deps::kanau::processor::Processor;
use wakuwaku::{CronJob, CronJobInterval, CronTick, ServiceCreation};

pub struct ExpireSessions {
    store: SessionStore,
}

impl ServiceCreation for ExpireSessions {
    type Dep = SessionStore;

    fn create_service(store: SessionStore) -> Self {
        Self { store }
    }
}

impl Processor<CronTick> for ExpireSessions {
    type Output = ();
    type Error = StoreError;

    async fn process(&self, _tick: CronTick) -> Result<(), StoreError> {
        let deleted = self.store.delete_expired(1_000).await?;
        tracing::info!(deleted, "expired sessions removed");
        Ok(())
    }
}

impl CronJob for ExpireSessions {
    const INTERVAL: CronJobInterval = CronJobInterval::AtLeast(Duration::from_secs(60));
}

/// データベースプールの代わり: clone は安価で、すべての clone がデータを共有する。
#[derive(Clone, Default)]
pub struct SessionStore {
    expired: Arc<AtomicU64>,
}

impl SessionStore {
    pub async fn delete_expired(&self, limit: u64) -> Result<u64, StoreError> {
        self.expired.fetch_add(1, Ordering::Relaxed);
        Ok(limit.min(1))
    }
}

#[derive(Debug)]
pub struct StoreError(String);

impl fmt::Display for StoreError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "session store: {}", self.0)
    }
}

impl Processor<CronTick> for ReindexSearch {
    type Output = ();
    type Error = StoreError;

    async fn process(&self, tick: CronTick) -> Result<(), StoreError> {
        // 24 回に一回(と初回)は全件を作り直し、
        // それ以外は変更分だけを取り込む
        let reindexed = if tick.run % 24 == 0 {
            self.index.reindex_all().await?
        } else {
            self.index.reindex_changed().await?
        };
        tracing::info!(run = tick.run, fired_at = ?tick.fired_at, reindexed, "search index refreshed");
        Ok(())
    }
}

impl CronJob for ReindexSearch {
    const INTERVAL: CronJobInterval = CronJobInterval::AfterFinish(Duration::from_secs(60 * 60));
}

let (services, cron) = ServiceBuilder::new(SessionStore::default())
    .push(SearchIndex::default())
    .cron_job::<ExpireSessions, _>()
    .cron_job::<ReindexSearch, _>()
    .start();

// `services` は返ってくる: ここから gRPC サーバーやコンシューマーを登録できる
let _store: &SessionStore = services.provide::<SessionStore, _>();

let _ = tokio::signal::ctrl_c().await;
// 実行中の処理を終わらせてから、すべてのジョブを止める
cron.shutdown().await;

impl CronJob for Busy {
    const INTERVAL: CronJobInterval = CronJobInterval::AtLeast(Duration::ZERO); // no schedule
}

#[tokio::main]
async fn main() {
    let (_services, cron) = ServiceBuilder::new(Config)
        .cron_job::<Busy, _>() // `cargo build` fails: INTERVAL is zero
        .start();
    cron.shutdown().await;
}

error[E0080]: evaluation panicked: CronJob::INTERVAL must not be zero
   --> wakuwaku/src/services/cron_job.rs:300:17
    |
300 |         const { assert_interval::<Job>() };
    |                 ^^^^^^^^^^^^^^^^^^^^^^^^ evaluation of `wakuwaku::services::cron_job::CronRegisterCenter::<Busy>::new::{constant#0}` failed inside this call
...
note: the above error was encountered while instantiating `fn CronRegisterCenter::<Busy>::new`

use std::time::Duration;

use wakuwaku::deps::kanau::processor::Processor;
use wakuwaku::deps::redis::RedisError;
use wakuwaku::entity::redis::RedisSource;
use wakuwaku::entity::redis::collection::lock::Lock;
use wakuwaku::entity::{Db, ReadWrite};
use wakuwaku::{CronJob, CronJobInterval, CronTick, ServiceCreation};

/// ジョブ名ごとに一つのロック。リースは一回の実行より長くする。
const CRON_LOCKS: Lock<str> = Lock::new("cron", Duration::from_secs(5 * 60));

pub struct SendDailyDigest {
    db: Db<RedisSource, ReadWrite>,
}

impl ServiceCreation for SendDailyDigest {
    type Dep = Db<RedisSource, ReadWrite>;

    fn create_service(db: Db<RedisSource, ReadWrite>) -> Self {
        Self { db }
    }
}

impl Processor<CronTick> for SendDailyDigest {
    type Output = ();
    type Error = RedisError;

    async fn process(&self, _tick: CronTick) -> Result<(), RedisError> {
        // すべてのレプリカが試し、トークンを得るのは多くても一つ
        let Some(token) = self.db.process(CRON_LOCKS.try_lock("send-daily-digest")).await? else {
            return Ok(()); // 別のレプリカがちょうど実行中
        };

        // ... ダイジェストを作って送る。長い処理なら `token.extend()` を呼ぶ ...

        self.db.process(token.release()).await?;
        Ok(())
    }
}

impl CronJob for SendDailyDigest {
    const INTERVAL: CronJobInterval = CronJobInterval::AtLeast(Duration::from_secs(60 * 60));
}

-- $1 = ジョブ名、$2 = ティックの fired_at
-- この UPDATE が行を変更した場合だけ実行を進める
-- (各ジョブの行は古い last_fired_at で用意しておく)
UPDATE cron_claims
SET    last_fired_at = $2
WHERE  job = $1
  AND  last_fired_at <= $2 - INTERVAL '1 hour';

/// 二つの登録値が必要なので、compose で組み立てる。
pub struct ReconcileIndex {
    store: SessionStore,
    index: SearchIndex,
}

let (services, cron) = ServiceBuilder::new(SessionStore::default())
    .push(SearchIndex::default())
    .cron_job::<ExpireSessions, _>()
    .compose_cron_job(|s| ReconcileIndex {
        store: s.provide::<SessionStore, _>().clone(),
        index: s.provide::<SearchIndex, _>().clone(),
    })
    .start();

use wakuwaku::services::cron_job::CronRegisterCenter;

let store = SessionStore::default();
let cron = CronRegisterCenter::new(ExpireSessions::create_service(store.clone()))
    .push(ReconcileIndex { store, index: SearchIndex::default() })
    .start();
cron.shutdown().await;

if tokio::time::timeout(Duration::from_secs(30), cron.shutdown()).await.is_err() {
    // タイムアウトすると `CronRuntime` が drop され、止まらない実行は中断される
    tracing::warn!("cron jobs did not stop within 30s; aborted");
}

let (services, cron) = ServiceBuilder::new(AccountStore::default())
    .push(SessionStore::default())
    .cron_job::<ExpireSessions, _>()
    .start();

let (_services, register) = services
    .grpc_service::<AccountsServer<_>, AccountsService, _>(AccountsServer::new)
    .into_router_fn();

let router = Server::builder().add_routes(Routes::default());
register(router)
    .serve_with_shutdown("[::1]:50051".parse()?, async {
        let _ = tokio::signal::ctrl_c().await;
    })
    .await?;

// サーバーは処理を終えた。次にジョブの実行中の処理を終わらせる
cron.shutdown().await;

#[cfg(test)]
mod tests {
    use std::time::SystemTime;

    use super::*;

    #[tokio::test]
    async fn expire_sessions_deletes_from_the_store() {
        let store = SessionStore::default();
        let job = ExpireSessions::create_service(store.clone());

        let tick = CronTick { run: 0, fired_at: SystemTime::now() };
        job.process(tick).await.unwrap();

        assert_eq!(store.expired.load(Ordering::Relaxed), 1);
    }
}

#[tokio::test(start_paused = true)]
async fn the_first_run_happens_at_start() {
    let store = SessionStore::default();
    let (_services, cron) = ServiceBuilder::new(store.clone())
        .cron_job::<ExpireSessions, _>()
        .start();

    // 60 秒の間隔の半分: 最初の実行だけが済んでいる
    tokio::time::sleep(Duration::from_secs(30)).await;
    cron.shutdown().await;

    assert_eq!(store.expired.load(Ordering::Relaxed), 1);
}