use crate::KvError;
use serde::{Deserialize, Serialize};
use std::fs;

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct ServerConfig {
    pub general: GeneralConfig,
    pub storage: StorageConfig,
    pub tls: ServerTlsConfig,
    pub log: LogConfig,
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct ClientConfig {
    pub general: GeneralConfig,
    pub tls: ClientTlsConfig,
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct GeneralConfig {
    pub addr: String,
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct LogConfig {
    pub path: String,
    pub rotation: RotationConfig,
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub enum RotationConfig {
    Hourly,
    Daily,
    Never,
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
#[serde(tag = "type", content = "args")]
pub enum StorageConfig {
    MemTable,
    SledDb(String),
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct ServerTlsConfig {
    pub cert: String,
    pub key: String,
    pub ca: Option<String>,
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct ClientTlsConfig {
    pub domain: String,
    pub identity: Option<(String, String)>,
    pub ca: Option<String>,
}

impl ServerConfig {
    pub fn load(path: &str) -> Result<Self, KvError> {
        let config = fs::read_to_string(path)?;
        let config: Self = toml::from_str(&config)?;
        Ok(config)
    }
}

impl ClientConfig {
    pub fn load(path: &str) -> Result<Self, KvError> {
        let config = fs::read_to_string(path)?;
        let config: Self = toml::from_str(&config)?;
        Ok(config)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn server_config_should_be_loaded() {
        let result: Result<ServerConfig, toml::de::Error> =
            toml::from_str(include_str!("../fixtures/server.conf"));
        assert!(result.is_ok());
    }

    #[test]
    fn client_config_should_be_loaded() {
        let result: Result<ClientConfig, toml::de::Error> =
            toml::from_str(include_str!("../fixtures/client.conf"));
        assert!(result.is_ok());
    }
}

mod config;
mod error;
mod network;
mod pb;
mod service;
mod storage;

pub use config::*;
pub use error::KvError;
pub use network::*;
pub use pb::abi::*;
pub use service::*;
pub use storage::*;

use anyhow::Result;
use tokio::net::{TcpListener, TcpStream};
use tokio_rustls::client;
use tokio_util::compat::FuturesAsyncReadCompatExt;
use tracing::{info, instrument, span};

/// 通过配置创建 KV 服务器
#[instrument(skip_all)]
pub async fn start_server_with_config(config: &ServerConfig) -> Result<()> {
    let acceptor =
        TlsServerAcceptor::new(&config.tls.cert, &config.tls.key, config.tls.ca.as_deref())?;

    let addr = &config.general.addr;
    match &config.storage {
        StorageConfig::MemTable => start_tls_server(addr, MemTable::new(), acceptor).await?,
        StorageConfig::SledDb(path) => start_tls_server(addr, SledDb::new(path), acceptor).await?,
    };

    Ok(())
}

/// 通过配置创建 KV 客户端
#[instrument(skip_all)]
pub async fn start_client_with_config(
    config: &ClientConfig,
) -> Result<YamuxCtrl<client::TlsStream<TcpStream>>> {
    let addr = &config.general.addr;
    let tls = &config.tls;

    let identity = tls.identity.as_ref().map(|(c, k)| (c.as_str(), k.as_str()));
    let connector = TlsClientConnector::new(&tls.domain, identity, tls.ca.as_deref())?;
    let stream = TcpStream::connect(addr).await?;
    let stream = connector.connect(stream).await?;

    // 打开一个 stream
    Ok(YamuxCtrl::new_client(stream, None))
}

async fn start_tls_server<Store: Storage>(
    addr: &str,
    store: Store,
    acceptor: TlsServerAcceptor,
) -> Result<()> {
    let service: Service<Store> = ServiceInner::new(store).into();
    let listener = TcpListener::bind(addr).await?;
    info!("Start listening on {}", addr);
    loop {
        let root = span!(tracing::Level::INFO, "server_process");
        let _enter = root.enter();
        let tls = acceptor.clone();
        let (stream, addr) = listener.accept().await?;
        info!("Client {:?} connected", addr);

        let svc = service.clone();
        tokio::spawn(async move {
            let stream = tls.accept(stream).await.unwrap();
            YamuxCtrl::new_server(stream, None, move |stream| {
                let svc1 = svc.clone();
                async move {
                    let stream = ProstServerStream::new(stream.compat(), svc1.clone());
                    stream.process().await.unwrap();
                    Ok(())
                }
            });
        });
    }
}

use std::env;

use anyhow::Result;
use kv6::{start_server_with_config, RotationConfig, ServerConfig};
use tokio::fs;
use tracing::span;
use tracing_subscriber::{
    fmt::{self, format},
    layer::SubscriberExt,
    prelude::*,
    EnvFilter,
};

#[tokio::main]
async fn main() -> Result<()> {
    let config = match env::var("KV_SERVER_CONFIG") {
        Ok(path) => fs::read_to_string(&path).await?,
        Err(_) => include_str!("../fixtures/server.conf").to_string(),
    };
    let config: ServerConfig = toml::from_str(&config)?;

    let tracer = opentelemetry_jaeger::new_pipeline()
        .with_service_name("kv-server")
        .install_simple()?;
    let opentelemetry = tracing_opentelemetry::layer().with_tracer(tracer);

    // 添加
    let log = &config.log;
    let file_appender = match log.rotation {
        RotationConfig::Hourly => tracing_appender::rolling::hourly(&log.path, "server.log"),
        RotationConfig::Daily => tracing_appender::rolling::daily(&log.path, "server.log"),
        RotationConfig::Never => tracing_appender::rolling::never(&log.path, "server.log"),
    };

    let (non_blocking, _guard1) = tracing_appender::non_blocking(file_appender);
    let fmt_layer = fmt::layer()
        .event_format(format().compact())
        .with_writer(non_blocking);

    tracing_subscriber::registry()
        .with(EnvFilter::from_default_env())
        .with(fmt_layer)
        .with(opentelemetry)
        .init();

    let root = span!(tracing::Level::INFO, "app_start", work_units = 2);
    let _enter = root.enter();

    start_server_with_config(&config).await?;

    Ok(())
}

use anyhow::Result;
use kv6::{
    start_client_with_config, start_server_with_config, ClientConfig, CommandRequest, ServerConfig,
    StorageConfig,
};
use std::time::Duration;
use tokio::time;

#[tokio::test]
async fn yamux_server_client_full_tests() -> Result<()> {
    let addr = "127.0.0.1:10086";

    let mut config: ServerConfig = toml::from_str(include_str!("../fixtures/server.conf"))?;
    config.general.addr = addr.into();
    config.storage = StorageConfig::MemTable;

    // 启动服务器
    tokio::spawn(async move {
        start_server_with_config(&config).await.unwrap();
    });

    time::sleep(Duration::from_millis(10)).await;
    let mut config: ClientConfig = toml::from_str(include_str!("../fixtures/client.conf"))?;
    config.general.addr = addr.into();

    let mut ctrl = start_client_with_config(&config).await.unwrap();
    let mut stream = ctrl.open_stream().await?;

    // 生成一个 HSET 命令
    let cmd = CommandRequest::new_hset("table1", "hello", "world".to_string().into());
    stream.execute_unary(&cmd).await?;

    // 生成一个 HGET 命令
    let cmd = CommandRequest::new_hget("table1", "hello");
    let data = stream.execute_unary(&cmd).await?;

    assert_eq!(data.status, 200);
    assert_eq!(data.values, &["world".into()]);

    Ok(())
}

可以看到，集成测试的写法和单元测试其实很类似，只不过我们不需要再使用 #[cfg(test)] 来做条件编译。

1. 在 tests 下 cargo new 出一个项目
2. 然后在那个项目里撰写辅助代码和测试代码
3. 如果你对此感兴趣，可以看[ tonic 的集成测试](https://github.com/hyperium/tonic/tree/master/tests)。
4. 不过注意了,集成测试和你的 crate 用同样的条件编译，所以在集成测试里，无法使用单元测试中构建的辅助代码。

</div>
</details>

## 性能测试

在之前不断完善 KV server 的过程中，你一定会好奇：我们的 KV server 性能究竟如何呢？那来写一个关于 Pub/Sub 的性能测试吧。

基本的想法是我们连上 100 个 subscriber 作为背景，然后看 publisher publish 的速度。

因为 BROADCAST_CAPACITY 有限，是 128，当 publisher 速度太快，而导致 server 不能及时往 subscriber 发送时，server 接收 client
数据的速度就会降下来，无法接收新的
client，整体的 publish 的速度也会降下来，所以这个测试能够了解 server 处理 publish 的速度。


<details id="admonition-为了确认这一点我们在-start_tls_server-函数中在-process-之前再加个-100ms-的延时人为减缓系统的处理速度" class="admonition note">
<summary class="admonition-title">

为了确认这一点，我们在 start_tls_server() 函数中，在 process() 之前，再加个 100ms 的延时，人为减缓系统的处理速度：

<a class="admonition-anchor-link" href="#admonition-为了确认这一点我们在-start_tls_server-函数中在-process-之前再加个-100ms-的延时人为减缓系统的处理速度"></a>
</summary>
<div>

```rust, editable

async move {
    let stream = ProstServerStream::new(stream.compat(), svc1.clone());
    // 延迟 100ms 处理
    time::sleep(Duration::from_millis(100)).await;
    stream.process().await.unwrap();
    Ok(())
}
```

</div>
</details>

好，现在可以写性能测试了。


<details id="admonition-在-rust-下我们可以用-criterion-库httpsgithubcombheislercriterionrs它可以处理基本的性能测试并生成漂亮的报告" class="admonition note">
<summary class="admonition-title">

在 Rust 下，我们可以用 [criterion 库](https://github.com/bheisler/criterion.rs)。它可以处理基本的性能测试，并生成漂亮的报告 

<a class="admonition-anchor-link" href="#admonition-在-rust-下我们可以用-criterion-库httpsgithubcombheislercriterionrs它可以处理基本的性能测试并生成漂亮的报告"></a>
</summary>
<div>

所以在 Cargo.toml 中加入：

```toml

[dev-dependencies]
...
criterion = { version = "0.3", features = ["async_futures", "async_tokio", "html_reports"] } # benchmark
...
rand = "0.8" # 随机数处理
...

[[bench]]
name = "pubsub"
harness = false
```

> 最后这个 bench section，描述了性能测试的名字，它对应 benches 目录下的同名文件。

我们创建和 src 平级的 benches，然后再创建 benches/pubsub.rs，添入如下代码：

~~~admonish info title="examples/gen_config.rs" collapsible=true
```rust, editable
use anyhow::Result;
use criterion::{criterion_group, criterion_main, Criterion};
use futures::StreamExt;
use kv6::{
    start_client_with_config, start_server_with_config, ClientConfig, CommandRequest, ServerConfig,
    StorageConfig, YamuxCtrl,
};
use rand::prelude::SliceRandom;
use std::time::Duration;
use tokio::net::TcpStream;
use tokio::runtime::Builder;
use tokio::time;
use tokio_rustls::client::TlsStream;
use tracing::{info, span};
use tracing_subscriber::{layer::SubscriberExt, prelude::*, EnvFilter};

async fn start_server() -> Result<()> {
    let addr = "127.0.0.1:9999";
    let mut config: ServerConfig = toml::from_str(include_str!("../fixtures/server.conf"))?;
    config.general.addr = addr.into();
    config.storage = StorageConfig::MemTable;

    tokio::spawn(async move {
        start_server_with_config(&config).await.unwrap();
    });

    Ok(())
}

async fn connect() -> Result<YamuxCtrl<TlsStream<TcpStream>>> {
    let addr = "127.0.0.1:9999";
    let mut config: ClientConfig = toml::from_str(include_str!("../fixtures/client.conf"))?;
    config.general.addr = addr.into();

    Ok(start_client_with_config(&config).await?)
}

async fn start_subscribers(topic: &'static str) -> Result<()> {
    let mut ctrl = connect().await?;
    let stream = ctrl.open_stream().await?;
    info!("C(subscriber): stream opened");
    let cmd = CommandRequest::new_subscribe(topic.to_string());
    tokio::spawn(async move {
        let mut stream = stream.execute_streaming(&cmd).await.unwrap();
        while let Some(Ok(data)) = stream.next().await {
            drop(data);
        }
    });

    Ok(())
}

async fn start_publishers(topic: &'static str, values: &'static [&'static str]) -> Result<()> {
    let mut rng = rand::thread_rng();
    let v = values.choose(&mut rng).unwrap();

    let mut ctrl = connect().await.unwrap();
    let mut stream = ctrl.open_stream().await.unwrap();
    info!("C(publisher): stream opened");

    let cmd = CommandRequest::new_publish(topic.to_string(), vec![(*v).into()]);
    stream.execute_unary(&cmd).await.unwrap();

    Ok(())
}

fn pubsub(c: &mut Criterion) {
    let tracer = opentelemetry_jaeger::new_pipeline()
        .with_service_name("kv-bench")
        .install_simple()
        .unwrap();
    let opentelemetry = tracing_opentelemetry::layer().with_tracer(tracer);

    tracing_subscriber::registry()
        .with(EnvFilter::from_default_env())
        .with(opentelemetry)
        .init();

    let root = span!(tracing::Level::INFO, "app_start", work_units = 2);
    let _enter = root.enter();
    // 创建 Tokio runtime
    let runtime = Builder::new_multi_thread()
        .worker_threads(4)
        .thread_name("pubsub")
        .enable_all()
        .build()
        .unwrap();

    let base_str = include_str!("../fixtures/server.conf"); // 891 bytes

    let values: &'static [&'static str] = Box::leak(
        vec![
            &base_str[..64],
            &base_str[..128],
            &base_str[..256],
            &base_str[..512],
        ]
        .into_boxed_slice(),
    );
    let topic = "lobby";

    // 运行服务器和 100 个 subscriber，为测试准备
    runtime.block_on(async {
        eprint!("preparing server and subscribers");
        start_server().await.unwrap();
        time::sleep(Duration::from_millis(50)).await;
        for _ in 0..1000 {
            start_subscribers(topic).await.unwrap();
            eprint!(".");
        }
        eprintln!("Done!");
    });

    // 进行 benchmark
    c.bench_function("publishing", move |b| {
        b.to_async(&runtime)
            .iter(|| async { start_publishers(topic, values).await })
    });
}

criterion_group! {
    name = benches;
    config = Criterion::default().sample_size(10);
    targets = pubsub
}
criterion_main!(benches);
```

</div>
</details>

大部分的代码都很好理解，就是创建服务器和客户端，为测试做准备。

说一下这里面核心的 benchmark 代码：

```rust, editable

c.bench_function("publishing", move |b| {
    b.to_async(&runtime)
        .iter(|| async { start_publishers(topic, values).await })
});
```

对于要测试的代码，我们可以封装成一个函数进行测试。

- 这里因为要做 async 函数的测试，需要使用 runtime。
- 普通的函数不需要调用 to_async。对于更多有关 criterion 的用法，可以参考它的文档。

````admonish success title="运行 cargo bench 后，会见到如下打印（如果你的代码无法通过，可以参考 repo 里的 diff_benchmark，我顺便做了一点小重构）： " collapsible=true

preparing server and subscribers....................................................................................................Done!
publishing              time:   [419.73 ms 426.84 ms 434.20 ms]                     
                        change: [-1.6712% +1.0499% +3.6586%] (p = 0.48 > 0.05)
                        No change in performance detected.

可以看到，单个 publish 的处理速度要 426ms，好慢！我们把之前在 start_tls_server() 里加的延迟去掉，再次测试：

preparing server and subscribers....................................................................................................Done!
publishing              time:   [318.61 ms 324.48 ms 329.81 ms]                     
                        change: [-25.854% -23.980% -22.144%] (p = 0.00 < 0.05)
                        Performance has improved.

 > 
 > 嗯，这下 324ms，正好是减去刚才加的 100ms。可是这个速度依旧不合理，凭直觉我们感觉一下这个速度，是 Python 这样的语言还正常，如果是
 > Rust 也太慢了吧？

## 测量和监控

 > 
 > 工业界有句名言：如果你无法测量，那你就无法改进（If you can’t measure it, you can’t improve it）。

现在知道了 KV server 性能有问题，但并不知道问题出在哪里。我们需要使用合适的测量方式。

目前，比较好的端对端的性能监控和测量工具是 jaeger，我们可以在 KV server/client 侧收集监控信息，发送给 jaeger
来查看在服务器和客户端的整个处理流程中，时间都花费到哪里去了。

之前我们在 KV server 里使用的日志工具是
tracing，不过日志只是它的诸多功能之一，它还能做[ instrument](https://docs.rs/tracing/0.1.28/tracing/attr.instrument.html)
，然后配合 [opentelemetry 库](https://github.com/open-telemetry/opentelemetry-rust)，我们就可以把 instrument 的结果发送给
jaeger 了。

<details id="admonition-1-加入jaeger" class="admonition note">
<summary class="admonition-title">

1. 加入jaeger 

<a class="admonition-anchor-link" href="#admonition-1-加入jaeger"></a>

</summary>
<div>

好，在 Cargo.toml 里添加新的依赖：

````toml

[dependencies]
...
opentelemetry-jaeger = "0.15" # opentelemetry jaeger 支持
...
tracing-appender = "0.1" # 文件日志
tracing-opentelemetry = "0.15" # opentelemetry 支持
tracing-subscriber = { version = "0.2", features = ["json", "chrono"] } # 日志处理

有了这些依赖后，在 benches/pubsub.rs 里，我们可以在初始化 tracing_subscriber 时，使用 jaeger 和 opentelemetry tracer：

fn pubsub(c: &mut Criterion) {
    let tracer = opentelemetry_jaeger::new_pipeline()
        .with_service_name("kv-bench")
        .install_simple()
        .unwrap();
    let opentelemetry = tracing_opentelemetry::layer().with_tracer(tracer);

    tracing_subscriber::registry()
        .with(EnvFilter::from_default_env())
        .with(opentelemetry)
        .init();

    let root = span!(tracing::Level::INFO, "app_start", work_units = 2);
    let _enter = root.enter();
    // 创建 Tokio runtime
    ...
}

设置好 tracing 后，就在系统的主流程上添加相应的 instrument：

新添加的代码你可以看 repo 中的 diff_telemetry。

注意 instrument 可以用不同的名称:

比如，对于 TlsConnector::new() 函数，可以用 #[instrument(name = “tls_connector_new”)]，这样它的名字辨识度高一些。

如果说 TlsClientConnector::connect() 花不少时间还情有可原，因为这是整个 TLS 协议的握手过程，涉及到网络调用、包的加解密等。

但 TlsClientConnector::new() 就是加载一些证书、创建 TlsConnector 这个数据结构而已，为何这么慢？

仔细阅读 TlsClientConnector::new() 的代码，你可以对照注释看：

#[instrument(name = "tls_connector_new", skip_all)]
pub fn new(
    domain: impl Into<String> + std::fmt::Debug,
    identity: Option<(&str, &str)>,
    server_ca: Option<&str>,
) -> Result<Self, KvError> {
    let mut config = ClientConfig::new();

    // 如果有客户端证书，加载之
    if let Some((cert, key)) = identity {
        let certs = load_certs(cert)?;
        let key = load_key(key)?;
        config.set_single_client_cert(certs, key)?;
    }

    // 加载本地信任的根证书链
    config.root_store = match rustls_native_certs::load_native_certs() {
        Ok(store) | Err((Some(store), _)) => store,
        Err((None, error)) => return Err(error.into()),
    };

    // 如果有签署服务器的 CA 证书，则加载它，这样服务器证书不在根证书链
    // 但是这个 CA 证书能验证它，也可以
    if let Some(cert) = server_ca {
        let mut buf = Cursor::new(cert);
        config.root_store.add_pem_file(&mut buf).unwrap();
    }

    Ok(Self {
        config: Arc::new(config),
        domain: Arc::new(domain.into()),
    })
}

可以发现，它的代码唯一可能影响性能的就是加载本地信任的根证书链的部分。这个代码会和操作系统交互，获取信任的根证书链。也许，这就是影响性能的原因之一？

那我们将其简单重构一下。因为根证书链，只有在客户端没有提供用于验证服务器证书的 CA 证书时，才需要，所以可以在没有 CA 证书时，才加载本地的根证书链：

#[instrument(name = "tls_connector_new", skip_all)]
pub fn new(
    domain: impl Into<String> + std::fmt::Debug,
    identity: Option<(&str, &str)>,
    server_ca: Option<&str>,
) -> Result<Self, KvError> {
    let mut config = ClientConfig::new();

    // 如果有客户端证书，加载之
    if let Some((cert, key)) = identity {
        let certs = load_certs(cert)?;
        let key = load_key(key)?;
        config.set_single_client_cert(certs, key)?;
    }

    // 如果有签署服务器的 CA 证书，则加载它，这样服务器证书不在根证书链
    // 但是这个 CA 证书能验证它，也可以
    if let Some(cert) = server_ca {
        let mut buf = Cursor::new(cert);
        config.root_store.add_pem_file(&mut buf).unwrap();
    } else {
        // 加载本地信任的根证书链
        config.root_store = match rustls_native_certs::load_native_certs() {
            Ok(store) | Err((Some(store), _)) => store,
            Err((None, error)) => return Err(error.into()),
        };
    }

    Ok(Self {
        config: Arc::new(config),
        domain: Arc::new(domain.into()),
    })
}

完成这个修改后，我们再运行 RUST_LOG=info cargo bench，现在的性能达到了 1.64ms，相比之前的 1.95ms，提升了 16%。

你可以用浏览器打开 ./target/criterion/publishing/report/index.html 查看（名字是 publishing ，因为 benchmark ID 是 publishing）：

• PR merge 到 master，在 build 完成后，触发 dev 服务器的部署，团队内部可以尝试；

• 如果 release tag 包含 alpha，在 build 完成后，触发 staging 服务器的部署，公司内部可以使用；

• 如果 release tag 包含 beta，在 build 完成后，触发 beta 服务器的部署，beta 用户可以使用；

• 正式的 release tag 会触发生产环境的滚动升级，升级覆盖到的用户可以使用。