2 files changed, 51 insertions, 8 deletions

diff --git a/src/sim/drain.cc b/src/sim/drain.cc
index bb8abd735..e920ee7ae 100644
--- a/src/sim/drain.cc
+++ b/src/sim/drain.cc
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2012, 2015 ARM Limited
+ * Copyright (c) 2012, 2015, 2017 ARM Limited
  * All rights reserved
  *
  * The license below extends only to copyright in the software and shall
@@ -101,14 +101,33 @@ DrainManager::resume()
             "Resuming a system that isn't fully drained, this is untested and "
             "likely to break\n");
 
+    panic_if(_state == DrainState::Resuming,
+             "Resuming a system that is already trying to resume. This should "
+             "never happen.\n");
+
     panic_if(_count != 0,
              "Resume called in the middle of a drain cycle. %u objects "
              "left to drain.\n", _count);
 
-    DPRINTF(Drain, "Resuming %u objects.\n", drainableCount());
+    // At this point in time the DrainManager and all objects will be
+    // in the the Drained state. New objects (i.e., objects created
+    // while resuming) will inherit the Resuming state from the
+    // DrainManager, which means we have to resume objects until all
+    // objects are in the Running state.
+    _state = DrainState::Resuming;
+
+    do {
+        DPRINTF(Drain, "Resuming %u objects.\n", drainableCount());
+        for (auto *obj : _allDrainable) {
+            if (obj->drainState() != DrainState::Running) {
+                assert(obj->drainState() == DrainState::Drained ||
+                       obj->drainState() == DrainState::Resuming);
+                obj->dmDrainResume();
+            }
+        }
+    } while (!allInState(DrainState::Running));
+
     _state = DrainState::Running;
-    for (auto *obj : _allDrainable)
-        obj->dmDrainResume();
 }
 
 void
@@ -154,6 +173,17 @@ DrainManager::unregisterDrainable(Drainable *obj)
     _allDrainable.erase(o);
 }
 
+bool
+DrainManager::allInState(DrainState state) const
+{
+    for (const auto *obj : _allDrainable) {
+        if (obj->drainState() != state)
+            return false;
+    }
+
+    return true;
+}
+
 size_t
 DrainManager::drainableCount() const
 {
@@ -189,7 +219,8 @@ Drainable::dmDrain()
 void
 Drainable::dmDrainResume()
 {
-    panic_if(_drainState != DrainState::Drained,
+    panic_if(_drainState != DrainState::Drained &&
+             _drainState != DrainState::Resuming,
              "Trying to resume an object that hasn't been drained\n");
 
     _drainState = DrainState::Running;
diff --git a/src/sim/drain.hh b/src/sim/drain.hh
index 663946616..7ff1d6e2a 100644
--- a/src/sim/drain.hh
+++ b/src/sim/drain.hh
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2012, 2015 ARM Limited
+ * Copyright (c) 2012, 2015, 2017 ARM Limited
  * All rights reserved
  *
  * The license below extends only to copyright in the software and shall
@@ -58,7 +58,11 @@ class Drainable;
  * all objects have entered the Drained state.
  *
  * Before resuming simulation, the simulator calls resume() to
- * transfer the object to the Running state.
+ * transfer the object to the Running state. This in turn results in a
+ * call to drainResume() for all Drainable objects in the
+ * simulator. New Drainable objects may be created while resuming. In
+ * such cases, the new objects will be created in the Resuming state
+ * and later resumed.
  *
  * \note Even though the state of an object (visible to the rest of
  * the world through Drainable::getState()) could be used to determine
@@ -68,7 +72,8 @@ class Drainable;
 enum class DrainState {
     Running,  /** Running normally */
     Draining, /** Draining buffers pending serialization/handover */
-    Drained   /** Buffers drained, ready for serialization/handover */
+    Drained,  /** Buffers drained, ready for serialization/handover */
+    Resuming, /** Transient state while the simulator is resuming */
 };
 #endif
 
@@ -153,6 +158,12 @@ class DrainManager
 
   private:
     /**
+     * Helper function to check if all Drainable objects are in a
+     * specific state.
+     */
+    bool allInState(DrainState state) const;
+
+    /**
      * Thread-safe helper function to get the number of Drainable
      * objects in a system.
      */
@@ -261,6 +272,7 @@ class Drainable
         switch (_drainState) {
           case DrainState::Running:
           case DrainState::Drained:
+          case DrainState::Resuming:
             return;
           case DrainState::Draining:
             _drainState = DrainState::Drained;