// Copyright 2026 PingCAP, Inc.

use std::{any::Any, fmt, panic::AssertUnwindSafe, sync::Arc, time::Duration};

use futures::FutureExt;
use kvproto::kvrpcpb;
use rfstore::store::{
    ResolveTxnFileLocksTask, TxnFileLockResolveLimiter, TxnFileLockResolveScheduler,
    metrics::{
        TXN_FILE_ACTIVE_RESOLVE_LOCKS_TOTAL, TXN_FILE_ACTIVE_RESOLVE_RUNNING_TASKS,
        TXN_FILE_ACTIVE_RESOLVE_TASK_DURATION_HISTOGRAM,
    },
};
use security::SecurityConfig;
use tikv_client::{
    Config as TxnClientConfig, Timestamp as TxnTimestamp, TimestampExt as _, TransactionClient,
    backoff::OPTIMISTIC_BACKOFF, transaction::ProtoLockInfo,
};
use tikv_util::{box_try, error, info, time::Instant, warn, worker::Runnable};
use tokio::{
    runtime::{Builder, Runtime},
    sync::{OnceCell, OwnedSemaphorePermit, Semaphore},
};
use txn_types::TimeStamp;

type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;

const DEFAULT_GRPC_MAX_DECODING_MESSAGE_SIZE: usize = 128 * 1024 * 1024; // 128MB
const RESOLVER_RUNTIME_WORKER_THREADS: usize = 2;
const RESOLVER_RUNTIME_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(3);

#[derive(Clone, Copy)]
struct ResolveTaskMetadata {
    tag: rfstore::store::PeerTag,
    total_locks: usize,
    min_age_ms: u64,
    max_age_ms: u64,
}

impl ResolveTaskMetadata {
    fn from_task(task: &ResolveTxnFileLocksTask) -> Self {
        let physical_now = TimeStamp::physical_now();
        let (first_lock, other_locks) = task.locks.split_first().unwrap();
        let first_age_ms = txn_file_lock_age_ms(first_lock.get_lock_version(), physical_now);
        let (min_age_ms, max_age_ms) = other_locks.iter().fold(
            (first_age_ms, first_age_ms),
            |(min_age_ms, max_age_ms), lock| {
                let age_ms = txn_file_lock_age_ms(lock.get_lock_version(), physical_now);
                (min_age_ms.min(age_ms), max_age_ms.max(age_ms))
            },
        );
        Self {
            tag: task.tag,
            total_locks: task.locks.len(),
            min_age_ms,
            max_age_ms,
        }
    }
}

struct RunningTaskGuard;

impl RunningTaskGuard {
    fn new() -> Self {
        TXN_FILE_ACTIVE_RESOLVE_RUNNING_TASKS.inc();
        Self
    }
}

impl Drop for RunningTaskGuard {
    fn drop(&mut self) {
        TXN_FILE_ACTIVE_RESOLVE_RUNNING_TASKS.dec();
    }
}

enum ResolveExecutor {
    Client(TransactionClient),
    #[cfg(test)]
    Stub(Arc<dyn Fn(ResolveTxnFileLocksTask) -> ResolveTaskFuture + Send + Sync>),
}

#[cfg(test)]
type ResolveTaskFuture = std::pin::Pin<Box<dyn std::future::Future<Output = Result<usize>> + Send>>;

#[cfg(test)]
type InitTaskFuture = std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send>>;

#[cfg(test)]
#[derive(Clone, Default)]
struct TxnFileLockResolverTestHooks {
    init_client: Option<Arc<dyn Fn() -> InitTaskFuture + Send + Sync>>,
    resolve_task: Option<Arc<dyn Fn(ResolveTxnFileLocksTask) -> ResolveTaskFuture + Send + Sync>>,
    on_wait_for_limiter: Option<Arc<dyn Fn(u64) + Send + Sync>>,
    before_task_start: Option<Arc<dyn Fn(u64) + Send + Sync>>,
}

struct ResolverShared {
    scheduler: TxnFileLockResolveScheduler,
    limiter: TxnFileLockResolveLimiter,
    admission: Arc<Semaphore>,
    txn_client: OnceCell<TransactionClient>,
    pd_endpoints: Vec<String>,
    client_config: TxnClientConfig,
    #[cfg(test)]
    test_hooks: TxnFileLockResolverTestHooks,
    #[cfg(test)]
    stub_init: OnceCell<()>,
}

/// Background runner that reuses the normal online resolve-lock path.
pub(crate) struct TxnFileLockResolverRunner {
    runtime: Option<Runtime>,
    shared: Arc<ResolverShared>,
}

