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

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

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