use std::time::Duration;

use anyhow::{Ok, Result};
use futures::future::BoxFuture;
use hyper::{Body, Client, Method, Request, StatusCode, Uri};
use scopeguard::defer;
use storage_backend::storage::{BackendStorage, create_backend_storage_for_test};
use tici::{
    HandleRange,
    config::{make_default_test_meta_config, make_default_test_worker_config},
    meta::reader::UnifiedShardInfo,
    worker::mocks::make_frag_meta,
};
use tici_common::{
    codec::{previous_key_helper, table_end_key, table_start_key},
    index::IndexKey,
};
use tici_grpc::{KeyRange, ReaderNodeHeartbeatRequest, ReaderNodeStatus};
use tici_system::readable::ReadableDuration;
use tici_test_utils::test_utils::{test_suite::E2ETestSuite, *};
use uuid::Uuid;

#[test]
fn test_manual_compaction() -> Result<()> {
    fail::cfg("worker_node_client_fail_fast", "return(2)").unwrap();
    defer!({
        fail::remove("worker_node_client_fail_fast");
    });
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    rt.block_on(async {
        let mut config = make_default_test_meta_config(true);
        config.shard.auto_compaction = false;
        config.shard.enable_split = false;
        config.shard.enable_merge = false;
        let status_addr = config.server.get_advertise_status_addr().to_string();

        let mut test_suite = E2ETestSuite::new_with_mock_cfg(
            config,
            make_default_test_worker_config(true),
        )
        .await
        .unwrap();

        let wh = test_suite.get_one_worker_node_handle();
        let mock = wh.shard_write_manager.as_mock();
        mock.mock_compact_process
            .store(true, std::sync::atomic::Ordering::Relaxed);
        // Make the compaction stay "running" for a while, so the endpoint has
        // time to observe it.
        mock.respond_compact_delay
            .store(1000, std::sync::atomic::Ordering::Relaxed);

        let table_id = 1;
        let index_id = 1001;
        let index_key = IndexKey::new_nullspace(table_id, index_id);

        test_suite
            .must_create_index(&index_key)
            .await
            .unwrap();

        test_suite
            .validate_add_index_result(&index_key)
            .await;

        let shard = test_suite.get_first_shard().await;
        let shard_id = shard.shard_id;
        let start_key = hex::encode(shard.range.encoded_start());
        let end_key = hex::encode(shard.range.encoded_end());

        for i in 1..20 {
            test_suite.must_append_delta_by_shard(&shard, &format!("A{}_{}", 222, i), 1, AppendExtraParam {
                do_manifest_check: true,
                upload_to_s3: false,
                ..Default::default()
            }).await;

            test_suite
                .must_append_base_frag(
                    &test_suite.get_one_worker_node_addr(),
                    shard.shard_id,
                    &index_key,
                    vec![make_frag_meta(
                        &index_key,
                        &mock_frag_path(format!("BASEBASE_{}", i)),
                    )],
                )
                .await
                .unwrap();
        }
        // A manual compaction request without required arguments will be reject.
        {
            let url = format!(
                "http://{}/op/meta/compact_range?range=delta&start_key={}&end_key={}&table_id={}",
                status_addr,
                start_key,
                end_key,
                table_id + 1
            );
            let resp = reqwest::get(&url).await.unwrap();
            assert!(
                resp.status() == StatusCode::BAD_REQUEST,
                "Manual compaction request not failed: {:?}",
                resp
            );
        }

        // A manual compaction request with inconsistent arguments will be reject.
        {
            let url = format!(
                "http://{}/op/meta/compact_range?range=delta&start_key={}&end_key={}&table_id={}&index_id={}",
                status_addr,
                start_key,
                end_key,
                table_id + 1,
                index_id
            );
            let resp = reqwest::get(&url).await.unwrap();
            assert!(
                resp.status() == StatusCode::INTERNAL_SERVER_ERROR,
                "Manual compaction request should fail: {:?}",
                resp
            );
        }

        {
            let url = format!(
                "http://{}/op/meta/compact_range?range=delta&start_key={}&end_key={}&table_id={}&index_id={}",
                status_addr, start_key, end_key, table_id, index_id
            );
            let resp = reqwest::get(&url).await.unwrap();
            assert!(
                resp.status().is_success(),
                "Manual compaction request should succeess: {:?}",
                resp
            );
        }

        {
            let url = format!(
                "http://{}/op/meta/compactions?start_key={}&end_key={}&table_id={}&index_id={}",
                status_addr, start_key, end_key, table_id, index_id
            );
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);

            let body = resp.json::<serde_json::Value>().await.unwrap();
            let shard_id_str = shard_id.to_string();
            let shard_json = body.get(shard_id_str.as_str()).unwrap();
            assert_eq!(
                shard_json.get("shard_id").unwrap().as_u64().unwrap(),
                shard_id
            );
            assert_eq!(
                shard_json.get("start_key").unwrap().as_str().unwrap(),
                start_key
            );
            assert_eq!(shard_json.get("end_key").unwrap().as_str().unwrap(), end_key);

            let compactions = shard_json.get("compactions").unwrap().as_array().unwrap();
            assert!(!compactions.is_empty(), "compactions={:?}", compactions);

            let c0 = &compactions[0];
            assert_eq!(c0.get("reason").unwrap().as_str().unwrap(), "Manual");
            assert_eq!(
                c0.get("key")
                    .unwrap()
                    .get("shard_id")
                    .unwrap()
                    .as_u64()
                    .unwrap(),
                shard_id
            );
        }

        // Compaction must finish.
        must_finally(async || {
            let shard = test_suite.get_first_shard().await;
            shard.progress.compaction_state.is_empty()
        })
        .await;

        mock.respond_compact_delay
            .store(200, std::sync::atomic::Ordering::Relaxed);

        {
            let url = format!(
                "http://{}/op/meta/compact_range?range=base&start_key={}&end_key={}&table_id={}&index_id={}",
                status_addr, start_key, end_key, table_id, index_id
            );
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(resp.status(), StatusCode::BAD_REQUEST, "{:?}", resp);
        }
        must_finally(async || {
            let shard = test_suite.get_first_shard().await;
            shard.progress.compaction_state.is_empty()
        }).await;

        must_finally(async || {
            let shard = test_suite.get_first_shard().await;
            shard.progress.compaction_state.is_empty()
        })
        .await;
        test_suite.check_meta_subscriber().await;
        test_suite.stop().await.unwrap();
    });
    Ok(())
}

