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.

59 files changed, 344 insertions, 536 deletions

diff --git a/src/arch/arm/table_walker.cc b/src/arch/arm/table_walker.cc
index 297054131..3e61a4bd6 100644
--- a/src/arch/arm/table_walker.cc
+++ b/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();
-        drainManager = NULL;
+        signalDrainDone();
     }
 }
 
-unsigned int
-TableWalker::drain(DrainManager *dm)
+DrainState
+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 port drain count plus the table walker itself needs to drain
-        return 1;
+        return DrainState::Draining;
     } else {
-        setDrainState(DrainState::Drained);
         DPRINTF(Drain, "TableWalker free, no need to drain\n");
-
-        // table walker is drained, but its ports may still need to be drained
-        return 0;
+        return DrainState::Drained;
     }
 }
 
 void
 TableWalker::drainResume()
 {
-    Drainable::drainResume();
     if (params()->sys->isTimingMode() && currState) {
         delete currState;
         currState = NULL;
diff --git a/src/arch/arm/table_walker.hh b/src/arch/arm/table_walker.hh
index a5327cd95..e973e9a74 100644
--- a/src/arch/arm/table_walker.hh
+++ b/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);
-    virtual void drainResume();
+    DrainState drain() M5_ATTR_OVERRIDE;
+    virtual void drainResume() M5_ATTR_OVERRIDE;
 
     virtual BaseMasterPort& getMasterPort(const std::string &if_name,
                                           PortID idx = InvalidPortID);
diff --git a/src/arch/arm/tlb.hh b/src/arch/arm/tlb.hh
index 28b99a8e0..63707dba2 100644
--- a/src/arch/arm/tlb.hh
+++ b/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;
diff --git a/src/cpu/kvm/base.cc b/src/cpu/kvm/base.cc
index 30e984366..47cff5917 100644
--- a/src/cpu/kvm/base.cc
+++ b/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
-BaseKvmCPU::drain(DrainManager *dm)
+DrainState
+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()) {
-            drainManager = dm;
-            return 1;
-        }
+        if (!archIsDrained())
+            return DrainState::Draining;
 
         // 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();
-        drainManager = NULL;
+        signalDrainDone();
         return true;
     } else {
         DPRINTF(Drain, "tryDrain: CPU not ready.\n");
diff --git a/src/cpu/kvm/base.hh b/src/cpu/kvm/base.hh
index a8e429dfa..df6a67808 100644
--- a/src/cpu/kvm/base.hh
+++ b/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);
-    void drainResume();
+    DrainState drain() M5_ATTR_OVERRIDE;
+    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;
diff --git a/src/cpu/minor/cpu.cc b/src/cpu/minor/cpu.cc
index ac1a18bf9..a93d0037d 100644
--- a/src/cpu/minor/cpu.cc
+++ b/src/cpu/minor/cpu.cc
@@ -47,8 +47,7 @@
 #include "debug/Quiesce.hh"
 
 MinorCPU::MinorCPU(MinorCPUParams *params) :
-    BaseCPU(params),
-    drainManager(NULL)
+    BaseCPU(params)
 {
     /* This is only written for one thread at the moment */
     Minor::MinorThread *thread;
@@ -194,39 +193,33 @@ MinorCPU::startup()
     activateContext(0);
 }
 
