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.

src/arch/arm/table_walker.cc

@@ -57,7 +57,7 @@
 using namespace ArmISA;
 
 TableWalker::TableWalker(const Params *p)
-    : MemObject(p), drainManager(NULL),
+    : MemObject(p),
       stage2Mmu(NULL), port(NULL), masterId(Request::invldMasterId),
       isStage2(p->is_stage2), tlb(NULL),
       currState(NULL), pending(false),
@@ -137,17 +137,17 @@ TableWalker::WalkerState::WalkerState() :
 void
 TableWalker::completeDrain()
 {
-    if (drainManager && stateQueues[L1].empty() && stateQueues[L2].empty() &&
+    if (drainState() == DrainState::Draining &&
+        stateQueues[L1].empty() && stateQueues[L2].empty() &&
         pendingQueue.empty()) {
-        setDrainState(DrainState::Drained);
+
         DPRINTF(Drain, "TableWalker done draining, processing drain event\n");
-        drainManager->signalDrainDone();
+        signalDrainDone();
-        drainManager = NULL;
     }
 }
 
-unsigned int
+DrainState
-TableWalker::drain(DrainManager *dm)
+TableWalker::drain()
 {
     bool state_queues_not_empty = false;
 
@@ -159,25 +159,17 @@ TableWalker::drain(DrainManager *dm)
     }
 
     if (state_queues_not_empty || pendingQueue.size()) {
-        drainManager = dm;
-        setDrainState(DrainState::Draining);
         DPRINTF(Drain, "TableWalker not drained\n");
-
+        return DrainState::Draining;
-        // return port drain count plus the table walker itself needs to drain
-        return 1;
     } else {
-        setDrainState(DrainState::Drained);
         DPRINTF(Drain, "TableWalker free, no need to drain\n");
-
+        return DrainState::Drained;
-        // table walker is drained, but its ports may still need to be drained
-        return 0;
     }
 }
 
 void
 TableWalker::drainResume()
 {
-    Drainable::drainResume();
     if (params()->sys->isTimingMode() && currState) {
         delete currState;
         currState = NULL;

src/arch/arm/table_walker.hh

@@ -819,9 +819,6 @@ class TableWalker : public MemObject
      * currently busy. */
     std::list<WalkerState *> pendingQueue;
 
-    /** If we're draining keep the drain event around until we're drained */
-    DrainManager *drainManager;
-
     /** The MMU to forward second stage look upts to */
     Stage2MMU *stage2Mmu;
 
@@ -894,8 +891,8 @@ class TableWalker : public MemObject
     bool haveLargeAsid64() const { return _haveLargeAsid64; }
     /** Checks if all state is cleared and if so, completes drain */
     void completeDrain();
-    unsigned int drain(DrainManager *dm);
+    DrainState drain() M5_ATTR_OVERRIDE;
-    virtual void drainResume();
+    virtual void drainResume() M5_ATTR_OVERRIDE;
 
     virtual BaseMasterPort& getMasterPort(const std::string &if_name,
                                           PortID idx = InvalidPortID);

src/arch/arm/tlb.hh

@@ -284,7 +284,7 @@ class TLB : public BaseTLB
             bool callFromS2);
     Fault finalizePhysical(RequestPtr req, ThreadContext *tc, Mode mode) const;
 
-    void drainResume();
+    void drainResume() M5_ATTR_OVERRIDE;
 
     // Checkpointing
     void serialize(CheckpointOut &cp) const M5_ATTR_OVERRIDE;

src/cpu/kvm/base.cc

@@ -78,7 +78,6 @@ BaseKvmCPU::BaseKvmCPU(BaseKvmCPUParams *params)
       activeInstPeriod(0),
       perfControlledByTimer(params->usePerfOverflow),
       hostFactor(params->hostFactor),
-      drainManager(NULL),
       ctrInsts(0)
 {
     if (pageSize == -1)
@@ -282,11 +281,11 @@ BaseKvmCPU::unserializeThread(CheckpointIn &cp, ThreadID tid)
     threadContextDirty = true;
 }
 
-unsigned int
+DrainState
-BaseKvmCPU::drain(DrainManager *dm)
+BaseKvmCPU::drain()
 {
     if (switchedOut())
-        return 0;
+        return DrainState::Drained;
 
     DPRINTF(Drain, "BaseKvmCPU::drain\n");
     switch (_status) {
@@ -296,10 +295,8 @@ BaseKvmCPU::drain(DrainManager *dm)
         // of a different opinion. This may happen when the CPU been
         // notified of an event that hasn't been accepted by the vCPU
         // yet.
-        if (!archIsDrained()) {
+        if (!archIsDrained())
-            drainManager = dm;
+            return DrainState::Draining;
-            return 1;
-        }
 
         // The state of the CPU is consistent, so we don't need to do
         // anything special to drain it. We simply de-schedule the
@@ -318,7 +315,7 @@ BaseKvmCPU::drain(DrainManager *dm)
         // switch CPUs or checkpoint the CPU.
         syncThreadContext();
 
-        return 0;
+        return DrainState::Drained;
 
       case RunningServiceCompletion:
         // The CPU has just requested a service that was handled in
@@ -327,22 +324,18 @@ BaseKvmCPU::drain(DrainManager *dm)
         // update the register state ourselves instead of letting KVM
         // handle it, but that would be tricky. Instead, we enter KVM
         // and let it do its stuff.
-        drainManager = dm;
-
         DPRINTF(Drain, "KVM CPU is waiting for service completion, "
                 "requesting drain.\n");
-        return 1;
+        return DrainState::Draining;
 
       case RunningService:
         // We need to drain since the CPU is waiting for service (e.g., MMIOs)
-        drainManager = dm;
-
         DPRINTF(Drain, "KVM CPU is waiting for service, requesting drain.\n");
-        return 1;
+        return DrainState::Draining;
 
       default:
         panic("KVM: Unhandled CPU state in drain()\n");
-        return 0;
+        return DrainState::Drained;
     }
 }
 
@@ -551,7 +544,7 @@ BaseKvmCPU::tick()
               setupInstStop();
 
           DPRINTF(KvmRun, "Entering KVM...\n");
-          if (drainManager) {
+          if (drainState() == DrainState::Draining) {
               // Force an immediate exit from KVM after completing
               // pending operations. The architecture-specific code
               // takes care to run until it is in a state where it can
@@ -1198,7 +1191,7 @@ BaseKvmCPU::setupCounters()
 bool
 BaseKvmCPU::tryDrain()
 {
-    if (!drainManager)
+    if (drainState() != DrainState::Draining)
         return false;
 
     if (!archIsDrained()) {
@@ -1209,8 +1202,7 @@ BaseKvmCPU::tryDrain()
     if (_status == Idle || _status == Running) {
         DPRINTF(Drain,
                 "tryDrain: CPU transitioned into the Idle state, drain done\n");
-        drainManager->signalDrainDone();
+        signalDrainDone();
-        drainManager = NULL;
         return true;
     } else {
         DPRINTF(Drain, "tryDrain: CPU not ready.\n");

src/cpu/kvm/base.hh

@@ -89,8 +89,8 @@ class BaseKvmCPU : public BaseCPU
     void unserializeThread(CheckpointIn &cp,
                            ThreadID tid) M5_ATTR_OVERRIDE;
 
-    unsigned int drain(DrainManager *dm);
+    DrainState drain() M5_ATTR_OVERRIDE;
-    void drainResume();
+    void drainResume() M5_ATTR_OVERRIDE;
 
     void switchOut();
     void takeOverFrom(BaseCPU *cpu);
@@ -749,13 +749,6 @@ class BaseKvmCPU : public BaseCPU
     /** Host factor as specified in the configuration */
     float hostFactor;
 
-    /**
-     * Drain manager to use when signaling drain completion
-     *
-     * This pointer is non-NULL when draining and NULL otherwise.
-     */
-    DrainManager *drainManager;
-
   public:
     /* @{ */
     Stats::Scalar numInsts;

src/cpu/minor/cpu.cc

@@ -47,8 +47,7 @@
 #include "debug/Quiesce.hh"
 
 MinorCPU::MinorCPU(MinorCPUParams *params) :
-    BaseCPU(params),
+    BaseCPU(params)
-    drainManager(NULL)
 {
     /* This is only written for one thread at the moment */
     Minor::MinorThread *thread;
@@ -194,39 +193,33 @@ MinorCPU::startup()
     activateContext(0);
 }
 
-unsigned int
+DrainState
-MinorCPU::drain(DrainManager *drain_manager)
+MinorCPU::drain()
 {
     DPRINTF(Drain, "MinorCPU drain\n");
 
-    drainManager = drain_manager;
-
     /* Need to suspend all threads and wait for Execute to idle.
      * Tell Fetch1 not to fetch */
-    unsigned int ret = pipeline->drain(drain_manager);
+    if (pipeline->drain()) {
-
-    if (ret == 0)
         DPRINTF(Drain, "MinorCPU drained\n");
-    else
+        return DrainState::Drained;
+    } else {
         DPRINTF(Drain, "MinorCPU not finished draining\n");
-
+        return DrainState::Draining;
-    return ret;
+    }
 }
 
 void
 MinorCPU::signalDrainDone()
 {
     DPRINTF(Drain, "MinorCPU drain done\n");
-    setDrainState(DrainState::Drained);
+    signalDrainDone();
-    drainManager->signalDrainDone();
-    drainManager = NULL;
 }
 
 void
 MinorCPU::drainResume()
 {
-    assert(getDrainState() == DrainState::Drained ||
+    assert(drainState() == DrainState::Drained);
-        getDrainState() == DrainState::Running);
 
     if (switchedOut()) {
         DPRINTF(Drain, "drainResume while switched out.  Ignoring\n");
@@ -242,8 +235,6 @@ MinorCPU::drainResume()
 
     wakeup();
     pipeline->drainResume();
-
-    setDrainState(DrainState::Running);
 }
 
 void

src/cpu/minor/cpu.hh

@@ -112,10 +112,6 @@ class MinorCPU : public BaseCPU
         virtual void recvTimingSnoopReq(PacketPtr pkt) { }
     };
 
-    /** The DrainManager passed into drain that needs be signalled when
-     *  draining is complete */
-    DrainManager *drainManager;
-
   protected:
      /** Return a reference to the data port. */
     MasterPort &getDataPort();
@@ -155,10 +151,10 @@ class MinorCPU : public BaseCPU
     void unserialize(CheckpointIn &cp);
 
     /** Drain interface */
-    unsigned int drain(DrainManager *drain_manager);
+    DrainState drain() M5_ATTR_OVERRIDE;
-    void drainResume();
+    void drainResume() M5_ATTR_OVERRIDE;
-    /** Signal from Pipeline that MinorCPU should signal the DrainManager
+    /** Signal from Pipeline that MinorCPU should signal that a drain
-     *  that a drain is complete and set its drainState */
+     *  is complete and set its drainState */
     void signalDrainDone();
     void memWriteback();
 

src/cpu/minor/pipeline.cc

@@ -192,8 +192,8 @@ Pipeline::wakeupFetch()
     execute.wakeupFetch();
 }
 
-unsigned int
+bool
-Pipeline::drain(DrainManager *manager)
+Pipeline::drain()
 {
     DPRINTF(MinorCPU, "Draining pipeline by halting inst fetches. "
         " Execution should drain naturally\n");
@@ -205,7 +205,7 @@ Pipeline::drain(DrainManager *manager)
     bool drained = isDrained();
     needToSignalDrained = !drained;
 
-    return (drained ? 0 : 1);
+    return drained;
 }
 
 void

src/cpu/minor/pipeline.hh

@@ -115,7 +115,7 @@ class Pipeline : public Ticked
     void wakeupFetch();
 
     /** Try to drain the CPU */
-    unsigned int drain(DrainManager *manager);
+    bool drain();
 
     void drainResume();
 

src/cpu/o3/cpu.cc

@@ -196,7 +196,6 @@ FullO3CPU<Impl>::FullO3CPU(DerivO3CPUParams *params)
 
       globalSeqNum(1),
       system(params->system),
-      drainManager(NULL),
       lastRunningCycle(curCycle())
 {
     if (!params->switched_out) {
@@ -539,7 +538,7 @@ FullO3CPU<Impl>::tick()
 {
     DPRINTF(O3CPU, "\n\nFullO3CPU: Ticking main, FullO3CPU.\n");
     assert(!switchedOut());
-    assert(getDrainState() != DrainState::Drained);
+    assert(drainState() != DrainState::Drained);
 
     ++numCycles;
     ppCycles->notify(1);
@@ -712,7 +711,7 @@ FullO3CPU<Impl>::activateContext(ThreadID tid)
     // We don't want to wake the CPU if it is drained. In that case,
     // we just want to flag the thread as active and schedule the tick
     // event from drainResume() instead.
-    if (getDrainState() == DrainState::Drained)
+    if (drainState() == DrainState::Drained)
         return;
 
     // If we are time 0 or if the last activation time is in the past,
@@ -999,17 +998,14 @@ FullO3CPU<Impl>::unserializeThread(CheckpointIn &cp, ThreadID tid)
 }
 
 template <class Impl>
-unsigned int
+DrainState
-FullO3CPU<Impl>::drain(DrainManager *drain_manager)
+FullO3CPU<Impl>::drain()
 {
     // If the CPU isn't doing anything, then return immediately.
-    if (switchedOut()) {
+    if (switchedOut())
-        setDrainState(DrainState::Drained);
+        return DrainState::Drained;
-        return 0;
-    }
 
     DPRINTF(Drain, "Draining...\n");
-    setDrainState(DrainState::Draining);
 
     // We only need to signal a drain to the commit stage as this
     // initiates squashing controls the draining. Once the commit
@@ -1022,16 +1018,13 @@ FullO3CPU<Impl>::drain(DrainManager *drain_manager)
     // Wake the CPU and record activity so everything can drain out if
     // the CPU was not able to immediately drain.
     if (!isDrained())  {
-        drainManager = drain_manager;
-
         wakeCPU();
         activityRec.activity();
 
         DPRINTF(Drain, "CPU not drained\n");
 
-        return 1;
+        return DrainState::Draining;
     } else {
-        setDrainState(DrainState::Drained);
         DPRINTF(Drain, "CPU is already drained\n");
         if (tickEvent.scheduled())
             deschedule(tickEvent);
@@ -1049,7 +1042,7 @@ FullO3CPU<Impl>::drain(DrainManager *drain_manager)
         }
 
         drainSanityCheck();
-        return 0;
+        return DrainState::Drained;
     }
 }
 
@@ -1057,15 +1050,14 @@ template <class Impl>
 bool
 FullO3CPU<Impl>::tryDrain()
 {
-    if (!drainManager || !isDrained())
+    if (drainState() != DrainState::Draining || !isDrained())
         return false;
 
     if (tickEvent.scheduled())
         deschedule(tickEvent);
 
     DPRINTF(Drain, "CPU done draining, processing drain event\n");
-    drainManager->signalDrainDone();
+    signalDrainDone();
-    drainManager = NULL;
 
     return true;
 }
@@ -1132,7 +1124,6 @@ template <class Impl>
 void
 FullO3CPU<Impl>::drainResume()
 {
-    setDrainState(DrainState::Running);
     if (switchedOut())
         return;
 

src/cpu/o3/cpu.hh

@@ -232,7 +232,7 @@ class FullO3CPU : public BaseO3CPU
     }
 
     /**
-     * Check if the pipeline has drained and signal the DrainManager.
+     * Check if the pipeline has drained and signal drain done.
      *
      * This method checks if a drain has been requested and if the CPU
      * has drained successfully (i.e., there are no instructions in
@@ -336,7 +336,7 @@ class FullO3CPU : public BaseO3CPU
     void updateThreadPriority();
 
     /** Is the CPU draining? */
-    bool isDraining() const { return getDrainState() == DrainState::Draining; }
+    bool isDraining() const { return drainState() == DrainState::Draining; }
 
     void serializeThread(CheckpointOut &cp,
                          ThreadID tid) const M5_ATTR_OVERRIDE;
@@ -350,10 +350,10 @@ class FullO3CPU : public BaseO3CPU
 
     /** Starts draining the CPU's pipeline of all instructions in
      * order to stop all memory accesses. */
-    unsigned int drain(DrainManager *drain_manager);
+    DrainState drain() M5_ATTR_OVERRIDE;
 
     /** Resumes execution after a drain. */
-    void drainResume();
+    void drainResume() M5_ATTR_OVERRIDE;
 
     /**
      * Commit has reached a safe point to drain a thread.
@@ -665,9 +665,6 @@ class FullO3CPU : public BaseO3CPU
     /** Pointer to the system. */
     System *system;
 
-    /** DrainManager to notify when draining has completed. */
-    DrainManager *drainManager;
-
     /** Pointers to all of the threads in the CPU. */
     std::vector<Thread *> thread;
 

src/cpu/simple/atomic.cc

@@ -108,7 +108,6 @@ AtomicSimpleCPU::AtomicSimpleCPU(AtomicSimpleCPUParams *p)
     : BaseSimpleCPU(p), tickEvent(this), width(p->width), locked(false),
       simulate_data_stalls(p->simulate_data_stalls),
       simulate_inst_stalls(p->simulate_inst_stalls),
-      drain_manager(NULL),
       icachePort(name() + ".icache_port", this),
       dcachePort(name() + ".dcache_port", this),
       fastmem(p->fastmem), dcache_access(false), dcache_latency(0),
@@ -125,23 +124,21 @@ AtomicSimpleCPU::~AtomicSimpleCPU()
     }
 }
 
-unsigned int
+DrainState
-AtomicSimpleCPU::drain(DrainManager *dm)
+AtomicSimpleCPU::drain()
 {
-    assert(!drain_manager);
     if (switchedOut())
-        return 0;
+        return DrainState::Drained;
 
     if (!isDrained()) {
         DPRINTF(Drain, "Requesting drain: %s\n", pcState());
-        drain_manager = dm;
+        return DrainState::Draining;
-        return 1;
     } else {
         if (tickEvent.scheduled())
             deschedule(tickEvent);
 
         DPRINTF(Drain, "Not executing microcode, no need to drain.\n");
-        return 0;
+        return DrainState::Drained;
     }
 }
 
@@ -149,7 +146,6 @@ void
 AtomicSimpleCPU::drainResume()
 {
     assert(!tickEvent.scheduled());
-    assert(!drain_manager);
     if (switchedOut())
         return;
 
@@ -173,7 +169,7 @@ AtomicSimpleCPU::drainResume()
 bool
 AtomicSimpleCPU::tryCompleteDrain()
 {
-    if (!drain_manager)
+    if (drainState() != DrainState::Draining)
         return false;
 
     DPRINTF(Drain, "tryCompleteDrain: %s\n", pcState());
@@ -181,8 +177,7 @@ AtomicSimpleCPU::tryCompleteDrain()
         return false;
 
     DPRINTF(Drain, "CPU done draining, processing drain event\n");
-    drain_manager->signalDrainDone();
+    signalDrainDone();
-    drain_manager = NULL;
 
     return true;
 }

src/cpu/simple/atomic.hh

@@ -74,13 +74,6 @@ class AtomicSimpleCPU : public BaseSimpleCPU
     const bool simulate_data_stalls;
     const bool simulate_inst_stalls;
 
-    /**
-     * Drain manager to use when signaling drain completion
-     *
-     * This pointer is non-NULL when draining and NULL otherwise.
-     */
-    DrainManager *drain_manager;
-
     // main simulation loop (one cycle)
     void tick();
 
@@ -192,8 +185,8 @@ class AtomicSimpleCPU : public BaseSimpleCPU
 
   public:
 
-    unsigned int drain(DrainManager *drain_manager);
+    DrainState drain() M5_ATTR_OVERRIDE;
-    void drainResume();
+    void drainResume() M5_ATTR_OVERRIDE;
 
     void switchOut();
     void takeOverFrom(BaseCPU *oldCPU);

src/cpu/simple/timing.cc

@@ -91,7 +91,7 @@ TimingSimpleCPU::TimingCPUPort::TickEvent::schedule(PacketPtr _pkt, Tick t)
 TimingSimpleCPU::TimingSimpleCPU(TimingSimpleCPUParams *p)
     : BaseSimpleCPU(p), fetchTranslation(this), icachePort(this),
       dcachePort(this), ifetch_pkt(NULL), dcache_pkt(NULL), previousCycle(0),
-      fetchEvent(this), drainManager(NULL)
+      fetchEvent(this)
 {
     _status = Idle;
 }
@@ -102,19 +102,17 @@ TimingSimpleCPU::~TimingSimpleCPU()
 {
 }
 
-unsigned int
+DrainState
-TimingSimpleCPU::drain(DrainManager *drain_manager)
+TimingSimpleCPU::drain()
 {
-    assert(!drainManager);
     if (switchedOut())
-        return 0;
+        return DrainState::Drained;
 
     if (_status == Idle ||
         (_status == BaseSimpleCPU::Running && isDrained())) {
         DPRINTF(Drain, "No need to drain.\n");
-        return 0;
+        return DrainState::Drained;
     } else {
-        drainManager = drain_manager;
         DPRINTF(Drain, "Requesting drain: %s\n", pcState());
 
         // The fetch event can become descheduled if a drain didn't
@@ -123,7 +121,7 @@ TimingSimpleCPU::drain(DrainManager *drain_manager)
         if (_status == BaseSimpleCPU::Running && !fetchEvent.scheduled())
             schedule(fetchEvent, clockEdge());
 
-        return 1;
+        return DrainState::Draining;
     }
 }
 
@@ -131,7 +129,6 @@ void
 TimingSimpleCPU::drainResume()
 {
     assert(!fetchEvent.scheduled());
-    assert(!drainManager);
     if (switchedOut())
         return;
 
@@ -155,7 +152,7 @@ TimingSimpleCPU::drainResume()
 bool
 TimingSimpleCPU::tryCompleteDrain()
 {
-    if (!drainManager)
+    if (drainState() != DrainState::Draining)
         return false;
 
     DPRINTF(Drain, "tryCompleteDrain: %s\n", pcState());
@@ -163,8 +160,7 @@ TimingSimpleCPU::tryCompleteDrain()
         return false;
 
     DPRINTF(Drain, "CPU done draining, processing drain event\n");
-    drainManager->signalDrainDone();
+    signalDrainDone();
-    drainManager = NULL;
 
     return true;
 }

src/cpu/simple/timing.hh

@@ -270,8 +270,8 @@ class TimingSimpleCPU : public BaseSimpleCPU
 
   public:
 
-    unsigned int drain(DrainManager *drain_manager);
+    DrainState drain() M5_ATTR_OVERRIDE;
-    void drainResume();
+    void drainResume() M5_ATTR_OVERRIDE;
 
     void switchOut();
     void takeOverFrom(BaseCPU *oldCPU);
@@ -351,13 +351,6 @@ class TimingSimpleCPU : public BaseSimpleCPU
      * @returns true if the CPU is drained, false otherwise.
      */
     bool tryCompleteDrain();
-
-    /**
-     * Drain manager to use when signaling drain completion
-     *
-     * This pointer is non-NULL when draining and NULL otherwise.
-     */
-    DrainManager *drainManager;
 };
 
 #endif // __CPU_SIMPLE_TIMING_HH__

src/cpu/testers/traffic_gen/traffic_gen.cc

@@ -63,8 +63,7 @@ TrafficGen::TrafficGen(const TrafficGenParams* p)
       port(name() + ".port", *this),
       retryPkt(NULL),
       retryPktTick(0),
-      updateEvent(this),
+      updateEvent(this)
-      drainManager(NULL)
 {
 }
 
@@ -118,12 +117,12 @@ TrafficGen::initState()
     }
 }
 
-unsigned int
+DrainState
-TrafficGen::drain(DrainManager *dm)
+TrafficGen::drain()
 {
     if (!updateEvent.scheduled()) {
         // no event has been scheduled yet (e.g. switched from atomic mode)
-        return 0;
+        return DrainState::Drained;
     }
 
     if (retryPkt == NULL) {
@@ -131,10 +130,9 @@ TrafficGen::drain(DrainManager *dm)
         nextPacketTick = MaxTick;
         nextTransitionTick = MaxTick;
         deschedule(updateEvent);
-        return 0;
+        return DrainState::Drained;
     } else {
-        drainManager = dm;
+        return DrainState::Draining;
-        return 1;
     }
 }
 
@@ -488,7 +486,7 @@ TrafficGen::recvReqRetry()
         retryPktTick = 0;
         retryTicks += delay;
 
-        if (drainManager == NULL) {
+        if (drainState() != DrainState::Draining) {
             // packet is sent, so find out when the next one is due
             nextPacketTick = states[currState]->nextPacketTick(elasticReq,
                                                                delay);
@@ -498,9 +496,7 @@ TrafficGen::recvReqRetry()
             // shut things down
             nextPacketTick = MaxTick;
             nextTransitionTick = MaxTick;
-            drainManager->signalDrainDone();
+            signalDrainDone();
-            // Clear the drain event once we're done with it.
-            drainManager = NULL;
         }
     }
 }

src/cpu/testers/traffic_gen/traffic_gen.hh

@@ -176,9 +176,6 @@ class TrafficGen : public MemObject
     /** Event for scheduling updates */
     EventWrapper<TrafficGen, &TrafficGen::update> updateEvent;
 
-    /** Manager to signal when drained */
-    DrainManager* drainManager;
-
     /** Count the number of generated packets. */
     Stats::Scalar numPackets;
 
@@ -201,7 +198,7 @@ class TrafficGen : public MemObject
 
     void initState();
 
-    unsigned int drain(DrainManager *dm);
+    DrainState drain() M5_ATTR_OVERRIDE;
 
     void serialize(CheckpointOut &cp) const M5_ATTR_OVERRIDE;
     void unserialize(CheckpointIn &cp) M5_ATTR_OVERRIDE;

src/dev/arm/flash_device.cc

@@ -86,7 +86,6 @@ FlashDevice::FlashDevice(const FlashDeviceParams* p):
     pagesPerDisk(0),
     blocksPerDisk(0),
     planeMask(numPlanes - 1),
-    drainManager(NULL),
     planeEventQueue(numPlanes),
     planeEvent(this)
 {
@@ -587,26 +586,16 @@ FlashDevice::unserialize(CheckpointIn &cp)
  * Drain; needed to enable checkpoints
  */
 
-unsigned int
+DrainState
-FlashDevice::drain(DrainManager *dm)
+FlashDevice::drain()
 {
-    unsigned int count = 0;
-
     if (planeEvent.scheduled()) {
-        count = 1;
+        DPRINTF(Drain, "Flash device is draining...\n");
-        drainManager = dm;
+        return DrainState::Draining;
     } else {
         DPRINTF(Drain, "Flash device in drained state\n");
+        return DrainState::Drained;
     }
-
-    if (count) {
-        DPRINTF(Drain, "Flash device is draining...\n");
-        setDrainState(DrainState::Draining);
-    } else {
-        DPRINTF(Drain, "Flash device drained\n");
-        setDrainState(DrainState::Drained);
-    }
-    return count;
 }
 
 /**
@@ -616,15 +605,13 @@ FlashDevice::drain(DrainManager *dm)
 void
 FlashDevice::checkDrain()
 {
-    if (drainManager == NULL) {
+    if (drainState() == DrainState::Draining)
         return;
-    }
 
     if (planeEvent.when() > curTick()) {
         DPRINTF(Drain, "Flash device is still draining\n");
     } else {
         DPRINTF(Drain, "Flash device is done draining\n");
-        drainManager->signalDrainDone();
+        signalDrainDone();
-        drainManager = NULL;
     }
 }

src/dev/arm/flash_device.hh

@@ -62,7 +62,7 @@ class FlashDevice : public AbstractNVM
     ~FlashDevice();
 
     /** Checkpoint functions*/
-    unsigned int drain(DrainManager *dm);
+    DrainState drain() M5_ATTR_OVERRIDE;
     void checkDrain();
 
     void serialize(CheckpointOut &cp) const M5_ATTR_OVERRIDE;
@@ -175,13 +175,6 @@ class FlashDevice : public AbstractNVM
 
     uint32_t planeMask;
 
-    /**
-     * drain manager
-     * Needed to be able to implement checkpoint functionality
-     */
-
-    DrainManager *drainManager;
-
     /**
      * when the disk is first started we are unsure of the number of
      * used pages, this variable will help determining what we do know.

src/dev/arm/ufs_device.cc

@@ -733,7 +733,6 @@ UFSHostDevice::UFSHostDevice(const UFSHostDeviceParams* p) :
     transferTrack(0),
     taskCommandTrack(0),
     idlePhaseStart(0),
-    drainManager(NULL),
     SCSIResumeEvent(this),
     UTPEvent(this)
 {
@@ -2316,18 +2315,15 @@ UFSHostDevice::unserialize(CheckpointIn &cp)
  * Drain; needed to enable checkpoints
  */
 
-unsigned int
+DrainState
-UFSHostDevice::drain(DrainManager *dm)
+UFSHostDevice::drain()
 {
     if (UFSHCIMem.TRUTRLDBR) {
-        drainManager = dm;
         DPRINTF(UFSHostDevice, "UFSDevice is draining...\n");
-        setDrainState(DrainState::Draining);
+        return DrainState::Draining;
-        return 1;
     } else {
         DPRINTF(UFSHostDevice, "UFSDevice drained\n");
-        setDrainState(DrainState::Drained);
+        return DrainState::Drained;
-        return 0;
     }
 }
 
@@ -2338,16 +2334,14 @@ UFSHostDevice::drain(DrainManager *dm)
 void
 UFSHostDevice::checkDrain()
 {
-    if (drainManager == NULL) {
+    if (drainState() != DrainState::Draining)
         return;
-    }
 
     if (UFSHCIMem.TRUTRLDBR) {
         DPRINTF(UFSHostDevice, "UFSDevice is still draining; with %d active"
             " doorbells\n", activeDoorbells);
     } else {
         DPRINTF(UFSHostDevice, "UFSDevice is done draining\n");
-        drainManager->signalDrainDone();
+        signalDrainDone();
-        drainManager = NULL;
     }
 }

src/dev/arm/ufs_device.hh

@@ -173,7 +173,7 @@ class UFSHostDevice : public DmaDevice
 
     UFSHostDevice(const UFSHostDeviceParams* p);
 
-    unsigned int drain(DrainManager *dm);
+    DrainState drain() M5_ATTR_OVERRIDE;
     void checkDrain();
     void serialize(CheckpointOut &cp) const M5_ATTR_OVERRIDE;
     void unserialize(CheckpointIn &cp) M5_ATTR_OVERRIDE;
@@ -1052,13 +1052,6 @@ class UFSHostDevice : public DmaDevice
     Tick transactionStart[32];
     Tick idlePhaseStart;
 
-    /**
-     * drain manager
-     * Needed to be able to implement checkpoint functionality
-     */
-
-    DrainManager *drainManager;
-
     /**
      * logic units connected to the UFS Host device
      * Note again that the "device" as such is represented by one or multiple

src/dev/copy_engine.cc

@@ -82,7 +82,7 @@ CopyEngine::CopyEngineChannel::CopyEngineChannel(CopyEngine *_ce, int cid)
       ce(_ce), channelId(cid), busy(false), underReset(false),
     refreshNext(false), latBeforeBegin(ce->params()->latBeforeBegin),
     latAfterCompletion(ce->params()->latAfterCompletion),
-    completionDataReg(0), nextState(Idle), drainManager(NULL),
+    completionDataReg(0), nextState(Idle),
     fetchCompleteEvent(this), addrCompleteEvent(this),
     readCompleteEvent(this), writeCompleteEvent(this),
     statusCompleteEvent(this)
@@ -140,12 +140,12 @@ CopyEngine::CopyEngineChannel::recvCommand()
         cr.status.dma_transfer_status(0);
         nextState = DescriptorFetch;
         fetchAddress = cr.descChainAddr;
-        if (ce->getDrainState() == DrainState::Running)
+        if (ce->drainState() == DrainState::Running)
             fetchDescriptor(cr.descChainAddr);
     } else if (cr.command.append_dma()) {
         if (!busy) {
             nextState = AddressFetch;
-            if (ce->getDrainState() == DrainState::Running)
+            if (ce->drainState() == DrainState::Running)
                 fetchNextAddr(lastDescriptorAddr);
         } else
             refreshNext = true;
@@ -635,25 +635,23 @@ CopyEngine::CopyEngineChannel::fetchAddrComplete()
 bool
 CopyEngine::CopyEngineChannel::inDrain()
 {
-    if (ce->getDrainState() == DrainState::Draining) {
+    if (drainState() == DrainState::Draining) {
         DPRINTF(Drain, "CopyEngine done draining, processing drain event\n");
-        assert(drainManager);
+        signalDrainDone();
-        drainManager->signalDrainDone();
-        drainManager = NULL;
     }
 
-    return ce->getDrainState() != DrainState::Running;
+    return ce->drainState() != DrainState::Running;
 }
 
-unsigned int
+DrainState
-CopyEngine::CopyEngineChannel::drain(DrainManager *dm)
+CopyEngine::CopyEngineChannel::drain()
 {
-    if (nextState == Idle || ce->getDrainState() != DrainState::Running)
+    if (nextState == Idle || ce->drainState() != DrainState::Running) {
-        return 0;
+        return DrainState::Drained;
-
+    } else {
         DPRINTF(Drain, "CopyEngineChannel not drained\n");
-    this->drainManager = dm;
+        return DrainState::Draining;
-    return 1;
+    }
 }
 
 void

src/dev/copy_engine.hh

@@ -92,7 +92,6 @@ class CopyEngine : public PciDevice
 
         ChannelState nextState;
 
-        DrainManager *drainManager;
       public:
         CopyEngineChannel(CopyEngine *_ce, int cid);
         virtual ~CopyEngineChannel();
@@ -107,8 +106,8 @@ class CopyEngine : public PciDevice
         void channelRead(PacketPtr pkt, Addr daddr, int size);
         void channelWrite(PacketPtr pkt, Addr daddr, int size);
 
-        unsigned int drain(DrainManager *drainManger);
+        DrainState drain() M5_ATTR_OVERRIDE;
-        void drainResume();
+        void drainResume() M5_ATTR_OVERRIDE;
 
         void serialize(CheckpointOut &cp) const M5_ATTR_OVERRIDE;
         void unserialize(CheckpointIn &cp) M5_ATTR_OVERRIDE;

src/dev/dma_device.cc

@@ -51,8 +51,7 @@
 DmaPort::DmaPort(MemObject *dev, System *s)
     : MasterPort(dev->name() + ".dma", dev), device(dev), sendEvent(this),
       sys(s), masterId(s->getMasterId(dev->name())),
-      pendingCount(0), drainManager(NULL),
+      pendingCount(0), inRetry(false)
-      inRetry(false)
 { }
 
 void
@@ -95,10 +94,8 @@ DmaPort::handleResp(PacketPtr pkt, Tick delay)
     delete pkt;
 
     // we might be drained at this point, if so signal the drain event
-    if (pendingCount == 0 && drainManager) {
+    if (pendingCount == 0)
-        drainManager->signalDrainDone();
+        signalDrainDone();
-        drainManager = NULL;
-    }
 }
 
 bool
@@ -125,14 +122,15 @@ DmaDevice::init()
     PioDevice::init();
 }
 
-unsigned int
+DrainState
-DmaPort::drain(DrainManager *dm)
+DmaPort::drain()
 {
-    if (pendingCount == 0)
+    if (pendingCount == 0) {
-        return 0;
+        return DrainState::Drained;
-    drainManager = dm;
+    } else {
         DPRINTF(Drain, "DmaPort not drained\n");
-    return 1;
+        return DrainState::Draining;
+    }
 }
 
 void

src/dev/dma_device.hh

@@ -123,10 +123,6 @@ class DmaPort : public MasterPort, public Drainable
     /** Number of outstanding packets the dma port has. */
     uint32_t pendingCount;
 
-    /** If we need to drain, keep the drain event around until we're done
-     * here.*/
-    DrainManager *drainManager;
-
     /** If the port is currently waiting for a retry before it can
      * send whatever it is that it's sending. */
     bool inRetry;
@@ -147,7 +143,7 @@ class DmaPort : public MasterPort, public Drainable
 
     bool dmaPending() const { return pendingCount > 0; }
 
-    unsigned int drain(DrainManager *drainManger);
+    DrainState drain() M5_ATTR_OVERRIDE;
 };
 
 class DmaDevice : public PioDevice

src/dev/i8254xGBe.cc

@@ -58,7 +58,7 @@ using namespace iGbReg;
 using namespace Net;
 
 IGbE::IGbE(const Params *p)
-    : EtherDevice(p), etherInt(NULL), cpa(NULL), drainManager(NULL),
+    : EtherDevice(p), etherInt(NULL), cpa(NULL),
       rxFifo(p->rx_fifo_size), txFifo(p->tx_fifo_size), rxTick(false),
       txTick(false), txFifoTick(false), rxDmaPacket(false), pktOffset(0),
       fetchDelay(p->fetch_delay), wbDelay(p->wb_delay), 
@@ -586,7 +586,7 @@ IGbE::write(PacketPtr pkt)
       case REG_RDT:
         regs.rdt = val;
         DPRINTF(EthernetSM, "RXS: RDT Updated.\n");
-        if (getDrainState() == DrainState::Running) {
+        if (drainState() == DrainState::Running) {
             DPRINTF(EthernetSM, "RXS: RDT Fetching Descriptors!\n");
             rxDescCache.fetchDescriptors();
         } else {
@@ -626,7 +626,7 @@ IGbE::write(PacketPtr pkt)
       case REG_TDT:
         regs.tdt = val;
         DPRINTF(EthernetSM, "TXS: TX Tail pointer updated\n");
-        if (getDrainState() == DrainState::Running) {
+        if (drainState() == DrainState::Running) {
             DPRINTF(EthernetSM, "TXS: TDT Fetching Descriptors!\n");
             txDescCache.fetchDescriptors();
         } else {
@@ -905,7 +905,7 @@ void
 IGbE::DescCache<T>::writeback1()
 {
     // If we're draining delay issuing this DMA
-    if (igbe->getDrainState() != DrainState::Running) {
+    if (igbe->drainState() != DrainState::Running) {
         igbe->schedule(wbDelayEvent, curTick() + igbe->wbDelay);
         return;
     }
@@ -986,7 +986,7 @@ void
 IGbE::DescCache<T>::fetchDescriptors1()
 {
     // If we're draining delay issuing this DMA
-    if (igbe->getDrainState() != DrainState::Running) {
+    if (igbe->drainState() != DrainState::Running) {
         igbe->schedule(fetchDelayEvent, curTick() + igbe->fetchDelay);
         return;
     }
@@ -1492,7 +1492,7 @@ IGbE::RxDescCache::pktComplete()
 void
 IGbE::RxDescCache::enableSm()
 {
-    if (!igbe->drainManager) {
+    if (igbe->drainState() != DrainState::Draining) {
         igbe->rxTick = true;
         igbe->restartClock();
     }
@@ -2031,7 +2031,7 @@ IGbE::TxDescCache::packetAvailable()
 void
 IGbE::TxDescCache::enableSm()
 {
-    if (!igbe->drainManager) {
+    if (igbe->drainState() != DrainState::Draining) {
         igbe->txTick = true;
         igbe->restartClock();
     }
@@ -2051,18 +2051,17 @@ void
 IGbE::restartClock()
 {
     if (!tickEvent.scheduled() && (rxTick || txTick || txFifoTick) &&
-        getDrainState() == DrainState::Running)
+        drainState() == DrainState::Running)
         schedule(tickEvent, clockEdge(Cycles(1)));
 }
 
-unsigned int
+DrainState
-IGbE::drain(DrainManager *dm)
+IGbE::drain()
 {
     unsigned int count(0);
     if (rxDescCache.hasOutstandingEvents() ||
         txDescCache.hasOutstandingEvents()) {
         count++;
-        drainManager = dm;
     }
 
     txFifoTick = false;
@@ -2074,11 +2073,9 @@ IGbE::drain(DrainManager *dm)
 
     if (count) {
         DPRINTF(Drain, "IGbE not drained\n");
-        setDrainState(DrainState::Draining);
+        return DrainState::Draining;
     } else
-        setDrainState(DrainState::Drained);
+        return DrainState::Drained;
-
-    return count;
 }
 
 void
@@ -2097,7 +2094,7 @@ IGbE::drainResume()
 void
 IGbE::checkDrain()
 {
-    if (!drainManager)
+    if (drainState() != DrainState::Draining)
         return;
 
     txFifoTick = false;
@@ -2106,8 +2103,7 @@ IGbE::checkDrain()
     if (!rxDescCache.hasOutstandingEvents() &&
         !txDescCache.hasOutstandingEvents()) {
         DPRINTF(Drain, "IGbE done draining, processing drain event\n");
-        drainManager->signalDrainDone();
+        signalDrainDone();
-        drainManager = NULL;
     }
 }
 
@@ -2131,7 +2127,7 @@ IGbE::txStateMachine()
         bool success =
 #endif
             txFifo.push(txPacket);
-        txFifoTick = true && !drainManager;
+        txFifoTick = true && drainState() != DrainState::Draining;
         assert(success);
         txPacket = NULL;
         anBegin("TXS", "Desc Writeback");
@@ -2230,7 +2226,7 @@ IGbE::ethRxPkt(EthPacketPtr pkt)
     }
 
     // restart the state machines if they are stopped
-    rxTick = true && !drainManager;
+    rxTick = true && drainState() != DrainState::Draining;
     if ((rxTick || txTick) && !tickEvent.scheduled()) {
         DPRINTF(EthernetSM,
                 "RXS: received packet into fifo, starting ticking\n");
@@ -2443,8 +2439,8 @@ IGbE::ethTxDone()
     // restart the tx state machines if they are stopped
     // fifo to send another packet
     // tx sm to put more data into the fifo
-    txFifoTick = true && !drainManager;
+    txFifoTick = true && drainState() != DrainState::Draining;
-    if (txDescCache.descLeft() != 0 && !drainManager)
+    if (txDescCache.descLeft() != 0 && drainState() != DrainState::Draining)
         txTick = true;
 
     restartClock();

src/dev/i8254xGBe.hh

@@ -67,9 +67,6 @@ class IGbE : public EtherDevice
     uint8_t eeOpcode, eeAddr;
     uint16_t flash[iGbReg::EEPROM_SIZE];
 
-    // The drain event if we have one
-    DrainManager *drainManager;
-
     // packet fifos
     PacketFifo rxFifo;
     PacketFifo txFifo;
@@ -352,7 +349,7 @@ class IGbE : public EtherDevice
         virtual void updateHead(long h) { igbe->regs.rdh(h); }
         virtual void enableSm();
         virtual void fetchAfterWb() {
-            if (!igbe->rxTick && igbe->getDrainState() == DrainState::Running)
+            if (!igbe->rxTick && igbe->drainState() == DrainState::Running)
                 fetchDescriptors();
         }
 
@@ -414,7 +411,7 @@ class IGbE : public EtherDevice
         virtual void enableSm();
         virtual void actionAfterWb();
         virtual void fetchAfterWb() {
-            if (!igbe->txTick && igbe->getDrainState() == DrainState::Running)
+            if (!igbe->txTick && igbe->drainState() == DrainState::Running)
                 fetchDescriptors();
         }
         
@@ -541,8 +538,8 @@ class IGbE : public EtherDevice
     void serialize(CheckpointOut &cp) const M5_ATTR_OVERRIDE;
     void unserialize(CheckpointIn &cp) M5_ATTR_OVERRIDE;
 
-    unsigned int drain(DrainManager *dm);
+    DrainState drain() M5_ATTR_OVERRIDE;
-    void drainResume();
+    void drainResume() M5_ATTR_OVERRIDE;
 
 };
 

src/dev/ide_disk.cc

@@ -342,7 +342,7 @@ IdeDisk::doDmaTransfer()
         panic("Inconsistent DMA transfer state: dmaState = %d devState = %d\n",
               dmaState, devState);
 
-    if (ctrl->dmaPending() || ctrl->getDrainState() != DrainState::Running) {
+    if (ctrl->dmaPending() || ctrl->drainState() != DrainState::Running) {
         schedule(dmaTransferEvent, curTick() + DMA_BACKOFF_PERIOD);
         return;
     } else
@@ -436,7 +436,7 @@ IdeDisk::doDmaRead()
                 curPrd.getByteCount(), TheISA::PageBytes);
 
     }
-    if (ctrl->dmaPending() || ctrl->getDrainState() != DrainState::Running) {
+    if (ctrl->dmaPending() || ctrl->drainState() != DrainState::Running) {
         schedule(dmaReadWaitEvent, curTick() + DMA_BACKOFF_PERIOD);
         return;
     } else if (!dmaReadCG->done()) {
@@ -518,7 +518,7 @@ IdeDisk::doDmaWrite()
         dmaWriteCG = new ChunkGenerator(curPrd.getBaseAddr(),
                 curPrd.getByteCount(), TheISA::PageBytes);
     }
-    if (ctrl->dmaPending() || ctrl->getDrainState() != DrainState::Running) {
+    if (ctrl->dmaPending() || ctrl->drainState() != DrainState::Running) {
         schedule(dmaWriteWaitEvent, curTick() + DMA_BACKOFF_PERIOD);
         DPRINTF(IdeDisk, "doDmaWrite: rescheduling\n");
         return;

src/dev/ns_gige.cc

@@ -1068,7 +1068,7 @@ NSGigE::doRxDmaRead()
     assert(rxDmaState == dmaIdle || rxDmaState == dmaReadWaiting);
     rxDmaState = dmaReading;
 
-    if (dmaPending() || getDrainState() != DrainState::Running)
+    if (dmaPending() || drainState() != DrainState::Running)
         rxDmaState = dmaReadWaiting;
     else
         dmaRead(rxDmaAddr, rxDmaLen, &rxDmaReadEvent, (uint8_t*)rxDmaData);
@@ -1099,7 +1099,7 @@ NSGigE::doRxDmaWrite()
     assert(rxDmaState == dmaIdle || rxDmaState == dmaWriteWaiting);
     rxDmaState = dmaWriting;
 
-    if (dmaPending() || getDrainState() != DrainState::Running)
+    if (dmaPending() || drainState() != DrainState::Running)
         rxDmaState = dmaWriteWaiting;
     else
         dmaWrite(rxDmaAddr, rxDmaLen, &rxDmaWriteEvent, (uint8_t*)rxDmaData);
@@ -1515,7 +1515,7 @@ NSGigE::doTxDmaRead()
     assert(txDmaState == dmaIdle || txDmaState == dmaReadWaiting);
     txDmaState = dmaReading;
 
-    if (dmaPending() || getDrainState() != DrainState::Running)
+    if (dmaPending() || drainState() != DrainState::Running)
         txDmaState = dmaReadWaiting;
     else
         dmaRead(txDmaAddr, txDmaLen, &txDmaReadEvent, (uint8_t*)txDmaData);
@@ -1546,7 +1546,7 @@ NSGigE::doTxDmaWrite()
     assert(txDmaState == dmaIdle || txDmaState == dmaWriteWaiting);
     txDmaState = dmaWriting;
 
-    if (dmaPending() || getDrainState() != DrainState::Running)
+    if (dmaPending() || drainState() != DrainState::Running)
         txDmaState = dmaWriteWaiting;
     else
         dmaWrite(txDmaAddr, txDmaLen, &txDmaWriteEvent, (uint8_t*)txDmaData);

src/dev/ns_gige.hh

@@ -369,7 +369,7 @@ class NSGigE : public EtherDevBase
     void serialize(CheckpointOut &cp) const M5_ATTR_OVERRIDE;
     void unserialize(CheckpointIn &cp) M5_ATTR_OVERRIDE;
 
-    void drainResume();
+    void drainResume() M5_ATTR_OVERRIDE;
 };
 
 /*

src/dev/sinic.cc

@@ -868,7 +868,7 @@ Device::rxKick()
         break;
 
       case rxBeginCopy:
-        if (dmaPending() || getDrainState() != DrainState::Running)
+        if (dmaPending() || drainState() != DrainState::Running)
             goto exit;
 
         rxDmaAddr = params()->platform->pciToDma(
@@ -1068,7 +1068,7 @@ Device::txKick()
         break;
 
       case txBeginCopy:
-        if (dmaPending() || getDrainState() != DrainState::Running)
+        if (dmaPending() || drainState() != DrainState::Running)
             goto exit;
 
         txDmaAddr = params()->platform->pciToDma(

src/dev/sinic.hh

@@ -271,7 +271,7 @@ class Device : public Base
   public:
     virtual Tick read(PacketPtr pkt);
     virtual Tick write(PacketPtr pkt);
-    virtual void drainResume();
+    virtual void drainResume() M5_ATTR_OVERRIDE;
 
     void prepareIO(int cpu, int index);
     void prepareRead(int cpu, int index);

src/mem/cache/mshr_queue.cc

@@ -56,7 +56,7 @@ MSHRQueue::MSHRQueue(const std::string &_label,
                      int _index)
     : label(_label), numEntries(num_entries + reserve - 1),
       numReserve(reserve), demandReserve(demand_reserve),
-      registers(numEntries), drainManager(NULL), allocated(0),
+      registers(numEntries), allocated(0),
       inServiceEntries(0), index(_index)
 {
     for (int i = 0; i < numEntries; ++i) {
@@ -180,13 +180,11 @@ MSHRQueue::deallocateOne(MSHR *mshr)
         readyList.erase(mshr->readyIter);
     }
     mshr->deallocate();
-    if (drainManager && allocated == 0) {
+    if (drainState() == DrainState::Draining && allocated == 0) {
         // Notify the drain manager that we have completed draining if
         // there are no other outstanding requests in this MSHR queue.
         DPRINTF(Drain, "MSHRQueue now empty, signalling drained\n");
-        drainManager->signalDrainDone();
+        signalDrainDone();
-        drainManager = NULL;
-        setDrainState(DrainState::Drained);
     }
     return retval;
 }
@@ -265,15 +263,8 @@ MSHRQueue::squash(int threadNum)
     }
 }
 
-unsigned int
+DrainState
-MSHRQueue::drain(DrainManager *dm)
+MSHRQueue::drain()
 {
-    if (allocated == 0) {
+    return allocated == 0 ? DrainState::Drained : DrainState::Draining;
-        setDrainState(DrainState::Drained);
-        return 0;
-    } else {
-        drainManager = dm;
-        setDrainState(DrainState::Draining);
-        return 1;
-    }
 }

src/mem/cache/mshr_queue.hh

@@ -92,9 +92,6 @@ class MSHRQueue : public Drainable
     /** Holds non allocated entries. */
     MSHR::List freeList;
 
-    /** Drain manager to inform of a completed drain */
-    DrainManager *drainManager;
-
     MSHR::Iterator addToReadyList(MSHR *mshr);
 
 
@@ -258,7 +255,7 @@ class MSHRQueue : public Drainable
         return readyList.empty() ? MaxTick : readyList.front()->readyTime;
     }
 
-    unsigned int drain(DrainManager *dm);
+    DrainState drain() M5_ATTR_OVERRIDE;
 };
 
 #endif //__MEM_CACHE_MSHR_QUEUE_HH__

src/mem/dram_ctrl.cc

@@ -61,7 +61,6 @@ DRAMCtrl::DRAMCtrl(const DRAMCtrlParams* p) :
     retryRdReq(false), retryWrReq(false),
     busState(READ),
     nextReqEvent(this), respondEvent(this),
-    drainManager(NULL),
     deviceSize(p->device_size),
     deviceBusWidth(p->device_bus_width), burstLength(p->burst_length),
     deviceRowBufferSize(p->device_rowbuffer_size),
@@ -694,11 +693,11 @@ DRAMCtrl::processRespondEvent()
         schedule(respondEvent, respQueue.front()->readyTime);
     } else {
         // if there is nothing left in any queue, signal a drain
-        if (writeQueue.empty() && readQueue.empty() &&
+        if (drainState() == DrainState::Draining &&
-            drainManager) {
+            writeQueue.empty() && readQueue.empty()) {
+
             DPRINTF(Drain, "DRAM controller done draining\n");
-            drainManager->signalDrainDone();
+            signalDrainDone();
-            drainManager = NULL;
         }
     }
 
@@ -1296,15 +1295,17 @@ DRAMCtrl::processNextReqEvent()
             // trigger writes if we have passed the low threshold (or
             // if we are draining)
             if (!writeQueue.empty() &&
-                (drainManager || writeQueue.size() > writeLowThreshold)) {
+                (drainState() == DrainState::Draining ||
+                 writeQueue.size() > writeLowThreshold)) {
 
                 switch_to_writes = true;
             } else {
                 // check if we are drained
-                if (respQueue.empty () && drainManager) {
+                if (drainState() == DrainState::Draining &&
+                    respQueue.empty()) {
+
                     DPRINTF(Drain, "DRAM controller done draining\n");
-                    drainManager->signalDrainDone();
+                    signalDrainDone();
-                    drainManager = NULL;
                 }
 
                 // nothing to do, not even any point in scheduling an
@@ -1416,7 +1417,7 @@ DRAMCtrl::processNextReqEvent()
         // writes, then switch to reads.
         if (writeQueue.empty() ||
             (writeQueue.size() + minWritesPerSwitch < writeLowThreshold &&
-             !drainManager) ||
+             drainState() != DrainState::Draining) ||
             (!readQueue.empty() && writesThisTime >= minWritesPerSwitch)) {
             // turn the bus back around for reads again
             busState = WRITE_TO_READ;
@@ -2166,28 +2167,24 @@ DRAMCtrl::getSlavePort(const string &if_name, PortID idx)
     }
 }
 
-unsigned int
+DrainState
-DRAMCtrl::drain(DrainManager *dm)
+DRAMCtrl::drain()
 {
     // if there is anything in any of our internal queues, keep track
     // of that as well
-    if (!(writeQueue.empty() && readQueue.empty() &&
+    if (!(writeQueue.empty() && readQueue.empty() && respQueue.empty())) {
-          respQueue.empty())) {
         DPRINTF(Drain, "DRAM controller not drained, write: %d, read: %d,"
                 " resp: %d\n", writeQueue.size(), readQueue.size(),
                 respQueue.size());
-        drainManager = dm;
 
         // the only part that is not drained automatically over time
         // is the write queue, thus kick things into action if needed
         if (!writeQueue.empty() && !nextReqEvent.scheduled()) {
             schedule(nextReqEvent, curTick());
         }
-        setDrainState(DrainState::Draining);
+        return DrainState::Draining;
-        return 1;
     } else {
-        setDrainState(DrainState::Drained);
+        return DrainState::Drained;
-        return 0;
     }
 }
 

src/mem/dram_ctrl.hh

@@ -670,12 +670,6 @@ class DRAMCtrl : public AbstractMemory
      */
     std::deque<DRAMPacket*> respQueue;
 
-    /**
-     * If we need to drain, keep the drain manager around until we're
-     * done here.
-     */
-    DrainManager *drainManager;
-
     /**
      * Vector of ranks
      */
@@ -878,7 +872,7 @@ class DRAMCtrl : public AbstractMemory
 
     DRAMCtrl(const DRAMCtrlParams* p);
 
-    unsigned int drain(DrainManager* dm);
+    DrainState drain() M5_ATTR_OVERRIDE;
 
     virtual BaseSlavePort& getSlavePort(const std::string& if_name,
                                         PortID idx = InvalidPortID);

src/mem/dramsim2.cc

@@ -52,7 +52,6 @@ DRAMSim2::DRAMSim2(const Params* p) :
             p->traceFile, p->range.size() / 1024 / 1024, p->enableDebug),
     retryReq(false), retryResp(false), startTick(0),
     nbrOutstandingReads(0), nbrOutstandingWrites(0),
-    drainManager(NULL),
     sendResponseEvent(this), tickEvent(this)
 {
     DPRINTF(DRAMSim2,
@@ -118,11 +117,8 @@ DRAMSim2::sendResponse()
         if (!responseQueue.empty() && !sendResponseEvent.scheduled())
             schedule(sendResponseEvent, curTick());
 
-        // check if we were asked to drain and if we are now done
+        if (nbrOutstanding() == 0)
-        if (drainManager && nbrOutstanding() == 0) {
+            signalDrainDone();
-            drainManager->signalDrainDone();
-            drainManager = NULL;
-        }
     } else {
         retryResp = true;
 
@@ -339,11 +335,8 @@ void DRAMSim2::writeComplete(unsigned id, uint64_t addr, uint64_t cycle)
     assert(nbrOutstandingWrites != 0);
     --nbrOutstandingWrites;
 
-    // check if we were asked to drain and if we are now done
+    if (nbrOutstanding() == 0)
-    if (drainManager && nbrOutstanding() == 0) {
+        signalDrainDone();
-        drainManager->signalDrainDone();
-        drainManager = NULL;
-    }
 }
 
 BaseSlavePort&
@@ -357,18 +350,11 @@ DRAMSim2::getSlavePort(const std::string &if_name, PortID idx)
 }
 
 unsigned int
-DRAMSim2::drain(DrainManager* dm)
+DRAMSim2::drain()
 {
     // check our outstanding reads and writes and if any they need to
     // drain
-    if (nbrOutstanding() != 0) {
+    return nbrOutstanding() != 0 ? DrainState::Draining : DrainState::Drained;
-        setDrainState(DrainState::Draining);
-        drainManager = dm;
-        return 1;
-    } else {
-        setDrainState(DrainState::Drained);
-        return 0;
-    }
 }
 
 DRAMSim2::MemoryPort::MemoryPort(const std::string& _name,

src/mem/dramsim2.hh

@@ -132,12 +132,6 @@ class DRAMSim2 : public AbstractMemory
      */
     std::deque<PacketPtr> responseQueue;
 
-    /**
-     * If we need to drain, keep the drain manager around until we're
-     * done here.
-     */
-    DrainManager *drainManager;
-
     unsigned int nbrOutstanding() const;
 
     /**
@@ -195,7 +189,7 @@ class DRAMSim2 : public AbstractMemory
      */
     void writeComplete(unsigned id, uint64_t addr, uint64_t cycle);
 
-    unsigned int drain(DrainManager* dm);
+    DrainState drain() M5_ATTR_OVERRIDE;
 
     virtual BaseSlavePort& getSlavePort(const std::string& if_name,
                                         PortID idx = InvalidPortID);

src/mem/packet_queue.cc

@@ -49,7 +49,7 @@
 using namespace std;
 
 PacketQueue::PacketQueue(EventManager& _em, const std::string& _label)
-    : em(_em), sendEvent(this), drainManager(NULL), label(_label),
+    : em(_em), sendEvent(this), label(_label),
       waitingOnRetry(false)
 {
 }
@@ -198,11 +198,12 @@ PacketQueue::schedSendEvent(Tick when)
     } else {
         // we get a MaxTick when there is no more to send, so if we're
         // draining, we may be done at this point
-        if (drainManager && transmitList.empty() && !sendEvent.scheduled()) {
+        if (drainState() == DrainState::Draining &&
+            transmitList.empty() && !sendEvent.scheduled()) {
+
             DPRINTF(Drain, "PacketQueue done draining,"
                     "processing drain event\n");
-            drainManager->signalDrainDone();
+            signalDrainDone();
-            drainManager = NULL;
         }
     }
 }
@@ -244,14 +245,15 @@ PacketQueue::processSendEvent()
     sendDeferredPacket();
 }
 
-unsigned int
+DrainState
-PacketQueue::drain(DrainManager *dm)
+PacketQueue::drain()
 {
-    if (transmitList.empty())
+    if (transmitList.empty()) {
-        return 0;
+        return DrainState::Drained;
+    } else {
         DPRINTF(Drain, "PacketQueue not drained\n");
-    drainManager = dm;
+        return DrainState::Draining;
-    return 1;
+    }
 }
 
 ReqPacketQueue::ReqPacketQueue(EventManager& _em, MasterPort& _masterPort,

src/mem/packet_queue.hh

@@ -89,10 +89,6 @@ class PacketQueue : public Drainable
     /** Event used to call processSendEvent. */
     EventWrapper<PacketQueue, &PacketQueue::processSendEvent> sendEvent;
 
-    /** If we need to drain, keep the drain manager around until we're done
-     * here.*/
-    DrainManager *drainManager;
-
   protected:
 
     /** Label to use for print request packets label stack. */
@@ -192,7 +188,7 @@ class PacketQueue : public Drainable
      */
     void retry();
 
-    unsigned int drain(DrainManager *dm);
+    DrainState drain() M5_ATTR_OVERRIDE;
 };
 
 class ReqPacketQueue : public PacketQueue

src/mem/ruby/structures/RubyMemoryControl.cc

@@ -640,14 +640,14 @@ RubyMemoryControl::executeCycle()
     }
 }
 
-unsigned int
+DrainState
-RubyMemoryControl::drain(DrainManager *dm)
+RubyMemoryControl::drain()
 {
     DPRINTF(RubyMemory, "MemoryController drain\n");
     if(m_event.scheduled()) {
         deschedule(m_event);
     }
-    return 0;
+    return DrainState::Drained;
 }
 
 // wakeup:  This function is called once per memory controller clock cycle.

src/mem/ruby/structures/RubyMemoryControl.hh

@@ -61,7 +61,7 @@ class RubyMemoryControl : public AbstractMemory, public Consumer
 
     virtual BaseSlavePort& getSlavePort(const std::string& if_name,
                                         PortID idx = InvalidPortID);
-    unsigned int drain(DrainManager *dm);
+    DrainState drain() M5_ATTR_OVERRIDE;
     void wakeup();
 
     void setDescription(const std::string& name) { m_description = name; };

src/mem/ruby/system/DMASequencer.cc

@@ -42,7 +42,7 @@ DMASequencer::DMASequencer(const Params *p)
       m_mandatory_q_ptr(NULL), m_usingRubyTester(p->using_ruby_tester),
       slave_port(csprintf("%s.slave", name()), this, 0, p->ruby_system,
                  p->ruby_system->getAccessBackingStore()),
-      drainManager(NULL), system(p->system), retry(false)
+      system(p->system), retry(false)
 {
     assert(m_version != -1);
 }
@@ -148,43 +148,34 @@ void
 DMASequencer::testDrainComplete()
 {
     //If we weren't able to drain before, we might be able to now.
-    if (drainManager != NULL) {
+    if (drainState() == DrainState::Draining) {
         unsigned int drainCount = outstandingCount();
         DPRINTF(Drain, "Drain count: %u\n", drainCount);
         if (drainCount == 0) {
             DPRINTF(Drain, "DMASequencer done draining, signaling drain done\n");
-            drainManager->signalDrainDone();
+            signalDrainDone();
-            // Clear the drain manager once we're done with it.
-            drainManager = NULL;
         }
     }
 }
 
-unsigned int
+DrainState
-DMASequencer::drain(DrainManager *dm)
+DMASequencer::drain()
 {
     if (isDeadlockEventScheduled()) {
         descheduleDeadlockEvent();
     }
 
     // If the DMASequencer is not empty, then it needs to clear all outstanding
-    // requests before it should call drainManager->signalDrainDone()
+    // requests before it should call signalDrainDone()
     DPRINTF(Config, "outstanding count %d\n", outstandingCount());
-    bool need_drain = outstandingCount() > 0;
-
 
     // Set status
-    if (need_drain) {
+    if (outstandingCount() > 0) {
-        drainManager = dm;
-
         DPRINTF(Drain, "DMASequencer not drained\n");
-        setDrainState(DrainState::Draining);
+        return DrainState::Draining;
-        return 1;
+    } else {
+        return DrainState::Drained;
     }
-
-    drainManager = NULL;
-    setDrainState(DrainState::Drained);
-    return 0;
 }
 
 void

src/mem/ruby/system/DMASequencer.hh

@@ -107,7 +107,7 @@ class DMASequencer : public MemObject
     // A pointer to the controller is needed for atomic support.
     void setController(AbstractController* _cntrl) { m_controller = _cntrl; }
     uint32_t getId() { return m_version; }
-    unsigned int drain(DrainManager *dm);
+    DrainState drain() M5_ATTR_OVERRIDE;
 
     /* SLICC callback */
     void dataCallback(const DataBlock & dblk);
@@ -129,7 +129,7 @@ class DMASequencer : public MemObject
      * @return Whether successfully sent
      */
     bool recvTimingResp(PacketPtr pkt, PortID master_port_id);
-    unsigned int getChildDrainCount(DrainManager *dm);
+    unsigned int getChildDrainCount();
 
   private:
     uint32_t m_version;
@@ -139,7 +139,6 @@ class DMASequencer : public MemObject
 
     MemSlavePort slave_port;
 
-    DrainManager *drainManager;
     System* system;
 
     bool retry;

src/mem/ruby/system/RubyPort.cc

@@ -59,7 +59,7 @@ RubyPort::RubyPort(const Params *p)
       memMasterPort(csprintf("%s.mem-master-port", name()), this),
       memSlavePort(csprintf("%s-mem-slave-port", name()), this,
           p->ruby_system, p->ruby_system->getAccessBackingStore(), -1),
-      gotAddrRanges(p->port_master_connection_count), drainManager(NULL)
+      gotAddrRanges(p->port_master_connection_count)
 {
     assert(m_version != -1);
 
@@ -387,20 +387,18 @@ void
 RubyPort::testDrainComplete()
 {
     //If we weren't able to drain before, we might be able to now.
-    if (drainManager != NULL) {
+    if (drainState() == DrainState::Draining) {
         unsigned int drainCount = outstandingCount();
         DPRINTF(Drain, "Drain count: %u\n", drainCount);
         if (drainCount == 0) {
             DPRINTF(Drain, "RubyPort done draining, signaling drain done\n");
-            drainManager->signalDrainDone();
+            signalDrainDone();
-            // Clear the drain manager once we're done with it.
-            drainManager = NULL;
         }
     }
 }
 
-unsigned int
+DrainState
-RubyPort::drain(DrainManager *dm)
+RubyPort::drain()
 {
     if (isDeadlockEventScheduled()) {
         descheduleDeadlockEvent();
@@ -408,23 +406,15 @@ RubyPort::drain(DrainManager *dm)
 
     //
     // If the RubyPort is not empty, then it needs to clear all outstanding
-    // requests before it should call drainManager->signalDrainDone()
+    // requests before it should call signalDrainDone()
     //
     DPRINTF(Config, "outstanding count %d\n", outstandingCount());
-    bool need_drain = outstandingCount() > 0;
+    if (outstandingCount() > 0) {
-
-    // Set status
-    if (need_drain) {
-        drainManager = dm;
-
         DPRINTF(Drain, "RubyPort not drained\n");
-        setDrainState(DrainState::Draining);
+        return DrainState::Draining;
-        return 1;
+    } else {
+        return DrainState::Drained;
     }
-
-    drainManager = NULL;
-    setDrainState(DrainState::Drained);
-    return 0;
 }
 
 void

src/mem/ruby/system/RubyPort.hh

@@ -162,7 +162,7 @@ class RubyPort : public MemObject
     //
     void setController(AbstractController* _cntrl) { m_controller = _cntrl; }
     uint32_t getId() { return m_version; }
-    unsigned int drain(DrainManager *dm);
+    DrainState drain() M5_ATTR_OVERRIDE;
 
   protected:
     void ruby_hit_callback(PacketPtr pkt);
@@ -204,8 +204,6 @@ class RubyPort : public MemObject
     std::vector<MemSlavePort *> slave_ports;
     std::vector<PioMasterPort *> master_ports;
 
-    DrainManager *drainManager;
-
     //
     // Based on similar code in the M5 bus.  Stores pointers to those ports
     // that should be called when the Sequencer becomes available after a stall.

src/mem/ruby/system/Sequencer.cc

@@ -77,7 +77,7 @@ Sequencer::~Sequencer()
 void
 Sequencer::wakeup()
 {
-    assert(getDrainState() != DrainState::Draining);
+    assert(drainState() != DrainState::Draining);
 
     // Check for deadlock of any of the requests
     Cycles current_time = curCycle();
@@ -215,7 +215,7 @@ Sequencer::insertRequest(PacketPtr pkt, RubyRequestType request_type)
 
     // See if we should schedule a deadlock check
     if (!deadlockCheckEvent.scheduled() &&
-        getDrainState() != DrainState::Draining) {
+        drainState() != DrainState::Draining) {
         schedule(deadlockCheckEvent, clockEdge(m_deadlock_threshold));
     }
 

src/mem/simple_mem.cc

@@ -53,7 +53,7 @@ SimpleMemory::SimpleMemory(const SimpleMemoryParams* p) :
     port(name() + ".port", *this), latency(p->latency),
     latency_var(p->latency_var), bandwidth(p->bandwidth), isBusy(false),
     retryReq(false), retryResp(false),
-    releaseEvent(this), dequeueEvent(this), drainManager(NULL)
+    releaseEvent(this), dequeueEvent(this)
 {
 }
 
@@ -200,10 +200,9 @@ SimpleMemory::dequeue()
             // already have an event scheduled, so use re-schedule
             reschedule(dequeueEvent,
                        std::max(packetQueue.front().tick, curTick()), true);
-        } else if (drainManager) {
+        } else if (drainState() == DrainState::Draining) {
-            DPRINTF(Drain, "Drainng of SimpleMemory complete\n");
+            DPRINTF(Drain, "Draining of SimpleMemory complete\n");
-            drainManager->signalDrainDone();
+            signalDrainDone();
-            drainManager = NULL;
         }
     }
 }
@@ -233,23 +232,15 @@ SimpleMemory::getSlavePort(const std::string &if_name, PortID idx)
     }
 }
 
-unsigned int
+DrainState
-SimpleMemory::drain(DrainManager *dm)
+SimpleMemory::drain()
 {
-    int count = 0;
-
-    // also track our internal queue
     if (!packetQueue.empty()) {
-        count += 1;
-        drainManager = dm;
         DPRINTF(Drain, "SimpleMemory Queue has requests, waiting to drain\n");
+        return DrainState::Draining;
+    } else {
+        return DrainState::Drained;
     }
-
-    if (count)
-        setDrainState(DrainState::Draining);
-    else
-        setDrainState(DrainState::Drained);
-    return count;
 }
 
 SimpleMemory::MemoryPort::MemoryPort(const std::string& _name,

src/mem/simple_mem.hh

@@ -181,17 +181,11 @@ class SimpleMemory : public AbstractMemory
      */
     std::vector<PacketPtr> pendingDelete;
 
-    /**
-     * If we need to drain, keep the drain manager around until we're
-     * done here.
-     */
-    DrainManager *drainManager;
-
   public:
 
     SimpleMemory(const SimpleMemoryParams *p);
 
-    unsigned int drain(DrainManager *dm);
+    DrainState drain() M5_ATTR_OVERRIDE;
 
     BaseSlavePort& getSlavePort(const std::string& if_name,
                                 PortID idx = InvalidPortID);

src/mem/xbar.cc

@@ -145,7 +145,7 @@ BaseXBar::calcPacketTiming(PacketPtr pkt, Tick header_delay)
 template <typename SrcType, typename DstType>
 BaseXBar::Layer<SrcType,DstType>::Layer(DstType& _port, BaseXBar& _xbar,
                                        const std::string& _name) :
-    port(_port), xbar(_xbar), _name(_name), state(IDLE), drainManager(NULL),
+    port(_port), xbar(_xbar), _name(_name), state(IDLE),
     waitingForPeer(NULL), releaseEvent(this)
 {
 }
@@ -252,12 +252,10 @@ BaseXBar::Layer<SrcType,DstType>::releaseLayer()
         // waiting for the peer
         if (waitingForPeer == NULL)
             retryWaiting();
-    } else if (waitingForPeer == NULL && drainManager) {
+    } else if (waitingForPeer == NULL && drainState() == DrainState::Draining) {
         DPRINTF(Drain, "Crossbar done draining, signaling drain manager\n");
         //If we weren't able to drain before, do it now.
-        drainManager->signalDrainDone();
+        signalDrainDone();
-        // Clear the drain event once we're done with it.
-        drainManager = NULL;
     }
 }
 
@@ -587,18 +585,18 @@ BaseXBar::regStats()
 }
 
 template <typename SrcType, typename DstType>
-unsigned int
+DrainState
-BaseXBar::Layer<SrcType,DstType>::drain(DrainManager *dm)
+BaseXBar::Layer<SrcType,DstType>::drain()
 {
     //We should check that we're not "doing" anything, and that noone is
     //waiting. We might be idle but have someone waiting if the device we
     //contacted for a retry didn't actually retry.
     if (state != IDLE) {
         DPRINTF(Drain, "Crossbar not drained\n");
-        drainManager = dm;
+        return DrainState::Draining;
-        return 1;
+    } else {
+        return DrainState::Drained;
     }
-    return 0;
 }
 
 template <typename SrcType, typename DstType>

src/mem/xbar.hh

@@ -114,7 +114,7 @@ class BaseXBar : public MemObject
          *
          * @return 1 if busy or waiting to retry, or 0 if idle
          */
-        unsigned int drain(DrainManager *dm);
+        DrainState drain() M5_ATTR_OVERRIDE;
 
         /**
          * Get the crossbar layer's name
@@ -217,9 +217,6 @@ class BaseXBar : public MemObject
         /** track the state of the layer */
         State state;
 
-        /** manager to signal when drained */
-        DrainManager *drainManager;
-
         /**
          * A deque of ports that retry should be called on because
          * the original send was delayed due to a busy layer.

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);
 }
 
-void
+DrainState
-Drainable::drainResume()
+Drainable::dmDrain()
 {
-    _drainState = DrainState::Running;
+    _drainState = DrainState::Draining;
+    _drainState = drain();
+    assert(_drainState == DrainState::Draining ||
+           _drainState == DrainState::Drained);
+
+    return _drainState;
+}
+
+void
+Drainable::dmDrainResume()
+{
+    panic_if(_drainState != DrainState::Drained,
+             "Trying to resume an object that hasn't been drained\n");
+
+    _drainState = DrainState::Running;
+    drainResume();
 }

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
+ * <li>DrainManager::tryDrain() calls Drainable::drain() for every
- *     system. Draining has completed if all of them return
+ *     object in the system. Draining has completed if all of them
- *     zero. Otherwise, the sum of the return values is loaded into
+ *     return true. Otherwise, the drain manager keeps track of the
- *     the counter of the DrainManager. A pointer to the drain
+ *     objects that requested draining and waits for them to signal
- *     manager is passed as an argument to the drain() method.
+ *     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
+ *     manager. The drain manager keeps track of the objects that
- *     simulation stops.
+ *     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
+ * <li>Check if any object still needs draining
- *     process above.
+ *     (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
+     * If the object does not need further simulation to drain
-     * with a pointer to a drain manager. If the object does not need
+     * internal buffers, it returns true and automatically switches to
-     * further simulation to drain internal buffers, it switched to
+     * the Drained state, otherwise it switches to the Draining state.
-     * 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.
      *
      * @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
+     * being drained. These perturbations are not visible in the
-     * therefore repeats the draining process until all objects return
+     * drain state. The simulator therefore repeats the draining
-     * 0 on the first call to drain().
+     * process until all objects return DrainState::Drained on the
+     * first call to drain().
      *
-     * @param drainManager DrainManager to use to inform the simulator
+     * @return DrainState::Drained if the object is ready for
-     * when draining has completed.
+     * serialization now, DrainState::Draining if it needs further
-     *
+     * simulation.
-     * @return 0 if the object is ready for serialization now, >0 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:
+  public:
-    void setDrainState(DrainState new_state) { _drainState = new_state; }
+    /** 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

src/sim/process.cc

@@ -259,11 +259,11 @@ Process::initState()
     pTable->initState(tc);
 }
 
-unsigned int
+DrainState
-Process::drain(DrainManager *dm)
+Process::drain()
 {
     find_file_offsets();
-    return 0;
+    return DrainState::Drained;
 }
 
 // map simulator fd sim_fd to target fd tgt_fd

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:
 

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)
 {

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
+     * Provide a default implementation of the drain interface for
-     * simply returns 0 (draining completed) and sets the drain state
+     * objects that don't need draining.
-     * to Drained.
      */
-    unsigned int drain(DrainManager *drainManger);
+    DrainState drain() M5_ATTR_OVERRIDE { return DrainState::Drained; }
 
     /**
      * Write back dirty buffers to memory using functional writes.

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;
 }
 

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() M5_ATTR_OVERRIDE;
-    void drainResume();
 
   public:
     Counter totalNumInsts;