#[test]
fn test_abort_compaction_via_status_server() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    rt.block_on(async {
        let mut config = make_default_test_meta_config(true);
        config.shard.auto_compaction = false;
        config.shard.enable_split = false;
        config.shard.enable_merge = false;
        let status_addr = config.server.get_advertise_status_addr().to_string();

        let mut test_suite = E2ETestSuite::new_with_mock_cfg(
            config,
            make_default_test_worker_config(true),
        )
        .await
        .unwrap();

        let wh = test_suite.get_one_worker_node_handle();
        let mock = wh.shard_write_manager.as_mock();
        mock.mock_compact_process
            .store(true, std::sync::atomic::Ordering::Relaxed);
        mock.respond_compact_delay
            .store(1000, std::sync::atomic::Ordering::Relaxed);

        let table_id = 1;
        let index_id = 1001;
        let index_key = IndexKey::new_nullspace(table_id, index_id);

        test_suite
            .must_create_index(&index_key)
            .await
            .unwrap();
        test_suite
            .validate_add_index_result(&index_key)
            .await;

        let shard = test_suite.get_first_shard().await;
        let shard_id = shard.shard_id;
        let start_key = hex::encode(shard.range.encoded_start());
        let end_key = hex::encode(shard.range.encoded_end());

        for i in 1..20 {
            test_suite
                .must_append_delta_by_shard(
                    &shard,
                    &format!("A{}_{}", 333, i),
                    1,
                    AppendExtraParam {
                        do_manifest_check: true,
                        upload_to_s3: false,
                        ..Default::default()
                    },
                )
                .await;

            test_suite
                .must_append_base_frag(
                    &test_suite.get_one_worker_node_addr(),
                    shard.shard_id,
                    &index_key,
                    vec![make_frag_meta(
                        &index_key,
                        &mock_frag_path(format!("BASEABORT_{}", i)),
                    )],
                )
                .await
                .unwrap();
        }

        {
            let url = format!(
                "http://{}/op/meta/compact_range?range=delta&start_key={}&end_key={}&table_id={}&index_id={}",
                status_addr, start_key, end_key, table_id, index_id
            );
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
        }

        let (_index, shard) = test_suite
            .get_shard_and_index(&index_key, shard_id)
            .await
            .unwrap();
        must_finally(|| {
            let shard = shard.clone();
            async move { !shard.read().await.progress.compaction_state.is_empty() }
        })
        .await;

        let compaction_aborted = shard.read().await.stats.get_compaction_abort();

        {
            let url = format!(
                "http://{}/op/meta/compactions?start_key={}&end_key={}&table_id={}&index_id={}",
                status_addr, start_key, end_key, table_id, index_id
            );
            let resp = reqwest::Client::new().delete(&url).send().await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);

            let body = resp.json::<serde_json::Value>().await.unwrap();
            assert_eq!(body.get("aborted").unwrap().as_u64().unwrap(), 1);

            let shard_key = shard_id.to_string();
            let shard_json = body.get("shards").unwrap().get(shard_key.as_str()).unwrap();
            assert_eq!(shard_json.get("shard_id").unwrap().as_u64().unwrap(), shard_id);
            assert_eq!(shard_json.get("aborted").unwrap().as_u64().unwrap(), 1);
        }

        must_finally(|| {
            let shard = shard.clone();
            async move {
                let guard = shard.read().await;
                guard.progress.compaction_state.is_empty()
                    && guard.stats.get_compaction_abort() == compaction_aborted + 1
            }
        })
        .await;

        test_suite.check_meta_subscriber().await;
        test_suite.stop().await.unwrap();
    });
    Ok(())
}

#[test]
fn test_manual_reschedule() -> Result<()> {
    fail::cfg("worker_node_client_fail_fast", "return(2)").unwrap();
    defer!({
        fail::remove("worker_node_client_fail_fast");
    });
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    rt.block_on(async {
        let mut config = make_default_test_meta_config(true);
        config.shard.auto_compaction = false;
        config.shard.enable_split = false;
        config.shard.enable_merge = false;
        let status_addr = config.server.get_advertise_status_addr().to_string();

        let mut test_suite = E2ETestSuite::new_with_mock_cfg(
            config,
            make_default_test_worker_config(true),
        )
        .await
        .unwrap();

        let table_id = 1;
        let index_id = 1001;
        let index_key = IndexKey::new_nullspace(table_id, index_id);

        test_suite
            .must_create_index(&index_key)
            .await
            .unwrap();

        test_suite
            .validate_add_index_result(&index_key)
            .await;

        let shard = test_suite.get_first_shard().await;
        let start_key = shard.range.encoded_start();
        let end_key = shard.range.encoded_end();
        let start_key = hex::encode(start_key);
        let end_key = hex::encode(end_key);

        let new_wn_addr = "127.0.0.1:8512".to_string();
        test_suite
            .add_mocked_worker_node(Some(new_wn_addr.clone()))
            .await;

        test_suite.wait_heartbeat().await;
        must_finally(async || {
            test_suite
                .get_meta_service()
                .worker_node_mgr
                .get_worker_nodes_len()
                .await
                == 2
        })
        .await;

        {
            let url = format!(
                "http://{}/op/meta/reschedule_range?writer=abc.baba%20.z%3A9999&start_key={}&end_key={}&table_id={}&index_id={}",
                status_addr,
                start_key,
                end_key,
                table_id,
                index_id,
            );
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(
                resp.status(),
                StatusCode::OK,
                "{:?}",
                resp,
            );
        }

        {
            let url = format!(
                "http://{}/op/meta/reschedule_range?writer=127.0.0.1:8512&start_key={}&end_key={}&table_id={}&index_id={}",
                status_addr,
                start_key,
                end_key,
                table_id,
                index_id,
            );
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(
                resp.status(),
                StatusCode::OK,
                "{:?}",
                resp,
            );

            assert_eq!(test_suite.collect_shards_from_restore().await.first().unwrap().writer, "127.0.0.1:8512");
        }

        {
            // Let meta service choose the writer
            let url = format!(
                "http://{}/op/meta/reschedule_range?start_key={}&end_key={}&table_id={}&index_id={}",
                status_addr, start_key, end_key, table_id, index_id,
            );
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
            assert_eq!(
                test_suite
                    .collect_shards_from_restore()
                    .await
                    .first()
                    .unwrap()
                    .writer,
                "127.0.0.1:8510"
            );

            let url = format!(
                "http://{}/op/meta/reschedule_range?start_key={}&end_key={}&table_id={}&index_id={}",
                status_addr, start_key, end_key, table_id, index_id,
            );
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
            assert_eq!(
                test_suite
                    .collect_shards_from_restore()
                    .await
                    .first()
                    .unwrap()
                    .writer,
                "127.0.0.1:8512"
            );
        }

        test_suite.check_meta_subscriber().await;
        test_suite.stop().await.unwrap();
    });
    Ok(())
}

