From fbef7f3f74ae20d3266e62323efb5b42c9197ad6 Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Thu, 18 Jul 2024 20:42:59 +0200
Subject: [PATCH] chore: track_caller for functions that panic (#9626)

---
 crates/stages/api/src/pipeline/set.rs | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/crates/stages/api/src/pipeline/set.rs b/crates/stages/api/src/pipeline/set.rs
index 99de4a06b2..13dd86e128 100644
--- a/crates/stages/api/src/pipeline/set.rs
+++ b/crates/stages/api/src/pipeline/set.rs
@@ -190,20 +190,23 @@ where
     /// # Panics
     ///
     /// Panics if the stage is not in this set.
+    #[track_caller]
     pub fn disable(mut self, stage_id: StageId) -> Self {
-        let entry =
-            self.stages.get_mut(&stage_id).expect("Cannot disable a stage that is not in the set.");
+        let entry = self
+            .stages
+            .get_mut(&stage_id)
+            .unwrap_or_else(|| panic!("Cannot disable a stage that is not in the set: {stage_id}"));
         entry.enabled = false;
         self
     }
 
     /// Disables all given stages. See [`disable`](Self::disable).
+    #[track_caller]
     pub fn disable_all(mut self, stages: &[StageId]) -> Self {
         for stage_id in stages {
-            let entry = self
-                .stages
-                .get_mut(stage_id)
-                .expect("Cannot disable a stage that is not in the set.");
+            let entry = self.stages.get_mut(stage_id).unwrap_or_else(|| {
+                panic!("Cannot disable a stage that is not in the set: {stage_id}")
+            });
             entry.enabled = false;
         }
         self
@@ -212,6 +215,7 @@ where
     /// Disables the given stage if the given closure returns true.
     ///
     /// See [`Self::disable`]
+    #[track_caller]
     pub fn disable_if<F>(self, stage_id: StageId, f: F) -> Self
     where
         F: FnOnce() -> bool,
@@ -225,6 +229,7 @@ where
     /// Disables all given stages if the given closure returns true.
     ///
     /// See [`Self::disable`]
+    #[track_caller]
     pub fn disable_all_if<F>(self, stages: &[StageId], f: F) -> Self
     where
         F: FnOnce() -> bool,