-unsigned int
-MinorCPU::drain(DrainManager *drain_manager)
+DrainState
+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 (ret == 0)
+    if (pipeline->drain()) {
         DPRINTF(Drain, "MinorCPU drained\n");
-    else
+        return DrainState::Drained;
+    } else {
         DPRINTF(Drain, "MinorCPU not finished draining\n");
-
-    return ret;
+        return DrainState::Draining;
+    }
 }
 
 void
 MinorCPU::signalDrainDone()
 {
     DPRINTF(Drain, "MinorCPU drain done\n");
-    setDrainState(DrainState::Drained);
-    drainManager->signalDrainDone();
-    drainManager = NULL;
+    signalDrainDone();
 }
 
 void
 MinorCPU::drainResume()
 {
-    assert(getDrainState() == DrainState::Drained ||
-        getDrainState() == DrainState::Running);
+    assert(drainState() == DrainState::Drained);
 
     if (switchedOut()) {
         DPRINTF(Drain, "drainResume while switched out.  Ignoring\n");
@@ -242,8 +235,6 @@ MinorCPU::drainResume()
 
     wakeup();
     pipeline->drainResume();
-
-    setDrainState(DrainState::Running);
 }
 
 void
diff --git a/src/cpu/minor/cpu.hh b/src/cpu/minor/cpu.hh
index fba54b515..2e877d786 100644
--- a/src/cpu/minor/cpu.hh
+++ b/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);
-    void drainResume();
-    /** Signal from Pipeline that MinorCPU should signal the DrainManager
-     *  that a drain is complete and set its drainState */
+    DrainState drain() M5_ATTR_OVERRIDE;
+    void drainResume() M5_ATTR_OVERRIDE;
+    /** Signal from Pipeline that MinorCPU should signal that a drain
+     *  is complete and set its drainState */
     void signalDrainDone();
     void memWriteback();
 
diff --git a/src/cpu/minor/pipeline.cc b/src/cpu/minor/pipeline.cc
index 9d802234b..39b7f31f9 100644
--- a/src/cpu/minor/pipeline.cc
+++ b/src/cpu/minor/pipeline.cc
@@ -192,8 +192,8 @@ Pipeline::wakeupFetch()
     execute.wakeupFetch();
 }
 
-unsigned int
-Pipeline::drain(DrainManager *manager)
+bool
+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
diff --git a/src/cpu/minor/pipeline.hh b/src/cpu/minor/pipeline.hh
index 355a3c6c2..bf2071b02 100644
--- a/src/cpu/minor/pipeline.hh
+++ b/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();
 
diff --git a/src/cpu/o3/cpu.cc b/src/cpu/o3/cpu.cc
index 89256a7f0..026907a94 100644
--- a/src/cpu/o3/cpu.cc
+++ b/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
-FullO3CPU<Impl>::drain(DrainManager *drain_manager)
+DrainState
+FullO3CPU<Impl>::drain()
 {
     // If the CPU isn't doing anything, then return immediately.
-    if (switchedOut()) {
-        setDrainState(DrainState::Drained);
-        return 0;
-    }
+    if (switchedOut())
+        return DrainState::Drained;
 
     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();
-    drainManager = NULL;
+    signalDrainDone();
 
     return true;
 }
@@ -1132,7 +1124,6 @@ template <class Impl>
 void
 FullO3CPU<Impl>::drainResume()
 {
-    setDrainState(DrainState::Running);
     if (switchedOut())
         return;
 
diff --git a/src/cpu/o3/cpu.hh b/src/cpu/o3/cpu.hh
index 44400542a..aa02ee2ea 100644
--- a/src/cpu/o3/cpu.hh
+++ b/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;
 
diff --git a/src/cpu/simple/atomic.cc b/src/cpu/simple/atomic.cc
index 4c1c45355..3777ddee9 100644
--- a/src/cpu/simple/atomic.cc
+++ b/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
-AtomicSimpleCPU::drain(DrainManager *dm)
+DrainState
+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 1;
+        return DrainState::Draining;
     } 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();
-    drain_manager = NULL;
+    signalDrainDone();
 
     return true;
 }
diff --git a/src/cpu/simple/atomic.hh b/src/cpu/simple/atomic.hh
index 5ad5c4305..3f587e671 100644
--- a/src/cpu/simple/atomic.hh
+++ b/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);
-    void drainResume();
+    DrainState drain() M5_ATTR_OVERRIDE;
+    void drainResume() M5_ATTR_OVERRIDE;
 
     void switchOut();
     void takeOverFrom(BaseCPU *oldCPU);
