【Rust投稿】从零实现消息中间件(3)-SUBLIST

简化实现

上一篇中提到subject是一种Trie树,在实现的过程中发现非常琐碎,也不好讲解,因此考虑先实现一个简单版本的.具体来说就是

1. 不支持*和>两种模糊匹配
   

2. 所有的主题就是一个普通字符串
   

3. 为了简化,先不考虑cache
   但是支持负载均衡,也就是SUB中的Queue.
   
   

思路

由于不支持*和>,那么就意味着所有的主题就是一个普通字符串,没有了层次结构. 所以直接使用Map即可.

实现

数据结构的定义

开工之前,先想好相关数据结构,这样事半功倍.

Subscription

订阅中要包含Client的信息也就是ClientMessageSender,他是为了找到相关订阅以后可以发送消息.
暂时不用处理里面的内容.

pub struct Subscription {
    pub msg_sender: Arc<Mutex<ClientMessageSender>>,
    pub subject: String,
    pub queue: Option<String>,
    pub sid: String,
}

SubResult

查找结果中为什么是ArcSubscription,而不是Subscription,是因为Subscription一旦定义,就不会修改了.
因此使用Arc方便传递,同时满足tokio多线程的要求.

#[derive(Debug, Default)]
pub struct SubResult {
    pub subs: Vec<ArcSubscription>,
    pub qsubs: Vec<Vec<ArcSubscription>>,
}
pub type ArcSubscription = Arc<Subscription>;
pub type ArcSubResult = Arc<SubResult>;

SubListTrait

SubListTrait是他对外提供的服务接口,主要是

1. 新增订阅
   这个是当一个Client 发送sub消息到服务端的时候要处理的
   

2. 删除订阅
   这个是当一个Client发送 unsub消息到服务端的时候要处理的,不过因为我们不支持unsub,那就是连接断开的时候处理的.
   

3. 查找相关订阅
   这个是当一个client发送pub消息到服务端后,服务端要查找所有相关的订阅,然后把消息逐一转发给他们.
   rust pub trait SubListTrait { fn insert(&mut self, sub: ArcSubscription) -> Result<()>; fn remove(&mut self, sub: ArcSubscription) -> Result<()>; fn match_subject(&mut self, subject: &str) -> Result<ArcSubResult>; }
   
   

SimpleSubList

这是我们的核心,也是我们要实现的简单的基于Map的操作.
考虑到增删改查的需求,这里面都用Map来保存,而不用Vector. 牺牲一些空间来换取时间.

#[derive(Debug, Default)]
pub struct SimpleSubList {
    subs: HashMap<String, BTreeSet<ArcSubscriptionWrapper>>,
    qsubs: HashMap<String, HashMap<String,
                 BTreeSet<ArcSubscriptionWrapper>>>,
}

ArcSubscriptionWrapper

注意到前面的SimpleSubList中,保存在BTreeSeet中的存放的是ArcSubscriptionWrapper,而不是ArcSubscriptionWrapper.
这是有意为之的,因为我们在向BTreeSet中插入新的Sub的时候不需要关心他们真实的顺序,只是需要关心他们是否相同.
所以我们比较的对象是他们的地址而不是内容. 但是因为孤儿原则的限制,我们不能为Arc实现Ord这个trait,只能再多一次wrapper,
相信我们代码中有不少为孤儿原则做出的让步.

#[derive(Debug, Clone)]
pub struct ArcSubscriptionWrapper(ArcSubscription);

代码实现

ArcSubscriptionWrapper

因为要放入BTreeSet中,所以必须实现自定义的Ord,而要实现这个Ord要做不少额外的工作.

 impl std::cmp::Ord for ArcSubscriptionWrapper {
    fn cmp(&self, other: &Self) -> Ordering {
        let a = self.0.as_ref() as *const Subscription as usize;
        let b = other.0.as_ref()
                as *const Subscription as usize;
        a.cmp(&b)
    }
}

除了上面这个trait,还需要额外其他的trait实现.

impl std::cmp::PartialEq for ArcSubscriptionWrapper {
    fn eq(&self, other: &Self) -> bool {
        self.cmp(other) == Ordering::Equal
    }
}
 
impl std::cmp::Eq for ArcSubscriptionWrapper {}
impl std::cmp::PartialOrd for ArcSubscriptionWrapper {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

SimpleSubList

思路理清后,实现起来反而简单了.关键代码总共也就一百行不到.

impl SubListTrait for SimpleSubList {
    fn insert(&mut self, sub: Arc<Subscription>) -> Result<()> {
        if let Some(ref q) = sub.queue {
            let entry = self
                .qsubs
                .entry(sub.subject.clone())
                .or_insert(Default::default());
            let queue = entry.entry(q.clone()).or_insert(Default::default());
            queue.insert(ArcSubscriptionWrapper(sub));
        } else {
            let subs = self
                .subs
                .entry(sub.subject.clone())
                .or_insert(Default::default());
            subs.insert(ArcSubscriptionWrapper(sub));
        }
        Ok(())
    }

    fn remove(&mut self, sub: Arc<Subscription>) -> Result<()> {
        if let Some(ref q) = sub.queue {
            if let Some(subs) = self.qsubs.get_mut(&sub.subject) {
                if let Some(qsubs) = subs.get_mut(q) {
                    qsubs.remove(&ArcSubscriptionWrapper(sub.clone()));
                    if qsubs.is_empty() {
                        subs.remove(q);
                    }
                } else {
                    return Err(NError::new(ERROR_SUBSCRIBTION_NOT_FOUND));
                }
                if subs.is_empty() {
                    self.qsubs.remove(&sub.subject);
                }
            } else {
                return Err(NError::new(ERROR_SUBSCRIBTION_NOT_FOUND));
            }
        } else {
            if let Some(subs) = self.subs.get_mut(&sub.subject) {
                subs.remove(&ArcSubscriptionWrapper(sub.clone()));
                if subs.is_empty() {
                    self.subs.remove(&sub.subject);
                }
            }
        }
        Ok(())
    }

    fn match_subject(&mut self, subject: &str) -> Result<ArcSubResult> {
        let mut r = SubResult::default();
        if let Some(subs) = self.subs.get(subject) {
            for s in subs {
                r.subs.push(s.0.clone());
            }
        }
        if let Some(qsubs) = self.qsubs.get(subject) {
            for (_, qsub) in qsubs {
                let mut v = Vec::with_capacity(qsub.len());
                for s in qsub {
                    v.push(s.0.clone());
                }
                r.qsubs.push(v);
            }
        }
        Ok(Arc::new(r))
    }
}

其他

相关代码都在我的github rnats 欢迎围观

https://github.com/nkbai/learnrustbynats