etcd_detector/

detector.rs

use std::collections::HashMap;
use std::error::Error;
use std::sync::Arc;
use std::sync::atomic::Ordering::Relaxed;
use std::sync::atomic::{AtomicI64, AtomicU32};
use std::time::Duration;

use crate::trace::debug;
use crate::util;
use detector::{
    OneselfRx, OneselfTx, Service, ServiceEvent, ServiceStatus, Services, WatchRx, build_oneself,
    build_watch,
};

#[derive(Clone)]
pub(crate) struct Inner {
    pub(crate) service: Service,
    pub(crate) client: etcd_client::Client,
    pub(crate) status: Arc<AtomicU32>,
    pub(crate) oneself_tx: OneselfTx,
    pub(crate) lease_id: Arc<AtomicI64>,
}

impl Inner {
    pub(crate) fn change_status(&self, status: ServiceStatus) {
        debug!("{:?} status change:{:?}", self.service, status);
        self.status.store(status as u32, Relaxed);
        let _ = self.oneself_tx.send(status);
    }

    pub(crate) fn get_lease_id(&self) -> i64 {
        self.lease_id.load(Relaxed)
    }

    pub(crate) fn set_lease_id(&self, id: i64) {
        self.lease_id.store(id, Relaxed)
    }
}

pub struct Detector {
    inner: Option<Inner>,
}

impl Detector {
    /// Connect to etcd servers from given endpoints.
    pub async fn connect<E: AsRef<str>, S: AsRef<[E]>>(
        service: Service,
        endpoints: S,
    ) -> Result<Detector, etcd_client::Error> {
        let options = Some(
            etcd_client::ConnectOptions::new()
                .with_keep_alive(Duration::from_secs(5), Duration::from_secs(10)),
        );
        let client = etcd_client::Client::connect(endpoints, options).await?;
        Ok(Detector::new(client, service))
    }

    pub fn new(client: etcd_client::Client, service: Service) -> Self {
        let (tx, _) = build_oneself(ServiceStatus::Unregistered);
        Detector {
            inner: Some(Inner {
                service,
                client,
                status: Arc::new(AtomicU32::new(ServiceStatus::Unregistered as u32)),
                oneself_tx: tx,
                lease_id: Arc::new(AtomicI64::new(0)),
            }),
        }
    }

    fn inner_mut(&mut self) -> &mut Inner {
        self.inner.as_mut().unwrap()
    }

    fn inner(&self) -> &Inner {
        self.inner.as_ref().unwrap()
    }

    async fn real_fetch(
        &mut self,
    ) -> Result<etcd_client::GetResponse, Box<dyn Error + 'static + Send>> {
        let inner = self.inner_mut();
        let resp = inner
            .client
            .get(
                inner.service.key.parent_path(),
                Some(etcd_client::GetOptions::new().with_prefix()),
            )
            .await
            .map_err(|e| Box::new(e) as Box<dyn Error + 'static + Send>);
        resp
    }

    async fn real_fetch_all(
        &mut self,
    ) -> Result<etcd_client::GetResponse, Box<dyn Error + 'static + Send>> {
        let inner = self.inner_mut();
        let resp = inner
            .client
            .get(
                inner.service.key.root_path(),
                Some(etcd_client::GetOptions::new().with_prefix()),
            )
            .await
            .map_err(|e| Box::new(e) as Box<dyn Error + 'static + Send>);
        resp
    }

    async fn watch_inner(
        &mut self,
        resp: etcd_client::GetResponse,
        key: String,
    ) -> Result<WatchRx<ServiceEvent>, Box<dyn Error + 'static + Send>> {
        // 2. 获取 revision
        let revision = resp.header().unwrap().revision();

        // 3. 创建 watch
        let inner = self.inner_mut();
        let watch = inner
            .client
            .watch(
                key,
                Some(
                    etcd_client::WatchOptions::new()
                        .with_prefix()
                        .with_start_revision(revision + 1),
                ),
            )
            .await
            .map_err(|e| Box::new(e) as Box<dyn Error + 'static + Send>)?;

        // 4. 建通道
        let (tx, rx) = build_watch();

        // 5. 把当前数据推送给 rx
        for m in util::parse_service_response(resp) {
            for s in m.1.services {
                let _ = tx.send(ServiceEvent::Changed(s));
            }
        }

        // 6. 异步监听 watch
        util::watch_service(watch, tx);

        Ok(rx)
    }
}

impl detector::Detector for Detector {
    fn service(&self) -> &Service {
        &self.inner().service
    }

    fn status(&self) -> ServiceStatus {
        ServiceStatus::from(self.inner().status.load(Relaxed))
    }

    /// 第一次注册如果失败则返回,如果成功也返回但内部会有掉线重试机制
    async fn register(&mut self) -> Result<(), Box<dyn Error + 'static + Send>> {
        let inner = self.inner_mut();

        // 已经执行过注册了
        if inner.service.key.has_id() {
            return Ok(());
        }

        match util::register(inner.clone()).await {
            Err(e) => Err(e),
            Ok(id) => {
                // 换上id
                inner.service.key.id = Some(id);
                Ok(())
            }
        }
    }

    /// 获取本类型服所有实例
    async fn fetch(&mut self) -> Result<Services, Box<dyn Error + 'static + Send>> {
        let name = &self.inner().service.key.name.clone();
        let mut resp = util::parse_service_response(self.real_fetch().await?);
        match resp.remove(name) {
            None => Ok(Services::new()),
            Some(s) => Ok(s),
        }
    }

    /// 获取所有注册的服务实例
    async fn fetch_all(
        &mut self,
    ) -> Result<HashMap<String, Services>, Box<dyn Error + 'static + Send>> {
        Ok(util::parse_service_response(self.real_fetch_all().await?))
    }

    /// 监听本类型服务实例
    async fn watch(&mut self) -> Result<WatchRx<ServiceEvent>, Box<dyn Error + 'static + Send>> {
        let resp = self.real_fetch().await?;
        let parent_path = self.inner().service.key.parent_path();
        self.watch_inner(resp, parent_path).await
    }

    /// 监听所有注册的服务实例
    async fn watch_all(
        &mut self,
    ) -> Result<WatchRx<ServiceEvent>, Box<dyn Error + 'static + Send>> {
        let resp = self.real_fetch_all().await?;
        let root_path = self.inner().service.key.root_path();
        self.watch_inner(resp, root_path).await
    }

    /// 监听服务自己本身状态, 需要在第一次调用register成功后才能使用它
    fn oneself(&mut self) -> Option<OneselfRx> {
        let inner = self.inner();
        if inner.service.key.has_id() {
            Some(inner.oneself_tx.oneself_rx())
        } else {
            None
        }
    }
}

impl Drop for Detector {
    fn drop(&mut self) {
        let inner = self.inner.take().unwrap();

        // 取消掉旧租约, 让注册的服务过段时间自己消亡掉
        let lease_id = inner.get_lease_id();
        if lease_id != 0 {
            let mut client = inner.client.clone();
            // 先释放掉inner再执行租约取消
            drop(inner);
            tokio::spawn(async move {
                let _ = client.lease_revoke(lease_id).await;
                debug!("lease revoke, lease id:{}", lease_id);
            });
        }
    }
}