#[test]
fn test_manual_split() -> Result<()> {
    fail::cfg("worker_node_client_fail_fast", "return(2)").unwrap();
    defer!({
        fail::remove("worker_node_client_fail_fast");
    });
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    rt.block_on(async {
        let mut config = make_default_test_meta_config(true);
        config.shard.auto_compaction = false;
        config.shard.enable_split = true;
        config.shard.enable_merge = false;
        let status_addr = config.server.get_advertise_status_addr().to_string();
        let max_size = config.shard.max_size.0;
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(config, make_default_test_worker_config(true))
                .await
                .unwrap();

        let table_id = 1;
        let index_id = 1001;
        let index_key = IndexKey::new_nullspace(table_id, index_id);

        test_suite.must_create_index(&index_key).await.unwrap();
        test_suite.validate_add_index_result(&index_key).await;
        let shard = test_suite.get_first_shard().await;

        test_suite
            .must_append_delta_by_shard(
                &shard,
                "CCC",
                15,
                AppendExtraParam {
                    do_manifest_check: true,
                    upload_to_s3: false,
                    size: Some(max_size),
                    count: Some(10000),
                    segs_count: Some(4),
                    ..Default::default()
                },
            )
            .await;

        let shard = test_suite.get_first_shard().await;
        let start_key = shard.range.encoded_start();
        let start_key = hex::encode(start_key);
        {
            let url = format!(
                "http://{}/op/meta/split_shard?start_key={}&table_id={}&index_id={}",
                status_addr, start_key, table_id, index_id,
            );
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp,);
        }

        test_suite.check_meta_subscriber().await;
        test_suite.stop().await.unwrap();
    });
    Ok(())
}

#[test]
fn test_manual_gc_frags() -> Result<()> {
    fail::cfg("worker_node_client_fail_fast", "return(2)").unwrap();
    defer!({
        fail::remove("worker_node_client_fail_fast");
    });
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    rt.block_on(async {
        let mut config = make_default_test_meta_config(true);
        config.shard.auto_compaction = false;
        config.shard.enable_split = false;
        config.shard.enable_merge = false;
        config.shard.gc_config.gc_worker_check_interval = ReadableDuration::days(1);
        let status_addr = config.server.get_advertise_status_addr().to_string();

        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(config, make_default_test_worker_config(true))
                .await
                .unwrap();

        let index_key = IndexKey::new_nullspace(1, 1001);
        test_suite.must_create_index(&index_key).await.unwrap();
        test_suite.validate_add_index_result(&index_key).await;

        let s3_cfg = test_suite.get_meta_service().config.load().s3.clone();
        let storage: std::sync::Arc<dyn BackendStorage> =
            create_backend_storage_for_test(&s3_cfg).await.unwrap();

        // Create an unreferenced fragment object; its "f_*" is intentionally not a
        // valid PD timestamp so gc_stale_fragment will delete it anyway.
        let stale_key = format!(
            "{}/{}/{}/f_gc_{}/multi.0",
            s3_cfg.prefix,
            tici::meta::writer::FRAG_FILE_PREFIX,
            IndexKey::new_nullspace(999, 999).to_s3_prefix(),
            Uuid::new_v4()
        );
        storage
            .put_data(&s3_cfg.bucket, &stale_key, b"gc-me".to_vec())
            .await
            .unwrap();

        let frags_prefix = format!("{}/{}", s3_cfg.prefix, tici::meta::writer::FRAG_FILE_PREFIX);
        assert!(
            storage
                .list_all(&s3_cfg.bucket, &frags_prefix)
                .await
                .unwrap()
                .iter()
                .any(|k| k == &stale_key),
            "stale object not uploaded, key={}",
            stale_key
        );

        {
            let url = format!("http://{}/op/meta/gc?frags=true", status_addr);
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
            let body = resp.text().await.unwrap();
            assert!(
                body.contains("frags:"),
                "manual gc response should mention frags, body={}",
                body
            );
        }

        must_finally(|| {
            let storage = storage.clone();
            let bucket = s3_cfg.bucket.clone();
            let frags_prefix = frags_prefix.clone();
            let stale_key = stale_key.clone();
            async move {
                !storage
                    .list_all(&bucket, &frags_prefix)
                    .await
                    .unwrap()
                    .iter()
                    .any(|k| k == &stale_key)
            }
        })
        .await;

        test_suite.check_meta_subscriber().await;
        test_suite.stop().await.unwrap();
    });
    Ok(())
}