impl TxnFileLockResolverRunner {
    pub(crate) fn new(
        pd_endpoints: Vec<String>,
        security_config: SecurityConfig,
        scheduler: TxnFileLockResolveScheduler,
        limiter: TxnFileLockResolveLimiter,
        admission_capacity: usize,
    ) -> Self {
        assert!(
            admission_capacity >= 1,
            "txn file lock resolve admission capacity must be >= 1"
        );
        let mut client_config = TxnClientConfig::default()
            .with_grpc_max_decoding_message_size(DEFAULT_GRPC_MAX_DECODING_MESSAGE_SIZE);
        if !security_config.ca_path.is_empty()
            || !security_config.cert_path.is_empty()
            || !security_config.key_path.is_empty()
        {
            let SecurityConfig {
                ca_path,
                cert_path,
                key_path,
                ..
            } = security_config;
            client_config = client_config.with_security(ca_path, cert_path, key_path);
        }

        let shared = Arc::new(ResolverShared {
            scheduler,
            limiter,
            admission: Arc::new(Semaphore::new(admission_capacity)),
            txn_client: OnceCell::new(),
            pd_endpoints,
            client_config,
            #[cfg(test)]
            test_hooks: TxnFileLockResolverTestHooks::default(),
            #[cfg(test)]
            stub_init: OnceCell::new(),
        });

        Self {
            runtime: None,
            shared,
        }
    }

    fn ensure_runtime(&mut self) -> Result<()> {
        if self.runtime.is_none() {
            self.runtime = Some(box_try!(
                Builder::new_multi_thread()
                    .worker_threads(RESOLVER_RUNTIME_WORKER_THREADS)
                    .thread_name("txn-file-lock-resolve")
                    .enable_all()
                    .build()
            ));
        }
        Ok(())
    }

    #[cfg(test)]
    fn test_hooks_mut(&mut self) -> &mut TxnFileLockResolverTestHooks {
        &mut Arc::get_mut(&mut self.shared)
            .expect("resolver runner test hooks must be configured before spawning tasks")
            .test_hooks
    }
}

impl Runnable for TxnFileLockResolverRunner {
    type Task = ResolveTxnFileLocksTask;

    fn run(&mut self, task: Self::Task) {
        if task.locks.is_empty() {
            self.shared.scheduler.clear(task.region_id);
            return;
        }

        let start = Instant::now_coarse();
        let region_id = task.region_id;

        if let Err(err) = self.ensure_runtime() {
            self.shared.scheduler.clear(region_id);
            record_resolve_error(ResolveTaskMetadata::from_task(&task), start, err);
            return;
        }

        // Hold admission for the full async lifetime so the runtime backlog
        // stays bounded by the worker's pending-capacity budget.
        let admission_permit = match self.shared.admission.clone().try_acquire_owned() {
            Ok(permit) => permit,
            Err(_) => {
                self.shared.scheduler.clear(region_id);
                record_resolve_dropped(ResolveTaskMetadata::from_task(&task), start);
                return;
            }
        };

        let metadata = ResolveTaskMetadata::from_task(&task);
        let attempt_guard = self.shared.scheduler.attempt_guard(region_id);
        let shared = self.shared.clone();
        let runtime = self.runtime.as_ref().unwrap();
        runtime.spawn(async move {
            // Keep the region inflight until the attempt outcome has been recorded.
            let _attempt_guard = attempt_guard;
            match AssertUnwindSafe(shared.run_attempt(task, admission_permit))
                .catch_unwind()
                .await
            {
                Ok(Ok(live_locks)) => record_resolve_completion(metadata, start, live_locks),
                Ok(Err(err)) => record_resolve_error(metadata, start, err),
                Err(panic_payload) => record_resolve_panic(metadata, start, panic_payload),
            }
        });
    }

    fn shutdown(&mut self) {
        if let Some(runtime) = self.runtime.take() {
            runtime.shutdown_timeout(RESOLVER_RUNTIME_SHUTDOWN_TIMEOUT);
        }
    }
}

impl ResolverShared {
    async fn run_attempt(
        self: Arc<Self>,
        task: ResolveTxnFileLocksTask,
        admission_permit: OwnedSemaphorePermit,
    ) -> Result<usize> {
        let region_id = task.region_id;
        let executor = self.resolve_executor().await?;
        self.on_wait_for_limiter(region_id);
        let permit = self.limiter.acquire().await?;
        self.before_task_start(region_id);

        let _admission_permit = admission_permit;
        let _running_guard = RunningTaskGuard::new();
        let _permit = permit;

        executor.resolve(task).await
    }

    async fn resolve_executor(&self) -> Result<ResolveExecutor> {
        #[cfg(test)]
        if let Some(resolve_task) = self.test_hooks.resolve_task.clone() {
            let init_client = self.test_hooks.init_client.clone();
            self.stub_init
                .get_or_try_init(move || async move {
                    if let Some(init_client) = init_client {
                        init_client().await?;
                    }
                    Ok::<(), Box<dyn std::error::Error + Send + Sync>>(())
                })
                .await?;
            return Ok(ResolveExecutor::Stub(resolve_task));
        }

        Ok(ResolveExecutor::Client(
            self.get_or_init_txn_client().await?,
        ))
    }

