sim: Refactor and simplify the drain API

The drain() call currently passes around a DrainManager pointer, which
is now completely pointless since there is only ever one global
DrainManager in the system. It also contains vestiges from the time
when SimObjects had to keep track of their child objects that needed
draining.

This changeset moves all of the DrainState handling to the Drainable
base class and changes the drain() and drainResume() calls to reflect
this. Particularly, the drain() call has been updated to take no
parameters (the DrainManager argument isn't needed) and return a
DrainState instead of an unsigned integer (there is no point returning
anything other than 0 or 1 any more). Drainable objects should return
either DrainState::Draining (equivalent to returning 1 in the old
system) if they need more time to drain or DrainState::Drained
(equivalent to returning 0 in the old system) if they are already in a
consistent state. Returning DrainState::Running is considered an
error.

Drain done signalling is now done through the signalDrainDone() method
in the Drainable class instead of using the DrainManager directly. The
new call checks if the state of the object is DrainState::Draining
before notifying the drain manager. This means that it is safe to call
signalDrainDone() without first checking if the simulator has
requested draining. The intention here is to reduce the code needed to
implement draining in simple objects.

8 files changed, 88 insertions, 65 deletions

diff --git a/src/sim/drain.cc b/src/sim/drain.cc
index 384374099..8549da261 100644
--- a/src/sim/drain.cc
+++ b/src/sim/drain.cc
@@ -68,7 +68,7 @@ DrainManager::tryDrain()
     DPRINTF(Drain, "Trying to drain %u objects.\n", drainableCount());
     _state = DrainState::Draining;
     for (auto *obj : _allDrainable)
-        _count += obj->drain(&_instance);
+        _count += obj->dmDrain() == DrainState::Drained ? 0 : 1;
 
     if (_count == 0) {
         DPRINTF(Drain, "Drain done.\n");
@@ -98,7 +98,7 @@ DrainManager::resume()
     DPRINTF(Drain, "Resuming %u objects.\n", drainableCount());
     _state = DrainState::Running;
     for (auto *obj : _allDrainable)
-        obj->drainResume();
+        obj->dmDrainResume();
 }
 
 void
@@ -160,8 +160,23 @@ Drainable::~Drainable()
     _drainManager.unregisterDrainable(this);
 }
 
+DrainState
+Drainable::dmDrain()
+{
+    _drainState = DrainState::Draining;
+    _drainState = drain();
+    assert(_drainState == DrainState::Draining ||
+           _drainState == DrainState::Drained);
+
+    return _drainState;
+}
+
 void
-Drainable::drainResume()
+Drainable::dmDrainResume()
 {
+    panic_if(_drainState != DrainState::Drained,
+             "Trying to resume an object that hasn't been drained\n");
+
     _drainState = DrainState::Running;
+    drainResume();
 }
diff --git a/src/sim/drain.hh b/src/sim/drain.hh
index a045bf169..f60a9978f 100644
--- a/src/sim/drain.hh
+++ b/src/sim/drain.hh
@@ -192,23 +192,26 @@ class DrainManager
  * follows (see simulate.py for details):
  *
  * <ol>
- * <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
- *     the counter of the DrainManager. A pointer to the drain
- *     manager is passed as an argument to the drain() method.
+ * <li>DrainManager::tryDrain() calls Drainable::drain() for every
+ *     object in the system. Draining has completed if all of them
+ *     return true. Otherwise, the drain manager keeps track of the
+ *     objects that requested draining and waits for them to signal
+ *     that they are done draining using the signalDrainDone() method.
  *
  * <li>Continue simulation. When an object has finished draining its
  *     internal state, it calls DrainManager::signalDrainDone() on the
- *     manager. When the counter in the manager reaches zero, the
- *     simulation stops.
+ *     manager. The drain manager keeps track of the objects that
+ *     haven't drained yet, simulation stops when the set of
+ *     non-drained objects becomes empty.
  *
- * <li>Check if any object still needs draining, if so repeat the
- *     process above.
+ * <li>Check if any object still needs draining
+ *     (DrainManager::tryDrain()), if so repeat the process above.
  *
  * <li>Serialize objects, switch CPU model, or change timing model.
  *
- * <li>Call Drainable::drainResume() and continue the simulation.
+ * <li>Call DrainManager::resume(), which intern calls
+ *     Drainable::drainResume() for all objects, and continue the
+ *     simulation.
  * </ol>
  *
  */
