From b737f4977938db33fe2e18b059db9471e1598d29 Mon Sep 17 00:00:00 2001
From: Chris Elion <chris.elion@unity3d.com>
Date: Mon, 12 Oct 2020 10:31:06 -0700
Subject: [PATCH] recursion checker proof-of-concept

---
 com.unity.ml-agents/Runtime/Academy.cs | 44 +++++++++++++-------------
 1 file changed, 22 insertions(+), 22 deletions(-)

diff --git a/com.unity.ml-agents/Runtime/Academy.cs b/com.unity.ml-agents/Runtime/Academy.cs
index 8ed370e350..991bdaee53 100644
--- a/com.unity.ml-agents/Runtime/Academy.cs
+++ b/com.unity.ml-agents/Runtime/Academy.cs
@@ -25,6 +25,26 @@
 
 namespace Unity.MLAgents
 {
+    internal class RecursionChecker : IDisposable
+    {
+        private bool m_IsRunning;
+
+        public IDisposable Start()
+        {
+            if (m_IsRunning)
+            {
+                throw new UnityAgentsException("Don't do this.");
+            }
+            m_IsRunning = true;
+            return this;
+        }
+
+        public void Dispose()
+        {
+            m_IsRunning = false;
+        }
+    }
+
     /// <summary>
     /// Helper class to step the Academy during FixedUpdate phase.
     /// </summary>
@@ -157,7 +177,7 @@ public bool IsCommunicatorOn
 
         // Whether the Academy is in the middle of a step. This is used to detect and Academy
         // step called by user code that is also called by the Academy.
-        bool m_IsStepping;
+        private RecursionChecker m_StepRecursionChecker = new RecursionChecker();
 
         // Random seed used for inference.
         int m_InferenceSeed;
@@ -535,22 +555,7 @@ void ForcedFullReset()
         /// </summary>
         public void EnvironmentStep()
         {
-            // Check whether we're already in the middle of a step.
-            // This shouldn't happen generally, but could happen if user code (e.g. CollectObservations)
-            // that is called by EnvironmentStep() also calls EnvironmentStep(). This would result
-            // in an infinite loop and/or stack overflow, so stop it before it happens.
-            if (m_IsStepping)
-            {
-                throw new UnityAgentsException(
-                    "Academy.EnvironmentStep() called recursively. " +
-                    "This might happen if you call EnvironmentStep() from custom code such as " +
-                    "CollectObservations() or OnActionReceived()."
-                );
-            }
-
-            m_IsStepping = true;
-
-            try
+            using (m_StepRecursionChecker.Start())
             {
                 if (!m_HadFirstReset)
                 {
@@ -584,11 +589,6 @@ public void EnvironmentStep()
                     AgentAct?.Invoke();
                 }
             }
-            finally
-            {
-                // Reset m_IsStepping when we're done (or if an exception occurred).
-                m_IsStepping = false;
-            }
         }
 
         /// <summary>