2 files changed, 7 insertions, 19 deletions
diff --git a/src/mem/cache/mshr_queue.cc b/src/mem/cache/mshr_queue.cc
index 3736c9f21..47f044d63 100644
--- a/src/mem/cache/mshr_queue.cc
+++ b/src/mem/cache/mshr_queue.cc
@@ -56,7 +56,7 @@ MSHRQueue::MSHRQueue(const std::string &_label,
                      int _index)
     : label(_label), numEntries(num_entries + reserve - 1),
       numReserve(reserve), demandReserve(demand_reserve),
-      registers(numEntries), drainManager(NULL), allocated(0),
+      registers(numEntries), allocated(0),
       inServiceEntries(0), index(_index)
 {
     for (int i = 0; i < numEntries; ++i) {
@@ -180,13 +180,11 @@ MSHRQueue::deallocateOne(MSHR *mshr)
         readyList.erase(mshr->readyIter);
     }
     mshr->deallocate();
-    if (drainManager && allocated == 0) {
+    if (drainState() == DrainState::Draining && allocated == 0) {
         // Notify the drain manager that we have completed draining if
         // there are no other outstanding requests in this MSHR queue.
         DPRINTF(Drain, "MSHRQueue now empty, signalling drained\n");
-        drainManager->signalDrainDone();
-        drainManager = NULL;
-        setDrainState(DrainState::Drained);
+        signalDrainDone();
     }
     return retval;
 }
@@ -265,15 +263,8 @@ MSHRQueue::squash(int threadNum)
     }
 }
 
-unsigned int
-MSHRQueue::drain(DrainManager *dm)
+DrainState
+MSHRQueue::drain()
 {
-    if (allocated == 0) {
-        setDrainState(DrainState::Drained);
-        return 0;
-    } else {
-        drainManager = dm;
-        setDrainState(DrainState::Draining);
-        return 1;
-    }
+    return allocated == 0 ? DrainState::Drained : DrainState::Draining;
 }
diff --git a/src/mem/cache/mshr_queue.hh b/src/mem/cache/mshr_queue.hh
index 4043bc565..308d371fe 100644
--- a/src/mem/cache/mshr_queue.hh
+++ b/src/mem/cache/mshr_queue.hh
@@ -92,9 +92,6 @@ class MSHRQueue : public Drainable
     /** Holds non allocated entries. */
     MSHR::List freeList;
 
-    /** Drain manager to inform of a completed drain */
-    DrainManager *drainManager;
-
     MSHR::Iterator addToReadyList(MSHR *mshr);
 
 
@@ -258,7 +255,7 @@ class MSHRQueue : public Drainable
         return readyList.empty() ? MaxTick : readyList.front()->readyTime;
     }
 
-    unsigned int drain(DrainManager *dm);
+    DrainState drain() M5_ATTR_OVERRIDE;
 };
 
 #endif //__MEM_CACHE_MSHR_QUEUE_HH__