sim: Refactor and simplify the drain API

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

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

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

14 files changed, 74 insertions, 144 deletions

diff --git a/src/cpu/kvm/base.cc b/src/cpu/kvm/base.cc
index 30e984366..47cff5917 100644
--- a/src/cpu/kvm/base.cc
+++ b/src/cpu/kvm/base.cc
@@ -78,7 +78,6 @@ BaseKvmCPU::BaseKvmCPU(BaseKvmCPUParams *params)
       activeInstPeriod(0),
       perfControlledByTimer(params->usePerfOverflow),
       hostFactor(params->hostFactor),
-      drainManager(NULL),
       ctrInsts(0)
 {
     if (pageSize == -1)
@@ -282,11 +281,11 @@ BaseKvmCPU::unserializeThread(CheckpointIn &cp, ThreadID tid)
     threadContextDirty = true;
 }
 
-unsigned int
-BaseKvmCPU::drain(DrainManager *dm)
+DrainState
+BaseKvmCPU::drain()
 {
     if (switchedOut())
-        return 0;
+        return DrainState::Drained;
 
     DPRINTF(Drain, "BaseKvmCPU::drain\n");
     switch (_status) {
@@ -296,10 +295,8 @@ BaseKvmCPU::drain(DrainManager *dm)
         // of a different opinion. This may happen when the CPU been
         // notified of an event that hasn't been accepted by the vCPU
         // yet.
-        if (!archIsDrained()) {
-            drainManager = dm;
-            return 1;
-        }
+        if (!archIsDrained())
+            return DrainState::Draining;
 
         // The state of the CPU is consistent, so we don't need to do
         // anything special to drain it. We simply de-schedule the
@@ -318,7 +315,7 @@ BaseKvmCPU::drain(DrainManager *dm)
         // switch CPUs or checkpoint the CPU.
         syncThreadContext();
 
-        return 0;
+        return DrainState::Drained;
 
       case RunningServiceCompletion:
         // The CPU has just requested a service that was handled in
@@ -327,22 +324,18 @@ BaseKvmCPU::drain(DrainManager *dm)
         // update the register state ourselves instead of letting KVM
         // handle it, but that would be tricky. Instead, we enter KVM
         // and let it do its stuff.
-        drainManager = dm;
-
         DPRINTF(Drain, "KVM CPU is waiting for service completion, "
                 "requesting drain.\n");
-        return 1;
+        return DrainState::Draining;
 
       case RunningService:
         // We need to drain since the CPU is waiting for service (e.g., MMIOs)
-        drainManager = dm;
-
         DPRINTF(Drain, "KVM CPU is waiting for service, requesting drain.\n");
-        return 1;
+        return DrainState::Draining;
 
       default:
         panic("KVM: Unhandled CPU state in drain()\n");
-        return 0;
+        return DrainState::Drained;
     }
 }
 
@@ -551,7 +544,7 @@ BaseKvmCPU::tick()
               setupInstStop();
 
           DPRINTF(KvmRun, "Entering KVM...\n");