#[test]
fn test_manual_reader_schedule() -> Result<()> {
    fail::cfg("worker_node_client_fail_fast", "return(2)").unwrap();
    defer!({
        fail::remove("worker_node_client_fail_fast");
    });
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    rt.block_on(async {
        let mut config = make_default_test_meta_config(true);
        config.shard.auto_compaction = false;
        config.shard.enable_split = false;
        config.shard.enable_merge = false;
        let status_addr = config.server.get_advertise_status_addr().to_string();
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(config, make_default_test_worker_config(true))
                .await
                .unwrap();

        test_suite.add_mocked_reader_node(
            "127.0.0.1:19410".to_string(),
            "127.0.0.1:19411".to_string(), Some(Duration::from_millis(400))).await;
        let index_key = IndexKey::new(111, 1, 1001);

        test_suite.must_create_index(&index_key).await.unwrap();
        test_suite.validate_add_index_result(&index_key).await;

        let range = HandleRange::new(
            table_start_key(&index_key).to_vec(),
            previous_key_helper(table_end_key(&index_key).as_slice(), 1),
        );

        {
            let url = format!(
                "http://{}/op/meta/reader/reschedule?start_key={}&end_key={}&table_id={}&index_id={}&keyspace_id={}",
                status_addr, hex::encode(range.encoded_start()), hex::encode(range.encoded_end()), index_key.table_id, index_key.index_id, index_key.keyspace_id
            );
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp,);
        }

        let res = test_suite
            .send_get_shard_local_cache_request(&index_key, vec![KeyRange {
                start_key: vec![],
                end_key: vec![],
            }])
            .await
            .unwrap();
        test_suite.check_valid_shard_allocation(&res, 0).unwrap();

        test_suite.check_meta_subscriber().await;
        test_suite.stop().await.unwrap();
    });
    Ok(())
}

#[test]
fn test_reader_allocation_by_node() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    rt.block_on(async {
        let mut config = make_default_test_meta_config(true);
        config.shard.auto_compaction = false;
        config.shard.enable_split = false;
        config.shard.enable_merge = false;
        let status_addr = config.server.get_advertise_status_addr().to_string();
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(config, make_default_test_worker_config(true))
                .await
                .unwrap();

        let reader_pool = test_suite.get_meta_service().get_reader_pool();
        let reader1_addr = "reader-1".to_string();
        let reader2_addr = "reader-2".to_string();

        let req1 = ReaderNodeHeartbeatRequest {
            addr: reader1_addr.clone(),
            status: Some(ReaderNodeStatus::default()),
            tiflash_server_addr: "tiflash-1".to_string(),
            cached_shard_ids: vec![],
            cached_shards: vec![],
        };
        let reader1 = reader_pool.register_or_update_reader(&req1).await.unwrap();

        let req2 = ReaderNodeHeartbeatRequest {
            addr: reader2_addr.clone(),
            status: Some(ReaderNodeStatus::default()),
            tiflash_server_addr: "tiflash-2".to_string(),
            cached_shard_ids: vec![],
            cached_shards: vec![],
        };
        let _reader2 = reader_pool.register_or_update_reader(&req2).await.unwrap();

        let mut shard = tici::ShardInfo::new_for_test(HandleRange::new(vec![1], vec![2]));
        shard.index_key = IndexKey::new_nullspace(1, 2);
        shard.shard_id = 42;
        shard.epoch = 7;
        shard.seq = 9;

        reader1
            .add_shard_cache(UnifiedShardInfo::from_shard_info(&shard))
            .await;

        {
            let url = format!(
                "http://{}/op/meta/reader/allocation/by_node?addr=missing",
                status_addr
            );
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
            let body = resp.json::<serde_json::Value>().await.unwrap();
            assert!(body.as_object().unwrap().is_empty());
        }

        {
            let url = format!(
                "http://{}/op/meta/reader/allocation/by_node?addr={}",
                status_addr, reader1_addr
            );
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
            let body = resp.json::<serde_json::Value>().await.unwrap();
            let shard_json = body.get("42").unwrap();
            assert_eq!(shard_json.get("shard_id").unwrap().as_u64().unwrap(), 42);
            assert_eq!(shard_json.get("epoch").unwrap().as_u64().unwrap(), 7);
            assert_eq!(shard_json.get("seq").unwrap().as_u64().unwrap(), 9);
            assert_eq!(shard_json.get("start_key").unwrap().as_str().unwrap(), "01");
            assert_eq!(shard_json.get("end_key").unwrap().as_str().unwrap(), "02");
        }

        {
            let url = format!("http://{}/op/meta/reader/allocation/by_node", status_addr);
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
            let body = resp.json::<serde_json::Value>().await.unwrap();
            assert!(body.get(&reader1_addr).is_some());
            assert!(body.get(&reader2_addr).is_some());
        }

        test_suite.stop().await.unwrap();
    });
    Ok(())
}