@@ -216,7 +219,7 @@ class Drainable
 {
     friend class DrainManager;
 
-  public:
+  protected:
     Drainable();
     virtual ~Drainable();
 
@@ -224,44 +227,67 @@ class Drainable
      * Determine if an object needs draining and register a
      * DrainManager.
      *
-     * When draining the state of an object, the simulator calls drain
-     * with a pointer to a drain manager. If the object does not need
-     * further simulation to drain internal buffers, it switched to
-     * the Drained state and returns 0, otherwise it switches to the
-     * Draining state and returns the number of times that it will
-     * call Event::process() on the drain event. Most objects are
-     * expected to return either 0 or 1.
+     * If the object does not need further simulation to drain
+     * internal buffers, it returns true and automatically switches to
+     * the Drained state, otherwise it switches to the Draining state.
      *
      * @note An object that has entered the Drained state can be
      * disturbed by other objects in the system and consequently be
-     * forced to enter the Draining state again. The simulator
-     * therefore repeats the draining process until all objects return
-     * 0 on the first call to drain().
+     * being drained. These perturbations are not visible in the
+     * drain state. The simulator therefore repeats the draining
+     * process until all objects return DrainState::Drained on the
+     * first call to drain().
      *
-     * @param drainManager DrainManager to use to inform the simulator
-     * when draining has completed.
-     *
-     * @return 0 if the object is ready for serialization now, >0 if
-     * it needs further simulation.
+     * @return DrainState::Drained if the object is ready for
+     * serialization now, DrainState::Draining if it needs further
+     * simulation.
      */
-    virtual unsigned int drain(DrainManager *drainManager) = 0;
+    virtual DrainState drain() = 0;
 
     /**
      * Resume execution after a successful drain.
-     *
-     * @note This method is normally only called from the simulation
-     * scripts.
      */
-    virtual void drainResume();
+    virtual void drainResume() {};
 
-    DrainState getDrainState() const { return _drainState; }
+    /**
+     * Signal that an object is drained
+     *
+     * This method is designed to be called whenever an object enters
+     * into a state where it is ready to be drained. The method is
+     * safe to call multiple times and there is no need to check that
+     * draining has been requested before calling this method.
+     */
+    void signalDrainDone() const {
+        switch (_drainState) {
+          case DrainState::Running:
+          case DrainState::Drained:
+            return;
+          case DrainState::Draining:
+            _drainState = DrainState::Drained;
+            _drainManager.signalDrainDone();
+            return;
+        }
+    }
 
-  protected:
-    void setDrainState(DrainState new_state) { _drainState = new_state; }
+  public:
+    /** Return the current drain state of an object. */
+    DrainState drainState() const { return _drainState; }
 
   private:
+    /** DrainManager interface to request a drain operation */
+    DrainState dmDrain();
+    /** DrainManager interface to request a resume operation */
+    void dmDrainResume();
+
+    /** Convenience reference to the drain manager */
     DrainManager &_drainManager;
-    DrainState _drainState;
+
+    /**
+     * Current drain state of the object. Needs to be mutable since
+     * objects need to be able to signal that they have transitioned
+     * into a Drained state even if the calling method is const.
+     */
+    mutable DrainState _drainState;
 };
 
 #endif
diff --git a/src/sim/process.cc b/src/sim/process.cc
index a820b0632..198f05adf 100644
--- a/src/sim/process.cc
+++ b/src/sim/process.cc
@@ -259,11 +259,11 @@ Process::initState()
     pTable->initState(tc);
 }
 
-unsigned int
-Process::drain(DrainManager *dm)
+DrainState
+Process::drain()
 {
     find_file_offsets();
-    return 0;
+    return DrainState::Drained;
 }
 
 // map simulator fd sim_fd to target fd tgt_fd
diff --git a/src/sim/process.hh b/src/sim/process.hh
index c1499ccf7..cdb3826ff 100644
--- a/src/sim/process.hh
+++ b/src/sim/process.hh
@@ -120,7 +120,7 @@ class Process : public SimObject
 
     virtual void initState();
 
-    unsigned int drain(DrainManager *dm) M5_ATTR_OVERRIDE;
+    DrainState drain() M5_ATTR_OVERRIDE;
 
   public:
 
diff --git a/src/sim/sim_object.cc b/src/sim/sim_object.cc
index 9ea51eb93..e87b7240f 100644
--- a/src/sim/sim_object.cc
+++ b/src/sim/sim_object.cc
@@ -180,14 +180,6 @@ debugObjectBreak(const char *objs)
 }
 #endif
 
-unsigned int
-SimObject::drain(DrainManager *drain_manager)
-{
-    setDrainState(DrainState::Drained);
-    return 0;
-}
-
-
 SimObject *
 SimObject::find(const char *name)
 {
diff --git a/src/sim/sim_object.hh b/src/sim/sim_object.hh
index 17714740b..c56dfd8bf 100644
--- a/src/sim/sim_object.hh
+++ b/src/sim/sim_object.hh
@@ -181,11 +181,10 @@ class SimObject : public EventManager, public Serializable, public Drainable
     virtual void startup();
 
     /**
-     * Provide a default implementation of the drain interface that
-     * simply returns 0 (draining completed) and sets the drain state
-     * to Drained.
+     * Provide a default implementation of the drain interface for
+     * objects that don't need draining.
      */
-    unsigned int drain(DrainManager *drainManger);
+    DrainState drain() M5_ATTR_OVERRIDE { return DrainState::Drained; }
 
     /**
      * Write back dirty buffers to memory using functional writes.
diff --git a/src/sim/system.cc b/src/sim/system.cc
index 3d4737617..c5e2e0b96 100644
--- a/src/sim/system.cc
+++ b/src/sim/system.cc
@@ -191,7 +191,7 @@ System::getMasterPort(const std::string &if_name, PortID idx)
 void
 System::setMemoryMode(Enums::MemoryMode mode)
 {
-    assert(getDrainState() == DrainState::Drained);
+    assert(drainState() == DrainState::Drained);
     memoryMode = mode;
 }
 
@@ -355,17 +355,9 @@ System::isMemAddr(Addr addr) const
     return physmem.isMemAddr(addr);
 }
 
-unsigned int
-System::drain(DrainManager *dm)
-{
-    setDrainState(DrainState::Drained);
-    return 0;
-}
-
 void
 System::drainResume()
 {
-    Drainable::drainResume();
     totalNumInsts = 0;
 }
 
diff --git a/src/sim/system.hh b/src/sim/system.hh
index b8114d0ca..97d271d3a 100644
--- a/src/sim/system.hh
+++ b/src/sim/system.hh
@@ -520,8 +520,7 @@ class System : public MemObject
     void serialize(CheckpointOut &cp) const M5_ATTR_OVERRIDE;
     void unserialize(CheckpointIn &cp) M5_ATTR_OVERRIDE;
 
-    unsigned int drain(DrainManager *dm);
-    void drainResume();
+    void drainResume() M5_ATTR_OVERRIDE;
 
   public:
     Counter totalNumInsts;