sim: Move the draining interface into a separate base class

This patch moves the draining interface from SimObject to a separate
class that can be used by any object needing draining. However,
objects not visible to the Python code (i.e., objects not deriving
from SimObject) still depend on their parents informing them when to
drain. This patch also gets rid of the CountedDrainEvent (which isn't
really an event) and replaces it with a DrainManager.

64 files changed, 689 insertions, 442 deletions

diff --git a/src/arch/arm/table_walker.cc b/src/arch/arm/table_walker.cc
index a10eb4a20..e4c6209d6 100644
--- a/src/arch/arm/table_walker.cc
+++ b/src/arch/arm/table_walker.cc
@@ -51,7 +51,7 @@
 using namespace ArmISA;
 
 TableWalker::TableWalker(const Params *p)
-    : MemObject(p), port(this, params()->sys), drainEvent(NULL),
+    : MemObject(p), port(this, params()->sys), drainManager(NULL),
       tlb(NULL), currState(NULL), pending(false),
       masterId(p->sys->getMasterId(name())),
       numSquashable(p->num_squash_per_cycle),
@@ -68,30 +68,30 @@ TableWalker::~TableWalker()
 void
 TableWalker::completeDrain()
 {
-    if (drainEvent && stateQueueL1.empty() && stateQueueL2.empty() &&
+    if (drainManager && stateQueueL1.empty() && stateQueueL2.empty() &&
         pendingQueue.empty()) {
-        changeState(Drained);
+        setDrainState(Drainable::Drained);
         DPRINTF(Drain, "TableWalker done draining, processing drain event\n");
-        drainEvent->process();
-        drainEvent = NULL;
+        drainManager->signalDrainDone();
+        drainManager = NULL;
     }
 }
 
 unsigned int
-TableWalker::drain(Event *de)
+TableWalker::drain(DrainManager *dm)
 {
-    unsigned int count = port.drain(de);
+    unsigned int count = port.drain(dm);
 
     if (stateQueueL1.empty() && stateQueueL2.empty() &&
         pendingQueue.empty()) {
-        changeState(Drained);
+        setDrainState(Drainable::Drained);
         DPRINTF(Drain, "TableWalker free, no need to drain\n");
 
         // table walker is drained, but its ports may still need to be drained
         return count;
     } else {
-        drainEvent = de;
-        changeState(Draining);
+        drainManager = dm;
+        setDrainState(Drainable::Draining);
         DPRINTF(Drain, "TableWalker not drained\n");
 
         // return port drain count plus the table walker itself needs to drain
@@ -101,9 +101,9 @@ TableWalker::drain(Event *de)
 }
 
 void
-TableWalker::resume()
+TableWalker::drainResume()
 {
-    MemObject::resume();
+    Drainable::drainResume();
     if ((params()->sys->getMemoryMode() == Enums::timing) && currState) {
         delete currState;
         currState = NULL;
diff --git a/src/arch/arm/table_walker.hh b/src/arch/arm/table_walker.hh
index 22c5d03b4..23464f56d 100644
--- a/src/arch/arm/table_walker.hh
+++ b/src/arch/arm/table_walker.hh
@@ -364,7 +364,7 @@ class TableWalker : public MemObject
     SnoopingDmaPort port;
 
     /** If we're draining keep the drain event around until we're drained */
-    Event *drainEvent;
+    DrainManager *drainManager;
 
     /** TLB that is initiating these table walks */
     TLB *tlb;
@@ -397,8 +397,8 @@ class TableWalker : public MemObject
 
     /** Checks if all state is cleared and if so, completes drain */
     void completeDrain();
-    virtual unsigned int drain(Event *de);
-    virtual void resume();
+    unsigned int drain(DrainManager *dm);
+    void drainResume();
     virtual BaseMasterPort& getMasterPort(const std::string &if_name,
                                           PortID idx = InvalidPortID);
 
diff --git a/src/cpu/o3/cpu.cc b/src/cpu/o3/cpu.cc
index fdd45fdda..bc5f096e6 100644
--- a/src/cpu/o3/cpu.cc
+++ b/src/cpu/o3/cpu.cc
@@ -619,7 +619,7 @@ FullO3CPU<Impl>::tick()
 
     if (!tickEvent.scheduled()) {
         if (_status == SwitchedOut ||
-            getState() == SimObject::Drained) {
+            getDrainState() == Drainable::Drained) {
             DPRINTF(O3CPU, "Switched out!\n");
             // increment stat
             lastRunningCycle = curCycle();
@@ -1077,7 +1077,7 @@ template <class Impl>
 void
 FullO3CPU<Impl>::serialize(std::ostream &os)
 {
-    SimObject::State so_state = SimObject::getState();
+    Drainable::State so_state(getDrainState());
     SERIALIZE_ENUM(so_state);
     BaseCPU::serialize(os);
     nameOut(os, csprintf("%s.tickEvent", name()));
@@ -1100,7 +1100,7 @@ template <class Impl>
 void
 FullO3CPU<Impl>::unserialize(Checkpoint *cp, const std::string &section)
 {
-    SimObject::State so_state;
+    Drainable::State so_state;
     UNSERIALIZE_ENUM(so_state);
     BaseCPU::unserialize(cp, section);
     tickEvent.unserialize(cp, csprintf("%s.tickEvent", section));
@@ -1120,7 +1120,7 @@ FullO3CPU<Impl>::unserialize(Checkpoint *cp, const std::string &section)
 
 template <class Impl>
 unsigned int
-FullO3CPU<Impl>::drain(Event *drain_event)
+FullO3CPU<Impl>::drain(DrainManager *drain_manager)
 {
     DPRINTF(O3CPU, "Switching out\n");
 
@@ -1137,12 +1137,12 @@ FullO3CPU<Impl>::drain(Event *drain_event)
 
     // Wake the CPU and record activity so everything can drain out if
     // the CPU was not able to immediately drain.
-    if (getState() != SimObject::Drained) {
-        // A bit of a hack...set the drainEvent after all the drain()
+    if (getDrainState() != Drainable::Drained) {
+        // A bit of a hack...set the drainManager after all the drain()
         // calls have been made, that way if all of the stages drain
         // immediately, the signalDrained() function knows not to call
         // process on the drain event.
-        drainEvent = drain_event;
+        drainManager = drain_manager;
 
         wakeCPU();
         activityRec.activity();
@@ -1157,7 +1157,7 @@ FullO3CPU<Impl>::drain(Event *drain_event)
 
 template <class Impl>
 void
-FullO3CPU<Impl>::resume()
+FullO3CPU<Impl>::drainResume()
 {
     fetch.resume();
     decode.resume();
@@ -1165,7 +1165,7 @@ FullO3CPU<Impl>::resume()
     iew.resume();
     commit.resume();
 
-    changeState(SimObject::Running);
+    setDrainState(Drainable::Running);
 
     if (_status == SwitchedOut)
         return;
@@ -1185,14 +1185,14 @@ FullO3CPU<Impl>::signalDrained()
         if (tickEvent.scheduled())
             tickEvent.squash();
 
-        changeState(SimObject::Drained);
+        setDrainState(Drainable::Drained);
 
         BaseCPU::switchOut();
 
-        if (drainEvent) {
+        if (drainManager) {
             DPRINTF(Drain, "CPU done draining, processing drain event\n");
-            drainEvent->process();
-            drainEvent = NULL;
+            drainManager->signalDrainDone();
+            drainManager = NULL;
         }
     }
     assert(drainCount <= 5);
diff --git a/src/cpu/o3/cpu.hh b/src/cpu/o3/cpu.hh
index 076cce0fb..1f9a8da6c 100644
--- a/src/cpu/o3/cpu.hh
+++ b/src/cpu/o3/cpu.hh
@@ -431,10 +431,10 @@ class FullO3CPU : public BaseO3CPU
 
     /** Starts draining the CPU's pipeline of all instructions in
      * order to stop all memory accesses. */
-    virtual unsigned int drain(Event *drain_event);
+    unsigned int drain(DrainManager *drain_manager);
 
     /** Resumes execution after a drain. */
-    virtual void resume();
+    void drainResume();
 
     /** Signals to this CPU that a stage has completed switching out. */
     void signalDrained();
@@ -730,8 +730,8 @@ class FullO3CPU : public BaseO3CPU
     /** Pointer to the system. */
     System *system;
 
-    /** Event to call process() on once draining has completed. */
-    Event *drainEvent;
+    /** DrainManager to notify when draining has completed. */
+    DrainManager *drainManager;
 
     /** Counter of how many stages have completed draining. */
     int drainCount;
diff --git a/src/cpu/simple/atomic.cc b/src/cpu/simple/atomic.cc
index 2d7afd221..e63d998a7 100644
--- a/src/cpu/simple/atomic.cc
+++ b/src/cpu/simple/atomic.cc
@@ -123,7 +123,7 @@ AtomicSimpleCPU::~AtomicSimpleCPU()
 void
 AtomicSimpleCPU::serialize(ostream &os)
 {
-    SimObject::State so_state = SimObject::getState();
+    Drainable::State so_state(getDrainState());
     SERIALIZE_ENUM(so_state);
     SERIALIZE_SCALAR(locked);
     BaseSimpleCPU::serialize(os);
@@ -134,15 +134,22 @@ AtomicSimpleCPU::serialize(ostream &os)
 void
 AtomicSimpleCPU::unserialize(Checkpoint *cp, const string &section)
 {
-    SimObject::State so_state;
+    Drainable::State so_state;
     UNSERIALIZE_ENUM(so_state);
     UNSERIALIZE_SCALAR(locked);
     BaseSimpleCPU::unserialize(cp, section);
     tickEvent.unserialize(cp, csprintf("%s.tickEvent", section));
 }
 
+unsigned int
+AtomicSimpleCPU::drain(DrainManager *drain_manager)
+{
+    setDrainState(Drainable::Drained);
+    return 0;
+}
+
 void
-AtomicSimpleCPU::resume()
+AtomicSimpleCPU::drainResume()
 {
     if (_status == Idle || _status == SwitchedOut)
         return;
@@ -150,7 +157,7 @@ AtomicSimpleCPU::resume()
     DPRINTF(SimpleCPU, "Resume\n");
     assert(system->getMemoryMode() == Enums::atomic);
 
-    changeState(SimObject::Running);
+    setDrainState(Drainable::Running);
     if (thread->status() == ThreadContext::Active) {
         if (!tickEvent.scheduled())
             schedule(tickEvent, nextCycle());
@@ -161,7 +168,7 @@ AtomicSimpleCPU::resume()
 void
 AtomicSimpleCPU::switchOut()
 {
-    assert(_status == Running || _status == Idle);
+    assert(_status == BaseSimpleCPU::Running || _status == Idle);
     _status = SwitchedOut;
 
     tickEvent.squash();
@@ -180,13 +187,14 @@ AtomicSimpleCPU::takeOverFrom(BaseCPU *oldCPU)
     ThreadID size = threadContexts.size();
     for (ThreadID i = 0; i < size; ++i) {
         ThreadContext *tc = threadContexts[i];
-        if (tc->status() == ThreadContext::Active && _status != Running) {
-            _status = Running;
+        if (tc->status() == ThreadContext::Active &&
+            _status != BaseSimpleCPU::Running) {
+            _status = BaseSimpleCPU::Running;
             schedule(tickEvent, nextCycle());
             break;
         }
     }
-    if (_status != Running) {
+    if (_status != BaseSimpleCPU::Running) {
         _status = Idle;
     }
     assert(threadContexts.size() == 1);
@@ -212,7 +220,7 @@ AtomicSimpleCPU::activateContext(ThreadID thread_num, Cycles delay)
 
     //Make sure ticks are still on multiples of cycles
     schedule(tickEvent, clockEdge(delay));
-    _status = Running;
+    _status = BaseSimpleCPU::Running;
 }
 
 
@@ -227,7 +235,7 @@ AtomicSimpleCPU::suspendContext(ThreadID thread_num)
     if (_status == Idle)
         return;
 
-    assert(_status == Running);
+    assert(_status == BaseSimpleCPU::Running);
 
     // tick event may not be scheduled if this gets called from inside
     // an instruction's execution, e.g. "quiesce"
diff --git a/src/cpu/simple/atomic.hh b/src/cpu/simple/atomic.hh
index d67ab67a5..94d2de081 100644
--- a/src/cpu/simple/atomic.hh
+++ b/src/cpu/simple/atomic.hh
@@ -122,7 +122,9 @@ class AtomicSimpleCPU : public BaseSimpleCPU
 
     virtual void serialize(std::ostream &os);
     virtual void unserialize(Checkpoint *cp, const std::string &section);
-    virtual void resume();
+
+    unsigned int drain(DrainManager *drain_manager);
+    void drainResume();
 
     void switchOut();
     void takeOverFrom(BaseCPU *oldCPU);
diff --git a/src/cpu/simple/timing.cc b/src/cpu/simple/timing.cc
index 15b277d53..41764302d 100644
--- a/src/cpu/simple/timing.cc
+++ b/src/cpu/simple/timing.cc
@@ -92,7 +92,7 @@ TimingSimpleCPU::TimingSimpleCPU(TimingSimpleCPUParams *p)
 {
     _status = Idle;
 
-    changeState(SimObject::Running);
+    setDrainState(Drainable::Running);
     system->totalNumInsts = 0;
 }
 
@@ -104,7 +104,7 @@ TimingSimpleCPU::~TimingSimpleCPU()
 void
 TimingSimpleCPU::serialize(ostream &os)
 {
-    SimObject::State so_state = SimObject::getState();
+    Drainable::State so_state(getDrainState());
     SERIALIZE_ENUM(so_state);
     BaseSimpleCPU::serialize(os);
 }
@@ -112,29 +112,31 @@ TimingSimpleCPU::serialize(ostream &os)
 void
 TimingSimpleCPU::unserialize(Checkpoint *cp, const string &section)
 {
-    SimObject::State so_state;
+    Drainable::State so_state;
     UNSERIALIZE_ENUM(so_state);
     BaseSimpleCPU::unserialize(cp, section);
 }
 
 unsigned int
-TimingSimpleCPU::drain(Event *drain_event)
+TimingSimpleCPU::drain(DrainManager *drain_manager)
 {
     // TimingSimpleCPU is ready to drain if it's not waiting for
     // an access to complete.
-    if (_status == Idle || _status == Running || _status == SwitchedOut) {
-        changeState(SimObject::Drained);
+    if (_status == Idle ||
+        _status == BaseSimpleCPU::Running ||
+        _status == SwitchedOut) {
+        setDrainState(Drainable::Drained);
         return 0;
     } else {
-        changeState(SimObject::Draining);
-        drainEvent = drain_event;
+        setDrainState(Drainable::Draining);
+        drainManager = drain_manager;
         DPRINTF(Drain, "CPU not drained\n");
         return 1;
     }
 }
 
 void
-TimingSimpleCPU::resume()
+TimingSimpleCPU::drainResume()
 {
     DPRINTF(SimpleCPU, "Resume\n");
     if (_status != SwitchedOut && _status != Idle) {
@@ -146,13 +148,13 @@ TimingSimpleCPU::resume()
         schedule(fetchEvent, nextCycle());
     }
 
-    changeState(SimObject::Running);
+    setDrainState(Drainable::Running);
 }
 
 void
 TimingSimpleCPU::switchOut()
 {
-    assert(_status == Running || _status == Idle);
+    assert(_status == BaseSimpleCPU::Running || _status == Idle);
     _status = SwitchedOut;
     numCycles += curCycle() - previousCycle;
 
@@ -172,13 +174,14 @@ TimingSimpleCPU::takeOverFrom(BaseCPU *oldCPU)
     // running and schedule its tick event.
     for (int i = 0; i < threadContexts.size(); ++i) {
         ThreadContext *tc = threadContexts[i];
-        if (tc->status() == ThreadContext::Active && _status != Running) {
-            _status = Running;
+        if (tc->status() == ThreadContext::Active &&
+            _status != BaseSimpleCPU::Running) {
+            _status = BaseSimpleCPU::Running;
             break;
         }
     }
 
-    if (_status != Running) {
+    if (_status != BaseSimpleCPU::Running) {
         _status = Idle;
     }
     assert(threadContexts.size() == 1);
@@ -197,7 +200,7 @@ TimingSimpleCPU::activateContext(ThreadID thread_num, Cycles delay)
     assert(_status == Idle);
 
     notIdleFraction++;
-    _status = Running;
+    _status = BaseSimpleCPU::Running;
 
     // kick things off by initiating the fetch of the next instruction
     schedule(fetchEvent, clockEdge(delay));
@@ -215,7 +218,7 @@ TimingSimpleCPU::suspendContext(ThreadID thread_num)
     if (_status == Idle)
         return;
 
-    assert(_status == Running);
+    assert(_status == BaseSimpleCPU::Running);
 
     // just change status to Idle... if status != Running,
     // completeInst() will not initiate fetch of next instruction.
@@ -330,7 +333,7 @@ TimingSimpleCPU::translationFault(Fault fault)
 
     postExecute();
 
-    if (getState() == SimObject::Draining) {
+    if (getDrainState() == Drainable::Draining) {
         advancePC(fault);
         completeDrain();
     } else {
@@ -511,7 +514,7 @@ TimingSimpleCPU::writeMem(uint8_t *data, unsigned size,
 void
 TimingSimpleCPU::finishTranslation(WholeTranslationState *state)
 {
-    _status = Running;
+    _status = BaseSimpleCPU::Running;
 
     if (state->getFault() != NoFault) {
         if (state->isPrefetch()) {
@@ -552,7 +555,7 @@ TimingSimpleCPU::fetch()
     bool needToFetch = !isRomMicroPC(pcState.microPC()) && !curMacroStaticInst;
 
     if (needToFetch) {
-        _status = Running;
+        _status = BaseSimpleCPU::Running;
         Request *ifetch_req = new Request();
         ifetch_req->setThreadContext(_cpuId, /* thread ID */ 0);
         setupFetchRequest(ifetch_req);
@@ -592,7 +595,7 @@ TimingSimpleCPU::sendFetch(Fault fault, RequestPtr req, ThreadContext *tc)
         DPRINTF(SimpleCPU, "Translation of addr %#x faulted\n", req->getVaddr());
         delete req;
         // fetch fault: advance directly to next instruction (fault handler)
-        _status = Running;
+        _status = BaseSimpleCPU::Running;
         advanceInst(fault);
     }
 
@@ -620,7 +623,7 @@ TimingSimpleCPU::advanceInst(Fault fault)
     if (!stayAtPC)
         advancePC(fault);
 
-    if (_status == Running) {
+    if (_status == BaseSimpleCPU::Running) {
         // kick off fetch of next instruction... callback from icache
         // response will cause that instruction to be executed,
         // keeping the CPU running.
@@ -641,12 +644,12 @@ TimingSimpleCPU::completeIfetch(PacketPtr pkt)
     assert(!pkt || !pkt->isError());
     assert(_status == IcacheWaitResponse);
 
-    _status = Running;
+    _status = BaseSimpleCPU::Running;
 
     numCycles += curCycle() - previousCycle;
     previousCycle = curCycle();
 
-    if (getState() == SimObject::Draining) {
+    if (getDrainState() == Drainable::Draining) {
         if (pkt) {
             delete pkt->req;
             delete pkt;
@@ -664,7 +667,7 @@ TimingSimpleCPU::completeIfetch(PacketPtr pkt)
         // If we're not running now the instruction will complete in a dcache
         // response callback or the instruction faulted and has started an
         // ifetch
-        if (_status == Running) {
+        if (_status == BaseSimpleCPU::Running) {
             if (fault != NoFault && traceData) {
                 // If there was a fault, we shouldn't trace this instruction.
                 delete traceData;
@@ -778,7 +781,7 @@ TimingSimpleCPU::completeDataAccess(PacketPtr pkt)
         }
     }
 
-    _status = Running;
+    _status = BaseSimpleCPU::Running;
 
     Fault fault = curStaticInst->completeAcc(pkt, this, traceData);
 
@@ -802,7 +805,7 @@ TimingSimpleCPU::completeDataAccess(PacketPtr pkt)
 
     postExecute();
 
-    if (getState() == SimObject::Draining) {
+    if (getDrainState() == Drainable::Draining) {
         advancePC(fault);
         completeDrain();
 
@@ -817,8 +820,8 @@ void
 TimingSimpleCPU::completeDrain()
 {
     DPRINTF(Drain, "CPU done draining, processing drain event\n");
-    changeState(SimObject::Drained);
-    drainEvent->process();
+    setDrainState(Drainable::Drained);
+    drainManager->signalDrainDone();
 }
 
 bool
diff --git a/src/cpu/simple/timing.hh b/src/cpu/simple/timing.hh
index a2570abe6..e7f5122d0 100644
--- a/src/cpu/simple/timing.hh
+++ b/src/cpu/simple/timing.hh
@@ -45,7 +45,7 @@ class TimingSimpleCPU : public BaseSimpleCPU
     virtual void init();
 
   public:
-    Event *drainEvent;
+    DrainManager *drainManager;
 
   private:
 
@@ -109,7 +109,7 @@ class TimingSimpleCPU : public BaseSimpleCPU
         void
         markDelayed()
         {
-            assert(cpu->_status == Running);
+            assert(cpu->_status == BaseSimpleCPU::Running);
             cpu->_status = ITBWaitResponse;
         }
 
@@ -249,8 +249,8 @@ class TimingSimpleCPU : public BaseSimpleCPU
     virtual void serialize(std::ostream &os);
     virtual void unserialize(Checkpoint *cp, const std::string &section);
 
-    virtual unsigned int drain(Event *drain_event);
-    virtual void resume();
+    unsigned int drain(DrainManager *drain_manager);
+    void drainResume();
 
     void switchOut();
     void takeOverFrom(BaseCPU *oldCPU);
diff --git a/src/cpu/testers/traffic_gen/traffic_gen.cc b/src/cpu/testers/traffic_gen/traffic_gen.cc
index af7ff89f4..e0657822c 100644
--- a/src/cpu/testers/traffic_gen/traffic_gen.cc
+++ b/src/cpu/testers/traffic_gen/traffic_gen.cc
@@ -110,11 +110,11 @@ TrafficGen::initState()
 }
 
 unsigned int
-TrafficGen::drain(Event* drain_event)
+TrafficGen::drain(DrainManager *dm)
 {
     // @todo we should also stop putting new requests in the queue and
     // either interrupt the current state or wait for a transition
-    return port.drain(drain_event);
+    return port.drain(dm);
 }
 
 void
diff --git a/src/cpu/testers/traffic_gen/traffic_gen.hh b/src/cpu/testers/traffic_gen/traffic_gen.hh
index 5f59be82c..4e94df548 100644
--- a/src/cpu/testers/traffic_gen/traffic_gen.hh
+++ b/src/cpu/testers/traffic_gen/traffic_gen.hh
@@ -604,7 +604,7 @@ class TrafficGen : public MemObject
 
     void initState();
 
-    unsigned int drain(Event *drain_event);
+    unsigned int drain(DrainManager *dm);
 
     void serialize(std::ostream &os);
 
diff --git a/src/dev/copy_engine.cc b/src/dev/copy_engine.cc
index 799e9f96a..d6162b689 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), drainEvent(NULL),
+    completionDataReg(0), nextState(Idle), drainManager(NULL),
     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->getState() == SimObject::Running)
+        if (ce->getDrainState() == Drainable::Running)
             fetchDescriptor(cr.descChainAddr);
     } else if (cr.command.append_dma()) {
         if (!busy) {
             nextState = AddressFetch;
-            if (ce->getState() == SimObject::Running)
+            if (ce->getDrainState() == Drainable::Running)
                 fetchNextAddr(lastDescriptorAddr);
         } else
             refreshNext = true;
@@ -637,41 +637,41 @@ CopyEngine::CopyEngineChannel::fetchAddrComplete()
 bool
 CopyEngine::CopyEngineChannel::inDrain()
 {
-    if (ce->getState() == SimObject::Draining) {
+    if (ce->getDrainState() == Drainable::Draining) {
         DPRINTF(Drain, "CopyEngine done draining, processing drain event\n");
-        assert(drainEvent);
-        drainEvent->process();
-        drainEvent = NULL;
+        assert(drainManager);
+        drainManager->signalDrainDone();
+        drainManager = NULL;
     }
 
-    return ce->getState() != SimObject::Running;
+    return ce->getDrainState() != Drainable::Running;
 }
 
 unsigned int
-CopyEngine::CopyEngineChannel::drain(Event *de)
+CopyEngine::CopyEngineChannel::drain(DrainManager *dm)
 {
-    if (nextState == Idle || ce->getState() != SimObject::Running)
+    if (nextState == Idle || ce->getDrainState() != Drainable::Running)
         return 0;
     unsigned int count = 1;
-    count += cePort.drain(de);
+    count += cePort.drain(dm);
 
     DPRINTF(Drain, "CopyEngineChannel not drained\n");
-    drainEvent = de;
+    this->drainManager = dm;
     return count;
 }
 
 unsigned int
-CopyEngine::drain(Event *de)
+CopyEngine::drain(DrainManager *dm)
 {
     unsigned int count;
-    count = pioPort.drain(de) + dmaPort.drain(de) + configPort.drain(de);
+    count = pioPort.drain(dm) + dmaPort.drain(dm) + configPort.drain(dm);
     for (int x = 0;x < chan.size(); x++)
-        count += chan[x]->drain(de);
+        count += chan[x]->drain(dm);
 
     if (count)
-        changeState(Draining);
+        setDrainState(Draining);
     else
-        changeState(Drained);
+        setDrainState(Drained);
 
     DPRINTF(Drain, "CopyEngine not drained\n");
     return count;
@@ -760,16 +760,16 @@ CopyEngine::CopyEngineChannel::restartStateMachine()
 }
 
 void
-CopyEngine::resume()
+CopyEngine::drainResume()
 {
-    SimObject::resume();
+    Drainable::drainResume();
     for (int x = 0;x < chan.size(); x++)
-        chan[x]->resume();
+        chan[x]->drainResume();
 }
 
 
 void
-CopyEngine::CopyEngineChannel::resume()
+CopyEngine::CopyEngineChannel::drainResume()
 {
     DPRINTF(DMACopyEngine, "Restarting state machine at state %d\n", nextState);
     restartStateMachine();
diff --git a/src/dev/copy_engine.hh b/src/dev/copy_engine.hh
index 9a0cb0628..c4b07c79d 100644
--- a/src/dev/copy_engine.hh
+++ b/src/dev/copy_engine.hh
@@ -55,11 +55,12 @@
 #include "dev/copy_engine_defs.hh"
 #include "dev/pcidev.hh"
 #include "params/CopyEngine.hh"
+#include "sim/drain.hh"
 #include "sim/eventq.hh"
 
 class CopyEngine : public PciDev
 {
-    class CopyEngineChannel
+    class CopyEngineChannel : public Drainable
     {
       private:
         DmaPort cePort;
@@ -91,7 +92,7 @@ class CopyEngine : public PciDev
 
         ChannelState nextState;
 
-        Event *drainEvent;
+        DrainManager *drainManager;
       public:
         CopyEngineChannel(CopyEngine *_ce, int cid);
         virtual ~CopyEngineChannel();
@@ -106,8 +107,9 @@ class CopyEngine : public PciDev
         void channelRead(PacketPtr pkt, Addr daddr, int size);
         void channelWrite(PacketPtr pkt, Addr daddr, int size);
 
-        unsigned int drain(Event *de);
-        void resume();
+        unsigned int drain(DrainManager *drainManger);
+        void drainResume();
+
         void serialize(std::ostream &os);
         void unserialize(Checkpoint *cp, const std::string &section);
 
@@ -205,8 +207,9 @@ class CopyEngine : public PciDev
 
     virtual void serialize(std::ostream &os);
     virtual void unserialize(Checkpoint *cp, const std::string &section);
-    virtual unsigned int drain(Event *de);
-    virtual void resume();
+
+    unsigned int drain(DrainManager *drainManger);
+    void drainResume();
 };
 
 #endif //__DEV_COPY_ENGINE_HH__
diff --git a/src/dev/dma_device.cc b/src/dev/dma_device.cc
index 1aa4a8647..952d6f622 100644
--- a/src/dev/dma_device.cc
+++ b/src/dev/dma_device.cc
@@ -51,7 +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), drainEvent(NULL),
+      pendingCount(0), drainManager(NULL),
       inRetry(false)
 { }
 
@@ -98,9 +98,9 @@ DmaPort::handleResp(PacketPtr pkt, Tick delay)
     delete pkt;
 
     // we might be drained at this point, if so signal the drain event
-    if (pendingCount == 0 && drainEvent) {
-        drainEvent->process();
-        drainEvent = NULL;
+    if (pendingCount == 0 && drainManager) {
+        drainManager->signalDrainDone();
+        drainManager = NULL;
     }
 }
 
@@ -128,22 +128,22 @@ DmaDevice::init()
 }
 
 unsigned int
-DmaDevice::drain(Event *de)
+DmaDevice::drain(DrainManager *dm)
 {
-    unsigned int count = pioPort.drain(de) + dmaPort.drain(de);
+    unsigned int count = pioPort.drain(dm) + dmaPort.drain(dm);
     if (count)
-        changeState(Draining);
+        setDrainState(Drainable::Draining);
     else
-        changeState(Drained);
+        setDrainState(Drainable::Drained);
     return count;
 }
 
 unsigned int
-DmaPort::drain(Event *de)
+DmaPort::drain(DrainManager *dm)
 {
     if (pendingCount == 0)
         return 0;
-    drainEvent = de;
+    drainManager = dm;
     DPRINTF(Drain, "DmaPort not drained\n");
     return 1;
 }
diff --git a/src/dev/dma_device.hh b/src/dev/dma_device.hh
index cd328f3d6..3b4bb522d 100644
--- a/src/dev/dma_device.hh
+++ b/src/dev/dma_device.hh
@@ -48,6 +48,7 @@
 
 #include "dev/io_device.hh"
 #include "params/DmaDevice.hh"
+#include "sim/drain.hh"
 
 class DmaPort : public MasterPort
 {
@@ -123,7 +124,7 @@ class DmaPort : public MasterPort
 
     /** If we need to drain, keep the drain event around until we're done
      * here.*/
-    Event *drainEvent;
+    DrainManager *drainManager;
 
     /** If the port is currently waiting for a retry before it can
      * send whatever it is that it's sending. */
@@ -146,7 +147,7 @@ class DmaPort : public MasterPort
     bool dmaPending() const { return pendingCount > 0; }
 
     unsigned cacheBlockSize() const { return peerBlockSize(); }
-    unsigned int drain(Event *de);
+    unsigned int drain(DrainManager *drainManger);
 };
 
 class DmaDevice : public PioDevice
@@ -175,7 +176,7 @@ class DmaDevice : public PioDevice
 
     virtual void init();
 
-    virtual unsigned int drain(Event *de);
+    unsigned int drain(DrainManager *drainManger);
 
     unsigned cacheBlockSize() const { return dmaPort.cacheBlockSize(); }
 
diff --git a/src/dev/i8254xGBe.cc b/src/dev/i8254xGBe.cc
index 159cc4726..9a3ddaeb7 100644
--- a/src/dev/i8254xGBe.cc
+++ b/src/dev/i8254xGBe.cc
@@ -57,7 +57,7 @@ using namespace iGbReg;
 using namespace Net;
 
 IGbE::IGbE(const Params *p)
-    : EtherDevice(p), etherInt(NULL),  drainEvent(NULL),
+    : EtherDevice(p), etherInt(NULL),  drainManager(NULL),
       useFlowControl(p->use_flow_control),
       rxFifo(p->rx_fifo_size), txFifo(p->tx_fifo_size), rxTick(false),
       txTick(false), txFifoTick(false), rxDmaPacket(false), pktOffset(0),
@@ -588,7 +588,7 @@ IGbE::write(PacketPtr pkt)
       case REG_RDT:
         regs.rdt = val;
         DPRINTF(EthernetSM, "RXS: RDT Updated.\n");
-        if (getState() == SimObject::Running) {
+        if (getDrainState() == Drainable::Running) {
             DPRINTF(EthernetSM, "RXS: RDT Fetching Descriptors!\n");
             rxDescCache.fetchDescriptors();
         } else {
@@ -628,7 +628,7 @@ IGbE::write(PacketPtr pkt)
       case REG_TDT:
         regs.tdt = val;
         DPRINTF(EthernetSM, "TXS: TX Tail pointer updated\n");
-        if (getState() == SimObject::Running) {
+        if (getDrainState() == Drainable::Running) {
             DPRINTF(EthernetSM, "TXS: TDT Fetching Descriptors!\n");
             txDescCache.fetchDescriptors();
         } else {
@@ -906,7 +906,7 @@ void
 IGbE::DescCache<T>::writeback1()
 {
     // If we're draining delay issuing this DMA
-    if (igbe->getState() != SimObject::Running) {
+    if (igbe->getDrainState() != Drainable::Running) {
         igbe->schedule(wbDelayEvent, curTick() + igbe->wbDelay);
         return;
     }
@@ -987,7 +987,7 @@ void
 IGbE::DescCache<T>::fetchDescriptors1()
 {
     // If we're draining delay issuing this DMA
-    if (igbe->getState() != SimObject::Running) {
+    if (igbe->getDrainState() != Drainable::Running) {
         igbe->schedule(fetchDelayEvent, curTick() + igbe->fetchDelay);
         return;
     }
@@ -1493,7 +1493,7 @@ IGbE::RxDescCache::pktComplete()
 void
 IGbE::RxDescCache::enableSm()
 {
-    if (!igbe->drainEvent) {
+    if (!igbe->drainManager) {
         igbe->rxTick = true;
         igbe->restartClock();
     }
@@ -2029,7 +2029,7 @@ IGbE::TxDescCache::packetAvailable()
 void
 IGbE::TxDescCache::enableSm()
 {
-    if (!igbe->drainEvent) {
+    if (!igbe->drainManager) {
         igbe->txTick = true;
         igbe->restartClock();
     }
@@ -2049,19 +2049,19 @@ void
 IGbE::restartClock()
 {
     if (!tickEvent.scheduled() && (rxTick || txTick || txFifoTick) &&
-        getState() == SimObject::Running)
+        getDrainState() == Drainable::Running)
         schedule(tickEvent, clockEdge(Cycles(1)));
 }
 
 unsigned int
-IGbE::drain(Event *de)
+IGbE::drain(DrainManager *dm)
 {
     unsigned int count;
-    count = pioPort.drain(de) + dmaPort.drain(de);
+    count = pioPort.drain(dm) + dmaPort.drain(dm);
     if (rxDescCache.hasOutstandingEvents() ||
         txDescCache.hasOutstandingEvents()) {
         count++;
-        drainEvent = de;
+        drainManager = dm;
     }
 
     txFifoTick = false;
@@ -2073,17 +2073,17 @@ IGbE::drain(Event *de)
 
     if (count) {
         DPRINTF(Drain, "IGbE not drained\n");
-        changeState(Draining);
+        setDrainState(Drainable::Draining);
     } else
-        changeState(Drained);
+        setDrainState(Drainable::Drained);
 
     return count;
 }
 
 void
-IGbE::resume()
+IGbE::drainResume()
 {
-    SimObject::resume();
+    Drainable::drainResume();
 
     txFifoTick = true;
     txTick = true;
@@ -2096,7 +2096,7 @@ IGbE::resume()
 void
 IGbE::checkDrain()
 {
-    if (!drainEvent)
+    if (!drainManager)
         return;
 
     txFifoTick = false;
@@ -2105,8 +2105,8 @@ IGbE::checkDrain()
     if (!rxDescCache.hasOutstandingEvents() &&
         !txDescCache.hasOutstandingEvents()) {
         DPRINTF(Drain, "IGbE done draining, processing drain event\n");
-        drainEvent->process();
-        drainEvent = NULL;
+        drainManager->signalDrainDone();
+        drainManager = NULL;
     }
 }
 
@@ -2130,7 +2130,7 @@ IGbE::txStateMachine()
         bool success =
 #endif
             txFifo.push(txPacket);
-        txFifoTick = true && !drainEvent;
+        txFifoTick = true && !drainManager;
         assert(success);
         txPacket = NULL;
         anBegin("TXS", "Desc Writeback");
@@ -2229,7 +2229,7 @@ IGbE::ethRxPkt(EthPacketPtr pkt)
     }
 
     // restart the state machines if they are stopped
-    rxTick = true && !drainEvent;
+    rxTick = true && !drainManager;
     if ((rxTick || txTick) && !tickEvent.scheduled()) {
         DPRINTF(EthernetSM,
                 "RXS: received packet into fifo, starting ticking\n");
@@ -2442,8 +2442,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 && !drainEvent;
-    if (txDescCache.descLeft() != 0 && !drainEvent)
+    txFifoTick = true && !drainManager;
+    if (txDescCache.descLeft() != 0 && !drainManager)
         txTick = true;
 
     restartClock();
diff --git a/src/dev/i8254xGBe.hh b/src/dev/i8254xGBe.hh
index 099cd0d11..b8099fb1c 100644
--- a/src/dev/i8254xGBe.hh
+++ b/src/dev/i8254xGBe.hh
@@ -68,7 +68,7 @@ class IGbE : public EtherDevice
     uint16_t flash[iGbReg::EEPROM_SIZE];
 
     // The drain event if we have one
-    Event *drainEvent;
+    DrainManager *drainManager;
 
     // cached parameters from params struct
     bool useFlowControl;
@@ -347,7 +347,7 @@ class IGbE : public EtherDevice
         virtual void updateHead(long h) { igbe->regs.rdh(h); }
         virtual void enableSm();
         virtual void fetchAfterWb() {
-            if (!igbe->rxTick && igbe->getState() == SimObject::Running)
+            if (!igbe->rxTick && igbe->getDrainState() == Drainable::Running)
                 fetchDescriptors();
         }
 
@@ -409,7 +409,7 @@ class IGbE : public EtherDevice
         virtual void enableSm();
         virtual void actionAfterWb();
         virtual void fetchAfterWb() {
-            if (!igbe->txTick && igbe->getState() == SimObject::Running)
+            if (!igbe->txTick && igbe->getDrainState() == Drainable::Running)
                 fetchDescriptors();
         }
         
@@ -535,8 +535,9 @@ class IGbE : public EtherDevice
 
     virtual void serialize(std::ostream &os);
     virtual void unserialize(Checkpoint *cp, const std::string &section);
-    virtual unsigned int drain(Event *de);
-    virtual void resume();
+
+    unsigned int drain(DrainManager *dm);
+    void drainResume();
 
 };
 
diff --git a/src/dev/ide_disk.cc b/src/dev/ide_disk.cc
index f0c8c8668..6c5ccdd86 100644
--- a/src/dev/ide_disk.cc
+++ b/src/dev/ide_disk.cc
@@ -323,7 +323,7 @@ IdeDisk::doDmaTransfer()
         panic("Inconsistent DMA transfer state: dmaState = %d devState = %d\n",
               dmaState, devState);
 
-    if (ctrl->dmaPending() || ctrl->getState() != SimObject::Running) {
+    if (ctrl->dmaPending() || ctrl->getDrainState() != Drainable::Running) {
         schedule(dmaTransferEvent, curTick() + DMA_BACKOFF_PERIOD);
         return;
     } else
@@ -404,7 +404,7 @@ IdeDisk::doDmaRead()
                 curPrd.getByteCount(), TheISA::PageBytes);
 
     }
-    if (ctrl->dmaPending() || ctrl->getState() != SimObject::Running) {
+    if (ctrl->dmaPending() || ctrl->getDrainState() != Drainable::Running) {
         schedule(dmaReadWaitEvent, curTick() + DMA_BACKOFF_PERIOD);
         return;
     } else if (!dmaReadCG->done()) {
@@ -481,7 +481,7 @@ IdeDisk::doDmaWrite()
         dmaWriteCG = new ChunkGenerator(curPrd.getBaseAddr(),
                 curPrd.getByteCount(), TheISA::PageBytes);
     }
-    if (ctrl->dmaPending() || ctrl->getState() != SimObject::Running) {
+    if (ctrl->dmaPending() || ctrl->getDrainState() != Drainable::Running) {
         schedule(dmaWriteWaitEvent, curTick() + DMA_BACKOFF_PERIOD);
         DPRINTF(IdeDisk, "doDmaWrite: rescheduling\n");
         return;
diff --git a/src/dev/io_device.cc b/src/dev/io_device.cc
index 0cc1324f5..988f8344a 100644
--- a/src/dev/io_device.cc
+++ b/src/dev/io_device.cc
@@ -89,14 +89,14 @@ PioDevice::getSlavePort(const std::string &if_name, PortID idx)
 }
 
 unsigned int
-PioDevice::drain(Event *de)
+PioDevice::drain(DrainManager *dm)
 {
     unsigned int count;
-    count = pioPort.drain(de);
+    count = pioPort.drain(dm);
     if (count)
-        changeState(Draining);
+        setDrainState(Drainable::Draining);
     else
-        changeState(Drained);
+        setDrainState(Drainable::Drained);
     return count;
 }
 
diff --git a/src/dev/io_device.hh b/src/dev/io_device.hh
index bd6a26d14..848b8f0ba 100644
--- a/src/dev/io_device.hh
+++ b/src/dev/io_device.hh
@@ -125,7 +125,7 @@ class PioDevice : public MemObject
 
     virtual void init();
 
-    virtual unsigned int drain(Event *de);
+    unsigned int drain(DrainManager *drainManger);
 
     virtual BaseSlavePort &getSlavePort(const std::string &if_name,
                                         PortID idx = InvalidPortID);
diff --git a/src/dev/ns_gige.cc b/src/dev/ns_gige.cc
index 0af9fbfc5..90eb14acd 100644
--- a/src/dev/ns_gige.cc
+++ b/src/dev/ns_gige.cc
@@ -1069,7 +1069,7 @@ NSGigE::doRxDmaRead()
     assert(rxDmaState == dmaIdle || rxDmaState == dmaReadWaiting);
     rxDmaState = dmaReading;
 
-    if (dmaPending() || getState() != Running)
+    if (dmaPending() || getDrainState() != Drainable::Running)
         rxDmaState = dmaReadWaiting;
     else
         dmaRead(rxDmaAddr, rxDmaLen, &rxDmaReadEvent, (uint8_t*)rxDmaData);
@@ -1100,7 +1100,7 @@ NSGigE::doRxDmaWrite()
     assert(rxDmaState == dmaIdle || rxDmaState == dmaWriteWaiting);
     rxDmaState = dmaWriting;
 
-    if (dmaPending() || getState() != Running)
+    if (dmaPending() || getDrainState() != Running)
         rxDmaState = dmaWriteWaiting;
     else
         dmaWrite(rxDmaAddr, rxDmaLen, &rxDmaWriteEvent, (uint8_t*)rxDmaData);
@@ -1518,7 +1518,7 @@ NSGigE::doTxDmaRead()
     assert(txDmaState == dmaIdle || txDmaState == dmaReadWaiting);
     txDmaState = dmaReading;
 
-    if (dmaPending() || getState() != Running)
+    if (dmaPending() || getDrainState() != Running)
         txDmaState = dmaReadWaiting;
     else
         dmaRead(txDmaAddr, txDmaLen, &txDmaReadEvent, (uint8_t*)txDmaData);
@@ -1549,7 +1549,7 @@ NSGigE::doTxDmaWrite()
     assert(txDmaState == dmaIdle || txDmaState == dmaWriteWaiting);
     txDmaState = dmaWriting;
 
-    if (dmaPending() || getState() != Running)
+    if (dmaPending() || getDrainState() != Running)
         txDmaState = dmaWriteWaiting;
     else
         dmaWrite(txDmaAddr, txDmaLen, &txDmaWriteEvent, (uint8_t*)txDmaData);
@@ -2112,9 +2112,9 @@ NSGigE::recvPacket(EthPacketPtr packet)
 
 
 void
-NSGigE::resume()
+NSGigE::drainResume()
 {
-    SimObject::resume();
+    Drainable::drainResume();
 
     // During drain we could have left the state machines in a waiting state and
     // they wouldn't get out until some other event occured to kick them.
diff --git a/src/dev/ns_gige.hh b/src/dev/ns_gige.hh
index f4d0171d6..6d5068a2b 100644
--- a/src/dev/ns_gige.hh
+++ b/src/dev/ns_gige.hh
@@ -369,7 +369,7 @@ class NSGigE : public EtherDevBase
     virtual void serialize(std::ostream &os);
     virtual void unserialize(Checkpoint *cp, const std::string &section);
 
-    virtual void resume();
+    void drainResume();
 };
 
 /*
diff --git a/src/dev/pcidev.cc b/src/dev/pcidev.cc
index fb4aaa799..592852e29 100644
--- a/src/dev/pcidev.cc
+++ b/src/dev/pcidev.cc
@@ -157,14 +157,14 @@ PciDev::init()
 }
 
 unsigned int
-PciDev::drain(Event *de)
+PciDev::drain(DrainManager *dm)
 {
     unsigned int count;
-    count = pioPort.drain(de) + dmaPort.drain(de) + configPort.drain(de);
+    count = pioPort.drain(dm) + dmaPort.drain(dm) + configPort.drain(dm);
     if (count)
-        changeState(Draining);
+        setDrainState(Drainable::Draining);
     else
-        changeState(Drained);
+        setDrainState(Drainable::Drained);
     return count;
 }
 
diff --git a/src/dev/pcidev.hh b/src/dev/pcidev.hh
index 9994d2a2d..df468d4c6 100644
--- a/src/dev/pcidev.hh
+++ b/src/dev/pcidev.hh
@@ -216,7 +216,7 @@ class PciDev : public DmaDevice
     virtual void unserialize(Checkpoint *cp, const std::string &section);
 
 
-    virtual unsigned int drain(Event *de);
+    virtual unsigned int drain(DrainManager *dm);
 
     virtual BaseSlavePort &getSlavePort(const std::string &if_name,
                                         PortID idx = InvalidPortID)
diff --git a/src/dev/sinic.cc b/src/dev/sinic.cc
index 2d109cbbb..cfb7548eb 100644
--- a/src/dev/sinic.cc
+++ b/src/dev/sinic.cc
@@ -870,7 +870,7 @@ Device::rxKick()
         break;
 
       case rxBeginCopy:
-        if (dmaPending() || getState() != Running)
+        if (dmaPending() || getDrainState() != Drainable::Running)
             goto exit;
 
         rxDmaAddr = params()->platform->pciToDma(
@@ -1070,7 +1070,7 @@ Device::txKick()
         break;
 
       case txBeginCopy:
-        if (dmaPending() || getState() != Running)
+        if (dmaPending() || getDrainState() != Drainable::Running)
             goto exit;
 
         txDmaAddr = params()->platform->pciToDma(
@@ -1246,9 +1246,9 @@ Device::recvPacket(EthPacketPtr packet)
 }
 
 void
-Device::resume()
+Device::drainResume()
 {
-    SimObject::resume();
+    Drainable::drainResume();
 
     // During drain we could have left the state machines in a waiting state and
     // they wouldn't get out until some other event occured to kick them.
diff --git a/src/dev/sinic.hh b/src/dev/sinic.hh
index 8189ce39a..58f9e7253 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 resume();
+    virtual void drainResume();
 
     void prepareIO(int cpu, int index);
     void prepareRead(int cpu, int index);
diff --git a/src/mem/bus.cc b/src/mem/bus.cc
index a0db6e52a..ddbdcb136 100644
--- a/src/mem/bus.cc
+++ b/src/mem/bus.cc
@@ -161,7 +161,8 @@ BaseBus::calcPacketTiming(PacketPtr pkt)
 template <typename PortClass>
 BaseBus::Layer<PortClass>::Layer(BaseBus& _bus, const std::string& _name,
                                  Tick _clock) :
-    bus(_bus), _name(_name), state(IDLE), clock(_clock), drainEvent(NULL),
+    Drainable(),
+    bus(_bus), _name(_name), state(IDLE), clock(_clock), drainManager(NULL),
     releaseEvent(this)
 {
 }
@@ -266,12 +267,12 @@ BaseBus::Layer<PortClass>::releaseLayer()
         // busy, and in the latter case the bus may be released before
         // we see a retry from the destination
         retryWaiting();
-    } else if (drainEvent) {
-        DPRINTF(Drain, "Bus done draining, processing drain event\n");
+    } else if (drainManager) {
+        DPRINTF(Drain, "Bus done draining, signaling drain manager\n");
         //If we weren't able to drain before, do it now.
-        drainEvent->process();
+        drainManager->signalDrainDone();
         // Clear the drain event once we're done with it.
-        drainEvent = NULL;
+        drainManager = NULL;
     }
 }
 
@@ -522,14 +523,14 @@ BaseBus::deviceBlockSize() const
 
 template <typename PortClass>
 unsigned int
-BaseBus::Layer<PortClass>::drain(Event * de)
+BaseBus::Layer<PortClass>::drain(DrainManager *dm)
 {
     //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 (!retryList.empty() || state != IDLE) {
         DPRINTF(Drain, "Bus not drained\n");
-        drainEvent = de;
+        drainManager = dm;
         return 1;
     }
     return 0;
diff --git a/src/mem/bus.hh b/src/mem/bus.hh
index f3cbc9d24..19ffa020c 100644
--- a/src/mem/bus.hh
+++ b/src/mem/bus.hh
@@ -94,7 +94,7 @@ class BaseBus : public MemObject
      * whereas a response layer holds master ports.
      */
     template <typename PortClass>
-    class Layer
+    class Layer : public Drainable
     {
 
       public:
@@ -118,7 +118,7 @@ class BaseBus : public MemObject
          *
          * @return 1 if busy or waiting to retry, or 0 if idle
          */
-        unsigned int drain(Event *de);
+        unsigned int drain(DrainManager *dm);
 
         /**
          * Get the bus layer's name
@@ -206,8 +206,8 @@ class BaseBus : public MemObject
         /** the clock speed for the bus layer */
         Tick clock;
 
-        /** event for signalling when drained */
-        Event * drainEvent;
+        /** manager to signal when drained */
+        DrainManager *drainManager;
 
         /**
          * An array of ports that retry should be called
@@ -366,7 +366,7 @@ class BaseBus : public MemObject
     BaseSlavePort& getSlavePort(const std::string& if_name,
                                 PortID idx = InvalidPortID);
 
-    virtual unsigned int drain(Event *de) = 0;
+    virtual unsigned int drain(DrainManager *dm) = 0;
 
 };
 
diff --git a/src/mem/cache/base.cc b/src/mem/cache/base.cc
index a88749627..ad1c751bc 100644
--- a/src/mem/cache/base.cc
+++ b/src/mem/cache/base.cc
@@ -77,7 +77,7 @@ BaseCache::BaseCache(const Params *p)
       blocked(0),
       noTargetMSHR(NULL),
       missCount(p->max_miss_count),
-      drainEvent(NULL),
+      drainManager(NULL),
       addrRanges(p->addr_ranges.begin(), p->addr_ranges.end()),
       system(p->system)
 {
@@ -749,19 +749,19 @@ BaseCache::regStats()
 }
 
 unsigned int
-BaseCache::drain(Event *de)
+BaseCache::drain(DrainManager *dm)
 {
-    int count = memSidePort->drain(de) + cpuSidePort->drain(de);
+    int count = memSidePort->drain(dm) + cpuSidePort->drain(dm);
 
     // Set status
     if (count != 0) {
-        drainEvent = de;
+        drainManager = dm;
 
-        changeState(SimObject::Draining);
+        setDrainState(Drainable::Draining);
         DPRINTF(Drain, "Cache not drained\n");
         return count;
     }
 
-    changeState(SimObject::Drained);
+    setDrainState(Drainable::Drained);
     return 0;
 }
diff --git a/src/mem/cache/base.hh b/src/mem/cache/base.hh
index 42ade9b0b..ab13be771 100644
--- a/src/mem/cache/base.hh
+++ b/src/mem/cache/base.hh
@@ -269,7 +269,7 @@ class BaseCache : public MemObject
     Counter missCount;
 
     /** The drain event. */
-    Event *drainEvent;
+    DrainManager *drainManager;
 
     /**
      * The address range to which the cache responds on the CPU side.
@@ -542,7 +542,7 @@ class BaseCache : public MemObject
         // interesting again.
     }
 
-    virtual unsigned int drain(Event *de);
+    virtual unsigned int drain(DrainManager *dm);
 
     virtual bool inCache(Addr addr) = 0;
 
diff --git a/src/mem/coherent_bus.cc b/src/mem/coherent_bus.cc
index 98d86f3f0..b1ac6dbcf 100644
--- a/src/mem/coherent_bus.cc
+++ b/src/mem/coherent_bus.cc
@@ -508,10 +508,10 @@ CoherentBus::forwardFunctional(PacketPtr pkt, PortID exclude_slave_port_id)
 }
 
 unsigned int
-CoherentBus::drain(Event *de)
+CoherentBus::drain(DrainManager *dm)
 {
     // sum up the individual layers
-    return reqLayer.drain(de) + respLayer.drain(de) + snoopRespLayer.drain(de);
+    return reqLayer.drain(dm) + respLayer.drain(dm) + snoopRespLayer.drain(dm);
 }
 
 CoherentBus *
diff --git a/src/mem/coherent_bus.hh b/src/mem/coherent_bus.hh
index 89a759546..61406608b 100644
--- a/src/mem/coherent_bus.hh
+++ b/src/mem/coherent_bus.hh
@@ -299,7 +299,7 @@ class CoherentBus : public BaseBus
 
     CoherentBus(const CoherentBusParams *p);
 
-    unsigned int drain(Event *de);
+    unsigned int drain(DrainManager *dm);
 };
 
 #endif //__MEM_COHERENT_BUS_HH__
diff --git a/src/mem/noncoherent_bus.cc b/src/mem/noncoherent_bus.cc
index 237e8726b..f14f6e3d6 100644
--- a/src/mem/noncoherent_bus.cc
+++ b/src/mem/noncoherent_bus.cc
@@ -212,10 +212,10 @@ NoncoherentBus::recvFunctional(PacketPtr pkt, PortID slave_port_id)
 }
 
 unsigned int
-NoncoherentBus::drain(Event *de)
+NoncoherentBus::drain(DrainManager *dm)
 {
     // sum up the individual layers
-    return reqLayer.drain(de) + respLayer.drain(de);
+    return reqLayer.drain(dm) + respLayer.drain(dm);
 }
 
 NoncoherentBus*
diff --git a/src/mem/noncoherent_bus.hh b/src/mem/noncoherent_bus.hh
index 16cf7deda..a42c26b2e 100644
--- a/src/mem/noncoherent_bus.hh
+++ b/src/mem/noncoherent_bus.hh
@@ -207,7 +207,7 @@ class NoncoherentBus : public BaseBus
 
     NoncoherentBus(const NoncoherentBusParams *p);
 
-    unsigned int drain(Event *de);
+    unsigned int drain(DrainManager *dm);
 
 };
 
diff --git a/src/mem/packet_queue.cc b/src/mem/packet_queue.cc
index 4a4543f61..eb94cc397 100644
--- a/src/mem/packet_queue.cc
+++ b/src/mem/packet_queue.cc
@@ -48,7 +48,7 @@
 using namespace std;
 
 PacketQueue::PacketQueue(EventManager& _em, const std::string& _label)
-    : em(_em), sendEvent(this), drainEvent(NULL), label(_label),
+    : em(_em), sendEvent(this), drainManager(NULL), label(_label),
       waitingOnRetry(false)
 {
 }
@@ -173,11 +173,11 @@ PacketQueue::scheduleSend(Tick time)
             em.schedule(&sendEvent, std::max(nextReady, curTick() + 1));
     } else {
         // no more to send, so if we're draining, we may be done
-        if (drainEvent && transmitList.empty() && !sendEvent.scheduled()) {
+        if (drainManager && transmitList.empty() && !sendEvent.scheduled()) {
             DPRINTF(Drain, "PacketQueue done draining,"
                     "processing drain event\n");
-            drainEvent->process();
-            drainEvent = NULL;
+            drainManager->signalDrainDone();
+            drainManager = NULL;
         }
     }
 }
@@ -204,12 +204,12 @@ PacketQueue::processSendEvent()
 }
 
 unsigned int
-PacketQueue::drain(Event *de)
+PacketQueue::drain(DrainManager *dm)
 {
     if (transmitList.empty() && !sendEvent.scheduled())
         return 0;
     DPRINTF(Drain, "PacketQueue not drained\n");
-    drainEvent = de;
+    drainManager = dm;
     return 1;
 }
 
diff --git a/src/mem/packet_queue.hh b/src/mem/packet_queue.hh
index 0171eb9a3..2321ec4f2 100644
--- a/src/mem/packet_queue.hh
+++ b/src/mem/packet_queue.hh
@@ -57,12 +57,13 @@
 
 #include "mem/port.hh"
 #include "sim/eventq.hh"
+#include "sim/drain.hh"
 
 /**
  * A packet queue is a class that holds deferred packets and later
  * sends them using the associated slave port or master port.
  */
-class PacketQueue
+class PacketQueue : public Drainable
 {
   private:
     /** A deferred packet, buffered to transmit later. */
@@ -95,9 +96,9 @@ class PacketQueue
      **/
     EventWrapper<PacketQueue, &PacketQueue::processSendEvent> sendEvent;
 
-    /** If we need to drain, keep the drain event around until we're done
+    /** If we need to drain, keep the drain manager around until we're done
      * here.*/
-    Event *drainEvent;
+    DrainManager *drainManager;
 
   protected:
 
@@ -207,13 +208,7 @@ class PacketQueue
      */
     void retry();
 
-    /**
-     * Hook for draining the packet queue.
-     *
-     * @param de An event which is used to signal back to the caller
-     * @return A number indicating how many times process will be called
-     */
-    unsigned int drain(Event *de);
+    unsigned int drain(DrainManager *dm);
 };
 
 class MasterPacketQueue : public PacketQueue
diff --git a/src/mem/qport.hh b/src/mem/qport.hh
index b771f6984..dd5caa084 100644
--- a/src/mem/qport.hh
+++ b/src/mem/qport.hh
@@ -97,13 +97,7 @@ class QueuedSlavePort : public SlavePort
      * functional request. */
     bool checkFunctional(PacketPtr pkt) { return queue.checkFunctional(pkt); }
 
-    /**
-     * Hook for draining the queued port.
-     *
-     * @param de an event which is used to signal back to the caller
-     * @returns a number indicating how many times process will be called
-     */
-    unsigned int drain(Event *de) { return queue.drain(de); }
+    unsigned int drain(DrainManager *dm) { return queue.drain(dm); }
 };
 
 class QueuedMasterPort : public MasterPort
@@ -156,13 +150,7 @@ class QueuedMasterPort : public MasterPort
      * functional request. */
     bool checkFunctional(PacketPtr pkt) { return queue.checkFunctional(pkt); }
 
-    /**
-     * Hook for draining the queued port.
-     *
-     * @param de an event which is used to signal back to the caller
-     * @returns a number indicating how many times process will be called
-     */
-    unsigned int drain(Event *de) { return queue.drain(de); }
+    unsigned int drain(DrainManager *dm) { return queue.drain(dm); }
 };
 
 #endif // __MEM_QPORT_HH__
diff --git a/src/mem/ruby/system/MemoryControl.hh b/src/mem/ruby/system/MemoryControl.hh
index 8d15b8dec..5c6adb0ab 100644
--- a/src/mem/ruby/system/MemoryControl.hh
+++ b/src/mem/ruby/system/MemoryControl.hh
@@ -56,8 +56,6 @@ class MemoryControl : public ClockedObject, public Consumer
 
     ~MemoryControl();
 
-    unsigned int drain(Event *de) = 0;
-
     virtual void wakeup() = 0;
 
     virtual void setConsumer(Consumer* consumer_ptr) = 0;
diff --git a/src/mem/ruby/system/RubyMemoryControl.cc b/src/mem/ruby/system/RubyMemoryControl.cc
index c0e91c28b..620113719 100644
--- a/src/mem/ruby/system/RubyMemoryControl.cc
+++ b/src/mem/ruby/system/RubyMemoryControl.cc
@@ -684,7 +684,7 @@ RubyMemoryControl::executeCycle()
 }
 
 unsigned int
-RubyMemoryControl::drain(Event *de)
+RubyMemoryControl::drain(DrainManager *dm)
 {
     DPRINTF(RubyMemory, "MemoryController drain\n");
     if(m_event.scheduled()) {
diff --git a/src/mem/ruby/system/RubyMemoryControl.hh b/src/mem/ruby/system/RubyMemoryControl.hh
index 1f3a8acf5..53e8fabef 100644
--- a/src/mem/ruby/system/RubyMemoryControl.hh
+++ b/src/mem/ruby/system/RubyMemoryControl.hh
@@ -62,7 +62,7 @@ class RubyMemoryControl : public MemoryControl
 
     ~RubyMemoryControl();
 
-    unsigned int drain(Event *de);
+    unsigned int drain(DrainManager *dm);
 
     void wakeup();
 
diff --git a/src/mem/ruby/system/RubyPort.cc b/src/mem/ruby/system/RubyPort.cc
index 1259f0f15..dd9e9676e 100644
--- a/src/mem/ruby/system/RubyPort.cc
+++ b/src/mem/ruby/system/RubyPort.cc
@@ -53,7 +53,7 @@ RubyPort::RubyPort(const Params *p)
       m_mandatory_q_ptr(NULL),
       pio_port(csprintf("%s-pio-port", name()), this),
       m_usingRubyTester(p->using_ruby_tester), m_request_cnt(0),
-      drainEvent(NULL), ruby_system(p->ruby_system), system(p->system),
+      drainManager(NULL), ruby_system(p->ruby_system), system(p->system),
       waitingOnSequencer(false), access_phys_mem(p->access_phys_mem)
 {
     assert(m_version != -1);
@@ -343,36 +343,36 @@ void
 RubyPort::testDrainComplete()
 {
     //If we weren't able to drain before, we might be able to now.
-    if (drainEvent != NULL) {
+    if (drainManager != NULL) {
         unsigned int drainCount = outstandingCount();
         DPRINTF(Drain, "Drain count: %u\n", drainCount);
         if (drainCount == 0) {
-            DPRINTF(Drain, "RubyPort done draining, processing drain event\n");
-            drainEvent->process();
-            // Clear the drain event once we're done with it.
-            drainEvent = NULL;
+            DPRINTF(Drain, "RubyPort done draining, signaling drain done\n");
+            drainManager->signalDrainDone();
+            // Clear the drain manager once we're done with it.
+            drainManager = NULL;
         }
     }
 }
 
 unsigned int
-RubyPort::getChildDrainCount(Event *de)
+RubyPort::getChildDrainCount(DrainManager *dm)
 {
     int count = 0;
 
     if (pio_port.isConnected()) {
-        count += pio_port.drain(de);
+        count += pio_port.drain(dm);
         DPRINTF(Config, "count after pio check %d\n", count);
     }
 
     for (CpuPortIter p = slave_ports.begin(); p != slave_ports.end(); ++p) {
-        count += (*p)->drain(de);
+        count += (*p)->drain(dm);
         DPRINTF(Config, "count after slave port check %d\n", count);
     }
 
     for (std::vector<PioPort*>::iterator p = master_ports.begin();
          p != master_ports.end(); ++p) {
-        count += (*p)->drain(de);
+        count += (*p)->drain(dm);
         DPRINTF(Config, "count after master port check %d\n", count);
     }
 
@@ -382,7 +382,7 @@ RubyPort::getChildDrainCount(Event *de)
 }
 
 unsigned int
-RubyPort::drain(Event *de)
+RubyPort::drain(DrainManager *dm)
 {
     if (isDeadlockEventScheduled()) {
         descheduleDeadlockEvent();
@@ -390,28 +390,28 @@ RubyPort::drain(Event *de)
 
     //
     // If the RubyPort is not empty, then it needs to clear all outstanding
-    // requests before it should call drainEvent->process()
+    // requests before it should call drainManager->signalDrainDone()
     //
     DPRINTF(Config, "outstanding count %d\n", outstandingCount());
     bool need_drain = outstandingCount() > 0;
 
     //
     // Also, get the number of child ports that will also need to clear
-    // their buffered requests before they call drainEvent->process()
+    // their buffered requests before they call drainManager->signalDrainDone()
     //
-    unsigned int child_drain_count = getChildDrainCount(de);
+    unsigned int child_drain_count = getChildDrainCount(dm);
 
     // Set status
     if (need_drain) {
-        drainEvent = de;
+        drainManager = dm;
 
         DPRINTF(Drain, "RubyPort not drained\n");
-        changeState(SimObject::Draining);
+        setDrainState(Drainable::Draining);
         return child_drain_count + 1;
     }
 
-    drainEvent = NULL;
-    changeState(SimObject::Drained);
+    drainManager = NULL;
+    setDrainState(Drainable::Drained);
     return child_drain_count;
 }
 
diff --git a/src/mem/ruby/system/RubyPort.hh b/src/mem/ruby/system/RubyPort.hh
index ab09bd90a..98bcede44 100644
--- a/src/mem/ruby/system/RubyPort.hh
+++ b/src/mem/ruby/system/RubyPort.hh
@@ -142,7 +142,7 @@ class RubyPort : public MemObject
     //
     void setController(AbstractController* _cntrl) { m_controller = _cntrl; }
     int getId() { return m_version; }
-    unsigned int drain(Event *de);
+    unsigned int drain(DrainManager *dm);
 
   protected:
     const std::string m_name;
@@ -166,7 +166,7 @@ class RubyPort : public MemObject
         }
     }
 
-    unsigned int getChildDrainCount(Event *de);
+    unsigned int getChildDrainCount(DrainManager *dm);
 
     uint16_t m_port_id;
     uint64_t m_request_cnt;
@@ -176,7 +176,7 @@ class RubyPort : public MemObject
     std::vector<M5Port*> slave_ports;
     std::vector<PioPort*> master_ports;
 
-    Event *drainEvent;
+    DrainManager *drainManager;
 
     RubySystem* ruby_system;
     System* system;
diff --git a/src/mem/ruby/system/Sequencer.cc b/src/mem/ruby/system/Sequencer.cc
index 9b6ef35cd..a45dfc98d 100644
--- a/src/mem/ruby/system/Sequencer.cc
+++ b/src/mem/ruby/system/Sequencer.cc
@@ -85,7 +85,7 @@ Sequencer::~Sequencer()
 void
 Sequencer::wakeup()
 {
-    assert(getState() != SimObject::Draining);
+    assert(getDrainState() != Drainable::Draining);
 
     // Check for deadlock of any of the requests
     Time current_time = g_system_ptr->getTime();
@@ -209,7 +209,8 @@ Sequencer::insertRequest(PacketPtr pkt, RubyRequestType request_type)
         (m_writeRequestTable.size() + m_readRequestTable.size()));
 
     // See if we should schedule a deadlock check
-    if (!deadlockCheckEvent.scheduled() && getState() != SimObject::Draining) {
+    if (!deadlockCheckEvent.scheduled() &&
+        getDrainState() != Drainable::Draining) {
         schedule(deadlockCheckEvent,
             g_system_ptr->clockPeriod() * m_deadlock_threshold + curTick());
     }
diff --git a/src/mem/simple_dram.cc b/src/mem/simple_dram.cc
index 0f6e9511c..42c97977a 100644
--- a/src/mem/simple_dram.cc
+++ b/src/mem/simple_dram.cc
@@ -51,7 +51,7 @@ SimpleDRAM::SimpleDRAM(const SimpleDRAMParams* p) :
     retryRdReq(false), retryWrReq(false),
     rowHitFlag(false), stopReads(false),
     writeEvent(this), respondEvent(this),
-    refreshEvent(this), nextReqEvent(this), drainEvent(NULL),
+    refreshEvent(this), nextReqEvent(this), drainManager(NULL),
     bytesPerCacheLine(0),
     linesPerRowBuffer(p->lines_per_rowbuffer),
     ranksPerChannel(p->ranks_per_channel),
@@ -346,9 +346,9 @@ SimpleDRAM::processWriteEvent()
 
     // if there is nothing left in any queue, signal a drain
     if (dramWriteQueue.empty() && dramReadQueue.empty() &&
-        dramRespQueue.empty () && drainEvent) {
-        drainEvent->process();
-        drainEvent = NULL;
+        dramRespQueue.empty () && drainManager) {
+        drainManager->signalDrainDone();
+        drainManager = NULL;
     }
 
     // Once you're done emptying the write queue, check if there's
@@ -595,9 +595,9 @@ SimpleDRAM::processRespondEvent()
      } else {
          // if there is nothing left in any queue, signal a drain
          if (dramWriteQueue.empty() && dramReadQueue.empty() &&
-             drainEvent) {
-             drainEvent->process();
-             drainEvent = NULL;
+             drainManager) {
+             drainManager->signalDrainDone();
+             drainManager = NULL;
          }
      }
 }
@@ -1197,22 +1197,22 @@ SimpleDRAM::getSlavePort(const string &if_name, PortID idx)
 }
 
 unsigned int
-SimpleDRAM::drain(Event *de)
+SimpleDRAM::drain(DrainManager *dm)
 {
-    unsigned int count = port.drain(de);
+    unsigned int count = port.drain(dm);
 
     // if there is anything in any of our internal queues, keep track
     // of that as well
     if (!(dramWriteQueue.empty() && dramReadQueue.empty() &&
           dramRespQueue.empty())) {
         ++count;
-        drainEvent = de;
+        drainManager = dm;
     }
 
     if (count)
-        changeState(Draining);
+        setDrainState(Drainable::Draining);
     else
-        changeState(Drained);
+        setDrainState(Drainable::Drained);
     return count;
 }
 
diff --git a/src/mem/simple_dram.hh b/src/mem/simple_dram.hh
index 74058afaa..373408c2a 100644
--- a/src/mem/simple_dram.hh
+++ b/src/mem/simple_dram.hh
@@ -341,10 +341,10 @@ class SimpleDRAM : public AbstractMemory
      */
     std::list<DRAMPacket*> dramRespQueue;
 
-    /** If we need to drain, keep the drain event around until we're done
+    /** If we need to drain, keep the drain manager around until we're done
      * here.
      */
-    Event *drainEvent;
+    DrainManager *drainManager;
 
     /**
      * Multi-dimensional vector of banks, first dimension is ranks,
@@ -459,7 +459,7 @@ class SimpleDRAM : public AbstractMemory
 
     SimpleDRAM(const SimpleDRAMParams* p);
 
-    unsigned int drain(Event* de);
+    unsigned int drain(DrainManager* dm);
 
     virtual BaseSlavePort& getSlavePort(const std::string& if_name,
                                         PortID idx = InvalidPortID);
diff --git a/src/mem/simple_mem.cc b/src/mem/simple_mem.cc
index c54e8e5ea..e78b57928 100644
--- a/src/mem/simple_mem.cc
+++ b/src/mem/simple_mem.cc
@@ -176,14 +176,14 @@ SimpleMemory::getSlavePort(const std::string &if_name, PortID idx)
 }
 
 unsigned int
-SimpleMemory::drain(Event *de)
+SimpleMemory::drain(DrainManager *dm)
 {
-    int count = port.drain(de);
+    int count = port.drain(dm);
 
     if (count)
-        changeState(Draining);
+        setDrainState(Drainable::Draining);
     else
-        changeState(Drained);
+        setDrainState(Drainable::Drained);
     return count;
 }
 
diff --git a/src/mem/simple_mem.hh b/src/mem/simple_mem.hh
index 7fd64db47..f1bad7d9f 100644
--- a/src/mem/simple_mem.hh
+++ b/src/mem/simple_mem.hh
@@ -123,7 +123,7 @@ class SimpleMemory : public AbstractMemory
     SimpleMemory(const SimpleMemoryParams *p);
     virtual ~SimpleMemory() { }
 
-    unsigned int drain(Event* de);
+    unsigned int drain(DrainManager *dm);
 
     virtual BaseSlavePort& getSlavePort(const std::string& if_name,
                                         PortID idx = InvalidPortID);
diff --git a/src/python/SConscript b/src/python/SConscript
index 751710665..d00432642 100644
--- a/src/python/SConscript
+++ b/src/python/SConscript
@@ -66,6 +66,7 @@ PySource('m5.util', 'm5/util/terminal.py')
 
 SwigSource('m5.internal', 'swig/core.i')
 SwigSource('m5.internal', 'swig/debug.i')
+SwigSource('m5.internal', 'swig/drain.i')
 SwigSource('m5.internal', 'swig/event.i')
 SwigSource('m5.internal', 'swig/pyobject.i')
 SwigSource('m5.internal', 'swig/range.i')
diff --git a/src/python/m5/SimObject.py b/src/python/m5/SimObject.py
index c01db2a80..3aea55f5f 100644
--- a/src/python/m5/SimObject.py
+++ b/src/python/m5/SimObject.py
@@ -123,7 +123,8 @@ class MetaSimObject(type):
                       'cxx_class' : str,
                       'cxx_type' : str,
                       'cxx_header' : str,
-                      'type' : str }
+                      'type' : str,
+                      'cxx_bases' : list }
     # Attributes that can be set any time
     keywords = { 'check' : FunctionType }
 
@@ -148,6 +149,8 @@ class MetaSimObject(type):
                 value_dict[key] = val
         if 'abstract' not in value_dict:
             value_dict['abstract'] = False
+        if 'cxx_bases' not in value_dict:
+            value_dict['cxx_bases'] = []
         cls_dict['_value_dict'] = value_dict
         cls = super(MetaSimObject, mcls).__new__(mcls, name, bases, cls_dict)
         if 'type' in value_dict:
@@ -414,6 +417,7 @@ class MetaSimObject(type):
         code('%module(package="m5.internal") param_$cls')
         code()
         code('%{')
+        code('#include "sim/sim_object.hh"')
         code('#include "params/$cls.hh"')
         for param in params:
             param.cxx_predecls(code)
@@ -458,7 +462,17 @@ using std::ptrdiff_t;
         code('%nodefault $classname;')
         code('class $classname')
         if cls._base:
-            code('    : public ${{cls._base.cxx_class}}')
+            bases = [ cls._base.cxx_class ] + cls.cxx_bases
+        else:
+            bases = cls.cxx_bases
+        base_first = True
+        for base in bases:
+            if base_first:
+                code('    : public ${{base}}')
+                base_first = False
+            else:
+                code('    , public ${{base}}')
+
         code('{')
         code('  public:')
         cls.export_methods(code)
@@ -581,30 +595,25 @@ class SimObject(object):
     abstract = True
     cxx_header = "sim/sim_object.hh"
 
+    cxx_bases = [ "Drainable" ]
+
     @classmethod
     def export_method_swig_predecls(cls, code):
         code('''
 %include <std_string.i>
+
+%import "python/swig/drain.i"
 ''')
 
     @classmethod
     def export_methods(cls, code):
         code('''
-    enum State {
-      Running,
-      Draining,
-      Drained
-    };
-
     void init();
     void loadState(Checkpoint *cp);
     void initState();
     void regStats();
     void resetStats();
     void startup();
-
-    unsigned int drain(Event *drain_event);
-    void resume();
 ''')
 
     # Initialize new instance.  For objects with SimObject-valued
diff --git a/src/python/m5/__init__.py b/src/python/m5/__init__.py
index 930609b6b..dc6c5a923 100644
--- a/src/python/m5/__init__.py
+++ b/src/python/m5/__init__.py
@@ -51,3 +51,4 @@ if internal:
     from event import *
     from main import main
     from simulate import *
+
diff --git a/src/python/m5/internal/__init__.py b/src/python/m5/internal/__init__.py
index ca09ab468..30090549a 100644
--- a/src/python/m5/internal/__init__.py
+++ b/src/python/m5/internal/__init__.py
@@ -31,3 +31,4 @@ import debug
 import event
 import stats
 import trace
+from drain import DrainManager, Drainable
diff --git a/src/python/m5/simulate.py b/src/python/m5/simulate.py
index 89b6b1e9d..df30f654a 100644
--- a/src/python/m5/simulate.py
+++ b/src/python/m5/simulate.py
@@ -169,19 +169,19 @@ def doDrain(root):
 # be drained.
 def drain(root):
     all_drained = False
-    drain_event = internal.event.createCountedDrain()
-    unready_objs = sum(obj.drain(drain_event) for obj in root.descendants())
+    dm = internal.drain.createDrainManager()
+    unready_objs = sum(obj.drain(dm) for obj in root.descendants())
     # If we've got some objects that can't drain immediately, then simulate
     if unready_objs > 0:
-        drain_event.setCount(unready_objs)
+        dm.setCount(unready_objs)
         simulate()
     else:
         all_drained = True
-    internal.event.cleanupCountedDrain(drain_event)
+    internal.drain.cleanupDrainManager(dm)
     return all_drained
 
 def resume(root):
-    for obj in root.descendants(): obj.resume()
+    for obj in root.descendants(): obj.drainResume()
 
 def checkpoint(dir):
     root = objects.Root.getInstance()
diff --git a/src/python/swig/drain.i b/src/python/swig/drain.i
new file mode 100644
index 000000000..4442db207
--- /dev/null
+++ b/src/python/swig/drain.i
@@ -0,0 +1,66 @@
+/*
+ * Copyright (c) 2012 ARM Limited
+ * All rights reserved
+ *
+ * The license below extends only to copyright in the software and shall
+ * not be construed as granting a license to any other intellectual
+ * property including but not limited to intellectual property relating
+ * to a hardware implementation of the functionality of the software
+ * licensed hereunder.  You may use the software subject to the license
+ * terms below provided that you ensure that this notice is replicated
+ * unmodified and in its entirety in all distributions of the software,
+ * modified or unmodified, in source code or in binary form.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met: redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer;
+ * redistributions in binary form must reproduce the above copyright
+ * notice, this list of conditions and the following disclaimer in the
+ * documentation and/or other materials provided with the distribution;
+ * neither the name of the copyright holders nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ * Authors: Andreas Sandberg
+ */
+
+%module(package="m5.internal") drain
+
+%{
+#include "sim/drain.hh"
+%}
+
+%nodefaultctor Drainable;
+
+%include "sim/drain.hh"
+
+%inline %{
+
+DrainManager *
+createDrainManager()
+{
+    return new DrainManager();
+}
+
+void
+cleanupDrainManager(DrainManager *drain_manager)
+{
+    assert(drain_manager);
+    assert(drain_manager->getCount() == 0);
+    delete drain_manager;
+}
+
+%}
diff --git a/src/python/swig/event.i b/src/python/swig/event.i
index 0af29e449..766dc2769 100644
--- a/src/python/swig/event.i
+++ b/src/python/swig/event.i
@@ -81,11 +81,6 @@
 // This must follow eventq.hh
 %include "python/swig/pyevent.hh"
 
-struct CountedDrainEvent : public Event
-{
-    void setCount(int _count);
-};
-
 // minimal definition of SimExitEvent interface to wrap
 class SimLoopExitEvent : public Event
 {
diff --git a/src/python/swig/pyevent.cc b/src/python/swig/pyevent.cc
index 0695ed2d3..4651d252b 100644
--- a/src/python/swig/pyevent.cc
+++ b/src/python/swig/pyevent.cc
@@ -65,22 +65,3 @@ PythonEvent::process()
     // reference count must be decremented.
     Py_DECREF(object);
 }
-
-CountedDrainEvent *
-createCountedDrain()
-{
-    return new CountedDrainEvent();
-}
-
-void
-cleanupCountedDrain(Event *counted_drain)
-{
-    CountedDrainEvent *event =
-        dynamic_cast<CountedDrainEvent *>(counted_drain);
-    if (event == NULL) {
-        fatal("Called cleanupCountedDrain() on an event that was not "
-              "a CountedDrainEvent.");
-    }
-    assert(event->getCount() == 0);
-    delete event;
-}
diff --git a/src/python/swig/pyevent.hh b/src/python/swig/pyevent.hh
index 9006a0404..f34fbd996 100644
--- a/src/python/swig/pyevent.hh
+++ b/src/python/swig/pyevent.hh
@@ -49,7 +49,4 @@ class PythonEvent : public Event
     virtual void process();
 };
 
-CountedDrainEvent *createCountedDrain();
-void cleanupCountedDrain(Event *counted_drain);
-
 #endif // __PYTHON_SWIG_PYEVENT_HH__
diff --git a/src/sim/SConscript b/src/sim/SConscript
index 16eeb36d3..42993b90f 100644
--- a/src/sim/SConscript
+++ b/src/sim/SConscript
@@ -44,6 +44,7 @@ Source('init.cc')
 Source('main.cc', main=True, skip_lib=True)
 Source('root.cc')
 Source('serialize.cc')
+Source('drain.cc')
 Source('sim_events.cc')
 Source('sim_object.cc')
 Source('simulate.cc')
diff --git a/src/sim/drain.cc b/src/sim/drain.cc
new file mode 100644
index 000000000..3daf762f6
--- /dev/null
+++ b/src/sim/drain.cc
@@ -0,0 +1,73 @@
+/*
+ * Copyright (c) 2012 ARM Limited
+ * All rights reserved
+ *
+ * The license below extends only to copyright in the software and shall
+ * not be construed as granting a license to any other intellectual
+ * property including but not limited to intellectual property relating
+ * to a hardware implementation of the functionality of the software
+ * licensed hereunder.  You may use the software subject to the license
+ * terms below provided that you ensure that this notice is replicated
+ * unmodified and in its entirety in all distributions of the software,
+ * modified or unmodified, in source code or in binary form.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met: redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer;
+ * redistributions in binary form must reproduce the above copyright
+ * notice, this list of conditions and the following disclaimer in the
+ * documentation and/or other materials provided with the distribution;
+ * neither the name of the copyright holders nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ * Authors: Andreas Sandberg
+ */
+
+#include "sim/drain.hh"
+#include "sim/sim_exit.hh"
+
+DrainManager::DrainManager()
+    : _count(0)
+{
+}
+
+DrainManager::~DrainManager()
+{
+}
+
+void
+DrainManager::drainCycleDone()
+{
+    exitSimLoop("Finished drain", 0);
+}
+
+
+
+Drainable::Drainable()
+    : _drainState(Running)
+{
+}
+
+Drainable::~Drainable()
+{
+}
+
+void
+Drainable::drainResume()
+{
+    _drainState = Running;
+}
diff --git a/src/sim/drain.hh b/src/sim/drain.hh
new file mode 100644
index 000000000..252022bc7
--- /dev/null
+++ b/src/sim/drain.hh
@@ -0,0 +1,216 @@
+/*
+ * Copyright (c) 2012 ARM Limited
+ * All rights reserved
+ *
+ * The license below extends only to copyright in the software and shall
+ * not be construed as granting a license to any other intellectual
+ * property including but not limited to intellectual property relating
+ * to a hardware implementation of the functionality of the software
+ * licensed hereunder.  You may use the software subject to the license
+ * terms below provided that you ensure that this notice is replicated
+ * unmodified and in its entirety in all distributions of the software,
+ * modified or unmodified, in source code or in binary form.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met: redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer;
+ * redistributions in binary form must reproduce the above copyright
+ * notice, this list of conditions and the following disclaimer in the
+ * documentation and/or other materials provided with the distribution;
+ * neither the name of the copyright holders nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ * Authors: Andreas Sandberg
+ */
+
+#ifndef __SIM_DRAIN_HH__
+#define __SIM_DRAIN_HH__
+
+#include <cassert>
+#include <vector>
+
+#include "base/flags.hh"
+
+class Event;
+
+/**
+ * This class coordinates draining of a System.
+ *
+ * When draining a System, we need to make sure that all SimObjects in
+ * that system have drained their state before declaring the operation
+ * to be successful. This class keeps track of how many objects are
+ * still in the process of draining their state. Once it determines
+ * that all objects have drained their state, it exits the simulation
+ * loop.
+ *
+ * @note A System might not be completely drained even though the
+ * DrainManager has caused the simulation loop to exit. Draining needs
+ * to be restarted until all Drainable objects declare that they don't
+ * need further simulation to be completely drained. See Drainable for
+ * more information.
+ */
+class DrainManager
+{
+  public:
+    DrainManager();
+    virtual ~DrainManager();
+
+    /**
+     * Get the number of objects registered with this DrainManager
+     * that are currently draining their state.
+     *
+     * @return Number of objects currently draining.
+     */
+    unsigned int getCount() const { return _count; }
+
+    void setCount(int count) { _count = count; }
+
+    /**
+     * Notify the DrainManager that a Drainable object has finished
+     * draining.
+     */
+    void signalDrainDone() {
+        assert(_count > 0);
+        if (--_count == 0)
+            drainCycleDone();
+    }
+
+  protected:
+    /**
+     * Callback when all registered Drainable objects have completed a
+     * drain cycle.
+     */
+    virtual void drainCycleDone();
+
+    /** Number of objects still draining. */
+    unsigned int _count;
+};
+
+/**
+ * Interface for objects that might require draining before
+ * checkpointing.
+ *
+ * An object's internal state needs to be drained when creating a
+ * checkpoint, switching between CPU models, or switching between
+ * timing models. Once the internal state has been drained from
+ * <i>all</i> objects in the system, the objects are serialized to
+ * disc or the configuration change takes place. The process works as
+ * follows (see simulate.py for details):
+ *
+ * <ol>
+ * <li>An instance of a DrainManager is created to keep track of how
+ *     many objects need to be drained. The object maintains an
+ *     internal counter that is decreased every time its
+ *     CountedDrainEvent::signalDrainDone() method is called. When the
+ *     counter reaches zero, the simulation is stopped.
+ *
+ * <li>Call Drainable::drain() for every object in the
+ *     system. Draining has completed if all of them return
+ *     zero. Otherwise, the sum of the return values is loaded into
+ *     the counter of the DrainManager. A pointer to the drain
+ *     manager is passed as an argument to the drain() method.
+ *
+ * <li>Continue simulation. When an object has finished draining its
+ *     internal state, it calls CountedDrainEvent::signalDrainDone()
+ *     on the manager. When the counter in the manager reaches zero,
+ *     the simulation stops.
+ *
+ * <li>Check if any object still needs draining, if so repeat the
+ *     process above.
+ *
+ * <li>Serialize objects, switch CPU model, or change timing model.
+ *
+ * <li>Call Drainable::drainResume() and continue the simulation.
+ * </ol>
+ *
+ */
+class Drainable
+{
+  public:
+    /**
+     * Object drain/handover states
+     *
+     * An object starts out in the Running state. When the simulator
+     * prepares to take a snapshot or prepares a CPU for handover, it
+     * calls the drain() method to transfer the object into the
+     * Draining or Drained state. If any object enters the Draining
+     * state (drain() returning >0), simulation continues until it all
+     * objects have entered the Drained state.
+     *
+     * Before resuming simulation, the simulator calls resume() to
+     * transfer the object to the Running state.
+     *
+     * \note Even though the state of an object (visible to the rest
+     * of the world through getState()) could be used to determine if
+     * all objects have entered the Drained state, the protocol is
+     * actually a bit more elaborate. See drain() for details.
+     */
+    enum State {
+        Running,  /** Running normally */
+        Draining, /** Draining buffers pending serialization/handover */
+        Drained   /** Buffers drained, ready for serialization/handover */
+    };
+
+    Drainable();
+    virtual ~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.
+     *
+     * @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().
+     *
+     * @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.
+     */
+    virtual unsigned int drain(DrainManager *drainManager) = 0;
+
+    /**
+     * Resume execution after a successful drain.
+     *
+     * @note This method is normally only called from the simulation
+     * scripts.
+     */
+    virtual void drainResume();
+
+    State getDrainState() const { return _drainState; }
+
+  protected:
+    void setDrainState(State new_state) { _drainState = new_state; }
+
+
+  private:
+    State _drainState;
+
+};
+
+#endif
diff --git a/src/sim/serialize.hh b/src/sim/serialize.hh
index 531b2e1cd..9ee34fe80 100644
--- a/src/sim/serialize.hh
+++ b/src/sim/serialize.hh
@@ -144,8 +144,14 @@ void fromSimObject(T &t, SimObject *s)
     fromSimObject(objptr, sptr);                        \
   } while (0)
 
-/*
+/**
  * Basic support for object serialization.
+ *
+ * @note Many objects that support serialization need to be put in a
+ * consistent state when serialization takes place. We refer to the
+ * action of forcing an object into a consistent state as
+ * 'draining'. Objects that need draining inherit from Drainable. See
+ * Drainable for more information.
  */
 class Serializable
 {
diff --git a/src/sim/sim_events.cc b/src/sim/sim_events.cc
index a77e8b103..2354e89f7 100644
--- a/src/sim/sim_events.cc
+++ b/src/sim/sim_events.cc
@@ -83,17 +83,6 @@ exitSimLoop(const std::string &message, int exit_code, Tick when, Tick repeat)
     mainEventQueue.schedule(event, when);
 }
 
-CountedDrainEvent::CountedDrainEvent()
-    : count(0)
-{ }
-
-void
-CountedDrainEvent::process()
-{
-    if (--count == 0)
-        exitSimLoop("Finished drain", 0);
-}
-
 //
 // constructor: automatically schedules at specified time
 //
diff --git a/src/sim/sim_object.cc b/src/sim/sim_object.cc
index 32e936ff2..345fb85cb 100644
--- a/src/sim/sim_object.cc
+++ b/src/sim/sim_object.cc
@@ -67,7 +67,6 @@ SimObject::SimObject(const Params *p)
 #endif
 
     simObjectList.push_back(this);
-    state = Running;
 }
 
 void
@@ -151,17 +150,12 @@ debugObjectBreak(const char *objs)
 #endif
 
 unsigned int
-SimObject::drain(Event *drain_event)
+SimObject::drain(DrainManager *drain_manager)
 {
-    state = Drained;
+    setDrainState(Drained);
     return 0;
 }
 
-void
-SimObject::resume()
-{
-    state = Running;
-}
 
 SimObject *
 SimObject::find(const char *name)
diff --git a/src/sim/sim_object.hh b/src/sim/sim_object.hh
index c1238e23f..f04289e2f 100644
--- a/src/sim/sim_object.hh
+++ b/src/sim/sim_object.hh
@@ -45,6 +45,7 @@
 
 #include "enums/MemoryMode.hh"
 #include "params/SimObject.hh"
+#include "sim/drain.hh"
 #include "sim/eventq.hh"
 #include "sim/serialize.hh"
 
@@ -72,40 +73,7 @@ class Event;
  *     </ul>
  * <li>SimObject::resetStats()
  * <li>SimObject::startup()
- * <li>SimObject::resume() if resuming from a checkpoint.
- * </ol>
- *
- * An object's internal state needs to be drained when creating a
- * checkpoint, switching between CPU models, or switching between
- * timing models. Once the internal state has been drained from
- * <i>all</i> objects in the system, the objects are serialized to
- * disc or the configuration change takes place. The process works as
- * follows (see simulate.py for details):
- *
- * <ol>
- * <li>An instance of a CountedDrainEvent is created to keep track of
- *     how many objects need to be drained. The object maintains an
- *     internal counter that is decreased every time its
- *     CountedDrainEvent::process() method is called. When the counter
- *     reaches zero, the simulation is stopped.
- *
- * <li>Call SimObject::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 CountedDrainEvent. A pointer of the drain
- *     event is passed as an argument to the drain() method.
- *
- * <li>Continue simulation. When an object has finished draining its
- *     internal state, it calls CountedDrainEvent::process() on the
- *     CountedDrainEvent. When counter in the CountedDrainEvent reaches
- *     zero, the simulation stops.
- *
- * <li>Check if any object still needs draining, if so repeat the
- *     process above.
- *
- * <li>Serialize objects, switch CPU model, or change timing model.
- *
- * <li>Call SimObject::resume() and continue the simulation.
+ * <li>Drainable::drainResume() if resuming from a checkpoint.
  * </ol>
  *
  * @note Whenever a method is called on all objects in the simulator's
@@ -114,42 +82,8 @@ class Event;
  * SimObject.py). This has the effect of calling the method on the
  * parent node <i>before</i> its children.
  */
-class SimObject : public EventManager, public Serializable
+class SimObject : public EventManager, public Serializable, public Drainable
 {
-  public:
-    /**
-     * Object drain/handover states
-     *
-     * An object starts out in the Running state. When the simulator
-     * prepares to take a snapshot or prepares a CPU for handover, it
-     * calls the drain() method to transfer the object into the
-     * Draining or Drained state. If any object enters the Draining
-     * state (drain() returning >0), simulation continues until it all
-     * objects have entered the Drained.
-     *
-     * The before resuming simulation, the simulator calls resume() to
-     * transfer the object to the Running state.
-     *
-     * \note Even though the state of an object (visible to the rest
-     * of the world through getState()) could be used to determine if
-     * all objects have entered the Drained state, the protocol is
-     * actually a bit more elaborate. See drain() for details.
-     */
-    enum State {
-        Running,  /** Running normally */
-        Draining, /** Draining buffers pending serialization/handover */
-        Drained   /** Buffers drained, ready for serialization/handover */
-    };
-
-  private:
-    State state;
-
-  protected:
-    void changeState(State new_state) { state = new_state; }
-
-  public:
-    State getState() { return state; }
-
   private:
     typedef std::vector<SimObject *> SimObjectList;
 
@@ -217,43 +151,16 @@ class SimObject : public EventManager, public Serializable
     virtual void startup();
 
     /**
-     * Serialize all SimObjects in the system.
+     * Provide a default implementation of the drain interface that
+     * simply returns 0 (draining completed) and sets the drain state
+     * to Drained.
      */
-    static void serializeAll(std::ostream &os);
+    unsigned int drain(DrainManager *drainManger);
 
     /**
-     * Determine if an object needs draining and register a drain
-     * event.
-     *
-     * When draining the state of an object, the simulator calls drain
-     * with a pointer to a drain event. 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.
-     *
-     * The default implementation simply switches to the Drained state
-     * and returns 0.
-     *
-     * @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().
-     *
-     * @param drain_event Event to use to inform the simulator when
-     * the draining has completed.
-     *
-     * @return 0 if the object is ready for serialization now, >0 if
-     * it needs further simulation.
-     */
-    virtual unsigned int drain(Event *drain_event);
-
-    /**
-     * Switch an object in the Drained stated into the Running state.
+     * Serialize all SimObjects in the system.
      */
-    virtual void resume();
+    static void serializeAll(std::ostream &os);
 
 #ifdef DEBUG
   public:
diff --git a/src/sim/system.cc b/src/sim/system.cc
index 5ec7f4b30..259ed3e88 100644
--- a/src/sim/system.cc
+++ b/src/sim/system.cc
@@ -181,7 +181,7 @@ System::getMasterPort(const std::string &if_name, PortID idx)
 void
 System::setMemoryMode(Enums::MemoryMode mode)
 {
-    assert(getState() == Drained);
+    assert(getDrainState() == Drainable::Drained);
     memoryMode = mode;
 }
 
@@ -328,10 +328,17 @@ System::isMemAddr(Addr addr) const
     return physmem.isMemAddr(addr);
 }
 
+unsigned int
+System::drain(DrainManager *dm)
+{
+    setDrainState(Drainable::Drained);
+    return 0;
+}
+
 void
-System::resume()
+System::drainResume()
 {
-    SimObject::resume();
+    Drainable::drainResume();
     totalNumInsts = 0;
 }
 
diff --git a/src/sim/system.hh b/src/sim/system.hh
index 645ef8af2..d1b79bbf4 100644
--- a/src/sim/system.hh
+++ b/src/sim/system.hh
@@ -382,7 +382,9 @@ class System : public MemObject
 
     void serialize(std::ostream &os);
     void unserialize(Checkpoint *cp, const std::string &section);
-    virtual void resume();
+
+    unsigned int drain(DrainManager *dm);
+    void drainResume();
 
   public:
     Counter totalNumInsts;