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

            
3
use bonsaidb_core::connection::SessionId;
4
use bonsaidb_core::pubsub::Receiver;
5

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

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

            
31
1328
        Subscriber {
32
1328
            id,
33
1328
            database,
34
1328
            subscriber,
35
1328
            receiver,
36
1328
        }
37
1328
    }
38

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