Skip to main content

steel_core/inventory/
lock.rs

1//! Container locking utilities for deadlock-free multi-container operations.
2//!
3//! This module provides types for safely locking multiple containers in a
4//! deterministic order to prevent deadlocks when performing operations that
5//! span multiple inventories (e.g., transferring items between containers).
6
7use parking_lot::ArcMutexGuard;
8use parking_lot::RawMutex;
9use rustc_hash::FxHashMap;
10use std::borrow::Borrow;
11use std::mem;
12use std::ops::{Deref, DerefMut};
13use std::sync::Arc;
14use steel_utils::locks::Shared;
15use steel_utils::{Downcast as _, DowncastType, locks::SyncMutex};
16
17use crate::{
18    block_entity::{BlockEntityBase, SharedBlockEntity},
19    inventory::container::Container,
20    player::{Player, player_inventory::PlayerInventory},
21};
22use steel_registry::item_stack::ItemStack;
23
24type ContainerChangedCallback = Arc<dyn Fn() + Send + Sync>;
25
26#[derive(Clone)]
27struct ContainerOwner {
28    block_entity: Arc<BlockEntityBase>,
29    after_changed: Option<ContainerChangedCallback>,
30}
31
32/// Thread-safe reference to an erased container.
33pub type SharedContainer = Shared<dyn Container>;
34
35struct LockedContainer(ArcMutexGuard<RawMutex, dyn Container>);
36
37impl Deref for LockedContainer {
38    type Target = dyn Container;
39
40    fn deref(&self) -> &Self::Target {
41        &*self.0
42    }
43}
44
45impl DerefMut for LockedContainer {
46    fn deref_mut(&mut self) -> &mut Self::Target {
47        &mut *self.0
48    }
49}
50
51/// A reference to a container that can be locked.
52///
53/// The optional owner is notified only after all container locks have been
54/// released at the corresponding Vanilla `setChanged` boundary.
55#[derive(Clone)]
56pub struct ContainerRef {
57    id: ContainerId,
58    source: SharedContainer,
59    owner: Option<ContainerOwner>,
60}
61
62impl<T> From<Shared<T>> for ContainerRef
63where
64    T: Container + 'static,
65{
66    fn from(container: Shared<T>) -> Self {
67        let id = ContainerId::from_arc(&container);
68        let container: SharedContainer = container;
69        Self {
70            id,
71            source: container,
72            owner: None,
73        }
74    }
75}
76
77impl<T> From<&Shared<T>> for ContainerRef
78where
79    T: Container + 'static,
80{
81    fn from(container: &Shared<T>) -> Self {
82        container.clone().into()
83    }
84}
85
86impl From<&ContainerRef> for ContainerRef {
87    fn from(r: &ContainerRef) -> Self {
88        r.clone()
89    }
90}
91
92impl From<&SharedContainer> for ContainerRef {
93    fn from(container: &SharedContainer) -> Self {
94        container.clone().into()
95    }
96}
97
98impl From<SharedContainer> for ContainerRef {
99    fn from(container: SharedContainer) -> Self {
100        Self {
101            id: ContainerId::from_arc(&container),
102            source: container,
103            owner: None,
104        }
105    }
106}
107
108impl ContainerRef {
109    /// Creates a `ContainerRef` from a block entity, if it implements Container.
110    ///
111    /// Returns `None` if the block entity has no container capability.
112    #[must_use]
113    pub fn from_block_entity(block_entity: SharedBlockEntity) -> Option<Self> {
114        block_entity.container_ref()
115    }
116
117    /// Creates a container capability owned by the block entity whose unique
118    /// base is returned from its [`crate::block_entity::BlockEntity::base`].
119    ///
120    /// The container implementation must obey [`Container`]'s storage-local
121    /// locking contract. This owner performs block-entity and comparator
122    /// notifications after all container locks are released.
123    #[must_use]
124    pub fn owned_by_block_entity(container: SharedContainer, owner: Arc<BlockEntityBase>) -> Self {
125        Self {
126            id: ContainerId::from_arc(&container),
127            source: container,
128            owner: Some(ContainerOwner {
129                block_entity: owner,
130                after_changed: None,
131            }),
132        }
133    }
134
135    /// Creates a block-entity container capability with an unlocked post-change callback.
136    ///
137    /// The callback runs after the container lock is released and after the owning
138    /// block entity is marked changed. Specialized containers can use this to
139    /// publish state derived from their contents without violating [`Container`]'s
140    /// storage-local locking contract.
141    #[must_use]
142    pub(crate) fn owned_by_block_entity_with_callback(
143        container: SharedContainer,
144        owner: Arc<BlockEntityBase>,
145        after_changed: ContainerChangedCallback,
146    ) -> Self {
147        Self {
148            id: ContainerId::from_arc(&container),
149            source: container,
150            owner: Some(ContainerOwner {
151                block_entity: owner,
152                after_changed: Some(after_changed),
153            }),
154        }
155    }
156
157    /// Returns a unique identifier for this container based on its Arc pointer address.
158    #[must_use]
159    pub const fn container_id(&self) -> ContainerId {
160        self.id
161    }
162
163    /// Checks container access without locking its item storage.
164    #[must_use]
165    pub fn still_valid(&self, player: &Player) -> bool {
166        self.owner
167            .as_ref()
168            .is_none_or(|owner| owner.block_entity.is_valid_container_for(player))
169    }
170
171    /// Locks this container and returns a guard.
172    fn lock(&self) -> LockedContainer {
173        LockedContainer(SyncMutex::lock_arc(&self.source))
174    }
175}
176
177/// A guard that holds locks on multiple containers in a deterministic order.
178///
179/// This struct ensures that when multiple containers need to be locked simultaneously,
180/// they are always locked in the same order (by pointer address) to prevent deadlocks.
181///
182/// # Example
183///
184/// ```ignore
185/// let player_inv = ContainerRef::from(player_inv_arc);
186/// let chest = ContainerRef::from(chest_arc);
187///
188/// let mut guard = ContainerLockGuard::lock_all(&[&player_inv, &chest]);
189///
190/// // Access containers by their IDs
191/// let player_id = player_inv.container_id();
192/// if let Some(inv) = guard.get_mut(player_id) {
193///     // Modify the player inventory
194/// }
195/// ```
196pub struct ContainerLockGuard {
197    // Stable sources retained while guards are temporarily released for an owner callback.
198    sources: Vec<(ContainerId, ContainerRef)>,
199    // Store locked guards in deterministic order
200    guards: Vec<(ContainerId, LockedContainer)>,
201    // For quick lookup
202    id_to_index: FxHashMap<ContainerId, usize>,
203}
204
205impl ContainerLockGuard {
206    /// Create a new lock guard and lock all containers in deterministic order.
207    ///
208    /// Containers are sorted by their pointer address before locking to ensure
209    /// a consistent lock order across all call sites, preventing deadlocks.
210    /// Duplicate containers (same Arc) are automatically deduplicated.
211    #[must_use]
212    pub fn lock_all<C>(containers: &[C]) -> Self
213    where
214        C: Borrow<ContainerRef>,
215    {
216        let mut sources: Vec<_> = containers
217            .iter()
218            .map(|c| (c.borrow().container_id(), c.borrow().clone()))
219            .collect();
220
221        // Sort by ID for deterministic lock order (prevents deadlocks)
222        sources.sort_by_key(|(id, _)| *id);
223
224        // Deduplicate (in case same container passed multiple times)
225        sources.dedup_by_key(|(id, _)| *id);
226
227        // Lock all in sorted order
228        let mut guards = Vec::with_capacity(sources.len());
229        for (id, container) in &sources {
230            let guard = container.lock();
231            guards.push((*id, guard));
232        }
233
234        // Build index map
235        let id_to_index = guards
236            .iter()
237            .enumerate()
238            .map(|(idx, (id, _))| (*id, idx))
239            .collect();
240
241        Self {
242            sources,
243            guards,
244            id_to_index,
245        }
246    }
247
248    /// Get mutable access to N locked containers simultaneously
249    ///
250    /// Returns `None` if any ID is not locked or if any IDs are duplicates
251    pub fn get_disjoint_mut<const N: usize>(
252        &mut self,
253        ids: [ContainerId; N],
254    ) -> Option<[&mut dyn Container; N]> {
255        let mut indices = [0usize; N];
256        for (i, id) in ids.iter().enumerate() {
257            indices[i] = *self.id_to_index.get(id)?;
258        }
259        let entries = self.guards.get_disjoint_mut(indices).ok()?;
260        Some(entries.map(|(_, locked)| &mut **locked as &mut dyn Container))
261    }
262
263    /// Unlock all containers and relock with a new set.
264    ///
265    /// This should only be called when you need to add more containers
266    /// mid-operation. All existing references from `get()`/`get_mut()` are invalidated.
267    #[must_use]
268    pub fn relock(self, containers: &[&ContainerRef]) -> Self {
269        // Drop self, releasing all locks
270        drop(self);
271        // Lock new set
272        Self::lock_all(containers)
273    }
274
275    /// Get immutable access to a locked container.
276    #[must_use]
277    pub fn get(&self, id: impl Into<ContainerId>) -> Option<&dyn Container> {
278        self.id_to_index
279            .get(&id.into())
280            .and_then(|&idx| self.guards.get(idx))
281            .map(|(_, guard)| &**guard as &dyn Container)
282    }
283
284    /// Get mutable access to a locked container.
285    ///
286    /// This bypasses owner notification and is only for deliberate no-update
287    /// mutation or callers that establish the notification boundary separately.
288    pub fn get_mut(&mut self, id: impl Into<ContainerId>) -> Option<&mut dyn Container> {
289        self.id_to_index
290            .get(&id.into())
291            .copied()
292            .and_then(|idx| self.guards.get_mut(idx))
293            .map(|(_, guard)| &mut **guard as &mut dyn Container)
294    }
295
296    /// Mirrors a container's own `setItem` call.
297    ///
298    /// Vanilla block-entity containers call `BlockEntity::setChanged` from
299    /// `setItem`, before a `Slot::set` performs its separate notification.
300    pub fn set_item(&mut self, id: impl Into<ContainerId>, slot: usize, stack: ItemStack) -> bool {
301        let id = id.into();
302        let Some(&index) = self.id_to_index.get(&id) else {
303            return false;
304        };
305        self.guards[index].1.set_item(slot, stack);
306        let owner = self.sources[index].1.owner.clone();
307        self.notify_owner(owner);
308        true
309    }
310
311    /// Mirrors a container's own conditional `removeItem` notification.
312    pub fn remove_item(
313        &mut self,
314        id: impl Into<ContainerId>,
315        slot: usize,
316        amount: i32,
317    ) -> Option<ItemStack> {
318        let id = id.into();
319        let index = *self.id_to_index.get(&id)?;
320        let removed = self.guards.get_mut(index)?.1.remove_item(slot, amount);
321        if !removed.is_empty() {
322            let owner = self.sources.get(index)?.1.owner.clone();
323            self.notify_owner(owner);
324        }
325        Some(removed)
326    }
327
328    /// Calls `Container::set_changed` and synchronously notifies its owner
329    /// after releasing every lock held by this guard.
330    pub fn set_changed(&mut self, id: impl Into<ContainerId>) -> bool {
331        let id = id.into();
332        let Some(&index) = self.id_to_index.get(&id) else {
333            return false;
334        };
335        self.guards[index].1.set_changed();
336        let owner = self.sources[index].1.owner.clone();
337        self.notify_owner(owner);
338        true
339    }
340
341    /// Runs a callback after releasing every container, then reacquires the
342    /// same sources in deterministic order before returning.
343    pub(crate) fn run_unlocked<R>(&mut self, callback: impl FnOnce() -> R) -> R {
344        drop(mem::take(&mut self.guards));
345        let result = callback();
346        self.guards = self
347            .sources
348            .iter()
349            .map(|(id, container)| (*id, container.lock()))
350            .collect();
351        result
352    }
353
354    fn notify_owner(&mut self, owner: Option<ContainerOwner>) {
355        let Some(owner) = owner else {
356            return;
357        };
358
359        self.run_unlocked(|| {
360            owner.block_entity.set_changed();
361            if let Some(after_changed) = owner.after_changed {
362                after_changed();
363            }
364        });
365    }
366
367    /// Gets immutable access when the locked container has concrete type `T`.
368    #[must_use]
369    pub fn get_typed<T>(&self, id: impl Into<ContainerId>) -> Option<&T>
370    where
371        T: Container + DowncastType,
372    {
373        self.get(id)?.downcast_ref::<T>()
374    }
375
376    /// Gets mutable access when the locked container has concrete type `T`.
377    ///
378    /// This bypasses owner notification and is only for deliberate no-update
379    /// mutation or callers that establish the notification boundary separately.
380    pub fn get_typed_mut<T>(&mut self, id: impl Into<ContainerId>) -> Option<&mut T>
381    where
382        T: Container + DowncastType,
383    {
384        self.get_mut(id)?.downcast_mut::<T>()
385    }
386
387    /// Gets mutable access to two distinct concrete containers.
388    ///
389    /// This bypasses owner notification and is only for deliberate no-update
390    /// mutation or callers that establish the notification boundary separately.
391    pub fn get_two_typed_mut<A, B>(
392        &mut self,
393        first: impl Into<ContainerId>,
394        second: impl Into<ContainerId>,
395    ) -> Option<(&mut A, &mut B)>
396    where
397        A: Container + DowncastType,
398        B: Container + DowncastType,
399    {
400        let first_index = *self.id_to_index.get(&first.into())?;
401        let second_index = *self.id_to_index.get(&second.into())?;
402        if first_index == second_index {
403            return None;
404        }
405
406        let (first, second): (&mut dyn Container, &mut dyn Container) =
407            if first_index < second_index {
408                let (before_second, from_second) = self.guards.split_at_mut(second_index);
409                (&mut *before_second[first_index].1, &mut *from_second[0].1)
410            } else {
411                let (before_first, from_first) = self.guards.split_at_mut(first_index);
412                (&mut *from_first[0].1, &mut *before_first[second_index].1)
413            };
414        Some((first.downcast_mut::<A>()?, second.downcast_mut::<B>()?))
415    }
416
417    /// Check if a container is locked.
418    #[must_use]
419    pub fn contains(&self, id: ContainerId) -> bool {
420        self.id_to_index.contains_key(&id)
421    }
422}
423
424/// Unique identifier for a container based on Arc pointer address.
425///
426/// This ID is used to establish a deterministic ordering when locking
427/// multiple containers, preventing deadlocks.
428#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
429pub struct ContainerId(usize);
430
431impl ContainerId {
432    /// Creates a container ID from an Arc's pointer address.
433    pub fn from_arc<T: ?Sized>(arc: &Arc<T>) -> Self {
434        Self(Arc::as_ptr(arc).cast::<()>() as usize)
435    }
436}
437
438impl From<&Shared<PlayerInventory>> for ContainerId {
439    fn from(value: &Shared<PlayerInventory>) -> Self {
440        Self::from_arc(value)
441    }
442}
443
444#[cfg(test)]
445mod tests {
446    use std::sync::{Arc, Weak};
447
448    use steel_registry::blocks::block_state_ext::BlockStateExt as _;
449    use steel_registry::blocks::properties::{BlockStateProperties, Direction};
450    use steel_registry::{
451        init_vanilla_registry, item_stack::ItemStack, vanilla_block_entity_types, vanilla_blocks,
452        vanilla_items,
453    };
454    use steel_utils::types::UpdateFlags;
455    use steel_utils::{BlockPos, ChunkPos, locks::SyncMutex};
456
457    use super::{ContainerId, ContainerLockGuard, ContainerRef};
458    use crate::behavior::{BLOCK_BEHAVIORS, init_behaviors};
459    use crate::block_entity::{
460        SharedBlockEntity,
461        entities::{BarrelBlockEntity, UnimplementedBlockEntity},
462        init_block_entities,
463    };
464    use crate::inventory::container::{Container, CraftingContainer, ResultContainer};
465    use crate::inventory::slots::{NormalSlot, Slot as _};
466    use crate::test_support::{fresh_test_world, insert_ready_full_chunk};
467
468    #[test]
469    fn erased_container_ref_preserves_id_and_typed_access() {
470        let crafting = Arc::new(SyncMutex::new(CraftingContainer::new(2, 2)));
471        let id = ContainerId::from_arc(&crafting);
472        let container_ref = ContainerRef::from(Arc::clone(&crafting));
473
474        assert_eq!(container_ref.container_id(), id);
475
476        let mut guard = ContainerLockGuard::lock_all(&[&container_ref]);
477        let Some(typed) = guard.get_typed::<CraftingContainer>(id) else {
478            panic!("erased crafting container should retain its concrete type");
479        };
480        assert_eq!((typed.width(), typed.height()), (2, 2));
481        assert!(guard.get_typed::<ResultContainer>(id).is_none());
482        assert!(guard.get_typed_mut::<CraftingContainer>(id).is_some());
483    }
484
485    #[test]
486    fn block_entity_container_capability_is_independently_lockable() {
487        init_vanilla_registry();
488        let barrel = Arc::new(BarrelBlockEntity::new(
489            Weak::new(),
490            BlockPos::new(1, 2, 3),
491            vanilla_blocks::BARREL.default_state(),
492        ));
493        let block_entity: SharedBlockEntity = barrel.clone();
494        let Some(container_ref) = ContainerRef::from_block_entity(block_entity) else {
495            panic!("barrel block entity should expose Container");
496        };
497        let id = container_ref.container_id();
498
499        let guard = ContainerLockGuard::lock_all(&[&container_ref]);
500        assert_eq!(guard.get(id).map(Container::get_container_size), Some(27));
501    }
502
503    #[test]
504    fn non_container_block_entity_ref_is_rejected() {
505        init_vanilla_registry();
506        let block_entity: SharedBlockEntity = Arc::new(UnimplementedBlockEntity::new(
507            &vanilla_block_entity_types::END_PORTAL,
508            Weak::new(),
509            BlockPos::new(1, 2, 3),
510            vanilla_blocks::END_PORTAL.default_state(),
511        ));
512
513        assert!(ContainerRef::from_block_entity(block_entity).is_none());
514    }
515
516    #[test]
517    fn unlocked_callback_releases_and_reacquires_every_container() {
518        let crafting = Arc::new(SyncMutex::new(CraftingContainer::new(2, 2)));
519        let result = Arc::new(SyncMutex::new(ResultContainer::new()));
520        let crafting_ref = ContainerRef::from(Arc::clone(&crafting));
521        let result_ref = ContainerRef::from(Arc::clone(&result));
522        let mut guard = ContainerLockGuard::lock_all(&[&crafting_ref, &result_ref]);
523
524        let both_unlocked = guard.run_unlocked(|| {
525            let crafting_guard = crafting.try_lock();
526            let result_guard = result.try_lock();
527            crafting_guard.is_some() && result_guard.is_some()
528        });
529
530        assert!(both_unlocked);
531        assert!(guard.get(crafting_ref.container_id()).is_some());
532        assert!(guard.get(result_ref.container_id()).is_some());
533    }
534
535    #[test]
536    fn barrel_change_reenters_analog_read_without_holding_container_lock() {
537        init_vanilla_registry();
538        init_behaviors();
539        init_block_entities();
540        let world = fresh_test_world("barrel_comparator_reentry");
541        let barrel_pos = BlockPos::new(8, 64, 8);
542        let comparator_pos = barrel_pos.west();
543        insert_ready_full_chunk(&world, ChunkPos::from_block_pos(barrel_pos));
544        assert!(world.set_block(
545            comparator_pos.below(),
546            vanilla_blocks::STONE.default_state(),
547            UpdateFlags::UPDATE_NONE,
548        ));
549        assert!(world.set_block(
550            barrel_pos,
551            vanilla_blocks::BARREL.default_state(),
552            UpdateFlags::UPDATE_NONE,
553        ));
554        let comparator_state = vanilla_blocks::COMPARATOR
555            .default_state()
556            .set_value(&BlockStateProperties::HORIZONTAL_FACING, Direction::East);
557        assert!(world.set_block(comparator_pos, comparator_state, UpdateFlags::UPDATE_NONE,));
558        assert!(!world.has_scheduled_block_tick(comparator_pos, &vanilla_blocks::COMPARATOR));
559
560        let container_ref = ContainerRef::from_block_entity(
561            world
562                .get_block_entity(barrel_pos)
563                .expect("barrel should create its block entity"),
564        )
565        .expect("barrel should expose a container capability");
566        let slot = NormalSlot::new(container_ref.clone(), 0);
567        let mut guard = ContainerLockGuard::lock_all(&[&container_ref]);
568        slot.set_item(&mut guard, ItemStack::new(&vanilla_items::STONE));
569        drop(guard);
570
571        let analog = BLOCK_BEHAVIORS
572            .get_behavior(&vanilla_blocks::BARREL)
573            .get_analog_output_signal(
574                world.get_block_state(barrel_pos),
575                world.as_ref(),
576                barrel_pos,
577                Direction::West,
578            );
579        assert_eq!(analog, 1);
580        assert!(
581            world.has_scheduled_block_tick(comparator_pos, &vanilla_blocks::COMPARATOR),
582            "barrel callback did not schedule the comparator despite analog output {analog}"
583        );
584    }
585}