1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
use crate::App;
use anyhow::Result;
use async_std::task;
use futures::channel::mpsc;
use futures::{Future, FutureExt};
use potatonet_client::{Client, SubscribeId};
use potatonet_common::{Context, LocalServiceId, Request, Response, ServiceId, Topic};
use serde::de::DeserializeOwned;
use serde::Serialize;
use std::sync::Arc;

/// 节点上下文
pub struct NodeContext<'a> {
    pub(crate) client: Arc<Client>,

    pub(crate) app: Arc<App>,

    /// 请求或者通知来源
    pub(crate) from: Option<ServiceId>,

    // 当前服务名
    pub(crate) service_name: &'a str,

    /// 当前服务id
    pub(crate) local_service_id: LocalServiceId,

    /// 停止节点用通道
    pub(crate) tx_abort: mpsc::Sender<()>,
}

impl<'a> NodeContext<'a> {
    /// 请求或者通知来源服务id
    pub fn from(&self) -> Option<ServiceId> {
        self.from
    }

    /// 当前服务名称
    pub fn service_name(&self) -> &str {
        self.service_name
    }

    /// 当前服务id
    pub fn service_id(&self) -> ServiceId {
        self.local_service_id.to_global(self.client.node_id())
    }

    /// 停止节点运行
    pub fn shutdown_node(&self) {
        self.tx_abort.clone().try_send(()).ok();
    }
}

impl<'a> NodeContext<'a> {
    /// 订阅指定主题的消息
    pub async fn subscribe_with_topic<T, F, R>(&self, topic: &str, mut handler: F) -> SubscribeId
    where
        T: Topic,
        F: FnMut(&NodeContext<'_>, T) -> R + Send + 'static,
        R: Future<Output = ()> + Send + 'static,
    {
        // 避免client->context->client的循环引用
        let client = Arc::downgrade(&self.client);
        let app = Arc::downgrade(&self.app);
        let service_name = self.service_name.to_string();
        let local_service_id = self.local_service_id;
        let tx_abort = self.tx_abort.clone();
        self.client
            .subscribe_with_topic(topic, move |msg| {
                let client = client.clone();
                let app = app.clone();
                let service_name = service_name.clone();
                let tx_abort = tx_abort.clone();

                if let (Some(client), Some(app)) = (client.upgrade(), app.upgrade()) {
                    let ctx = NodeContext {
                        client: client.clone(),
                        app: app.clone(),
                        from: None,
                        service_name: &service_name,
                        local_service_id,
                        tx_abort,
                    };
                    handler(&ctx, msg).boxed()
                } else {
                    Box::pin(futures::future::ready(()))
                }
            })
            .await
    }

    /// 订阅消息
    /// 如果重复订阅,则会按订阅顺序依次触发回调函数,并不会增加数据传输流量
    pub async fn subscribe<T, F, R>(&self, handler: F) -> SubscribeId
    where
        T: Topic,
        F: FnMut(&NodeContext<'_>, T) -> R + Send + 'static,
        R: Future<Output = ()> + Send + 'static,
    {
        self.subscribe_with_topic(T::name(), handler).await
    }

    /// 取消订阅
    pub async fn unsubscribe(&self, id: SubscribeId) {
        self.client.unsubscribe(id).await;
    }
}

#[async_trait::async_trait]
impl<'a> Context for NodeContext<'a> {
    async fn call<T, R>(&self, service_name: &str, request: Request<T>) -> Result<Response<R>>
    where
        T: Serialize + Send + 'static,
        R: DeserializeOwned + Send + 'static,
    {
        trace!("call. service={} method={}", service_name, request.method);

        if let Some(lid) = self.app.services_map.get(service_name).copied() {
            // 优先调用本地服务
            if let Some((_, service)) = self.app.services.get(lid.0 as usize) {
                let resp = service.call(self, request.to_bytes()).await?;
                return Ok(Response::from_bytes(resp));
            }
        }

        self.client.call(service_name, request).await
    }

    async fn notify<T: Serialize + Send + 'static>(&self, service_name: &str, request: Request<T>) {
        trace!("notify. service={} method={}", service_name, request.method);
        let request = request.to_bytes();

        // 通知本地服务
        if let Some(lid) = self.app.services_map.get(service_name).copied() {
            let client = self.client.clone();
            let app = self.app.clone();
            let service_name = service_name.to_string();
            let from = self.local_service_id.to_global(client.node_id());
            let request = request.clone();
            let tx_abort = self.tx_abort.clone();
            task::spawn(async move {
                if let Some((_, service)) = app.services.get(lid.0 as usize) {
                    service
                        .notify(
                            &NodeContext {
                                client,
                                app: app.clone(),
                                from: Some(from),
                                service_name: &service_name,
                                local_service_id: lid,
                                tx_abort,
                            },
                            request,
                        )
                        .await;
                }
            });
        }

        // 通知远程服务
        self.client.notify(service_name, request).await;
    }

    async fn notify_to<T: Serialize + Send + 'static>(&self, to: ServiceId, request: Request<T>) {
        trace!("notify to. to={} method={}", to, request.method);
        let request = request.to_bytes();

        if to.node_id == self.client.node_id() {
            // 是本地服务
            let client = self.client.clone();
            let app = self.app.clone();
            let from = self.local_service_id.to_global(self.client.node_id());
            let service_name = self.service_name.to_string();
            let tx_abort = self.tx_abort.clone();
            task::spawn(async move {
                if let Some((_, service)) = app.services.get(to.local_service_id.0 as usize) {
                    service.notify(
                        &NodeContext {
                            client,
                            app: app.clone(),
                            from: Some(from),
                            service_name: &service_name,
                            local_service_id: to.local_service_id,
                            tx_abort,
                        },
                        request,
                    );
                }
            });
            return;
        }

        // 通知远程服务
        self.client.notify_to(to, request).await;
    }

    async fn publish_with_topic<T: Topic>(&self, topic: &str, msg: T) {
        trace!("publish. topic={}", topic);
        self.client.publish_with_topic(topic, msg).await;
    }
}