[Merged by Bors] - Panic on dropping NonSend in non-origin thread.

crates/bevy_ecs/src/lib.rs

@@ -1267,6 +1267,15 @@ mod tests {
  assert_eq!(*world.non_send_resource_mut::<i64>(), 456);
  }
 
+ #[test]
+ fn non_send_resource_points_to_distinct_data() {
+ let mut world = World::default();
+ world.insert_resource(A(123));
+ world.insert_non_send_resource(A(456));
+ assert_eq!(*world.resource::<A>(), A(123));
+ assert_eq!(*world.non_send_resource::<A>(), A(456));
+ }
+
  #[test]
  #[should_panic]
  fn non_send_resource_panic() {
@@ -1406,38 +1415,32 @@ mod tests {
  assert_eq!(world.resource::<A>().0, 1);
  }
 
- #[test]
- fn non_send_resource_scope() {
- let mut world = World::default();
- world.insert_non_send_resource(NonSendA::default());
- world.resource_scope(|world: &mut World, mut value: Mut<NonSendA>| {
- value.0 += 1;
- assert!(!world.contains_resource::<NonSendA>());
- });
- assert_eq!(world.non_send_resource::<NonSendA>().0, 1);
- }
-
  #[test]
  #[should_panic(
- expected = "attempted to access NonSend resource bevy_ecs::tests::NonSendA off of the main thread"
+ expected = "Attempted to access or drop non-send resource bevy_ecs::tests::NonSendA from thread"
  )]
