From eb597c86d25850e1708815ec5185a63e96e3ce1f 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 | 24 +++++++++++++++++++
 core/src/stage.cpp                            | 11 +++++++++
 2 files changed, 35 insertions(+)

diff --git a/core/include/moveit/task_constructor/stage_p.h b/core/include/moveit/task_constructor/stage_p.h
index 5d359b5c6..ad2012dce 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
 {
@@ -69,6 +80,9 @@ class StagePrivate
 	StagePrivate(Stage* me, const std::string& name);
 	virtual ~StagePrivate() = default;
 
+	/// optional callback checking if this stage is preempted
+	using PreemptedCallback = std::function<bool()>;
+
 	/// actually configured interface of this stage (only valid after init())
 	InterfaceFlags interfaceFlags() const;
 
@@ -146,6 +160,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 +177,10 @@ class StagePrivate
 	/** compute cost for solution through configured CostTerm */
 	void computeCost(const InterfaceState& from, const InterfaceState& to, SolutionBase& solution);
 
+	void setPreemptedCallback(const PreemptedCallback& preempted_callback);
+	/// is the stage preempted ? defaults to false if callback is not set
+	bool preempted() const;
+
 protected:
 	StagePrivate& operator=(StagePrivate&& other);
 
@@ -197,6 +219,8 @@ class StagePrivate
 	InterfaceWeakPtr next_starts_;  // interface to be used for sendForward()
 
 	Introspection* introspection_;  // task's introspection instance
+
+	PreemptedCallback preempted_callback_;
 };
 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..7615c1c67 100644
--- a/core/src/stage.cpp
+++ b/core/src/stage.cpp
@@ -305,6 +305,17 @@ void StagePrivate::computeCost(const InterfaceState& from, const InterfaceState&
 	}
 }
 
+void StagePrivate::setPreemptedCallback(const PreemptedCallback& preempted_callback) {
+	preempted_callback_ = preempted_callback;
+}
+
+bool StagePrivate::preempted() const {
+	if (preempted_callback_)
+		return preempted_callback_();
+
+	return false;
+}
+
 Stage::Stage(StagePrivate* impl) : pimpl_(impl) {
 	assert(impl);
 	auto& p = properties();