Proper coroutine handle destruction for dangling handles

zompi2 · Aug 3, 2024 · de45049 · de45049
1 parent 53d8e26
commit de45049
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 36 deletions.
diff --git a/Source/EnhancedCodeFlow/EnhancedCodeFlow.Build.cs b/Source/EnhancedCodeFlow/EnhancedCodeFlow.Build.cs
@@ -56,8 +56,5 @@ public EnhancedCodeFlow(ReadOnlyTargetRules Target) : base(Target)
         {
             PublicDefinitions.Add("ECF_INSIGHT_PROFILING=0");
         }
-
-        // Enable when testing coroutines explicit destroy
-        PublicDefinitions.Add("ECF_USE_EXPLICIT_CORO_DESTROY=0");
     }
 }
diff --git a/Source/EnhancedCodeFlow/Public/Coroutines/ECFCoroutine.h b/Source/EnhancedCodeFlow/Public/Coroutines/ECFCoroutine.h
@@ -22,18 +22,10 @@ struct FECFCoroutinePromise
 {
 	FECFCoroutine get_return_object() { return { FECFCoroutine::from_promise(*this) }; }
 	std::suspend_never initial_suspend() noexcept { return {}; }
-#if ECF_USE_EXPLICIT_CORO_DESTROY
-	std::suspend_always final_suspend() noexcept { return {}; }
-#else
 	std::suspend_never final_suspend() noexcept { return {}; }
-#endif
-	void return_void() {}
+	void return_void() { bHasFinished = true; }
 	void unhandled_exception() {}
-
-#if ECF_USE_EXPLICIT_CORO_DESTROY
-	int32 HandleCounter = 0;
-	bool bDestroyed = false;
-#endif
+	bool bHasFinished = false;
 };
 
 #else
@@ -44,9 +36,18 @@ struct FECFCoroutinePromise
 
 using FECFCoroutine = void;
 
+struct FECFCoroutinePromise
+{
+	bool bHasFinished = false;
+};
+
 struct FECFCoroutineHandle 
 {
-	void resume() {};
+	void resume() {}
+	void destroy() {}
+
+	FECFCoroutinePromise CoroPromise;
+	FECFCoroutinePromise& promise() { return CoroPromise; }
 };
 
 #define co_await static_assert(false, "Trying to use co_await without coroutine support!")

diff --git a/Source/EnhancedCodeFlow/Public/Coroutines/ECFCoroutineActionBase.h b/Source/EnhancedCodeFlow/Public/Coroutines/ECFCoroutineActionBase.h
@@ -21,27 +21,20 @@ class ENHANCEDCODEFLOW_API UECFCoroutineActionBase : public UECFActionBase
 	// Coroutine handle used to control the coroutine inside the Action.
 	FECFCoroutineHandle CoroutineHandle;
 
-#ifdef __cpp_impl_coroutine
-#if ECF_USE_EXPLICIT_CORO_DESTROY
-	bool bHasValidCoroHandle = false;
+	// Flag indicating if the coroutine handle has been set.
+	bool bHasCoroutineHandle = false;
+
 	void BeginDestroy() override
 	{
-		if (bHasValidCoroHandle)
+		// Handling a case when the owner has beed destroyed before the coroutine has been fully finished.
+		// In such case the handle must be explicitly destroyed.
+		if (bHasCoroutineHandle && (HasValidOwner() == false) && (CoroutineHandle.promise().bHasFinished == false))
 		{
-			if (CoroutineHandle.promise().bDestroyed == false)
-			{
-				CoroutineHandle.promise().HandleCounter--;
-				if (CoroutineHandle.promise().HandleCounter <= 0)
-				{
-					CoroutineHandle.promise().bDestroyed = true;
-					CoroutineHandle.destroy();
-				}
-			}
+			CoroutineHandle.promise().bHasFinished = true;
+			CoroutineHandle.destroy();
 		}
 		Super::BeginDestroy();
 	}
-#endif
-#endif
 
 private:
 
@@ -50,13 +43,7 @@ class ENHANCEDCODEFLOW_API UECFCoroutineActionBase : public UECFActionBase
 	{
 		UECFActionBase::SetAction(InOwner, InHandleId, {}, InSettings);
 		CoroutineHandle = InCoroutineHandle;
-
-#ifdef __cpp_impl_coroutine
-#if ECF_USE_EXPLICIT_CORO_DESTROY
-		CoroutineHandle.promise().HandleCounter++;
-		bHasValidCoroHandle = true;
-#endif
-#endif
+		bHasCoroutineHandle = true;
 	}
 };