#[test]
fn test_reader_allocation_add() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    rt.block_on(async {
        let mut config = make_default_test_meta_config(true);
        config.shard.auto_compaction = false;
        config.shard.enable_split = false;
        config.shard.enable_merge = false;
        config.cache_scheduling.eager_schedule = false;
        let status_addr = config.server.get_advertise_status_addr().to_string();
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(config, make_default_test_worker_config(true))
                .await
                .unwrap();

        let reader1_addr = "127.0.0.1:19510".to_string();
        let reader2_addr = "127.0.0.1:19512".to_string();
        test_suite
            .add_mocked_reader_node(
                reader1_addr.clone(),
                "127.0.0.1:19511".to_string(),
                Some(Duration::from_millis(200)),
            )
            .await;
        test_suite
            .add_mocked_reader_node(
                reader2_addr.clone(),
                "127.0.0.1:19513".to_string(),
                Some(Duration::from_millis(200)),
            )
            .await;
        must_finally(|| {
            let reader_pool = test_suite.get_meta_service().get_reader_pool();
            async move { reader_pool.get_readers_len().await == 2 }
        })
        .await;

        let index_key_a = IndexKey::new(111, 10, 1010);
        let index_key_b = IndexKey::new(111, 11, 1011);
        test_suite.must_create_index(&index_key_a).await.unwrap();
        test_suite.validate_add_index_result(&index_key_a).await;
        test_suite.must_create_index(&index_key_b).await.unwrap();
        test_suite.validate_add_index_result(&index_key_b).await;

        let shard_a = test_suite
            .collect_shards_from_index(&index_key_a)
            .await
            .first()
            .unwrap()
            .clone();
        let shard_b = test_suite
            .collect_shards_from_index(&index_key_b)
            .await
            .first()
            .unwrap()
            .clone();

        let url_a_reader2 = format!(
            "http://{}/op/meta/reader/allocation/add?start_key={}&end_key={}&table_id={}&index_id={}&keyspace_id={}&addr={}",
            status_addr,
            hex::encode(shard_a.range.encoded_start()),
            hex::encode(shard_a.range.encoded_end()),
            index_key_a.table_id,
            index_key_a.index_id,
            index_key_a.keyspace_id,
            reader2_addr,
        );
        let resp = reqwest::get(&url_a_reader2).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
        assert!(resp.text().await.unwrap().contains("error_code=Success"));

        let shard_a_ref = test_suite
            .get_meta_service()
            .shard_cache_manager
            .get_shard_by_id(&shard_a.shard_id)
            .await
            .unwrap();
        must_finally(|| {
            let shard = shard_a_ref.clone();
            let reader2 = reader2_addr.clone();
            async move { shard.get_reader_addrs().await == vec![reader2] }
        })
        .await;

        let url_b = format!(
            "http://{}/op/meta/reader/allocation/add?start_key={}&end_key={}&table_id={}&index_id={}&keyspace_id={}",
            status_addr,
            hex::encode(shard_b.range.encoded_start()),
            hex::encode(shard_b.range.encoded_end()),
            index_key_b.table_id,
            index_key_b.index_id,
            index_key_b.keyspace_id,
        );
        let resp = reqwest::get(&url_b).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
        assert!(resp.text().await.unwrap().contains("error_code=Success"));

        let shard_b_ref = test_suite
            .get_meta_service()
            .shard_cache_manager
            .get_shard_by_id(&shard_b.shard_id)
            .await
            .unwrap();
        must_finally(|| {
            let shard = shard_b_ref.clone();
            async move { shard.get_reader_addrs().await.len() == 1 }
        })
        .await;
        let existing = shard_b_ref.get_reader_addrs().await.first().unwrap().clone();
        let designated_for_b = if existing == reader1_addr {
            reader2_addr.clone()
        } else {
            reader1_addr.clone()
        };

        let url_a_reader1 = format!(
            "http://{}/op/meta/reader/allocation/add?start_key={}&end_key={}&table_id={}&index_id={}&keyspace_id={}&addr={}",
            status_addr,
            hex::encode(shard_a.range.encoded_start()),
            hex::encode(shard_a.range.encoded_end()),
            index_key_a.table_id,
            index_key_a.index_id,
            index_key_a.keyspace_id,
            reader1_addr,
        );
        let resp = reqwest::get(&url_a_reader1).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
        assert!(resp.text().await.unwrap().contains("error_code=Success"));
        must_finally(|| {
            let shard = shard_a_ref.clone();
            let reader1 = reader1_addr.clone();
            let reader2 = reader2_addr.clone();
            async move {
                let mut addrs = shard.get_reader_addrs().await;
                addrs.sort();
                addrs == vec![reader1, reader2]
            }
        })
        .await;

        let url_b_designated = format!("{}&addr={}", url_b, designated_for_b);
        let resp = reqwest::get(&url_b_designated).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
        assert!(resp.text().await.unwrap().contains("error_code=Success"));
        must_finally(|| {
            let shard = shard_b_ref.clone();
            let reader1 = reader1_addr.clone();
            let reader2 = reader2_addr.clone();
            async move {
                let mut addrs = shard.get_reader_addrs().await;
                addrs.sort();
                addrs == vec![reader1, reader2]
            }
        })
        .await;

        test_suite.stop().await.unwrap();
    });
    Ok(())
}

