myko_server/
server_ownership.rs

1use std::{
2    collections::{HashMap, HashSet},
3    sync::Arc,
4};
5
6use myko::{
7    entities::server::{GetAllServers, ServerId},
8    relationship::iter_server_owned_registrations,
9    server::{CellServerCtx, PersistError},
10};
11
12pub struct ServerOwnershipManager;
13
14impl ServerOwnershipManager {
15    /// Get IDs of all currently live servers.
16    fn live_server_ids(ctx: &CellServerCtx) -> Vec<ServerId> {
17        use hyphae::Gettable;
18        let req = ctx.new_server_transaction();
19        ctx.query_map(GetAllServers {}, req)
20            .items()
21            .get()
22            .iter()
23            .map(|s| s.id.clone())
24            .collect()
25    }
26
27    /// Count how many server_owned items each server currently owns
28    /// across ALL registered #[server_owned] types.
29    fn count_distribution(ctx: &CellServerCtx) -> HashMap<Arc<str>, usize> {
30        let mut counts: HashMap<Arc<str>, usize> = HashMap::new();
31
32        for reg in iter_server_owned_registrations() {
33            let store = ctx.registry.get_or_create(reg.entity_type);
34            for (_, item) in store.snapshot() {
35                if let Some(owner) = item.server_owner() {
36                    *counts.entry(Arc::from(owner)).or_default() += 1;
37                }
38            }
39        }
40        counts
41    }
42
43    /// Pick the server with the lowest item count from the given set.
44    fn least_loaded(live_ids: &[ServerId], counts: &HashMap<Arc<str>, usize>) -> Option<ServerId> {
45        live_ids
46            .iter()
47            .min_by_key(|id| counts.get(id.0.as_ref()).copied().unwrap_or(0))
48            .cloned()
49    }
50
51    /// Scan all server_owned items and reassign any referencing dead/empty servers.
52    pub fn claim_orphaned(ctx: &CellServerCtx) -> Result<(), PersistError> {
53        let live_ids = Self::live_server_ids(ctx);
54        if live_ids.is_empty() {
55            log::warn!("[ServerOwnership] No live servers found, skipping orphan claim");
56            return Ok(());
57        }
58
59        let live_set: HashSet<&str> = live_ids.iter().map(|id| id.0.as_ref()).collect();
60        let mut counts = Self::count_distribution(ctx);
61        counts.retain(|k, _| live_set.contains(k.as_ref()));
62
63        let mut reassigned = 0usize;
64
65        for reg in iter_server_owned_registrations() {
66            let store = ctx.registry.get_or_create(reg.entity_type);
67            let items: Vec<_> = store.snapshot().into_iter().map(|(_, item)| item).collect();
68
69            for item in &items {
70                let current_owner = item.server_owner().unwrap_or("");
71
72                if !current_owner.is_empty() && live_set.contains(current_owner) {
73                    continue; // healthy
74                }
75
76                let Some(new_owner) = Self::least_loaded(&live_ids, &counts) else {
77                    continue;
78                };
79
80                if let Some(patched) = item.bake_server_owner(&new_owner.0) {
81                    // Server-ownership rebakes are Local (per the event-bus design):
82                    // re-emit the item normally rather than suppressing relationships.
83                    ctx.set_dyn(patched)?;
84                    *counts.entry(new_owner.0.clone()).or_default() += 1;
85                    reassigned += 1;
86                }
87            }
88        }
89
90        if reassigned > 0 {
91            log::info!(
92                "[ServerOwnership] Reassigned {} orphaned item(s)",
93                reassigned
94            );
95        }
96        Ok(())
97    }
98
99    /// Watch for Server entity removals and redistribute orphaned items.
100    /// Returns a SubscriptionGuard that must be kept alive.
101    pub fn watch_peer_deaths(ctx: &CellServerCtx) -> hyphae::SubscriptionGuard {
102        use hyphae::{Gettable, Signal, Watchable};
103
104        let req = ctx.new_server_transaction();
105        let servers_cell = ctx.query_map(GetAllServers {}, req).items();
106
107        let prev_ids: std::sync::Mutex<HashSet<Arc<str>>> =
108            std::sync::Mutex::new(servers_cell.get().iter().map(|s| s.id.0.clone()).collect());
109
110        let ctx = ctx.clone();
111        servers_cell.subscribe(move |signal| {
112            let Signal::Value(servers) = signal else {
113                return;
114            };
115
116            let current_ids: HashSet<Arc<str>> = servers.iter().map(|s| s.id.0.clone()).collect();
117
118            let mut prev = prev_ids.lock().unwrap();
119            let removed: Vec<Arc<str>> = prev.difference(&current_ids).cloned().collect();
120            *prev = current_ids;
121            drop(prev);
122
123            if removed.is_empty() {
124                return;
125            }
126
127            for id in &removed {
128                log::warn!(
129                    "[ServerOwnership] Server {} left cluster, redistributing",
130                    id
131                );
132            }
133
134            if let Err(e) = ServerOwnershipManager::claim_orphaned(&ctx) {
135                log::error!("[ServerOwnership] Failed to redistribute: {}", e);
136            }
137        })
138    }
139}
myko_server/server_ownership.rs

myko_server/
server_ownership.rs
