sim: Decouple draining from the SimObject hierarchy

Draining is currently done by traversing the SimObject graph and
calling drain()/drainResume() on the SimObjects. This is not ideal
when non-SimObjects (e.g., ports) need draining since this means that
SimObjects owning those objects need to be aware of this.

This changeset moves the responsibility for finding objects that need
draining from SimObjects and the Python-side of the simulator to the
DrainManager. The DrainManager now maintains a set of all objects that
need draining. To reduce the overhead in classes owning non-SimObjects
that need draining, objects inheriting from Drainable now
automatically register with the DrainManager. If such an object is
destroyed, it is automatically unregistered. This means that drain()
and drainResume() should never be called directly on a Drainable
object.

While implementing the new functionality, the DrainManager has now
been made thread safe. In practice, this means that it takes a lock
whenever it manipulates the set of Drainable objects since SimObjects
in different threads may create Drainable objects
dynamically. Similarly, the drain counter is now an atomic_uint, which
ensures that it is manipulated correctly when objects signal that they
are done draining.

A nice side effect of these changes is that it makes the drain state
changes stricter, which the simulation scripts can exploit to avoid
redundant drains.

1 files changed, 87 insertions, 39 deletions

diff --git a/src/sim/drain.hh b/src/sim/drain.hh
index 4ed6074dd..a045bf169 100644
--- a/src/sim/drain.hh
+++ b/src/sim/drain.hh
@@ -40,8 +40,9 @@
 #ifndef __SIM_DRAIN_HH__
 #define __SIM_DRAIN_HH__
 
-#include <cassert>
-#include <vector>
+#include <atomic>
+#include <mutex>
+#include <unordered_set>
 
 #include "base/flags.hh"
 
@@ -76,12 +77,12 @@ enum class DrainState {
 /**
  * This class coordinates draining of a System.
  *
- * When draining a System, we need to make sure that all SimObjects in
- * that system have drained their state before declaring the operation
- * to be successful. This class keeps track of how many objects are
- * still in the process of draining their state. Once it determines
- * that all objects have drained their state, it exits the simulation
- * loop.
+ * When draining the simulator, we need to make sure that all
+ * Drainable objects within the system have ended up in the drained
+ * state before declaring the operation to be successful. This class
+ * keeps track of how many objects are still in the process of
+ * draining. Once it determines that all objects have drained their
+ * state, it exits the simulation loop.
  *
  * @note A System might not be completely drained even though the
  * DrainManager has caused the simulation loop to exit. Draining needs
@@ -91,39 +92,92 @@ enum class DrainState {
  */
 class DrainManager
 {
-  public:
+  private:
     DrainManager();
-    virtual ~DrainManager();
+#ifndef SWIG
+    DrainManager(DrainManager &) = delete;
+#endif
+    ~DrainManager();
+
+  public:
+    /** Get the singleton DrainManager instance */
+    static DrainManager &instance() { return _instance; }
+
+    /**
+     * Try to drain the system.
+     *
+     * Try to drain the system and return true if all objects are in a
+     * the Drained state at which point the whole simulator is in a
+     * consistent state and ready for checkpointing or CPU
+     * handover. The simulation script must continue simulating until
+     * the simulation loop returns "Finished drain", at which point
+     * this method should be called again. This cycle should continue
+     * until this method returns true.
+     *
+     * @return true if all objects were drained successfully, false if
+     * more simulation is needed.
+     */
+    bool tryDrain();
 
     /**
-     * Get the number of objects registered with this DrainManager
-     * that are currently draining their state.
+     * Resume normal simulation in a Drained system.
+     */
+    void resume();
+
+    /**
+     * Run state fixups before a checkpoint restore operation
      *
-     * @return Number of objects currently draining.
+     * The drain state of an object isn't stored in a checkpoint since
+     * the whole system is always going to be in the Drained state
+     * when the checkpoint is created. When the checkpoint is restored
+     * at a later stage, recreated objects will be in the Running
+     * state since the state isn't stored in checkpoints. This method
+     * performs state fixups on all Drainable objects and the
+     * DrainManager itself.
      */
-    unsigned int getCount() const { return _count; }
+    void preCheckpointRestore();
+
+    /** Check if the system is drained */
+    bool isDrained() { return _state == DrainState::Drained; }
 
-    void setCount(int count) { _count = count; }
+    /** Get the simulators global drain state */
+    DrainState state() { return _state; }
 
     /**
      * Notify the DrainManager that a Drainable object has finished
      * draining.
      */
-    void signalDrainDone() {
-        assert(_count > 0);
-        if (--_count == 0)
-            drainCycleDone();
-    }
+    void signalDrainDone();
 
-  protected:
+  public:
+    void registerDrainable(Drainable *obj);
+    void unregisterDrainable(Drainable *obj);
+
+  private:
     /**
-     * Callback when all registered Drainable objects have completed a
-     * drain cycle.
+     * Thread-safe helper function to get the number of Drainable
+     * objects in a system.
      */
-    virtual void drainCycleDone();
+    size_t drainableCount() const;
 
-    /** Number of objects still draining. */
-    unsigned int _count;
+    /** Lock protecting the set of drainable objects */
+    mutable std::mutex globalLock;
+
+    /** Set of all drainable objects */
+    std::unordered_set<Drainable *> _allDrainable;
+
+    /**
+     * Number of objects still draining. This is flagged atomic since
+     * it can be manipulated by SimObjects living in different
+     * threads.
+     */
+    std::atomic_uint _count;
+
+    /** Global simulator drain state */
+    DrainState _state;
+
+    /** Singleton instance of the drain manager */
+    static DrainManager _instance;
 };
 
 /**
@@ -133,17 +187,11 @@ class DrainManager
  * An object's internal state needs to be drained when creating a
  * checkpoint, switching between CPU models, or switching between
  * timing models. Once the internal state has been drained from
- * <i>all</i> objects in the system, the objects are serialized to
+ * <i>all</i> objects in the simulator, the objects are serialized to
  * disc or the configuration change takes place. The process works as
  * follows (see simulate.py for details):
  *
  * <ol>
- * <li>An instance of a DrainManager is created to keep track of how
- *     many objects need to be drained. The object maintains an
- *     internal counter that is decreased every time its
- *     CountedDrainEvent::signalDrainDone() method is called. When the
- *     counter reaches zero, the simulation is stopped.
- *
  * <li>Call Drainable::drain() for every object in the
  *     system. Draining has completed if all of them return
  *     zero. Otherwise, the sum of the return values is loaded into
@@ -151,9 +199,9 @@ class DrainManager
  *     manager is passed as an argument to the drain() method.
  *
  * <li>Continue simulation. When an object has finished draining its
- *     internal state, it calls CountedDrainEvent::signalDrainDone()
- *     on the manager. When the counter in the manager reaches zero,
- *     the simulation stops.
+ *     internal state, it calls DrainManager::signalDrainDone() on the
+ *     manager. When the counter in the manager reaches zero, the
+ *     simulation stops.
  *
  * <li>Check if any object still needs draining, if so repeat the
  *     process above.
@@ -166,6 +214,8 @@ class DrainManager
  */
 class Drainable
 {
+    friend class DrainManager;
+
   public:
     Drainable();
     virtual ~Drainable();
@@ -210,10 +260,8 @@ class Drainable
     void setDrainState(DrainState new_state) { _drainState = new_state; }
 
   private:
+    DrainManager &_drainManager;
     DrainState _drainState;
 };
 
-DrainManager *createDrainManager();
-void cleanupDrainManager(DrainManager *drain_manager);
-
 #endif