Reduce preempt stop time in stages

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
+
+	const std::atomic<bool>* preempt_requested_;
 };
 PIMPL_FUNCTIONS(Stage)
 std::ostream& operator<<(std::ostream& os, const StagePrivate& stage);

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(const std::atomic<bool>* 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();