    async fn get_or_init_txn_client(&self) -> Result<TransactionClient> {
        let pd_endpoints = self.pd_endpoints.clone();
        let client_config = self.client_config.clone();
        let txn_client = self
            .txn_client
            .get_or_try_init(move || async move {
                TransactionClient::new_with_config_api_v2_no_prefix(pd_endpoints, client_config)
                    .await
                    .map_err(|err| -> Box<dyn std::error::Error + Send + Sync> { Box::new(err) })
            })
            .await?;
        Ok(txn_client.clone())
    }

    #[cfg(test)]
    fn on_wait_for_limiter(&self, region_id: u64) {
        if let Some(on_wait_for_limiter) = self.test_hooks.on_wait_for_limiter.clone() {
            on_wait_for_limiter(region_id);
        }
    }

    #[cfg(not(test))]
    fn on_wait_for_limiter(&self, _region_id: u64) {}

    #[cfg(test)]
    fn before_task_start(&self, region_id: u64) {
        if let Some(before_task_start) = self.test_hooks.before_task_start.clone() {
            before_task_start(region_id);
        }
    }

    #[cfg(not(test))]
    fn before_task_start(&self, _region_id: u64) {}
}

impl ResolveExecutor {
    async fn resolve(self, task: ResolveTxnFileLocksTask) -> Result<usize> {
        match self {
            Self::Client(txn_client) => resolve_task(task.locks, txn_client).await,
            #[cfg(test)]
            Self::Stub(resolve_fn) => resolve_fn(task).await,
        }
    }
}

fn txn_file_lock_age_ms(lock_version: u64, physical_now: u64) -> u64 {
    physical_now.saturating_sub(TimeStamp::new(lock_version).physical())
}

async fn resolve_task(
    locks: Vec<kvrpcpb::LockInfo>,
    txn_client: TransactionClient,
) -> Result<usize> {
    let client_locks = locks.into_iter().map(convert_lock_info).collect::<Vec<_>>();

    // Use a zero start_ts to avoid pushing min_commit_ts, matching the
    // server-side resolve-lock flow used elsewhere in the repo.
    let live_locks = box_try!(
        txn_client
            .resolve_locks(
                client_locks,
                TxnTimestamp::from_version(0),
                OPTIMISTIC_BACKOFF,
            )
            .await
    );

    Ok(live_locks.len())
}

fn record_resolve_completion(metadata: ResolveTaskMetadata, start: Instant, live_locks: usize) {
    let resolved_locks = metadata.total_locks.saturating_sub(live_locks);
    if resolved_locks > 0 {
        TXN_FILE_ACTIVE_RESOLVE_LOCKS_TOTAL
            .with_label_values(&["resolved"])
            .inc_by(resolved_locks as u64);
    }
    if live_locks > 0 {
        TXN_FILE_ACTIVE_RESOLVE_LOCKS_TOTAL
            .with_label_values(&["live"])
            .inc_by(live_locks as u64);
        warn!(
            "{} active resolve txn file locks completed with live locks",
            metadata.tag;
            "total_locks" => metadata.total_locks,
            "resolved_locks" => resolved_locks,
            "live_locks" => live_locks,
            "min_age_ms" => metadata.min_age_ms,
            "max_age_ms" => metadata.max_age_ms,
        );
    } else {
        info!("{} active resolve txn file locks completed", metadata.tag;
            "total_locks" => metadata.total_locks,
            "resolved_locks" => resolved_locks,
            "live_locks" => live_locks,
            "min_age_ms" => metadata.min_age_ms,
            "max_age_ms" => metadata.max_age_ms,
        );
    }
    observe_resolve_duration(start);
}

fn record_resolve_error<E>(metadata: ResolveTaskMetadata, start: Instant, err: E)
where
    E: fmt::Debug,
{
    TXN_FILE_ACTIVE_RESOLVE_LOCKS_TOTAL
        .with_label_values(&["error"])
        .inc_by(metadata.total_locks as u64);
    error!("{} active resolve txn file locks failed", metadata.tag;
        "total_locks" => metadata.total_locks,
        "min_age_ms" => metadata.min_age_ms,
        "max_age_ms" => metadata.max_age_ms,
        "err" => ?err,
    );
    observe_resolve_duration(start);
}

fn record_resolve_dropped(metadata: ResolveTaskMetadata, start: Instant) {
    TXN_FILE_ACTIVE_RESOLVE_LOCKS_TOTAL
        .with_label_values(&["dropped"])
        .inc_by(metadata.total_locks as u64);
    warn!(
        "{} active resolve txn file locks dropped before start because admission is full",
        metadata.tag;
        "total_locks" => metadata.total_locks,
        "min_age_ms" => metadata.min_age_ms,
        "max_age_ms" => metadata.max_age_ms,
    );
    observe_resolve_duration(start);
}

