// Copyright 2016 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.
//! Storage online config manager.

use std::{
    convert::TryInto,
    sync::{Arc, Mutex},
};

use concurrency_manager::ConcurrencyManager;
use file_system::{IoPriority, IoRateLimiter, IoType};
use online_config::{ConfigChange, ConfigManager, ConfigValue, Result as CfgResult};
use strum::IntoEnumIterator;
use tikv_kv::Engine;
use tikv_util::config::{ReadableDuration, ReadableSize};

use crate::storage::{
    TxnScheduler,
    config::{BlockCacheConfig, ExtraCfScanMode, StorageRuntimeConfig},
    lock_manager::LockManager,
};

pub struct StorageConfigManager<E: Engine, L: LockManager> {
    engine: Mutex<E>,
    concurrency_manager: ConcurrencyManager,
    io_rate_limiter: Arc<IoRateLimiter>,
    scheduler: Mutex<TxnScheduler<E, L>>,
    runtime_config: StorageRuntimeConfig,
    // whether scheduler worker pool shares unified read pool.
    use_separated_scheduler_pool: bool,
}

impl<E: Engine, L: LockManager> StorageConfigManager<E, L> {
    pub fn new(
        engine: E,
        concurrency_manager: ConcurrencyManager,
        io_rate_limiter: Arc<IoRateLimiter>,
        scheduler: TxnScheduler<E, L>,
        runtime_config: StorageRuntimeConfig,
        use_separated_scheduler_pool: bool,
    ) -> Self {
        StorageConfigManager {
            engine: Mutex::new(engine),
            concurrency_manager,
            io_rate_limiter,
            scheduler: Mutex::new(scheduler),
            runtime_config,
            use_separated_scheduler_pool,
        }
    }

    fn set_block_cache_capacity(&self, capacity: ConfigValue) -> CfgResult<()> {
        let capacity = match capacity {
            ConfigValue::None => BlockCacheConfig::default_capacity(),
            capacity => capacity.into(),
        };
        let kvengine = self
            .engine
            .lock()
            .unwrap()
            .get_kvengine()
            .ok_or("storage engine does not expose kvengine block cache")?;
        kvengine.set_block_cache_capacity(capacity.0)?;
        Ok(())
    }
}

