1
use std::collections::hash_map::Entry;
2

            
3
use bonsaidb_core::{connection::SessionId, pubsub::Receiver};
4

            
5
use crate::{storage::SessionSubscriber, Database, Subscriber};
6

            
7
impl crate::storage::StorageInstance {
8
1112
    pub(crate) fn register_subscriber(
9
1112
        &self,
10
1112
        session_id: Option<SessionId>,
11
1112
        database: Database,
12
1112
    ) -> Subscriber {
13
1112
        let subscriber = self.relay().create_subscriber();
14
1112
        let mut data = self.data.subscribers.write();
15
1112
        let receiver = Receiver::new_stripping_prefixes(subscriber.receiver().clone());
16
1112
        let id = loop {
17
1112
            data.last_id = data.last_id.wrapping_add(1);
18
1112
            let id = data.last_id;
19
1112
            let entry = data.subscribers.entry(id);
20
1112
            if matches!(entry, Entry::Vacant(_)) {
21
1112
                entry.or_insert(SessionSubscriber {
22
1112
                    session_id,
23
1112
                    subscriber: subscriber.clone(),
24
1112
                });
25
1112
                break id;
26
            }
27
        };
28

            
29
1112
        Subscriber {
30
1112
            id,
31
1112
            database,
32
1112
            subscriber,
33
1112
            receiver,
34
1112
        }
35
1112
    }
36

            
37
692
    pub(crate) fn unregister_subscriber(&self, subscriber: &Subscriber) {
38
692
        let mut data = self.data.subscribers.write();
39
692
        data.unregister(subscriber.id);
40
692
    }
41
}