diff --git a/src/cpu/simple/timing.cc b/src/cpu/simple/timing.cc
index a3c4e27e8..5dc042f1e 100644
--- a/src/cpu/simple/timing.cc
+++ b/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
-TimingSimpleCPU::drain(DrainManager *drain_manager)
+DrainState
+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();
-    drainManager = NULL;
+    signalDrainDone();
 
     return true;
 }
diff --git a/src/cpu/simple/timing.hh b/src/cpu/simple/timing.hh
index 3ce596fc7..b6a1da4e2 100644
--- a/src/cpu/simple/timing.hh
+++ b/src/cpu/simple/timing.hh
@@ -270,8 +270,8 @@ class TimingSimpleCPU : public BaseSimpleCPU
 
   public:
 
-    unsigned int drain(DrainManager *drain_manager);
-    void drainResume();
+    DrainState drain() M5_ATTR_OVERRIDE;
+    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__
diff --git a/src/cpu/testers/traffic_gen/traffic_gen.cc b/src/cpu/testers/traffic_gen/traffic_gen.cc
index 0fc8848fb..984c9950d 100644
--- a/src/cpu/testers/traffic_gen/traffic_gen.cc
+++ b/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),
-      drainManager(NULL)
+      updateEvent(this)
 {
 }
 
@@ -118,12 +117,12 @@ TrafficGen::initState()
     }
 }
 
-unsigned int
-TrafficGen::drain(DrainManager *dm)
+DrainState
+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 1;
+        return DrainState::Draining;
     }
 }
 
@@ -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();
-            // Clear the drain event once we're done with it.
-            drainManager = NULL;
+            signalDrainDone();
         }
     }
 }
diff --git a/src/cpu/testers/traffic_gen/traffic_gen.hh b/src/cpu/testers/traffic_gen/traffic_gen.hh
index ba7fda7dd..8b71443f9 100644
--- a/src/cpu/testers/traffic_gen/traffic_gen.hh
+++ b/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;
diff --git a/src/dev/arm/flash_device.cc b/src/dev/arm/flash_device.cc
index e4bfcf5c9..b651a1eeb 100644
--- a/src/dev/arm/flash_device.cc
+++ b/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
-FlashDevice::drain(DrainManager *dm)
+DrainState
+FlashDevice::drain()
 {
-    unsigned int count = 0;
-
     if (planeEvent.scheduled()) {
-        count = 1;
-        drainManager = dm;
-    } else {
-        DPRINTF(Drain, "Flash device in drained state\n");
-    }
-
-    if (count) {
         DPRINTF(Drain, "Flash device is draining...\n");
-        setDrainState(DrainState::Draining);
+        return DrainState::Draining;
     } else {
-        DPRINTF(Drain, "Flash device drained\n");
-        setDrainState(DrainState::Drained);
+        DPRINTF(Drain, "Flash device in drained state\n");
+        return 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();
-        drainManager = NULL;
+        signalDrainDone();
     }
 }
diff --git a/src/dev/arm/flash_device.hh b/src/dev/arm/flash_device.hh
index 330299451..891217cbf 100644
--- a/src/dev/arm/flash_device.hh
+++ b/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;
@@ -176,13 +176,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.
      */
diff --git a/src/dev/arm/ufs_device.cc b/src/dev/arm/ufs_device.cc
index 696aeba6f..61b125ef5 100644
--- a/src/dev/arm/ufs_device.cc
+++ b/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
-UFSHostDevice::drain(DrainManager *dm)
+DrainState
+UFSHostDevice::drain()
 {
     if (UFSHCIMem.TRUTRLDBR) {
-        drainManager = dm;
         DPRINTF(UFSHostDevice, "UFSDevice is draining...\n");
-        setDrainState(DrainState::Draining);
-        return 1;
+        return DrainState::Draining;
     } else {
         DPRINTF(UFSHostDevice, "UFSDevice drained\n");
-        setDrainState(DrainState::Drained);
-        return 0;
+        return DrainState::Drained;
     }
 }
 
@@ -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();
-        drainManager = NULL;
+        signalDrainDone();
     }
 }