#[test]
fn test_block_reader_filters_active_readers() -> Result<()> {
    fail::cfg("worker_node_client_fail_fast", "return(2)").unwrap();
    defer!({
        fail::remove("worker_node_client_fail_fast");
    });
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    rt.block_on(async {
        let mut config = make_default_test_meta_config(true);
        config.shard.auto_compaction = false;
        config.shard.enable_split = false;
        config.shard.enable_merge = false;
        let status_addr = config.server.get_advertise_status_addr().to_string();
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(config, make_default_test_worker_config(true))
                .await
                .unwrap();

        let blocked_addr = "127.0.0.1:19410".to_string();
        let blocked_tiflash = "127.0.0.1:19411".to_string();
        let other_addr = "127.0.0.1:19412".to_string();
        let other_tiflash = "127.0.0.1:19413".to_string();
        let missing_addr = "127.0.0.1:19499".to_string();

        {
            let url = format!("http://{}/op/meta/reader/block", status_addr);
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
            let blocked_list = resp.json::<Vec<String>>().await.unwrap();
            assert!(
                blocked_list.is_empty(),
                "blocked list should be empty, list={:?}",
                blocked_list
            );
        }

        {
            let url = format!(
                "http://{}/op/meta/reader/block?addrs={},{}",
                status_addr, blocked_addr, missing_addr
            );
            let resp = reqwest::Client::new().post(&url).send().await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
        }

        must_finally(|| {
            let reader_pool = test_suite.get_meta_service().get_reader_pool();
            let blocked_addr = blocked_addr.clone();
            let missing_addr = missing_addr.clone();
            async move {
                reader_pool.is_reader_blocked(&blocked_addr).await
                    && reader_pool.is_reader_blocked(&missing_addr).await
            }
        })
        .await;

        {
            let url = format!("http://{}/op/meta/reader/block", status_addr);
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
            let blocked_list = resp.json::<Vec<String>>().await.unwrap();
            assert!(
                blocked_list.contains(&blocked_addr),
                "blocked list missing addr, list={:?}",
                blocked_list
            );
            assert!(
                blocked_list.contains(&missing_addr),
                "blocked list missing addr, list={:?}",
                blocked_list
            );
        }

        test_suite
            .add_mocked_reader_node(
                blocked_addr.clone(),
                blocked_tiflash,
                Some(Duration::from_millis(200)),
            )
            .await;
        test_suite
            .add_mocked_reader_node(
                other_addr.clone(),
                other_tiflash,
                Some(Duration::from_millis(200)),
            )
            .await;

        must_finally(|| {
            let reader_pool = test_suite.get_meta_service().get_reader_pool();
            async move { reader_pool.get_readers_len().await == 2 }
        })
        .await;

        must_finally(|| {
            let reader_pool = test_suite.get_meta_service().get_reader_pool();
            let blocked_addr = blocked_addr.clone();
            let other_addr = other_addr.clone();
            async move {
                let active = reader_pool.get_readers_for_schedule().await;
                active.len() == 1
                    && active[0].addr == other_addr
                    && !reader_pool.is_reader_blocked(&other_addr).await
                    && reader_pool.is_reader_blocked(&blocked_addr).await
            }
        })
        .await;

        {
            let url = format!(
                "http://{}/op/meta/reader/block?addr={}",
                status_addr, blocked_addr
            );
            let resp = reqwest::Client::new().delete(&url).send().await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
        }

        must_finally(|| {
            let reader_pool = test_suite.get_meta_service().get_reader_pool();
            let blocked_addr = blocked_addr.clone();
            let missing_addr = missing_addr.clone();
            async move {
                !reader_pool.is_reader_blocked(&blocked_addr).await
                    && reader_pool.is_reader_blocked(&missing_addr).await
            }
        })
        .await;

        {
            let url = format!("http://{}/op/meta/reader/block", status_addr);
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
            let blocked_list = resp.json::<Vec<String>>().await.unwrap();
            assert!(
                !blocked_list.contains(&blocked_addr),
                "blocked list should not include addr, list={:?}",
                blocked_list
            );
            assert_eq!(
                blocked_list,
                vec![missing_addr.clone()],
                "blocked list mismatch",
            );
        }

        must_finally(|| {
            let reader_pool = test_suite.get_meta_service().get_reader_pool();
            let blocked_addr = blocked_addr.clone();
            let other_addr = other_addr.clone();
            async move {
                let active = reader_pool.get_readers_for_schedule().await;
                active.len() == 2
                    && active.iter().any(|r| r.addr == blocked_addr)
                    && active.iter().any(|r| r.addr == other_addr)
            }
        })
        .await;

        test_suite.check_meta_subscriber().await;
        test_suite.stop().await.unwrap();
    });
    Ok(())
}

#[test]
fn test_config_update() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;

    let f = |persist: bool| -> BoxFuture<'static, ()> {
        let a = async move {
            let config = make_default_test_meta_config(true);
            let status_addr = config.server.get_advertise_status_addr().to_string();
            let mut test_suite =
                E2ETestSuite::new_with_mock_cfg(config, make_default_test_worker_config(true))
                    .await
                    .unwrap();

            let client = Client::new();
            let uri = if persist {
                Uri::builder()
                    .scheme("http")
                    .authority(status_addr)
                    .path_and_query("/config")
                    .build()
                    .unwrap()
            } else {
                Uri::builder()
                    .scheme("http")
                    .authority(status_addr)
                    .path_and_query("/config?persist=false")
                    .build()
                    .unwrap()
            };
            let mut req = Request::new(Body::from(
                "{\"shard.gc_config.gc_worker_check_interval\": \"1234567ms\"}",
            ));
            *req.method_mut() = Method::POST;
            *req.uri_mut() = uri.clone();
            let resp = client.request(req).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK);
            assert_eq!(
                test_suite
                    .get_meta_service()
                    .config
                    .load()
                    .shard
                    .gc_config
                    .gc_worker_check_interval,
                ReadableDuration::millis(1234567)
            );
            let mut req = Request::new(Body::from("{\"server.addr\": \"hahaha\"}"));
            *req.method_mut() = Method::POST;
            *req.uri_mut() = uri.clone();
            let resp = client.request(req).await.unwrap();
            assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);
            test_suite.stop().await.unwrap();
        };
        Box::pin(a)
    };
    rt.block_on(f(false));
    Ok(())
}

#[test]
fn test_metrics() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;

    rt.block_on(async {
        let config = make_default_test_meta_config(true);
        let status_addr = config.server.get_advertise_status_addr().to_string();
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(config, make_default_test_worker_config(true))
                .await
                .unwrap();

        {
            let url = format!("http://{}/metrics", status_addr);
            let resp = reqwest::get(&url).await.unwrap();
            assert!(resp.status().is_success());
            assert!(
                resp.text()
                    .await
                    .unwrap()
                    .contains("tici_grpc_handle_duration")
            )
        }

        test_suite.stop().await.unwrap();
    });
    Ok(())
}