impl<E: Engine, L: LockManager> ConfigManager for StorageConfigManager<E, L> {
    fn dispatch(&mut self, mut change: ConfigChange) -> CfgResult<()> {
        if let Some(ConfigValue::Module(mut block_cache)) = change.remove("block_cache") {
            if let Some(capacity) = block_cache.remove("capacity") {
                self.set_block_cache_capacity(capacity)?;
            }
        } else if let Some(ConfigValue::Module(_flow_control)) = change.remove("flow_control") {
            // TODO: support OnlineConfig for kvengine's flow_control
        }
        if let Some(ConfigValue::Module(mut io_rate_limit)) = change.remove("io_rate_limit") {
            if let Some(limit) = io_rate_limit.remove("max_bytes_per_sec") {
                let limit: ReadableSize = limit.into();
                self.io_rate_limiter.set_io_rate_limit(limit.0 as usize);
            }

            for t in IoType::iter() {
                if let Some(priority) = io_rate_limit.remove(&(t.as_str().to_owned() + "_priority"))
                {
                    let priority: IoPriority = priority.try_into()?;
                    self.io_rate_limiter.set_io_priority(t, priority);
                }
            }
        }
        if let Some(v) = change.get("scheduler_worker_pool_size") {
            if self.use_separated_scheduler_pool {
                let pool_size: usize = v.into();
                self.scheduler.lock().unwrap().scale_pool_size(pool_size);
            } else {
                warn!(
                    "cannot scale scheduler worker pool when config 'use_separated_scheduler_pool' is false."
                );
            }
        }
        if let Some(v) = change.remove("max_txn_file_locks_per_region") {
            let limit: usize = v.into();
            self.scheduler
                .lock()
                .unwrap()
                .set_max_txn_file_locks_per_region(limit);
        }
        if let Some(v) = change.remove("scheduler_async_resolve_lock_batch_interval") {
            let interval: ReadableDuration = v.into();
            self.scheduler
                .lock()
                .unwrap()
                .set_async_resolve_lock_batch_interval(interval.0);
        }
        if let Some(v) = change.remove("range_scan_reschedule_check_keys") {
            self.runtime_config
                .read_reschedule()
                .set_range_scan_reschedule_check_keys(v.into());
        }
        if let Some(v) = change.remove("range_scan_max_time_slice_us") {
            self.runtime_config
                .read_reschedule()
                .set_range_scan_max_time_slice_us(v.into());
        }
        if let Some(v) = change.remove("batch_get_fairness_checkpoint_keys") {
            self.runtime_config
                .read_reschedule()
                .set_batch_get_fairness_checkpoint_keys(v.into());
        }
        if let Some(v) = change.remove("cloud_store_fair_cpu_checkpoint_entries") {
            self.runtime_config
                .read_reschedule()
                .set_cloud_store_fair_cpu_checkpoint_entries(v.into());
        }
        if let Some(v) = change.remove("enable_get_in_read_pool") {
            self.runtime_config.set_enable_get_in_read_pool(v.into());
        }
        if let Some(ConfigValue::Module(mut max_ts)) = change.remove("max_ts") {
            if let Some(v) = max_ts.remove("action_on_invalid_update") {
                let str_v: String = v.into();
                let action: concurrency_manager::ActionOnInvalidMaxTs = str_v.try_into()?;
                self.concurrency_manager
                    .set_action_on_invalid_max_ts_update(action);
            }
            if let Some(v) = max_ts.remove("max_drift") {
                let dur_v: ReadableDuration = v.into();
                self.concurrency_manager.set_max_ts_drift_allowance(dur_v.0);
            }
        }
        if let Some(v) = change.remove("extra_cf_scan_mode") {
            let mode: ExtraCfScanMode = String::from(v).try_into()?;
            self.scheduler.lock().unwrap().set_extra_cf_scan_mode(mode);
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        config::{ConfigController, Module, TikvConfig},
        storage::{
            TestEngineBuilder, TestStorageBuilderApiV1,
            config::{Config, ExtraCfScanMode},
            lock_manager::MockLockManager,
        },
    };

    #[test]
    fn test_dispatch_updates_only_its_scheduler_extra_cf_scan_mode() {
        let storage = TestStorageBuilderApiV1::new(MockLockManager::new())
            .config(Config::default())
            .build()
            .unwrap();
        let scheduler = storage.get_scheduler();
        let other_storage = TestStorageBuilderApiV1::new(MockLockManager::new())
            .config(Config::default())
            .build()
            .unwrap();
        let other_scheduler = other_storage.get_scheduler();
        let mut mgr = StorageConfigManager::new(
            storage.get_engine(),
            storage.get_concurrency_manager(),
            Arc::new(IoRateLimiter::new_for_test()),
            scheduler.clone(),
            storage.get_runtime_config(),
            false,
        );

        assert_eq!(scheduler.extra_cf_scan_mode(), ExtraCfScanMode::Both);
        assert_eq!(other_scheduler.extra_cf_scan_mode(), ExtraCfScanMode::Both);

        let mut change = ConfigChange::default();
        change.insert(
            "extra_cf_scan_mode".to_owned(),
            ConfigValue::String("pruned".to_owned()),
        );
        mgr.dispatch(change).unwrap();
        assert_eq!(scheduler.extra_cf_scan_mode(), ExtraCfScanMode::Pruned);
        assert_eq!(other_scheduler.extra_cf_scan_mode(), ExtraCfScanMode::Both);

        let mut change = ConfigChange::default();
        change.insert(
            "extra_cf_scan_mode".to_owned(),
            ConfigValue::String("both".to_owned()),
        );
        mgr.dispatch(change).unwrap();
        assert_eq!(scheduler.extra_cf_scan_mode(), ExtraCfScanMode::Both);
        assert_eq!(other_scheduler.extra_cf_scan_mode(), ExtraCfScanMode::Both);

        let mut change = ConfigChange::default();
        change.insert(
            "extra_cf_scan_mode".to_owned(),
            ConfigValue::String("legacy".to_owned()),
        );
        mgr.dispatch(change).unwrap();
        assert_eq!(scheduler.extra_cf_scan_mode(), ExtraCfScanMode::Legacy);
        assert_eq!(other_scheduler.extra_cf_scan_mode(), ExtraCfScanMode::Both);
    }

    #[test]
    fn test_storage_cfg_update_block_cache() {
        let (mut cfg, _dir) = TikvConfig::with_tmp().unwrap();
        cfg.storage.block_cache.capacity = Some(ReadableSize::kb(1024));
        cfg.validate().unwrap();
        let use_separated_scheduler_pool = cfg.storage.use_separated_scheduler_pool;

        let storage = TestStorageBuilderApiV1::from_engine_and_lock_mgr(
            TestEngineBuilder::new().build().unwrap(),
            MockLockManager::new(),
        )
        .config(cfg.storage.clone())
        .build()
        .unwrap();
        let scheduler = storage.get_scheduler();

        let cfg_controller = ConfigController::new(cfg);
        cfg_controller.register(
            Module::Storage,
            Box::new(StorageConfigManager::new(
                storage.get_engine(),
                storage.get_concurrency_manager(),
                Arc::new(IoRateLimiter::new_for_test()),
                scheduler.clone(),
                storage.get_runtime_config(),
                use_separated_scheduler_pool,
            )),
        );

        if let Some(kvengine) = storage.get_engine().get_kvengine() {
            assert_eq!(kvengine.block_cache_capacity(), Some(1024 * 1024));
            cfg_controller
                .update_config("storage.block-cache.capacity", "2MB")
                .unwrap();
            assert_eq!(kvengine.block_cache_capacity(), Some(2 * 1024 * 1024));
        } else {
            // Verify that updating capacity on unsupported engines returns an error.
            assert!(
                cfg_controller
                    .update_config("storage.block-cache.capacity", "2MB")
                    .is_err()
            );
        }
    }
}