diff --git a/src/dev/arm/ufs_device.hh b/src/dev/arm/ufs_device.hh
index 15e983ad8..716b1bdcb 100644
--- a/src/dev/arm/ufs_device.hh
+++ b/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;
@@ -1053,13 +1053,6 @@ class UFSHostDevice : public DmaDevice
     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
      * logic units.
diff --git a/src/dev/copy_engine.cc b/src/dev/copy_engine.cc
index ed177cf7f..646393b8f 100644
--- a/src/dev/copy_engine.cc
+++ b/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);
-        drainManager->signalDrainDone();
-        drainManager = NULL;
+        signalDrainDone();
     }
 
-    return ce->getDrainState() != DrainState::Running;
+    return ce->drainState() != DrainState::Running;
 }
 
-unsigned int
-CopyEngine::CopyEngineChannel::drain(DrainManager *dm)
+DrainState
+CopyEngine::CopyEngineChannel::drain()
 {
-    if (nextState == Idle || ce->getDrainState() != DrainState::Running)
-        return 0;
-
-    DPRINTF(Drain, "CopyEngineChannel not drained\n");
-    this->drainManager = dm;
-    return 1;
+    if (nextState == Idle || ce->drainState() != DrainState::Running) {
+        return DrainState::Drained;
+    } else {
+        DPRINTF(Drain, "CopyEngineChannel not drained\n");
+        return DrainState::Draining;
+    }
 }
 
 void
diff --git a/src/dev/copy_engine.hh b/src/dev/copy_engine.hh
index d09a18dbd..db701d451 100644
--- a/src/dev/copy_engine.hh
+++ b/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);
-        void drainResume();
+        DrainState drain() M5_ATTR_OVERRIDE;
+        void drainResume() M5_ATTR_OVERRIDE;
 
         void serialize(CheckpointOut &cp) const M5_ATTR_OVERRIDE;
         void unserialize(CheckpointIn &cp) M5_ATTR_OVERRIDE;
diff --git a/src/dev/dma_device.cc b/src/dev/dma_device.cc
index 97b31e976..343f706b0 100644
--- a/src/dev/dma_device.cc
+++ b/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),
-      inRetry(false)
+      pendingCount(0), 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) {
-        drainManager->signalDrainDone();
-        drainManager = NULL;
-    }
+    if (pendingCount == 0)
+        signalDrainDone();
 }
 
 bool
@@ -125,14 +122,15 @@ DmaDevice::init()
     PioDevice::init();
 }
 
-unsigned int
-DmaPort::drain(DrainManager *dm)
+DrainState
+DmaPort::drain()
 {
-    if (pendingCount == 0)
-        return 0;
-    drainManager = dm;
-    DPRINTF(Drain, "DmaPort not drained\n");
-    return 1;
+    if (pendingCount == 0) {
+        return DrainState::Drained;
+    } else {
+        DPRINTF(Drain, "DmaPort not drained\n");
+        return DrainState::Draining;
+    }
 }
 
 void
diff --git a/src/dev/dma_device.hh b/src/dev/dma_device.hh
index 92bf8f72c..b0e35e44f 100644
--- a/src/dev/dma_device.hh
+++ b/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
diff --git a/src/dev/i8254xGBe.cc b/src/dev/i8254xGBe.cc
index cd96c9eca..50126e62a 100644
--- a/src/dev/i8254xGBe.cc
+++ b/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
-IGbE::drain(DrainManager *dm)
+DrainState
+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 count;
+        return DrainState::Drained;
 }
 
 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();
-        drainManager = NULL;
+        signalDrainDone();
     }
 }
 
@@ -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;
-    if (txDescCache.descLeft() != 0 && !drainManager)
+    txFifoTick = true && drainState() != DrainState::Draining;
+    if (txDescCache.descLeft() != 0 && drainState() != DrainState::Draining)
         txTick = true;
 
     restartClock();
