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

crates/bevy_ecs/src/lib.rs

@@ -1334,7 +1334,7 @@ mod tests {
 
  #[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() {
  let mut world = World::default();
@@ -1353,6 +1353,25 @@ mod tests {
  }
  }
 
+ #[test]
+ #[should_panic(
+ expected = "Attempted to access or drop non-send resource bevy_ecs::tests::NonSendA from 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 || {
+ // Dropping the non-send resource on a different thread
+ // Should result in a panic
+ drop(world);
+ });
+
+ if let Err(err) = thread.join() {
+ std::panic::resume_unwind(err);
+ }
+ }
+
  #[test]
  fn insert_overwrite_drop() {
  let (dropck1, dropped1) = DropCk::new_pair();

crates/bevy_ecs/src/storage/resource.rs

@@ -2,16 +2,54 @@ use crate::archetype::ArchetypeComponentId;
 use crate::component::{ComponentId, ComponentTicks, Components, TickCells};
 use crate::storage::{Column, SparseSet};
 use bevy_ptr::{OwningPtr, Ptr, UnsafeCellDeref};
+use std::cell::UnsafeCell;
+use std::mem::ManuallyDrop;
+use std::thread::ThreadId;
 
 /// The type-erased backing storage and metadata for a single resource within a [`World`].
 ///
 /// [`World`]: crate::world::World
 pub struct ResourceData {
- column: Column,
+ column: ManuallyDrop<Column>,
+ type_name: String,
  id: ArchetypeComponentId,
+ origin_thread_id: Option<ThreadId>,
+}
+
+impl Drop for ResourceData {
+ fn drop(&mut self) {
+ if self.is_present() {
+ 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.
+ unsafe {
+ ManuallyDrop::drop(&mut self.column);
+ }
+ }
 }
 
 impl ResourceData {
+ #[inline]
+ fn validate_access(&self) {
+ // Avoid aborting due to double panicking on the same thread.
+ if std::thread::panicking() {
+ return;
+ }
+ if let Some(origin_thread_id) = self.origin_thread_id {
+ if origin_thread_id != 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,
+ origin_thread_id,
+ std::thread::current().id()
+ );
+ }
+ }
+ }
+
  /// Returns true if the resource is populated.
  #[inline]
  pub fn is_present(&self) -> bool {
@@ -25,9 +63,16 @@ impl ResourceData {
  }
 
  /// Gets a read-only pointer to the underlying resource, if available.
+ ///
+ /// # Panics
+ /// This will panic if a value is present, the underlying type is
+ /// `!Send`, and is not accessed from the original thread it was inserted in.
  #[inline]
  pub fn get_data(&self) -> Option<Ptr<'_>> {
- self.column.get_data(0)
+ self.column.get_data(0).map(|res| {
+ self.validate_access();
+ res
+ })
  }
 
  /// Gets a read-only reference to the change ticks of the underlying resource, if available.
@@ -36,80 +81,97 @@ impl ResourceData {
  self.column.get_ticks(0)
  }
 
+ /// # Panics
+ /// This will panic if a value is present, the underlying type is
+ /// `!Send`, 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(0)
+ self.column.get(0).map(|res| {
+ self.validate_access();
+ res
+ })
  }
 
  /// Inserts a value into the resource. If a value is already present
  /// it will be replaced.
  ///
+ /// # Panics
+ /// This will panic if a value is present, the underlying type is
+ /// `!Send`, and is not accessed from the original thread it was inserted in.
+ ///
  /// # 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`].
