From 000d3be66ec5b5fe812798715aa22a2a53b0843e Mon Sep 17 00:00:00 2001
From: Captain Yoshi <captain.yoshisaur@gmail.com>
Date: Tue, 16 Jul 2024 15:20:15 -0400
Subject: [PATCH] [impl] don't compute stage if preempted

Checked at the start of the StagePrivate::runCompute(). A callback must
be issued, otherwise the stage/s cannot be preempted.
---
 .../include/moveit/task_constructor/stage_p.h | 21 +++++++++++++++++++
 core/src/stage.cpp                            | 14 ++++++++++++-
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/core/include/moveit/task_constructor/stage_p.h b/core/include/moveit/task_constructor/stage_p.h
index 5d359b5c6..36cd8558a 100644
--- a/core/include/moveit/task_constructor/stage_p.h
+++ b/core/include/moveit/task_constructor/stage_p.h
@@ -57,6 +57,17 @@
 namespace moveit {
 namespace task_constructor {
 
+/// exception thrown by StagePrivate::runCompute()
+class PreemptStageException : public std::exception
+{
+public:
+	explicit PreemptStageException() {}
+	const char* what() const noexcept override {
+		static const char* msg = "";
+		return msg;
+	}
+};
+
 class ContainerBase;
 class StagePrivate
 {
@@ -146,6 +157,10 @@ class StagePrivate
 	bool storeFailures() const { return introspection_ != nullptr; }
 	void runCompute() {
 		ROS_DEBUG_STREAM_NAMED("Stage", fmt::format("Computing stage '{}'", name()));
+
+		if (preempted())
+			throw PreemptStageException();
+
 		auto compute_start_time = std::chrono::steady_clock::now();
 		try {
 			compute();
@@ -159,6 +174,10 @@ class StagePrivate
 	/** compute cost for solution through configured CostTerm */
 	void computeCost(const InterfaceState& from, const InterfaceState& to, SolutionBase& solution);
 
+	void setPreemptedCheck(const std::atomic<bool>* preempt_requested);
+	/// is the stage preempted ? defaults to false
+	bool preempted() const;
+
 protected:
 	StagePrivate& operator=(StagePrivate&& other);
 
@@ -197,6 +216,8 @@ class StagePrivate
 	InterfaceWeakPtr next_starts_;  // interface to be used for sendForward()
 
 	Introspection* introspection_;  // task's introspection instance
+
+	std::atomic<bool> const* preempt_requested_;
 };
 PIMPL_FUNCTIONS(Stage)
 std::ostream& operator<<(std::ostream& os, const StagePrivate& stage);
diff --git a/core/src/stage.cpp b/core/src/stage.cpp
index 97c99826b..d53dfebb0 100644
--- a/core/src/stage.cpp
+++ b/core/src/stage.cpp
@@ -102,7 +102,8 @@ StagePrivate::StagePrivate(Stage* me, const std::string& name)
   , cost_term_{ std::make_unique<CostTerm>() }
   , total_compute_time_{}
   , parent_{ nullptr }
-  , introspection_{ nullptr } {}
+  , introspection_{ nullptr }
+  , preempt_requested_{ nullptr } {}
 
 StagePrivate& StagePrivate::operator=(StagePrivate&& other) {
 	assert(typeid(*this) == typeid(other));
@@ -305,6 +306,17 @@ void StagePrivate::computeCost(const InterfaceState& from, const InterfaceState&
 	}
 }
 
+void StagePrivate::setPreemptedCheck(std::atomic<bool> const* preempt_requested) {
+	preempt_requested_ = preempt_requested;
+}
+
+bool StagePrivate::preempted() const {
+	if (preempt_requested_)
+		return *preempt_requested_;
+
+	return false;
+}
+
 Stage::Stage(StagePrivate* impl) : pimpl_(impl) {
 	assert(impl);
 	auto& p = properties();