-          if (drainManager) {
+          if (drainState() == DrainState::Draining) {
               // Force an immediate exit from KVM after completing
               // pending operations. The architecture-specific code
               // takes care to run until it is in a state where it can
@@ -1198,7 +1191,7 @@ BaseKvmCPU::setupCounters()
 bool
 BaseKvmCPU::tryDrain()
 {
-    if (!drainManager)
+    if (drainState() != DrainState::Draining)
         return false;
 
     if (!archIsDrained()) {
@@ -1209,8 +1202,7 @@ BaseKvmCPU::tryDrain()
     if (_status == Idle || _status == Running) {
         DPRINTF(Drain,
                 "tryDrain: CPU transitioned into the Idle state, drain done\n");
-        drainManager->signalDrainDone();
-        drainManager = NULL;
+        signalDrainDone();
         return true;
     } else {
         DPRINTF(Drain, "tryDrain: CPU not ready.\n");
diff --git a/src/cpu/kvm/base.hh b/src/cpu/kvm/base.hh
index a8e429dfa..df6a67808 100644
--- a/src/cpu/kvm/base.hh
+++ b/src/cpu/kvm/base.hh
@@ -89,8 +89,8 @@ class BaseKvmCPU : public BaseCPU
     void unserializeThread(CheckpointIn &cp,
                            ThreadID tid) M5_ATTR_OVERRIDE;
 
-    unsigned int drain(DrainManager *dm);
-    void drainResume();
+    DrainState drain() M5_ATTR_OVERRIDE;
+    void drainResume() M5_ATTR_OVERRIDE;
 
     void switchOut();
     void takeOverFrom(BaseCPU *cpu);
@@ -749,13 +749,6 @@ class BaseKvmCPU : public BaseCPU
     /** Host factor as specified in the configuration */
     float hostFactor;
 
-    /**
-     * Drain manager to use when signaling drain completion
-     *
-     * This pointer is non-NULL when draining and NULL otherwise.
-     */
-    DrainManager *drainManager;
-
   public:
     /* @{ */
     Stats::Scalar numInsts;
diff --git a/src/cpu/minor/cpu.cc b/src/cpu/minor/cpu.cc
index ac1a18bf9..a93d0037d 100644
--- a/src/cpu/minor/cpu.cc
+++ b/src/cpu/minor/cpu.cc
@@ -47,8 +47,7 @@
 #include "debug/Quiesce.hh"
 
 MinorCPU::MinorCPU(MinorCPUParams *params) :
-    BaseCPU(params),
-    drainManager(NULL)
+    BaseCPU(params)
 {
     /* This is only written for one thread at the moment */
     Minor::MinorThread *thread;
@@ -194,39 +193,33 @@ MinorCPU::startup()
     activateContext(0);
 }
 
-unsigned int
-MinorCPU::drain(DrainManager *drain_manager)
+DrainState
+MinorCPU::drain()
 {
     DPRINTF(Drain, "MinorCPU drain\n");
 
-    drainManager = drain_manager;
-
     /* Need to suspend all threads and wait for Execute to idle.
      * Tell Fetch1 not to fetch */
-    unsigned int ret = pipeline->drain(drain_manager);
-
-    if (ret == 0)
+    if (pipeline->drain()) {
         DPRINTF(Drain, "MinorCPU drained\n");
-    else
+        return DrainState::Drained;
+    } else {
         DPRINTF(Drain, "MinorCPU not finished draining\n");
-
-    return ret;
+        return DrainState::Draining;
+    }
 }
 
 void
 MinorCPU::signalDrainDone()
 {
     DPRINTF(Drain, "MinorCPU drain done\n");
-    setDrainState(DrainState::Drained);
-    drainManager->signalDrainDone();
-    drainManager = NULL;
+    signalDrainDone();
 }
 
 void
 MinorCPU::drainResume()
 {
-    assert(getDrainState() == DrainState::Drained ||
-        getDrainState() == DrainState::Running);
+    assert(drainState() == DrainState::Drained);
 
     if (switchedOut()) {
         DPRINTF(Drain, "drainResume while switched out.  Ignoring\n");
@@ -242,8 +235,6 @@ MinorCPU::drainResume()
 
     wakeup();
     pipeline->drainResume();
-
-    setDrainState(DrainState::Running);
 }
 
 void
diff --git a/src/cpu/minor/cpu.hh b/src/cpu/minor/cpu.hh
index fba54b515..2e877d786 100644
--- a/src/cpu/minor/cpu.hh
+++ b/src/cpu/minor/cpu.hh
@@ -112,10 +112,6 @@ class MinorCPU : public BaseCPU
         virtual void recvTimingSnoopReq(PacketPtr pkt) { }
     };
 
-    /** The DrainManager passed into drain that needs be signalled when
-     *  draining is complete */
-    DrainManager *drainManager;
-
   protected:
      /** Return a reference to the data port. */
     MasterPort &getDataPort();
@@ -155,10 +151,10 @@ class MinorCPU : public BaseCPU
     void unserialize(CheckpointIn &cp);
 
     /** Drain interface */
-    unsigned int drain(DrainManager *drain_manager);
-    void drainResume();
-    /** Signal from Pipeline that MinorCPU should signal the DrainManager
-     *  that a drain is complete and set its drainState */
+    DrainState drain() M5_ATTR_OVERRIDE;
+    void drainResume() M5_ATTR_OVERRIDE;
+    /** Signal from Pipeline that MinorCPU should signal that a drain
+     *  is complete and set its drainState */
     void signalDrainDone();
     void memWriteback();
 
diff --git a/src/cpu/minor/pipeline.cc b/src/cpu/minor/pipeline.cc
index 9d802234b..39b7f31f9 100644
--- a/src/cpu/minor/pipeline.cc
+++ b/src/cpu/minor/pipeline.cc
@@ -192,8 +192,8 @@ Pipeline::wakeupFetch()
     execute.wakeupFetch();
 }
 
-unsigned int
-Pipeline::drain(DrainManager *manager)
+bool
+Pipeline::drain()
 {
     DPRINTF(MinorCPU, "Draining pipeline by halting inst fetches. "
         " Execution should drain naturally\n");
@@ -205,7 +205,7 @@ Pipeline::drain(DrainManager *manager)
     bool drained = isDrained();
     needToSignalDrained = !drained;
 
-    return (drained ? 0 : 1);
+    return drained;
 }
 
 void
diff --git a/src/cpu/minor/pipeline.hh b/src/cpu/minor/pipeline.hh
index 355a3c6c2..bf2071b02 100644
--- a/src/cpu/minor/pipeline.hh
+++ b/src/cpu/minor/pipeline.hh
@@ -115,7 +115,7 @@ class Pipeline : public Ticked
     void wakeupFetch();
 
     /** Try to drain the CPU */
-    unsigned int drain(DrainManager *manager);
+    bool drain();
 
     void drainResume();
 
diff --git a/src/cpu/o3/cpu.cc b/src/cpu/o3/cpu.cc
index 89256a7f0..026907a94 100644
--- a/src/cpu/o3/cpu.cc
+++ b/src/cpu/o3/cpu.cc
@@ -196,7 +196,6 @@ FullO3CPU<Impl>::FullO3CPU(DerivO3CPUParams *params)
 
       globalSeqNum(1),
       system(params->system),
-      drainManager(NULL),
       lastRunningCycle(curCycle())
 {
     if (!params->switched_out) {
@@ -539,7 +538,7 @@ FullO3CPU<Impl>::tick()
 {
     DPRINTF(O3CPU, "\n\nFullO3CPU: Ticking main, FullO3CPU.\n");
     assert(!switchedOut());
-    assert(getDrainState() != DrainState::Drained);
+    assert(drainState() != DrainState::Drained);
 
     ++numCycles;
     ppCycles->notify(1);
@@ -712,7 +711,7 @@ FullO3CPU<Impl>::activateContext(ThreadID tid)
     // We don't want to wake the CPU if it is drained. In that case,
     // we just want to flag the thread as active and schedule the tick
     // event from drainResume() instead.
-    if (getDrainState() == DrainState::Drained)
+    if (drainState() == DrainState::Drained)
         return;
 
     // If we are time 0 or if the last activation time is in the past,
@@ -999,17 +998,14 @@ FullO3CPU<Impl>::unserializeThread(CheckpointIn &cp, ThreadID tid)
 }
 
 template <class Impl>
-unsigned int
-FullO3CPU<Impl>::drain(DrainManager *drain_manager)
+DrainState
+FullO3CPU<Impl>::drain()
 {
     // If the CPU isn't doing anything, then return immediately.
-    if (switchedOut()) {
-        setDrainState(DrainState::Drained);
-        return 0;
-    }
+    if (switchedOut())
+        return DrainState::Drained;
 
     DPRINTF(Drain, "Draining...\n");
-    setDrainState(DrainState::Draining);
 
     // We only need to signal a drain to the commit stage as this
     // initiates squashing controls the draining. Once the commit
@@ -1022,16 +1018,13 @@ FullO3CPU<Impl>::drain(DrainManager *drain_manager)
     // Wake the CPU and record activity so everything can drain out if
     // the CPU was not able to immediately drain.
     if (!isDrained())  {
-        drainManager = drain_manager;
-
         wakeCPU();
         activityRec.activity();
 
         DPRINTF(Drain, "CPU not drained\n");
 
-        return 1;
+        return DrainState::Draining;
     } else {
-        setDrainState(DrainState::Drained);
         DPRINTF(Drain, "CPU is already drained\n");
         if (tickEvent.scheduled())
             deschedule(tickEvent);
@@ -1049,7 +1042,7 @@ FullO3CPU<Impl>::drain(DrainManager *drain_manager)
         }
 
         drainSanityCheck();
-        return 0;
+        return DrainState::Drained;
     }
 }
 
@@ -1057,15 +1050,14 @@ template <class Impl>
 bool
 FullO3CPU<Impl>::tryDrain()
 {
-    if (!drainManager || !isDrained())
+    if (drainState() != DrainState::Draining || !isDrained())
         return false;
 
     if (tickEvent.scheduled())
         deschedule(tickEvent);
 
     DPRINTF(Drain, "CPU done draining, processing drain event\n");
-    drainManager->signalDrainDone();
-    drainManager = NULL;
+    signalDrainDone();
 
     return true;
 }
@@ -1132,7 +1124,6 @@ template <class Impl>
 void
 FullO3CPU<Impl>::drainResume()
 {
-    setDrainState(DrainState::Running);
     if (switchedOut())
         return;
 
diff --git a/src/cpu/o3/cpu.hh b/src/cpu/o3/cpu.hh
index 44400542a..aa02ee2ea 100644
--- a/src/cpu/o3/cpu.hh
+++ b/src/cpu/o3/cpu.hh
@@ -232,7 +232,7 @@ class FullO3CPU : public BaseO3CPU
     }
 
     /**
-     * Check if the pipeline has drained and signal the DrainManager.
+     * Check if the pipeline has drained and signal drain done.
      *
      * This method checks if a drain has been requested and if the CPU
      * has drained successfully (i.e., there are no instructions in
@@ -336,7 +336,7 @@ class FullO3CPU : public BaseO3CPU
     void updateThreadPriority();
 
     /** Is the CPU draining? */
-    bool isDraining() const { return getDrainState() == DrainState::Draining; }
+    bool isDraining() const { return drainState() == DrainState::Draining; }
 
     void serializeThread(CheckpointOut &cp,
                          ThreadID tid) const M5_ATTR_OVERRIDE;
@@ -350,10 +350,10 @@ class FullO3CPU : public BaseO3CPU
 
     /** Starts draining the CPU's pipeline of all instructions in
      * order to stop all memory accesses. */
-    unsigned int drain(DrainManager *drain_manager);
+    DrainState drain() M5_ATTR_OVERRIDE;
 
     /** Resumes execution after a drain. */
-    void drainResume();
+    void drainResume() M5_ATTR_OVERRIDE;
 
     /**
      * Commit has reached a safe point to drain a thread.
@@ -665,9 +665,6 @@ class FullO3CPU : public BaseO3CPU
     /** Pointer to the system. */
     System *system;
 
-    /** DrainManager to notify when draining has completed. */
-    DrainManager *drainManager;
-
     /** Pointers to all of the threads in the CPU. */
     std::vector<Thread *> thread;
 
diff --git a/src/cpu/simple/atomic.cc b/src/cpu/simple/atomic.cc
index 4c1c45355..3777ddee9 100644
--- a/src/cpu/simple/atomic.cc
+++ b/src/cpu/simple/atomic.cc
@@ -108,7 +108,6 @@ AtomicSimpleCPU::AtomicSimpleCPU(AtomicSimpleCPUParams *p)
     : BaseSimpleCPU(p), tickEvent(this), width(p->width), locked(false),
       simulate_data_stalls(p->simulate_data_stalls),
       simulate_inst_stalls(p->simulate_inst_stalls),
-      drain_manager(NULL),
       icachePort(name() + ".icache_port", this),
       dcachePort(name() + ".dcache_port", this),
       fastmem(p->fastmem), dcache_access(false), dcache_latency(0),
@@ -125,23 +124,21 @@ AtomicSimpleCPU::~AtomicSimpleCPU()
     }
 }
 
-unsigned int
-AtomicSimpleCPU::drain(DrainManager *dm)
+DrainState
+AtomicSimpleCPU::drain()
 {
-    assert(!drain_manager);
     if (switchedOut())
-        return 0;
+        return DrainState::Drained;
 
     if (!isDrained()) {
         DPRINTF(Drain, "Requesting drain: %s\n", pcState());
-        drain_manager = dm;
-        return 1;
+        return DrainState::Draining;
     } else {
         if (tickEvent.scheduled())
             deschedule(tickEvent);
 
         DPRINTF(Drain, "Not executing microcode, no need to drain.\n");
-        return 0;
+        return DrainState::Drained;
     }
 }
 
@@ -149,7 +146,6 @@ void
 AtomicSimpleCPU::drainResume()
 {
     assert(!tickEvent.scheduled());
-    assert(!drain_manager);
     if (switchedOut())
         return;
 
@@ -173,7 +169,7 @@ AtomicSimpleCPU::drainResume()
 bool
 AtomicSimpleCPU::tryCompleteDrain()
 {
-    if (!drain_manager)
+    if (drainState() != DrainState::Draining)
         return false;
 
     DPRINTF(Drain, "tryCompleteDrain: %s\n", pcState());
@@ -181,8 +177,7 @@ AtomicSimpleCPU::tryCompleteDrain()
         return false;
 
     DPRINTF(Drain, "CPU done draining, processing drain event\n");
-    drain_manager->signalDrainDone();
-    drain_manager = NULL;
+    signalDrainDone();
 
     return true;
 }
diff --git a/src/cpu/simple/atomic.hh b/src/cpu/simple/atomic.hh
index 5ad5c4305..3f587e671 100644
--- a/src/cpu/simple/atomic.hh
+++ b/src/cpu/simple/atomic.hh
@@ -74,13 +74,6 @@ class AtomicSimpleCPU : public BaseSimpleCPU
     const bool simulate_data_stalls;
     const bool simulate_inst_stalls;
 
-    /**
-     * Drain manager to use when signaling drain completion
-     *
-     * This pointer is non-NULL when draining and NULL otherwise.
-     */
-    DrainManager *drain_manager;
-
     // main simulation loop (one cycle)
     void tick();
 
@@ -192,8 +185,8 @@ class AtomicSimpleCPU : public BaseSimpleCPU
 
   public:
 
-    unsigned int drain(DrainManager *drain_manager);
-    void drainResume();
+    DrainState drain() M5_ATTR_OVERRIDE;
+    void drainResume() M5_ATTR_OVERRIDE;
 
     void switchOut();
     void takeOverFrom(BaseCPU *oldCPU);
diff --git a/src/cpu/simple/timing.cc b/src/cpu/simple/timing.cc
index a3c4e27e8..5dc042f1e 100644
--- a/src/cpu/simple/timing.cc
+++ b/src/cpu/simple/timing.cc
@@ -91,7 +91,7 @@ TimingSimpleCPU::TimingCPUPort::TickEvent::schedule(PacketPtr _pkt, Tick t)
 TimingSimpleCPU::TimingSimpleCPU(TimingSimpleCPUParams *p)
     : BaseSimpleCPU(p), fetchTranslation(this), icachePort(this),
       dcachePort(this), ifetch_pkt(NULL), dcache_pkt(NULL), previousCycle(0),
-      fetchEvent(this), drainManager(NULL)
+      fetchEvent(this)
 {
     _status = Idle;
 }
@@ -102,19 +102,17 @@ TimingSimpleCPU::~TimingSimpleCPU()
 {
 }
 
-unsigned int
-TimingSimpleCPU::drain(DrainManager *drain_manager)
+DrainState
+TimingSimpleCPU::drain()
 {
-    assert(!drainManager);
     if (switchedOut())
-        return 0;
+        return DrainState::Drained;
 
     if (_status == Idle ||
         (_status == BaseSimpleCPU::Running && isDrained())) {
         DPRINTF(Drain, "No need to drain.\n");
-        return 0;
+        return DrainState::Drained;
     } else {
-        drainManager = drain_manager;
         DPRINTF(Drain, "Requesting drain: %s\n", pcState());
 
         // The fetch event can become descheduled if a drain didn't
@@ -123,7 +121,7 @@ TimingSimpleCPU::drain(DrainManager *drain_manager)
         if (_status == BaseSimpleCPU::Running && !fetchEvent.scheduled())
             schedule(fetchEvent, clockEdge());
 
-        return 1;
+        return DrainState::Draining;
     }
 }
 
@@ -131,7 +129,6 @@ void
 TimingSimpleCPU::drainResume()
 {
     assert(!fetchEvent.scheduled());
-    assert(!drainManager);
     if (switchedOut())
         return;
 
@@ -155,7 +152,7 @@ TimingSimpleCPU::drainResume()
 bool
 TimingSimpleCPU::tryCompleteDrain()
 {
-    if (!drainManager)
+    if (drainState() != DrainState::Draining)
         return false;
 
     DPRINTF(Drain, "tryCompleteDrain: %s\n", pcState());
@@ -163,8 +160,7 @@ TimingSimpleCPU::tryCompleteDrain()
         return false;
 
     DPRINTF(Drain, "CPU done draining, processing drain event\n");
-    drainManager->signalDrainDone();
-    drainManager = NULL;
+    signalDrainDone();
 
     return true;
 }
diff --git a/src/cpu/simple/timing.hh b/src/cpu/simple/timing.hh
index 3ce596fc7..b6a1da4e2 100644
--- a/src/cpu/simple/timing.hh
+++ b/src/cpu/simple/timing.hh
@@ -270,8 +270,8 @@ class TimingSimpleCPU : public BaseSimpleCPU
 
   public:
 
-    unsigned int drain(DrainManager *drain_manager);
-    void drainResume();
+    DrainState drain() M5_ATTR_OVERRIDE;
+    void drainResume() M5_ATTR_OVERRIDE;
 
     void switchOut();
     void takeOverFrom(BaseCPU *oldCPU);
@@ -351,13 +351,6 @@ class TimingSimpleCPU : public BaseSimpleCPU
      * @returns true if the CPU is drained, false otherwise.
      */
     bool tryCompleteDrain();
-
-    /**
-     * Drain manager to use when signaling drain completion
-     *
-     * This pointer is non-NULL when draining and NULL otherwise.
-     */
-    DrainManager *drainManager;
 };
 
 #endif // __CPU_SIMPLE_TIMING_HH__
diff --git a/src/cpu/testers/traffic_gen/traffic_gen.cc b/src/cpu/testers/traffic_gen/traffic_gen.cc
index 0fc8848fb..984c9950d 100644
--- a/src/cpu/testers/traffic_gen/traffic_gen.cc
+++ b/src/cpu/testers/traffic_gen/traffic_gen.cc
@@ -63,8 +63,7 @@ TrafficGen::TrafficGen(const TrafficGenParams* p)
       port(name() + ".port", *this),
       retryPkt(NULL),
       retryPktTick(0),
-      updateEvent(this),
-      drainManager(NULL)
+      updateEvent(this)
 {
 }
 
@@ -118,12 +117,12 @@ TrafficGen::initState()
     }
 }
 
-unsigned int
-TrafficGen::drain(DrainManager *dm)
+DrainState
+TrafficGen::drain()
 {
     if (!updateEvent.scheduled()) {
         // no event has been scheduled yet (e.g. switched from atomic mode)
-        return 0;
+        return DrainState::Drained;
     }
 
     if (retryPkt == NULL) {
@@ -131,10 +130,9 @@ TrafficGen::drain(DrainManager *dm)
         nextPacketTick = MaxTick;
         nextTransitionTick = MaxTick;
         deschedule(updateEvent);
-        return 0;
+        return DrainState::Drained;
     } else {
-        drainManager = dm;
-        return 1;
+        return DrainState::Draining;
     }
 }
 
@@ -488,7 +486,7 @@ TrafficGen::recvReqRetry()
         retryPktTick = 0;
         retryTicks += delay;
 
-        if (drainManager == NULL) {
+        if (drainState() != DrainState::Draining) {
             // packet is sent, so find out when the next one is due
             nextPacketTick = states[currState]->nextPacketTick(elasticReq,
                                                                delay);
@@ -498,9 +496,7 @@ TrafficGen::recvReqRetry()
             // shut things down
             nextPacketTick = MaxTick;
             nextTransitionTick = MaxTick;
-            drainManager->signalDrainDone();
-            // Clear the drain event once we're done with it.
-            drainManager = NULL;
+            signalDrainDone();
         }
     }
 }
diff --git a/src/cpu/testers/traffic_gen/traffic_gen.hh b/src/cpu/testers/traffic_gen/traffic_gen.hh
index ba7fda7dd..8b71443f9 100644
--- a/src/cpu/testers/traffic_gen/traffic_gen.hh
+++ b/src/cpu/testers/traffic_gen/traffic_gen.hh
@@ -176,9 +176,6 @@ class TrafficGen : public MemObject
     /** Event for scheduling updates */
     EventWrapper<TrafficGen, &TrafficGen::update> updateEvent;
 
-    /** Manager to signal when drained */
-    DrainManager* drainManager;
-
     /** Count the number of generated packets. */
     Stats::Scalar numPackets;
 
@@ -201,7 +198,7 @@ class TrafficGen : public MemObject
 
     void initState();
 
-    unsigned int drain(DrainManager *dm);
+    DrainState drain() M5_ATTR_OVERRIDE;
 
     void serialize(CheckpointOut &cp) const M5_ATTR_OVERRIDE;
     void unserialize(CheckpointIn &cp) M5_ATTR_OVERRIDE;