- ///
- /// [`World::validate_non_send_access_untyped`]: crate::world::World::validate_non_send_access_untyped
  #[inline]
  pub(crate) unsafe fn insert(&mut self, value: OwningPtr<'_>, change_tick: u32) {
  if self.is_present() {
+ self.validate_access();
  self.column.replace(0, value, change_tick);
  } else {
+ self.origin_thread_id = self.origin_thread_id.map(|_| 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.
  ///
+ /// # Panics
+ /// This will panic if a value is present, the underlying type is
+ /// `!Send`, and is not accessed from the original thread it was inserted in.
+ ///
  /// # 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`].
- ///
- /// [`World::validate_non_send_access_untyped`]: crate::world::World::validate_non_send_access_untyped
  #[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(0, value);
  *self.column.get_added_ticks_unchecked(0).deref_mut() = change_ticks.added;
  *self.column.get_changed_ticks_unchecked(0).deref_mut() = change_ticks.changed;
  } else {
+ self.origin_thread_id = self.origin_thread_id.map(|_| std::thread::current().id());
  self.column.push(value, change_ticks);
  }
  }
 
  /// Removes a value from the resource, if present.
  ///
+ /// # Panics
+ /// This will panic if a value is present, the underlying type is
+ /// `!Send`, and is not accessed from the original thread it was inserted in.
+ ///
  /// # 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 underlying type must be [`Send`] or be removed from the thread it was
+ /// inserted from.
  ///
  /// The removed value must be used or dropped.
- ///
- /// [`World::validate_non_send_access_untyped`]: crate::world::World::validate_non_send_access_untyped
  #[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(0)
+ self.is_present()
+ .then(|| self.validate_access())
+ .and_then(|_| self.column.swap_remove_and_forget(0))
  }
 
  /// 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`].
+ /// # Panics
+ /// This will panic if a value is present, the underlying type is
+ /// `!Send`, and is not accessed from the original thread it was inserted in.
  ///
- /// [`World::validate_non_send_access_untyped`]: crate::world::World::validate_non_send_access_untyped
+ /// # Safety
+ /// The underlying type must be [`Send`] or be removed from the thread it was
+ /// inserted from.
  #[inline]
  pub(crate) unsafe fn remove_and_drop(&mut self) {
- self.column.clear();
+ if self.is_present() {
+ self.validate_access();
+ self.column.clear();
+ }
  }
 }
 
@@ -161,17 +223,23 @@ impl Resources {
  ///
  /// # Panics
  /// Will panic if `component_id` is not valid for the provided `components`
- pub(crate) fn initialize_with(
+ ///
+ /// # Safety
+ /// `is_send` must be accurate for the Resource that is being initialized.
+ pub(crate) unsafe fn initialize_with(
  &mut self,
  component_id: ComponentId,
  components: &Components,
+ is_send: bool,
  f: impl FnOnce() -> ArchetypeComponentId,
  ) -> &mut ResourceData {
  self.resources.get_or_insert_with(component_id, || {
  let component_info = components.get_info(component_id).unwrap();
  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: (!is_send).then(|| std::thread::current().id()),
  }
  })
  }

crates/bevy_ecs/src/system/system_param.rs

@@ -1012,7 +1012,6 @@ impl<'w, 's, T: 'static> SystemParamFetch<'w, 's> for NonSendState<T> {
  world: &'w World,
  change_tick: u32,
  ) -> Self::Item {
- world.validate_non_send_access::<T>();
  let (ptr, ticks) = world
  .get_resource_with_ticks(state.component_id)
  .unwrap_or_else(|| {
@@ -1062,7 +1061,6 @@ impl<'w, 's, T: 'static> SystemParamFetch<'w, 's> for OptionNonSendState<T> {
  world: &'w World,
  change_tick: u32,
  ) -> Self::Item {
- world.validate_non_send_access::<T>();
  world
  .get_resource_with_ticks(state.0.component_id)
  .map(|(ptr, ticks)| NonSend {
@@ -1129,7 +1127,6 @@ impl<'w, 's, T: 'static> SystemParamFetch<'w, 's> for NonSendMutState<T> {
  world: &'w World,
  change_tick: u32,
  ) -> Self::Item {
- world.validate_non_send_access::<T>();
  let (ptr, ticks) = world
  .get_resource_with_ticks(state.component_id)
  .unwrap_or_else(|| {
@@ -1173,7 +1170,6 @@ impl<'w, 's, T: 'static> SystemParamFetch<'w, 's> for OptionNonSendMutState<T> {
  world: &'w World,
  change_tick: u32,
  ) -> Self::Item {
- world.validate_non_send_access::<T>();
  world
  .get_resource_with_ticks(state.0.component_id)
  .map(|(ptr, ticks)| NonSendMut {