diff --git a/src/dev/i8254xGBe.hh b/src/dev/i8254xGBe.hh
index 116fa5b95..3a3efb795 100644
--- a/src/dev/i8254xGBe.hh
+++ b/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);
-    void drainResume();
+    DrainState drain() M5_ATTR_OVERRIDE;
+    void drainResume() M5_ATTR_OVERRIDE;
 
 };
 
diff --git a/src/dev/ide_disk.cc b/src/dev/ide_disk.cc
index 2e6df3805..7298da521 100644
--- a/src/dev/ide_disk.cc
+++ b/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;
diff --git a/src/dev/ns_gige.cc b/src/dev/ns_gige.cc
index 7a8d202c5..c6c570749 100644
--- a/src/dev/ns_gige.cc
+++ b/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);
diff --git a/src/dev/ns_gige.hh b/src/dev/ns_gige.hh
index f8e9be853..08b02027a 100644
--- a/src/dev/ns_gige.hh
+++ b/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;
 };
 
 /*
diff --git a/src/dev/sinic.cc b/src/dev/sinic.cc
index 1ba076200..d4a3f19b3 100644
--- a/src/dev/sinic.cc
+++ b/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(
diff --git a/src/dev/sinic.hh b/src/dev/sinic.hh
index ac7af34af..69b81b1b2 100644
--- a/src/dev/sinic.hh
+++ b/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);
diff --git a/src/mem/cache/mshr_queue.cc b/src/mem/cache/mshr_queue.cc
index 3736c9f21..47f044d63 100644
--- a/src/mem/cache/mshr_queue.cc
+++ b/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();
-        drainManager = NULL;
-        setDrainState(DrainState::Drained);
+        signalDrainDone();
     }
     return retval;
 }
@@ -265,15 +263,8 @@ MSHRQueue::squash(int threadNum)
     }
 }
 
-unsigned int
-MSHRQueue::drain(DrainManager *dm)
+DrainState
+MSHRQueue::drain()
 {
-    if (allocated == 0) {
-        setDrainState(DrainState::Drained);
-        return 0;
-    } else {
-        drainManager = dm;
-        setDrainState(DrainState::Draining);
-        return 1;
-    }
+    return allocated == 0 ? DrainState::Drained : DrainState::Draining;
 }
diff --git a/src/mem/cache/mshr_queue.hh b/src/mem/cache/mshr_queue.hh
index 4043bc565..308d371fe 100644
--- a/src/mem/cache/mshr_queue.hh
+++ b/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__
diff --git a/src/mem/dram_ctrl.cc b/src/mem/dram_ctrl.cc
index 0b1509e2f..dc2c03332 100644
--- a/src/mem/dram_ctrl.cc
+++ b/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() &&
-            drainManager) {
+        if (drainState() == DrainState::Draining &&
+            writeQueue.empty() && readQueue.empty()) {
+
             DPRINTF(Drain, "DRAM controller done draining\n");
-            drainManager->signalDrainDone();
-            drainManager = NULL;
+            signalDrainDone();
         }
     }
 
@@ -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();
-                    drainManager = NULL;
+                    signalDrainDone();
                 }
 
                 // 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
-DRAMCtrl::drain(DrainManager *dm)
+DrainState
+DRAMCtrl::drain()
 {
     // if there is anything in any of our internal queues, keep track
     // of that as well
-    if (!(writeQueue.empty() && readQueue.empty() &&
-          respQueue.empty())) {
+    if (!(writeQueue.empty() && readQueue.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 1;
+        return DrainState::Draining;
     } else {
-        setDrainState(DrainState::Drained);
-        return 0;
+        return DrainState::Drained;
     }
 }
 
diff --git a/src/mem/dram_ctrl.hh b/src/mem/dram_ctrl.hh
index 1ea04fa56..bf0be04a7 100644
--- a/src/mem/dram_ctrl.hh
+++ b/src/mem/dram_ctrl.hh
@@ -671,12 +671,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
      */
     std::vector<Rank*> 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);
