lightyear_link/
server.rs

1use crate::{LinkPlugin, Linked, Linking, Unlink, Unlinked};
2use alloc::{format, string::String, vec::Vec};
3use bevy_app::{App, Plugin};
4use bevy_ecs::{
5    component::{Component, HookContext},
6    entity::Entity,
7    observer::Trigger,
8    query::{With, Without},
9    relationship::{
10        Relationship, RelationshipHookMode, RelationshipSourceCollection, RelationshipTarget,
11    },
12    system::{Commands, Query},
13    world::{DeferredWorld, OnAdd, OnInsert},
14};
15use bevy_reflect::Reflect;
16use lightyear_core::prelude::LocalTimeline;
17use tracing::{trace, warn};
18// TODO: should we also have a LinkId (remote addr/etc.) that uniquely identifies the link?
19
20#[derive(Component, Default, Debug, PartialEq, Eq, Reflect)]
21#[component(on_add = Server::on_add)]
22#[relationship_target(relationship = LinkOf, linked_spawn)]
23pub struct Server {
24    links: Vec<Entity>,
25}
26
27impl Server {
28    fn on_add(mut world: DeferredWorld, context: HookContext) {
29        let entity_ref = world.entity(context.entity);
30        if !entity_ref.contains::<Unlinked>()
31            && !entity_ref.contains::<Linked>()
32            && !entity_ref.contains::<Linking>()
33        {
34            trace!("Inserting Unlinked because Server was added");
35            world.commands().entity(context.entity).insert(Unlinked {
36                reason: String::new(),
37            });
38        };
39    }
40
41    fn unlinked(
42        trigger: Trigger<OnAdd, Unlinked>,
43        mut query: Query<(&Server, &Unlinked)>,
44        mut commands: Commands,
45    ) {
46        if let Ok((server_link, unlinked)) = query.get_mut(trigger.target()) {
47            for link_of in server_link.collection() {
48                if let Ok(mut c) = commands.get_entity(*link_of) {
49                    // cannot simply insert Unlinked because then we wouldn't close aeronet sessions...
50                    c.trigger(Unlink {
51                        reason: unlinked.reason.clone(),
52                    });
53                    c.despawn();
54                }
55            }
56        }
57    }
58}
59
60// We add our own relationship hooks instead of deriving relationship
61//  because we don't want to despawn Server if there are no more LinkOfs.
62#[derive(Component, Clone, Copy, PartialEq, Eq, Debug, Reflect)]
63#[component(on_insert = LinkOf::on_insert_hook)]
64#[component(on_replace = LinkOf::on_replace)]
65pub struct LinkOf {
66    pub server: Entity,
67}
68
69impl Relationship for LinkOf {
70    type RelationshipTarget = Server;
71    #[inline(always)]
72    fn get(&self) -> Entity {
73        self.server
74    }
75    #[inline]
76    fn from(entity: Entity) -> Self {
77        Self { server: entity }
78    }
79}
80
81impl LinkOf {
82    fn on_insert_hook(
83        mut world: DeferredWorld,
84        HookContext {
85            entity,
86            caller,
87            relationship_hook_mode,
88            ..
89        }: HookContext,
90    ) {
91        match relationship_hook_mode {
92            RelationshipHookMode::Run => {}
93            RelationshipHookMode::Skip => return,
94            RelationshipHookMode::RunIfNotLinked => return,
95        }
96        let target_entity = world.entity(entity).get::<Self>().unwrap().get();
97        if target_entity == entity {
98            warn!(
99                "{}The {}({target_entity:?}) relationship on entity {entity:?} points to itself. The invalid {} relationship has been removed.",
100                caller
101                    .map(|location| format!("{location}: "))
102                    .unwrap_or_default(),
103                core::any::type_name::<Self>(),
104                core::any::type_name::<Self>()
105            );
106            world.commands().entity(entity).remove::<Self>();
107            return;
108        }
109        if let Ok(mut target_entity_mut) = world.get_entity_mut(target_entity) {
110            if let Some(mut relationship_target) = target_entity_mut.get_mut::<Server>() {
111                relationship_target.collection_mut_risky().add(entity);
112            } else {
113                let mut target = <Server as RelationshipTarget>::with_capacity(1);
114                target.collection_mut_risky().add(entity);
115                world.commands().entity(target_entity).insert(target);
116            }
117        } else {
118            warn!(
119                "{}The {}({target_entity:?}) relationship on entity {entity:?} relates to an entity that does not exist. The invalid {} relationship has been removed.",
120                caller
121                    .map(|location| format!("{location}: "))
122                    .unwrap_or_default(),
123                core::any::type_name::<Self>(),
124                core::any::type_name::<Self>()
125            );
126            world.commands().entity(entity).remove::<Self>();
127        }
128    }
129
130    fn on_replace(
131        mut world: DeferredWorld,
132        HookContext {
133            entity,
134            relationship_hook_mode,
135            ..
136        }: HookContext,
137    ) {
138        match relationship_hook_mode {
139            RelationshipHookMode::Run => {}
140            RelationshipHookMode::Skip => return,
141            RelationshipHookMode::RunIfNotLinked => {
142                if <Server as RelationshipTarget>::LINKED_SPAWN {
143                    return;
144                }
145            }
146        }
147        let target_entity = world.entity(entity).get::<Self>().unwrap().get();
148        if let Ok(mut target_entity_mut) = world.get_entity_mut(target_entity) {
149            if let Some(mut relationship_target) = target_entity_mut.get_mut::<Server>() {
150                RelationshipSourceCollection::remove(
151                    relationship_target.collection_mut_risky(),
152                    entity,
153                );
154            }
155        }
156    }
157
158    pub(crate) fn on_insert(
159        trigger: Trigger<OnInsert, (LinkOf, LocalTimeline)>,
160        server: Query<&LocalTimeline, (Without<LinkOf>, With<Server>)>,
161        mut query: Query<(&mut LocalTimeline, &LinkOf)>,
162    ) {
163        if let Ok((mut timeline, link_of)) = query.get_mut(trigger.target()) {
164            if let Ok(server_timeline) = server.get(link_of.get()) {
165                *timeline = server_timeline.clone();
166            }
167        }
168    }
169}
170
171pub struct ServerLinkPlugin;
172impl Plugin for ServerLinkPlugin {
173    fn build(&self, app: &mut App) {
174        if !app.is_plugin_added::<LinkPlugin>() {
175            app.add_plugins(LinkPlugin);
176        }
177        app.register_required_components::<Server, LocalTimeline>();
178        app.add_observer(Server::unlinked);
179        app.add_observer(LinkOf::on_insert);
180    }
181}