fn record_resolve_panic(
    metadata: ResolveTaskMetadata,
    start: Instant,
    panic_payload: Box<dyn Any + Send>,
) {
    TXN_FILE_ACTIVE_RESOLVE_LOCKS_TOTAL
        .with_label_values(&["panic"])
        .inc_by(metadata.total_locks as u64);
    error!("{} active resolve txn file locks panicked", metadata.tag;
        "total_locks" => metadata.total_locks,
        "min_age_ms" => metadata.min_age_ms,
        "max_age_ms" => metadata.max_age_ms,
        "panic" => panic_payload_message(panic_payload.as_ref()),
    );
    observe_resolve_duration(start);
    debug_assert!(false, "txn file lock resolve task panicked");
}

fn observe_resolve_duration(start: Instant) {
    TXN_FILE_ACTIVE_RESOLVE_TASK_DURATION_HISTOGRAM
        .observe(start.saturating_elapsed().as_secs_f64());
}

fn panic_payload_message(panic_payload: &(dyn Any + Send)) -> String {
    if let Some(message) = panic_payload.downcast_ref::<&'static str>() {
        return (*message).to_owned();
    }
    if let Some(message) = panic_payload.downcast_ref::<String>() {
        return message.clone();
    }
    "Box<Any>".to_owned()
}

fn convert_lock_info(mut lock: kvrpcpb::LockInfo) -> ProtoLockInfo {
    ProtoLockInfo {
        primary_lock: lock.take_primary_lock(),
        lock_version: lock.get_lock_version(),
        key: lock.take_key(),
        lock_ttl: lock.get_lock_ttl(),
        txn_size: lock.get_txn_size(),
        lock_type: lock.get_lock_type() as i32,
        lock_for_update_ts: lock.get_lock_for_update_ts(),
        use_async_commit: lock.get_use_async_commit(),
        min_commit_ts: lock.get_min_commit_ts(),
        secondaries: lock.take_secondaries().into_iter().collect(),
        duration_to_last_update_ms: lock.get_duration_to_last_update_ms(),
        shared_lock_infos: lock
            .take_shared_lock_infos()
            .into_iter()
            .map(convert_lock_info)
            .collect(),
        is_txn_file: lock.get_is_txn_file(),
    }
}

#[cfg(test)]
mod tests {
    use std::{
        sync::{
            Arc, Mutex as StdMutex,
            atomic::{AtomicUsize, Ordering},
            mpsc,
        },
        thread,
        time::Duration as StdDuration,
    };

    use rfstore::store::ResolveScheduleCheck;
    use security::SecurityConfig;
    use tikv_util::worker::{LazyWorker, dummy_scheduler};
    use tokio::{runtime::Builder, sync::oneshot, time::timeout};

    use super::*;

    static TEST_GUARD: StdMutex<()> = StdMutex::new(());