- fn non_send_resource_scope_from_different_thread() {
+ fn non_send_resource_drop_from_different_thread() {
  let mut world = World::default();
  world.insert_non_send_resource(NonSendA::default());
 
  let thread = std::thread::spawn(move || {
- // Accessing the non-send resource on a different thread
+ // Dropping the non-send resource on a different thread
  // Should result in a panic
- world.resource_scope(|_: &mut World, mut value: Mut<NonSendA>| {
- value.0 += 1;
- });
+ drop(world);
  });
 
  if let Err(err) = thread.join() {
  std::panic::resume_unwind(err);
  }
  }
 
+ #[test]
+ fn non_send_resource_drop_from_same_thread() {
+ let mut world = World::default();
+ world.insert_non_send_resource(NonSendA::default());
+ drop(world);
+ }
+
  #[test]
  fn insert_overwrite_drop() {
  let (dropck1, dropped1) = DropCk::new_pair();

crates/bevy_ecs/src/schedule/ambiguity_detection.rs

@@ -335,6 +335,7 @@ mod tests {
  fn read_only() {
  let mut world = World::new();
  world.insert_resource(R);
+ world.insert_non_send_resource(R);
  world.spawn(A);
  world.init_resource::<Events<E>>();
 
@@ -394,6 +395,7 @@ mod tests {
  fn nonsend() {
  let mut world = World::new();
  world.insert_resource(R);
+ world.insert_non_send_resource(R);
 
  let mut test_stage = SystemStage::parallel();
  test_stage
@@ -497,6 +499,7 @@ mod tests {
  fn ignore_all_ambiguities() {
  let mut world = World::new();
  world.insert_resource(R);
+ world.insert_non_send_resource(R);
 
  let mut test_stage = SystemStage::parallel();
  test_stage
@@ -513,6 +516,7 @@ mod tests {
  fn ambiguous_with_label() {
  let mut world = World::new();
  world.insert_resource(R);
+ world.insert_non_send_resource(R);
 
  #[derive(SystemLabel)]
  struct IgnoreMe;

crates/bevy_ecs/src/storage/mod.rs

@@ -14,5 +14,6 @@ pub use table::*;
 pub struct Storages {
  pub sparse_sets: SparseSets,
  pub tables: Tables,
- pub resources: Resources,
+ pub resources: Resources<true>,
+ pub non_send_resources: Resources<false>,
 }

crates/bevy_ecs/src/storage/resource.rs

@@ -2,19 +2,61 @@ use crate::archetype::ArchetypeComponentId;
 use crate::component::{ComponentId, ComponentTicks, Components, TickCells};
 use crate::storage::{Column, SparseSet, TableRow};
 use bevy_ptr::{OwningPtr, Ptr, UnsafeCellDeref};
+use std::{mem::ManuallyDrop, thread::ThreadId};
 
 /// The type-erased backing storage and metadata for a single resource within a [`World`].
 ///
+/// If `SEND` is false, values of this type will panic if dropped from a different thread.
+///
 /// [`World`]: crate::world::World
-pub struct ResourceData {
- column: Column,
+pub struct ResourceData<const SEND: bool> {
+ column: ManuallyDrop<Column>,
+ type_name: String,
  id: ArchetypeComponentId,
+ origin_thread_id: Option<ThreadId>,
+}
+
+impl<const SEND: bool> Drop for ResourceData<SEND> {
+ fn drop(&mut self) {
+ if self.is_present() {
+ // If this thread is already panicking, panicking again will cause
+ // the entire process to abort. In this case we choose to avoid
+ // dropping or checking this altogether and just leak the column.
+ if std::thread::panicking() {
+ return;
+ }
+ self.validate_access();
+ }
+ // SAFETY: Drop is only called once upon dropping the ResourceData
+ // and is inaccessible after this as the parent ResourceData has
+ // been dropped. The validate_access call above will check that the
+ // data is dropped on the thread it was inserted from.
+ unsafe {
+ ManuallyDrop::drop(&mut self.column);
+ }
+ }
 }
 
-impl ResourceData {
+impl<const SEND: bool> ResourceData<SEND> {
  /// The only row in the underlying column.
  const ROW: TableRow = TableRow::new(0);
 
+ #[inline]
+ fn validate_access(&self) {
+ if SEND {
+ return;
+ }
+ if self.origin_thread_id != Some(std::thread::current().id()) {
+ // Panic in tests, as testing for aborting is nearly impossible
+ panic!(
+ "Attempted to access or drop non-send resource {} from thread {:?} on a thread {:?}. This is not allowed. Aborting.",
+ self.type_name,
+ self.origin_thread_id,
+ std::thread::current().id()
+ );
+ }
+ }
+
  /// Returns true if the resource is populated.
  #[inline]
  pub fn is_present(&self) -> bool {
@@ -28,9 +70,16 @@ impl ResourceData {
  }
 
  /// Gets a read-only pointer to the underlying resource, if available.
+ ///
+ /// # Panics
+ /// If `SEND` is false, this will panic if a value is present and is not accessed from the
+ /// original thread it was inserted from.
  #[inline]
  pub fn get_data(&self) -> Option<Ptr<'_>> {
- self.column.get_data(Self::ROW)
+ self.column.get_data(Self::ROW).map(|res| {
+ self.validate_access();
+ res
+ })
  }
 
  /// Gets a read-only reference to the change ticks of the underlying resource, if available.
@@ -39,83 +88,98 @@ impl ResourceData {
  self.column.get_ticks(Self::ROW)
  }
 
+ /// # Panics
+ /// If `SEND` is false, this will panic if a value is present and is not accessed from the
+ /// original thread it was inserted in.
  #[inline]
  pub(crate) fn get_with_ticks(&self) -> Option<(Ptr<'_>, TickCells<'_>)> {
- self.column.get(Self::ROW)
+ self.column.get(Self::ROW).map(|res| {
+ self.validate_access();
+ res
+ })
  }
 
  /// Inserts a value into the resource. If a value is already present
  /// it will be replaced.
  ///
- /// # Safety
- /// `value` must be valid for the underlying type for the resource.
- ///
- /// The underlying type must be [`Send`] or be inserted from the main thread.
- /// This can be validated with [`World::validate_non_send_access_untyped`].
+ /// # Panics
+ /// If `SEND` is false, this will panic if a value is present and is not replaced from
+ /// the original thread it was inserted in.
  ///
- /// [`World::validate_non_send_access_untyped`]: crate::world::World::validate_non_send_access_untyped
+ /// # Safety
+ /// - `value` must be valid for the underlying type for the resource.
  #[inline]
  pub(crate) unsafe fn insert(&mut self, value: OwningPtr<'_>, change_tick: u32) {
  if self.is_present() {
+ self.validate_access();
  self.column.replace(Self::ROW, value, change_tick);
  } else {
+ if !SEND {
+ self.origin_thread_id = Some(std::thread::current().id());
+ }
  self.column.push(value, ComponentTicks::new(change_tick));
  }
  }
 
  /// Inserts a value into the resource with a pre-existing change tick. If a
  /// value is already present it will be replaced.
  ///
- /// # Safety
- /// `value` must be valid for the underlying type for the resource.
- ///
- /// The underlying type must be [`Send`] or be inserted from the main thread.
- /// This can be validated with [`World::validate_non_send_access_untyped`].
+ /// # Panics
+ /// If `SEND` is false, this will panic if a value is present and is not replaced from
+ /// the original thread it was inserted in.
  ///
- /// [`World::validate_non_send_access_untyped`]: crate::world::World::validate_non_send_access_untyped
+ /// # Safety
+ /// - `value` must be valid for the underlying type for the resource.
  #[inline]
  pub(crate) unsafe fn insert_with_ticks(
  &mut self,
  value: OwningPtr<'_>,
  change_ticks: ComponentTicks,
  ) {
  if self.is_present() {
+ self.validate_access();
  self.column.replace_untracked(Self::ROW, value);
  *self.column.get_added_ticks_unchecked(Self::ROW).deref_mut() = change_ticks.added;
  *self
  .column
  .get_changed_ticks_unchecked(Self::ROW)
  .deref_mut() = change_ticks.changed;
  } else {
+ if !SEND {
+ self.origin_thread_id = Some(std::thread::current().id());
+ }
  self.column.push(value, change_ticks);
  }
  }
 
  /// Removes a value from the resource, if present.
  ///
- /// # Safety
- /// The underlying type must be [`Send`] or be removed from the main thread.
- /// This can be validated with [`World::validate_non_send_access_untyped`].
- ///
- /// The removed value must be used or dropped.
- ///
- /// [`World::validate_non_send_access_untyped`]: crate::world::World::validate_non_send_access_untyped
+ /// # Panics
+ /// If `SEND` is false, this will panic if a value is present and is not removed from the
+ /// original thread it was inserted from.
  #[inline]
  #[must_use = "The returned pointer to the removed component should be used or dropped"]
- pub(crate) unsafe fn remove(&mut self) -> Option<(OwningPtr<'_>, ComponentTicks)> {
- self.column.swap_remove_and_forget(Self::ROW)
+ pub(crate) fn remove(&mut self) -> Option<(OwningPtr<'_>, ComponentTicks)> {
+ if SEND {
+ self.column.swap_remove_and_forget(Self::ROW)
+ } else {
+ self.is_present()
+ .then(|| self.validate_access())
+ .and_then(|_| self.column.swap_remove_and_forget(Self::ROW))
+ }
  }
 
  /// Removes a value from the resource, if present, and drops it.
  ///
- /// # Safety
- /// The underlying type must be [`Send`] or be removed from the main thread.
- /// This can be validated with [`World::validate_non_send_access_untyped`].
- ///
- /// [`World::validate_non_send_access_untyped`]: crate::world::World::validate_non_send_access_untyped
+ /// # Panics
+ /// If `SEND` is false, this will panic if a value is present and is not
+ /// accessed from the original thread it was inserted in.
  #[inline]
- pub(crate) unsafe fn remove_and_drop(&mut self) {
- self.column.clear();
+ pub(crate) fn remove_and_drop(&mut self) {
+ if self.is_present() {
+ self.validate_access();
+ self.column.clear();
+ }
  }
 }
 
@@ -124,11 +188,11 @@ impl ResourceData {
 /// [`Resource`]: crate::system::Resource
 /// [`World`]: crate::world::World
 #[derive(Default)]
-pub struct Resources {
- resources: SparseSet<ComponentId, ResourceData>,
+pub struct Resources<const SEND: bool> {
+ resources: SparseSet<ComponentId, ResourceData<SEND>>,
 }
 
-impl Resources {
+impl<const SEND: bool> Resources<SEND> {
  /// The total number of resources stored in the [`World`]
  ///
  /// [`World`]: crate::world::World
@@ -138,7 +202,7 @@ impl Resources {
  }
 
  /// Iterate over all resources that have been initialized, i.e. given a [`ComponentId`]
- pub fn iter(&self) -> impl Iterator<Item = (ComponentId, &ResourceData)> {
+ pub fn iter(&self) -> impl Iterator<Item = (ComponentId, &ResourceData<SEND>)> {
  self.resources.iter().map(|(id, data)| (*id, data))
  }
 
@@ -153,31 +217,37 @@ impl Resources {
 
  /// Gets read-only access to a resource, if it exists.
  #[inline]
- pub fn get(&self, component_id: ComponentId) -> Option<&ResourceData> {
+ pub fn get(&self, component_id: ComponentId) -> Option<&ResourceData<SEND>> {
  self.resources.get(component_id)
  }
 
  /// Gets mutable access to a resource, if it exists.
  #[inline]
- pub(crate) fn get_mut(&mut self, component_id: ComponentId) -> Option<&mut ResourceData> {
+ pub(crate) fn get_mut(&mut self, component_id: ComponentId) -> Option<&mut ResourceData<SEND>> {
  self.resources.get_mut(component_id)
  }
 
  /// Fetches or initializes a new resource and returns back it's underlying column.
  ///
  /// # Panics
  /// Will panic if `component_id` is not valid for the provided `components`
+ /// If `SEND` is false, this will panic if `component_id`'s `ComponentInfo` is not registered as being `Send` + `Sync`.
  pub(crate) fn initialize_with(
  &mut self,
  component_id: ComponentId,
  components: &Components,
  f: impl FnOnce() -> ArchetypeComponentId,
- ) -> &mut ResourceData {
+ ) -> &mut ResourceData<SEND> {
  self.resources.get_or_insert_with(component_id, || {
  let component_info = components.get_info(component_id).unwrap();
+ if SEND {
+ assert!(component_info.is_send_and_sync());
+ }
  ResourceData {
- column: Column::with_capacity(component_info, 1),
+ column: ManuallyDrop::new(Column::with_capacity(component_info, 1)),
+ type_name: String::from(component_info.name()),
  id: f(),
+ origin_thread_id: None,
  }
  })
  }