diff --git a/src/mem/dramsim2.cc b/src/mem/dramsim2.cc
index dfeca0d47..6d0b84336 100644
--- a/src/mem/dramsim2.cc
+++ b/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 (drainManager && nbrOutstanding() == 0) {
-            drainManager->signalDrainDone();
-            drainManager = NULL;
-        }
+        if (nbrOutstanding() == 0)
+            signalDrainDone();
     } 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 (drainManager && nbrOutstanding() == 0) {
-        drainManager->signalDrainDone();
-        drainManager = NULL;
-    }
+    if (nbrOutstanding() == 0)
+        signalDrainDone();
 }
 
 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) {
-        setDrainState(DrainState::Draining);
-        drainManager = dm;
-        return 1;
-    } else {
-        setDrainState(DrainState::Drained);
-        return 0;
-    }
+    return nbrOutstanding() != 0 ? DrainState::Draining : DrainState::Drained;
 }
 
 DRAMSim2::MemoryPort::MemoryPort(const std::string& _name,
diff --git a/src/mem/dramsim2.hh b/src/mem/dramsim2.hh
index cd87fe73d..5d8e64282 100644
--- a/src/mem/dramsim2.hh
+++ b/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);
diff --git a/src/mem/packet_queue.cc b/src/mem/packet_queue.cc
index 86cf599ab..ff248d388 100644
--- a/src/mem/packet_queue.cc
+++ b/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();
-            drainManager = NULL;
+            signalDrainDone();
         }
     }
 }
@@ -244,14 +245,15 @@ PacketQueue::processSendEvent()
     sendDeferredPacket();
 }
 
-unsigned int
-PacketQueue::drain(DrainManager *dm)
+DrainState
+PacketQueue::drain()
 {
-    if (transmitList.empty())
-        return 0;
-    DPRINTF(Drain, "PacketQueue not drained\n");
-    drainManager = dm;
-    return 1;
+    if (transmitList.empty()) {
+        return DrainState::Drained;
+    } else {
+        DPRINTF(Drain, "PacketQueue not drained\n");
+        return DrainState::Draining;
+    }
 }
 
 ReqPacketQueue::ReqPacketQueue(EventManager& _em, MasterPort& _masterPort,
diff --git a/src/mem/packet_queue.hh b/src/mem/packet_queue.hh
index 6584fe997..4eabd1bc4 100644
--- a/src/mem/packet_queue.hh
+++ b/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
diff --git a/src/mem/ruby/structures/RubyMemoryControl.cc b/src/mem/ruby/structures/RubyMemoryControl.cc
index 7e07a82bc..6e40d415c 100644
--- a/src/mem/ruby/structures/RubyMemoryControl.cc
+++ b/src/mem/ruby/structures/RubyMemoryControl.cc
@@ -640,14 +640,14 @@ RubyMemoryControl::executeCycle()
     }
 }
 
-unsigned int
-RubyMemoryControl::drain(DrainManager *dm)
+DrainState
+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.
diff --git a/src/mem/ruby/structures/RubyMemoryControl.hh b/src/mem/ruby/structures/RubyMemoryControl.hh
index 78bc74ad5..d69e8f435 100644
--- a/src/mem/ruby/structures/RubyMemoryControl.hh
+++ b/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; };
diff --git a/src/mem/ruby/system/DMASequencer.cc b/src/mem/ruby/system/DMASequencer.cc
index 0dd28d91c..05592b231 100644
--- a/src/mem/ruby/system/DMASequencer.cc
+++ b/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();
-            // Clear the drain manager once we're done with it.
-            drainManager = NULL;
+            signalDrainDone();
         }
     }
 }
 
-unsigned int
-DMASequencer::drain(DrainManager *dm)
+DrainState
+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) {
-        drainManager = dm;
-
+    if (outstandingCount() > 0) {
         DPRINTF(Drain, "DMASequencer not drained\n");
-        setDrainState(DrainState::Draining);
-        return 1;
+        return DrainState::Draining;
+    } else {
+        return DrainState::Drained;
     }
-
-    drainManager = NULL;
-    setDrainState(DrainState::Drained);
-    return 0;
 }
 
 void