    fn test_guard() -> std::sync::MutexGuard<'static, ()> {
        TEST_GUARD
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
    }

    fn block_on<F>(future: F) -> F::Output
    where
        F: std::future::Future,
    {
        Builder::new_current_thread()
            .enable_all()
            .build()
            .unwrap()
            .block_on(future)
    }

    fn wait_until(timeout: StdDuration, mut condition: impl FnMut() -> bool) {
        let start = std::time::Instant::now();
        while start.elapsed() < timeout {
            if condition() {
                return;
            }
            thread::sleep(StdDuration::from_millis(10));
        }
        panic!("condition not met within {:?}", timeout);
    }

    fn wait_for_region_signal(
        receiver: &mpsc::Receiver<u64>,
        expected_region_id: u64,
        timeout: StdDuration,
    ) {
        wait_until(timeout, || {
            while let Ok(region_id) = receiver.try_recv() {
                if region_id == expected_region_id {
                    return true;
                }
            }
            false
        });
    }

    fn build_task(region_id: u64) -> ResolveTxnFileLocksTask {
        let mut lock = kvrpcpb::LockInfo::default();
        lock.set_primary_lock(b"primary".to_vec());
        lock.set_key(b"key".to_vec());
        lock.set_lock_version(1);
        ResolveTxnFileLocksTask {
            tag: Default::default(),
            region_id,
            locks: vec![lock],
        }
    }

    #[test]
    fn test_convert_lock_info_preserves_shared_lock_holders() {
        let mut holder = kvrpcpb::LockInfo::default();
        holder.set_primary_lock(b"primary".to_vec());
        holder.set_key(b"key".to_vec());
        holder.set_lock_version(42);

        let mut wrapper = kvrpcpb::LockInfo::default();
        wrapper.mut_shared_lock_infos().push(holder);

        let converted = convert_lock_info(wrapper);
        assert_eq!(converted.shared_lock_infos.len(), 1);
        assert_eq!(converted.shared_lock_infos[0].primary_lock, b"primary");
        assert_eq!(converted.shared_lock_infos[0].key, b"key");
        assert_eq!(converted.shared_lock_infos[0].lock_version, 42);
    }

    fn build_runner(
        scheduler: TxnFileLockResolveScheduler,
        limiter: TxnFileLockResolveLimiter,
    ) -> TxnFileLockResolverRunner {
        build_runner_with_admission(scheduler, limiter, 64)
    }

    fn build_runner_with_admission(
        scheduler: TxnFileLockResolveScheduler,
        limiter: TxnFileLockResolveLimiter,
        admission_capacity: usize,
    ) -> TxnFileLockResolverRunner {
        TxnFileLockResolverRunner::new(
            vec!["127.0.0.1:2379".to_owned()],
            SecurityConfig::default(),
            scheduler,
            limiter,
            admission_capacity,
        )
    }

    fn schedule_task(scheduler: &TxnFileLockResolveScheduler, region_id: u64) {
        assert_eq!(
            scheduler.schedule(build_task(region_id)).unwrap(),
            rfstore::store::ResolveScheduleOutcome::Scheduled
        );
    }

    #[test]
    fn test_run_cancels_region_on_init_failure() {
        let _guard = test_guard();
        let (raw_scheduler, _receiver) = dummy_scheduler();
        let scheduler = TxnFileLockResolveScheduler::new(raw_scheduler);
        let limiter = TxnFileLockResolveLimiter::new(1);
        let mut runner = build_runner(scheduler.clone(), limiter);
        runner.test_hooks_mut().resolve_task = Some(Arc::new(|_task| Box::pin(async { Ok(0) })));
        runner.test_hooks_mut().init_client = Some(Arc::new(|| {
            Box::pin(async {
                Err(tikv_util::box_err!(
                    "txn file lock resolver init failure (test)"
                ))
            })
        }));

        let task = build_task(1);
        schedule_task(&scheduler, task.region_id);
        let error_before = TXN_FILE_ACTIVE_RESOLVE_LOCKS_TOTAL
            .with_label_values(&["error"])
            .get();

        runner.run(task);

        wait_until(StdDuration::from_secs(1), || {
            scheduler.check(1) == ResolveScheduleCheck::Ready
        });
        wait_until(StdDuration::from_secs(1), || {
            TXN_FILE_ACTIVE_RESOLVE_LOCKS_TOTAL
                .with_label_values(&["error"])
                .get()
                == error_before + 1
        });
    }

    #[test]
    fn test_successful_completion_returns_region_to_ready() {
        let _guard = test_guard();
        let (raw_scheduler, _receiver) = dummy_scheduler();
        let scheduler = TxnFileLockResolveScheduler::new(raw_scheduler);
        let limiter = TxnFileLockResolveLimiter::new(1);
        let mut runner = build_runner(scheduler.clone(), limiter);

        let (started_tx, started_rx) = mpsc::channel::<u64>();
        let (finish_tx, finish_rx) = oneshot::channel::<()>();
        let finish_rx = Arc::new(StdMutex::new(Some(finish_rx)));
        runner.test_hooks_mut().resolve_task = Some(Arc::new(move |task| {
            let region_id = task.region_id;
            let started_tx = started_tx.clone();
            let finish_rx = finish_rx.clone();
            Box::pin(async move {
                started_tx.send(region_id).unwrap();
                let finish_rx = finish_rx.lock().unwrap().take().unwrap();
                finish_rx.await.unwrap();
                Ok(0)
            })
        }));

        let task = build_task(2);
        schedule_task(&scheduler, task.region_id);
        runner.run(task);
        wait_for_region_signal(&started_rx, 2, StdDuration::from_secs(1));
        assert_eq!(scheduler.check(2), ResolveScheduleCheck::Inflight);

        finish_tx.send(()).unwrap();
        wait_until(StdDuration::from_secs(1), || {
            scheduler.check(2) == ResolveScheduleCheck::Ready
        });
    }

    #[test]
    fn test_stop_worker_preempts_task_waiting_on_limiter() {
        let _guard = test_guard();
        let admission_capacity = 4;
        let mut worker = LazyWorker::<ResolveTxnFileLocksTask>::new_with_pending_capacity(
            "txn-file-lock-resolve-test",
            admission_capacity,
        );
        let scheduler = TxnFileLockResolveScheduler::new(worker.scheduler());
        let limiter = TxnFileLockResolveLimiter::new(1);
        let mut runner =
            build_runner_with_admission(scheduler.clone(), limiter, admission_capacity);

        let (started_tx, started_rx) = mpsc::channel::<u64>();
        let (waiting_tx, waiting_rx) = mpsc::channel::<u64>();
        let (finish_tx, finish_rx) = oneshot::channel::<()>();
        let finish_rx = Arc::new(StdMutex::new(Some(finish_rx)));
        runner.test_hooks_mut().resolve_task = Some(Arc::new(move |task| {
            let region_id = task.region_id;
            let started_tx = started_tx.clone();
            let finish_rx = finish_rx.clone();
            Box::pin(async move {
                started_tx.send(region_id).unwrap();
                if region_id == 3 {
                    let finish_rx = finish_rx.lock().unwrap().take().unwrap();
                    let _ = finish_rx.await;
                }
                Ok(0)
            })
        }));
        runner.test_hooks_mut().on_wait_for_limiter = Some(Arc::new(move |region_id| {
            if region_id == 4 {
                waiting_tx.send(region_id).unwrap();
            }
        }));
        assert!(worker.start(runner));

        schedule_task(&scheduler, 3);
        wait_for_region_signal(&started_rx, 3, StdDuration::from_secs(1));
        schedule_task(&scheduler, 4);
        wait_for_region_signal(&waiting_rx, 4, StdDuration::from_secs(1));

        let (stopped_tx, stopped_rx) = mpsc::channel();
        thread::spawn(move || {
            let stop_start = std::time::Instant::now();
            worker.stop_worker();
            stopped_tx.send(stop_start.elapsed()).unwrap();
        });
        let max_stop_time = RESOLVER_RUNTIME_SHUTDOWN_TIMEOUT + StdDuration::from_secs(1);
        let stop_elapsed = stopped_rx.recv_timeout(max_stop_time).unwrap();
        assert!(stop_elapsed <= max_stop_time);

        wait_until(StdDuration::from_secs(1), || {
            scheduler.check(3) == ResolveScheduleCheck::Ready
        });
        wait_until(StdDuration::from_secs(1), || {
            scheduler.check(4) == ResolveScheduleCheck::Ready
        });

        drop(finish_tx);
    }

    #[test]
    fn test_shutdown_before_task_start_returns_region_to_ready() {
        let _guard = test_guard();
        let (raw_scheduler, _receiver) = dummy_scheduler();
        let scheduler = TxnFileLockResolveScheduler::new(raw_scheduler);
        let limiter = TxnFileLockResolveLimiter::new(1);
        let mut runner = build_runner(scheduler.clone(), limiter);

        let (init_started_tx, init_started_rx) = mpsc::channel::<u64>();
        let (finish_tx, finish_rx) = oneshot::channel::<()>();
        let finish_rx = Arc::new(StdMutex::new(Some(finish_rx)));
        runner.test_hooks_mut().resolve_task = Some(Arc::new(|_task| Box::pin(async { Ok(0) })));
        runner.test_hooks_mut().init_client = Some(Arc::new(move || {
            let init_started_tx = init_started_tx.clone();
            let finish_rx = finish_rx.clone();
            Box::pin(async move {
                init_started_tx.send(5).unwrap();
                let finish_rx = finish_rx.lock().unwrap().take().unwrap();
                let _ = finish_rx.await;
                Ok(())
            })
        }));

        schedule_task(&scheduler, 5);
        runner.run(build_task(5));
        wait_for_region_signal(&init_started_rx, 5, StdDuration::from_secs(1));
        <TxnFileLockResolverRunner as Runnable>::shutdown(&mut runner);

        wait_until(StdDuration::from_secs(1), || {
            scheduler.check(5) == ResolveScheduleCheck::Ready
        });

        drop(finish_tx);
    }

    #[test]
    fn test_started_task_shutdown_returns_region_to_ready() {
        let _guard = test_guard();
        let (raw_scheduler, _receiver) = dummy_scheduler();
        let scheduler = TxnFileLockResolveScheduler::new(raw_scheduler);
        let limiter = TxnFileLockResolveLimiter::new(1);
        let mut runner = build_runner(scheduler.clone(), limiter);

        let (started_tx, started_rx) = mpsc::channel::<u64>();
        let (finish_tx, finish_rx) = oneshot::channel::<()>();
        let finish_rx = Arc::new(StdMutex::new(Some(finish_rx)));
        runner.test_hooks_mut().resolve_task = Some(Arc::new(move |task| {
            let region_id = task.region_id;
            let started_tx = started_tx.clone();
            let finish_rx = finish_rx.clone();
            Box::pin(async move {
                started_tx.send(region_id).unwrap();
                let finish_rx = finish_rx.lock().unwrap().take().unwrap();
                let _ = finish_rx.await;
                Ok(0)
            })
        }));

        schedule_task(&scheduler, 6);
        runner.run(build_task(6));
        wait_for_region_signal(&started_rx, 6, StdDuration::from_secs(1));
        <TxnFileLockResolverRunner as Runnable>::shutdown(&mut runner);

        wait_until(StdDuration::from_secs(1), || {
            scheduler.check(6) == ResolveScheduleCheck::Ready
        });

        drop(finish_tx);
    }

    #[test]
    fn test_running_task_gauge_counts_only_started_tasks() {
        let _guard = test_guard();
        let (raw_scheduler, _receiver) = dummy_scheduler();
        let scheduler = TxnFileLockResolveScheduler::new(raw_scheduler);
        let limiter = TxnFileLockResolveLimiter::new(1);
        let mut runner = build_runner(scheduler.clone(), limiter);
        let baseline = TXN_FILE_ACTIVE_RESOLVE_RUNNING_TASKS.get();

        let (started_tx, started_rx) = mpsc::channel::<u64>();
        let (waiting_tx, waiting_rx) = mpsc::channel::<u64>();
        let (finish_tx, finish_rx) = oneshot::channel::<()>();
        let finish_rx = Arc::new(StdMutex::new(Some(finish_rx)));
        runner.test_hooks_mut().resolve_task = Some(Arc::new(move |task| {
            let region_id = task.region_id;
            let started_tx = started_tx.clone();
            let finish_rx = finish_rx.clone();
            Box::pin(async move {
                started_tx.send(region_id).unwrap();
                if region_id == 7 {
                    let finish_rx = finish_rx.lock().unwrap().take().unwrap();
                    let _ = finish_rx.await;
                }
                Ok(0)
            })
        }));
        runner.test_hooks_mut().on_wait_for_limiter = Some(Arc::new(move |region_id| {
            if region_id == 8 {
                waiting_tx.send(region_id).unwrap();
            }
        }));

        schedule_task(&scheduler, 7);
        runner.run(build_task(7));
        wait_for_region_signal(&started_rx, 7, StdDuration::from_secs(1));
        schedule_task(&scheduler, 8);
        runner.run(build_task(8));
        wait_for_region_signal(&waiting_rx, 8, StdDuration::from_secs(1));
        assert_eq!(TXN_FILE_ACTIVE_RESOLVE_RUNNING_TASKS.get(), baseline + 1);

        <TxnFileLockResolverRunner as Runnable>::shutdown(&mut runner);
        wait_until(StdDuration::from_secs(1), || {
            TXN_FILE_ACTIVE_RESOLVE_RUNNING_TASKS.get() == baseline
        });
        wait_until(StdDuration::from_secs(1), || {
            scheduler.check(7) == ResolveScheduleCheck::Ready
        });
        wait_until(StdDuration::from_secs(1), || {
            scheduler.check(8) == ResolveScheduleCheck::Ready
        });

        drop(finish_tx);
    }

    #[test]
    fn test_async_init_retry_after_failure() {
        let _guard = test_guard();
        let (raw_scheduler, _receiver) = dummy_scheduler();
        let scheduler = TxnFileLockResolveScheduler::new(raw_scheduler);
        let limiter = TxnFileLockResolveLimiter::new(1);
        let mut runner = build_runner(scheduler.clone(), limiter);
        let init_attempts = Arc::new(AtomicUsize::new(0));
        let init_attempts_for_hook = init_attempts.clone();
        let (started_tx, started_rx) = mpsc::channel::<u64>();
        let error_before = TXN_FILE_ACTIVE_RESOLVE_LOCKS_TOTAL
            .with_label_values(&["error"])
            .get();

        runner.test_hooks_mut().resolve_task = Some(Arc::new(move |task| {
            let region_id = task.region_id;
            let started_tx = started_tx.clone();
            Box::pin(async move {
                started_tx.send(region_id).unwrap();
                Ok(0)
            })
        }));
        runner.test_hooks_mut().init_client = Some(Arc::new(move || {
            let init_attempts = init_attempts_for_hook.clone();
            Box::pin(async move {
                let attempt = init_attempts.fetch_add(1, Ordering::SeqCst);
                if attempt == 0 {
                    Err(tikv_util::box_err!(
                        "txn file lock resolver init failure (test)"
                    ))
                } else {
                    Ok(())
                }
            })
        }));

        schedule_task(&scheduler, 9);
        runner.run(build_task(9));
        wait_until(StdDuration::from_secs(1), || {
            scheduler.check(9) == ResolveScheduleCheck::Ready
        });
        assert_eq!(
            TXN_FILE_ACTIVE_RESOLVE_LOCKS_TOTAL
                .with_label_values(&["error"])
                .get(),
            error_before + 1
        );

        schedule_task(&scheduler, 10);
        runner.run(build_task(10));
        wait_for_region_signal(&started_rx, 10, StdDuration::from_secs(1));
        wait_until(StdDuration::from_secs(1), || {
            scheduler.check(10) == ResolveScheduleCheck::Ready
        });
        assert_eq!(init_attempts.load(Ordering::SeqCst), 2);
    }

    #[test]
    fn test_admission_capacity_drops_without_blocking_run() {
        let _guard = test_guard();
        let (raw_scheduler, _receiver) = dummy_scheduler();
        let scheduler = TxnFileLockResolveScheduler::new(raw_scheduler);
        let limiter = TxnFileLockResolveLimiter::new(1);
        let mut runner = build_runner_with_admission(scheduler.clone(), limiter, 1);
        let baseline = TXN_FILE_ACTIVE_RESOLVE_RUNNING_TASKS.get();
        let dropped_before = TXN_FILE_ACTIVE_RESOLVE_LOCKS_TOTAL
            .with_label_values(&["dropped"])
            .get();

        let (started_tx, started_rx) = mpsc::channel::<u64>();
        let (finish_tx, finish_rx) = oneshot::channel::<()>();
        let finish_rx = Arc::new(StdMutex::new(Some(finish_rx)));
        runner.test_hooks_mut().resolve_task = Some(Arc::new(move |task| {
            let region_id = task.region_id;
            let started_tx = started_tx.clone();
            let finish_rx = finish_rx.clone();
            Box::pin(async move {
                started_tx.send(region_id).unwrap();
                if region_id == 11 {
                    let finish_rx = finish_rx.lock().unwrap().take().unwrap();
                    let _ = finish_rx.await;
                }
                Ok(0)
            })
        }));

        schedule_task(&scheduler, 11);
        runner.run(build_task(11));
        wait_for_region_signal(&started_rx, 11, StdDuration::from_secs(1));
        assert_eq!(TXN_FILE_ACTIVE_RESOLVE_RUNNING_TASKS.get(), baseline + 1);

        schedule_task(&scheduler, 12);
        let run_start = std::time::Instant::now();
        runner.run(build_task(12));
        assert!(run_start.elapsed() < StdDuration::from_millis(500));
        assert_eq!(scheduler.check(12), ResolveScheduleCheck::Ready);
        assert_eq!(
            TXN_FILE_ACTIVE_RESOLVE_LOCKS_TOTAL
                .with_label_values(&["dropped"])
                .get(),
            dropped_before + 1
        );

        <TxnFileLockResolverRunner as Runnable>::shutdown(&mut runner);
        wait_until(StdDuration::from_secs(1), || {
            TXN_FILE_ACTIVE_RESOLVE_RUNNING_TASKS.get() == baseline
        });

        drop(finish_tx);
    }

    #[test]
    fn test_panic_in_spawned_task_counts_panic_and_returns_region_to_ready() {
        let _guard = test_guard();
        let (raw_scheduler, _receiver) = dummy_scheduler();
        let scheduler = TxnFileLockResolveScheduler::new(raw_scheduler);
        let limiter = TxnFileLockResolveLimiter::new(1);
        let mut runner = build_runner(scheduler.clone(), limiter.clone());
        let panic_before = TXN_FILE_ACTIVE_RESOLVE_LOCKS_TOTAL
            .with_label_values(&["panic"])
            .get();

        runner.test_hooks_mut().resolve_task = Some(Arc::new(move |_task| {
            Box::pin(async move {
                panic!("resolve task panic");
                #[allow(unreachable_code)]
                Ok(0)
            })
        }));

        schedule_task(&scheduler, 13);
        runner.run(build_task(13));

        wait_until(StdDuration::from_secs(1), || {
            scheduler.check(13) == ResolveScheduleCheck::Ready
        });
        assert_eq!(
            TXN_FILE_ACTIVE_RESOLVE_LOCKS_TOTAL
                .with_label_values(&["panic"])
                .get(),
            panic_before + 1
        );
        assert_eq!(TXN_FILE_ACTIVE_RESOLVE_RUNNING_TASKS.get(), 0);

        let _permit = block_on(async {
            timeout(StdDuration::from_secs(1), limiter.acquire())
                .await
                .unwrap()
                .unwrap()
        });
    }

    #[test]
    fn test_panic_before_task_start_counts_panic_and_returns_region_to_ready() {
        let _guard = test_guard();
        let (raw_scheduler, _receiver) = dummy_scheduler();
        let scheduler = TxnFileLockResolveScheduler::new(raw_scheduler);
        let limiter = TxnFileLockResolveLimiter::new(1);
        let mut runner = build_runner(scheduler.clone(), limiter.clone());
        let panic_before = TXN_FILE_ACTIVE_RESOLVE_LOCKS_TOTAL
            .with_label_values(&["panic"])
            .get();

        runner.test_hooks_mut().resolve_task = Some(Arc::new(|_task| Box::pin(async { Ok(0) })));
        runner.test_hooks_mut().before_task_start = Some(Arc::new(|region_id| {
            let _ = region_id;
            panic!("resolve task panic before start");
        }));

        schedule_task(&scheduler, 14);
        runner.run(build_task(14));

        wait_until(StdDuration::from_secs(1), || {
            scheduler.check(14) == ResolveScheduleCheck::Ready
        });
        assert_eq!(
            TXN_FILE_ACTIVE_RESOLVE_LOCKS_TOTAL
                .with_label_values(&["panic"])
                .get(),
            panic_before + 1
        );
        assert_eq!(TXN_FILE_ACTIVE_RESOLVE_RUNNING_TASKS.get(), 0);

        let _permit = block_on(async {
            timeout(StdDuration::from_secs(1), limiter.acquire())
                .await
                .unwrap()
                .unwrap()
        });
    }
}