#[test]
fn test_worker_allocation_shards() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;

    rt.block_on(async {
        let mut config = make_default_test_meta_config(true);
        config.shard.auto_compaction = false;
        config.shard.enable_split = false;
        config.shard.enable_merge = false;
        let status_addr = config.server.get_advertise_status_addr().to_string();
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(config, make_default_test_worker_config(true))
                .await
                .unwrap();

        let index_key = IndexKey::new_nullspace(1, 2001);
        test_suite.must_create_index(&index_key).await.unwrap();
        test_suite.validate_add_index_result(&index_key).await;

        let shard = test_suite.get_first_shard().await;
        let shard_id = shard.shard_id;
        let shard_id_str = shard_id.to_string();

        must_finally(|| {
            let test_suite = &test_suite;
            async move {
                for (_, worker) in test_suite
                    .get_meta_service()
                    .worker_node_mgr
                    .collect_worker_nodes()
                    .await
                {
                    if worker.get_shard_schedule(&shard_id).await.is_some() {
                        return true;
                    }
                }
                false
            }
        })
        .await;

        let url = format!("http://{}/op/meta/worker/allocation/shards", status_addr);
        let resp = reqwest::get(&url).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);

        let body = resp.json::<serde_json::Value>().await.unwrap();
        let all_workers = body.as_object().unwrap();
        assert!(
            !all_workers.is_empty(),
            "worker allocation should not be empty"
        );

        let mut found = false;
        for worker_shards in all_workers.values() {
            if let Some(shard_json) = worker_shards.get(&shard_id_str) {
                assert_eq!(
                    shard_json.get("shard_id").unwrap().as_u64().unwrap(),
                    shard_id
                );
                assert_eq!(
                    shard_json.get("keyspace_id").unwrap().as_u64().unwrap(),
                    index_key.keyspace_id as u64
                );
                assert_eq!(
                    shard_json.get("table_id").unwrap().as_i64().unwrap(),
                    index_key.table_id
                );
                assert_eq!(
                    shard_json.get("index_id").unwrap().as_i64().unwrap(),
                    index_key.index_id
                );
                found = true;
                break;
            }
        }
        assert!(
            found,
            "missing shard {} in endpoint body={}",
            shard_id, body
        );

        let mut scheduled_worker_addr = None;
        for (addr, worker) in test_suite
            .get_meta_service()
            .worker_node_mgr
            .collect_worker_nodes()
            .await
        {
            if worker.get_shard_schedule(&shard_id).await.is_some() {
                scheduled_worker_addr = Some(addr);
                break;
            }
        }
        let scheduled_worker_addr = scheduled_worker_addr.expect("shard should be scheduled");

        let mut url = reqwest::Url::parse(&format!(
            "http://{}/op/meta/worker/allocation/shards",
            status_addr
        ))
        .unwrap();
        url.query_pairs_mut()
            .append_pair("addr", &scheduled_worker_addr);
        let resp = reqwest::get(url).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);

        let filtered_body = resp.json::<serde_json::Value>().await.unwrap();
        let filtered_workers = filtered_body.as_object().unwrap();
        assert_eq!(
            filtered_workers.len(),
            1,
            "filtered worker allocation should have exactly one worker, body={}",
            filtered_body
        );
        assert!(
            filtered_workers.contains_key(&scheduled_worker_addr),
            "filtered worker allocation should contain addr={}, body={}",
            scheduled_worker_addr,
            filtered_body
        );
        let shard_json = filtered_workers
            .get(&scheduled_worker_addr)
            .unwrap()
            .get(&shard_id_str);
        assert!(
            shard_json.is_some(),
            "missing shard {} in filtered endpoint body={}",
            shard_id,
            filtered_body
        );

        test_suite.stop().await.unwrap();
    });
    Ok(())
}

#[test]
fn test_meta_shard_info() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;

    rt.block_on(async {
        let mut config = make_default_test_meta_config(true);
        config.shard.auto_compaction = false;
        config.shard.enable_split = false;
        config.shard.enable_merge = false;
        let status_addr = config.server.get_advertise_status_addr().to_string();
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(config, make_default_test_worker_config(true))
                .await
                .unwrap();

        let index_key = IndexKey::new_nullspace(1, 3001);
        test_suite.must_create_index(&index_key).await.unwrap();
        test_suite.validate_add_index_result(&index_key).await;

        let shard = test_suite.get_first_shard().await;
        let url = format!(
            "http://{}/op/meta/shard_info?shard_id={}",
            status_addr, shard.shard_id
        );
        let resp = reqwest::get(&url).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);

        let body = resp.json::<serde_json::Value>().await.unwrap();
        assert_eq!(
            body.get("shard_id").unwrap().as_u64().unwrap(),
            shard.shard_id
        );
        assert_eq!(
            body.get("index_key")
                .unwrap()
                .get("keyspace_id")
                .unwrap()
                .as_u64()
                .unwrap(),
            index_key.keyspace_id as u64
        );
        assert_eq!(
            body.get("index_key")
                .unwrap()
                .get("table_id")
                .unwrap()
                .as_i64()
                .unwrap(),
            index_key.table_id
        );
        assert_eq!(
            body.get("index_key")
                .unwrap()
                .get("index_id")
                .unwrap()
                .as_i64()
                .unwrap(),
            index_key.index_id
        );
        assert_eq!(
            body.get("range")
                .unwrap()
                .get("start_key")
                .unwrap()
                .as_str()
                .unwrap(),
            hex::encode(shard.range.encoded_start())
        );
        assert_eq!(
            body.get("range")
                .unwrap()
                .get("end_key")
                .unwrap()
                .as_str()
                .unwrap(),
            hex::encode(shard.range.encoded_end())
        );
        assert_eq!(
            body.get("writer").unwrap().as_str().unwrap(),
            shard.writer.as_str()
        );
        assert!(body.get("manifest").is_some(), "body={}", body);
        assert!(body.get("frag_upload_location").is_some(), "body={}", body);
        assert!(
            body.get("progress")
                .unwrap()
                .get("compaction_state")
                .unwrap()
                .as_array()
                .is_some(),
            "body={}",
            body
        );
        assert_eq!(
            body.get("stats")
                .unwrap()
                .get("size")
                .unwrap()
                .as_u64()
                .unwrap(),
            shard.stats.get_size()
        );

        let resp = reqwest::get(&format!(
            "http://{}/op/meta/shard_info?shard_id={}",
            status_addr,
            shard.shard_id + 100000
        ))
        .await
        .unwrap();
        assert_eq!(resp.status(), StatusCode::NOT_FOUND, "{:?}", resp);

        let resp = reqwest::get(&format!("http://{}/op/meta/shard_info", status_addr))
            .await
            .unwrap();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST, "{:?}", resp);

        let resp = reqwest::get(&format!(
            "http://{}/op/meta/shard_info?shard_id=abc",
            status_addr
        ))
        .await
        .unwrap();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST, "{:?}", resp);

        test_suite.stop().await.unwrap();
    });
    Ok(())
}