diff --git a/src/mem/ruby/system/DMASequencer.hh b/src/mem/ruby/system/DMASequencer.hh
index bcf586acf..ee5a4b698 100644
--- a/src/mem/ruby/system/DMASequencer.hh
+++ b/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;
diff --git a/src/mem/ruby/system/RubyPort.cc b/src/mem/ruby/system/RubyPort.cc
index 5818056e9..d34cb3c2f 100644
--- a/src/mem/ruby/system/RubyPort.cc
+++ b/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();
-            // Clear the drain manager once we're done with it.
-            drainManager = NULL;
+            signalDrainDone();
         }
     }
 }
 
-unsigned int
-RubyPort::drain(DrainManager *dm)
+DrainState
+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;
-
-    // Set status
-    if (need_drain) {
-        drainManager = dm;
-
+    if (outstandingCount() > 0) {
         DPRINTF(Drain, "RubyPort not drained\n");
-        setDrainState(DrainState::Draining);
-        return 1;
+        return DrainState::Draining;
+    } else {
+        return DrainState::Drained;
     }
-
-    drainManager = NULL;
-    setDrainState(DrainState::Drained);
-    return 0;
 }
 
 void
diff --git a/src/mem/ruby/system/RubyPort.hh b/src/mem/ruby/system/RubyPort.hh
index ff1b2af04..1fbaeba7b 100644
--- a/src/mem/ruby/system/RubyPort.hh
+++ b/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.
diff --git a/src/mem/ruby/system/Sequencer.cc b/src/mem/ruby/system/Sequencer.cc
index c33f5f819..0a48817ca 100644
--- a/src/mem/ruby/system/Sequencer.cc
+++ b/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));
     }
 
diff --git a/src/mem/simple_mem.cc b/src/mem/simple_mem.cc
index f36301809..8938b2a55 100644
--- a/src/mem/simple_mem.cc
+++ b/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) {
-            DPRINTF(Drain, "Drainng of SimpleMemory complete\n");
-            drainManager->signalDrainDone();
-            drainManager = NULL;
+        } else if (drainState() == DrainState::Draining) {
+            DPRINTF(Drain, "Draining of SimpleMemory complete\n");
+            signalDrainDone();
         }
     }
 }
@@ -233,23 +232,15 @@ SimpleMemory::getSlavePort(const std::string &if_name, PortID idx)
     }
 }
 
-unsigned int
-SimpleMemory::drain(DrainManager *dm)
+DrainState
+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");
-     }
-
-    if (count)
-        setDrainState(DrainState::Draining);
-    else
-        setDrainState(DrainState::Drained);
-    return count;
+        return DrainState::Draining;
+    } else {
+        return DrainState::Drained;
+    }
 }
 
 SimpleMemory::MemoryPort::MemoryPort(const std::string& _name,
diff --git a/src/mem/simple_mem.hh b/src/mem/simple_mem.hh
index f69633690..98c41623f 100644
--- a/src/mem/simple_mem.hh
+++ b/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);
diff --git a/src/mem/xbar.cc b/src/mem/xbar.cc
index bc649581b..7316fbc21 100644
--- a/src/mem/xbar.cc
+++ b/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();
-        // Clear the drain event once we're done with it.
-        drainManager = NULL;
+        signalDrainDone();
     }
 }
 
@@ -587,18 +585,18 @@ BaseXBar::regStats()
 }
 
 template <typename SrcType, typename DstType>
-unsigned int
-BaseXBar::Layer<SrcType,DstType>::drain(DrainManager *dm)
+DrainState
+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 1;
+        return DrainState::Draining;
+    } else {
+        return DrainState::Drained;
     }
-    return 0;
 }
 
 template <typename SrcType, typename DstType>
diff --git a/src/mem/xbar.hh b/src/mem/xbar.hh
index 547e138e9..79e9a5380 100644
--- a/src/mem/xbar.hh
+++ b/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.
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;