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

use std::collections::hash_map::Entry;

use bonsaidb_core::connection::SessionId;
use bonsaidb_core::pubsub::Receiver;

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

impl crate::storage::StorageInstance {
    pub(crate) fn register_subscriber(
        &self,
        session_id: Option<SessionId>,
        database: Database,
    ) -> Subscriber {
        let subscriber = self.relay().create_subscriber();
        let mut data = self.data.subscribers.write();
        let receiver = Receiver::new_stripping_prefixes(subscriber.receiver().clone());
        let id = loop {
            data.last_id = data.last_id.wrapping_add(1);
            let id = data.last_id;
            let entry = data.subscribers.entry(id);
            if matches!(entry, Entry::Vacant(_)) {
                entry.or_insert(SessionSubscriber {
                    session_id,
                    subscriber: subscriber.clone(),
                });
                break id;
            }
        };

        Subscriber {
            id,
            database,
            subscriber,
            receiver,
        }
    }

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