#[test]
fn test_block_worker_prevents_shard_scheduling() -> Result<()> {
    fail::cfg("worker_node_client_fail_fast", "return(2)").unwrap();
    defer!({
        fail::remove("worker_node_client_fail_fast");
    });

    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;

    rt.block_on(async {
        let mut config = make_default_test_meta_config(true);
        config.shard.auto_compaction = false;
        config.shard.enable_split = false;
        config.shard.enable_merge = false;
        let status_addr = config.server.get_advertise_status_addr().to_string();

        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(config, make_default_test_worker_config(true))
                .await
                .unwrap();

        // Default mocked worker is 127.0.0.1:8510; add another worker for rescheduling.
        let blocked_addr = "127.0.0.1:8510".to_string();
        let other_addr = "127.0.0.1:8512".to_string();
        test_suite
            .add_mocked_worker_node(Some(other_addr.clone()))
            .await;

        test_suite.wait_heartbeat().await;
        must_finally(async || {
            test_suite
                .get_meta_service()
                .worker_node_mgr
                .get_worker_nodes_len()
                .await
                == 2
        })
        .await;

        let index_key = IndexKey::new_nullspace(1, 1001);
        test_suite.must_create_index(&index_key).await.unwrap();
        test_suite.validate_add_index_result(&index_key).await;

        let shard = test_suite.get_first_shard().await;
        let shard_id = shard.shard_id;
        let start_key = hex::encode(shard.range.encoded_start());
        let end_key = hex::encode(shard.range.encoded_end());

        // Pin the shard to the worker that we're about to block.
        {
            let url = format!(
                "http://{}/op/meta/reschedule_range?writer={}&start_key={}&end_key={}&table_id={}&index_id={}",
                status_addr,
                blocked_addr,
                start_key,
                end_key,
                index_key.table_id,
                index_key.index_id,
            );
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
        }

        must_finally(|| {
            let test_suite = &test_suite;
            let blocked_addr = blocked_addr.clone();
            async move {
                test_suite
                    .get_shard_from_restore(shard_id)
                    .await
                    .unwrap()
                    .writer
                    == blocked_addr
            }
        })
        .await;

        // Block the worker and ensure the shard won't be scheduled onto it.
        {
            let url = format!(
                "http://{}/op/meta/worker/block?addr={}",
                status_addr, blocked_addr
            );
            let resp = reqwest::Client::new().post(&url).send().await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
        }

        {
            let url = format!("http://{}/op/meta/worker/block", status_addr);
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
            let blocked_list = resp.json::<Vec<String>>().await.unwrap();
            assert!(
                blocked_list.contains(&blocked_addr),
                "blocked list missing addr, list={:?}",
                blocked_list
            );
        }

        must_finally(|| {
            let test_suite = &test_suite;
            let blocked_addr = blocked_addr.clone();
            async move { test_suite.get_meta_service().worker_node_mgr.is_worker_blocked(&blocked_addr).await }
        })
        .await;

        must_finally(|| {
            let test_suite = &test_suite;
            let blocked_addr = blocked_addr.clone();
            async move {
                let writer = test_suite
                    .get_shard_from_restore(shard_id)
                    .await
                    .unwrap()
                    .writer;
                writer == blocked_addr
            }
        })
        .await;

        {
            let url = format!(
                "http://{}/op/meta/reschedule_range?writer={}&start_key={}&end_key={}&table_id={}&index_id={}",
                status_addr,
                blocked_addr,
                start_key,
                end_key,
                index_key.table_id,
                index_key.index_id,
            );
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
            let body = resp.text().await.unwrap();
            assert!(
                body.contains("error_code=TryAgain"),
                "reschedule to blocked worker should fail, body={}",
                body
            );
        }

        // Create another index; its initial shard must not be scheduled to the blocked worker.
        let index_key2 = IndexKey::new_nullspace(1, 1002);
        test_suite.must_create_index(&index_key2).await.unwrap();
        test_suite.validate_add_index_result(&index_key2).await;
        let shards2 = test_suite.collect_shards_from_index(&index_key2).await;
        assert_eq!(shards2.len(), 1);
        assert_eq!(shards2[0].writer, other_addr);

        // Unblock the worker and ensure rescheduling is allowed again.
        {
            let url = format!(
                "http://{}/op/meta/worker/block?addr={}",
                status_addr, blocked_addr
            );
            let resp = reqwest::Client::new().delete(&url).send().await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
        }

        must_finally(|| {
            let test_suite = &test_suite;
            let blocked_addr = blocked_addr.clone();
            async move {
                !test_suite
                    .get_meta_service()
                    .worker_node_mgr
                    .is_worker_blocked(&blocked_addr)
                    .await
            }
        })
        .await;

        {
            let url = format!("http://{}/op/meta/worker/block", status_addr);
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
            let blocked_list = resp.json::<Vec<String>>().await.unwrap();
            assert!(
                !blocked_list.contains(&blocked_addr),
                "blocked list should not include addr, list={:?}",
                blocked_list
            );
        }

        {
            let url = format!(
                "http://{}/op/meta/reschedule_range?writer={}&start_key={}&end_key={}&table_id={}&index_id={}",
                status_addr,
                blocked_addr,
                start_key,
                end_key,
                index_key.table_id,
                index_key.index_id,
            );
            let resp = reqwest::get(&url).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK, "{:?}", resp);
            let body = resp.text().await.unwrap();
            assert!(
                !body.contains("error_code=TryAgain"),
                "reschedule to unblocked worker should succeed, body={}",
                body
            );
        }

        must_finally(|| {
            let test_suite = &test_suite;
            let blocked_addr = blocked_addr.clone();
            async move {
                test_suite
                    .get_shard_from_restore(shard_id)
                    .await
                    .unwrap()
                    .writer
                    == blocked_addr
            }
        })
        .await;

        test_suite.check_meta_subscriber().await;
        test_suite.stop().await.unwrap();
    });
    Ok(())
}