4 files changed, 11 insertions, 4 deletions

diff --git a/src/arch/arm/table_walker.cc b/src/arch/arm/table_walker.cc
index ed669919c..adff57fdf 100644
--- a/src/arch/arm/table_walker.cc
+++ b/src/arch/arm/table_walker.cc
@@ -388,7 +388,6 @@ TableWalker::processWalkWrapper()
     // if we've still got pending translations schedule more work
     nextWalk(tc);
     currState = NULL;
-    completeDrain();
 }
 
 Fault
@@ -1677,7 +1676,6 @@ TableWalker::doL1DescriptorWrapper()
     doL1Descriptor();
 
     stateQueues[L1].pop_front();
-    completeDrain();
     // Check if fault was generated
     if (currState->fault != NoFault) {
         currState->transState->finish(currState->fault, currState->req,
@@ -1745,7 +1743,6 @@ TableWalker::doL2DescriptorWrapper()
 
 
     stateQueues[L2].pop_front();
-    completeDrain();
     pending = false;
     nextWalk(currState->tc);
 
@@ -1844,6 +1841,8 @@ TableWalker::nextWalk(ThreadContext *tc)
 {
     if (pendingQueue.size())
         schedule(doProcessEvent, clockEdge(Cycles(1)));
+    else
+        completeDrain();
 }
 
 bool
diff --git a/src/mem/cache/mshr_queue.cc b/src/mem/cache/mshr_queue.cc
index 7bfbb90f5..9146cddf7 100644
--- a/src/mem/cache/mshr_queue.cc
+++ b/src/mem/cache/mshr_queue.cc
@@ -45,7 +45,9 @@
  * Definition of MSHRQueue class functions.
  */
 
+#include "base/trace.hh"
 #include "mem/cache/mshr_queue.hh"
+#include "debug/Drain.hh"
 
 using namespace std;
 
@@ -191,6 +193,7 @@ MSHRQueue::deallocateOne(MSHR *mshr)
     if (drainManager && 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(Drainable::Drained);
diff --git a/src/mem/dram_ctrl.cc b/src/mem/dram_ctrl.cc
index f65f7e1dd..1beebdd01 100644
--- a/src/mem/dram_ctrl.cc
+++ b/src/mem/dram_ctrl.cc
@@ -737,6 +737,7 @@ DRAMCtrl::processRespondEvent()
         // if there is nothing left in any queue, signal a drain
         if (writeQueue.empty() && readQueue.empty() &&
             drainManager) {
+            DPRINTF(Drain, "DRAM controller done draining\n");
             drainManager->signalDrainDone();
             drainManager = NULL;
         }
@@ -1290,6 +1291,7 @@ DRAMCtrl::processNextReqEvent()
             } else {
                 // check if we are drained
                 if (respQueue.empty () && drainManager) {
+                    DPRINTF(Drain, "DRAM controller done draining\n");
                     drainManager->signalDrainDone();
                     drainManager = NULL;
                 }
diff --git a/src/mem/simple_mem.cc b/src/mem/simple_mem.cc
index 27d3f1186..4e1020de5 100644
--- a/src/mem/simple_mem.cc
+++ b/src/mem/simple_mem.cc
@@ -44,6 +44,7 @@
 
 #include "base/random.hh"
 #include "mem/simple_mem.hh"
+#include "debug/Drain.hh"
 
 using namespace std;
 
@@ -200,6 +201,7 @@ SimpleMemory::dequeue()
             reschedule(dequeueEvent,
                        std::max(packetQueue.front().tick, curTick()), true);
         } else if (drainManager) {
+            DPRINTF(Drain, "Drainng of SimpleMemory complete\n");
             drainManager->signalDrainDone();
             drainManager = NULL;
         }
@@ -240,7 +242,8 @@ SimpleMemory::drain(DrainManager *dm)
     if (!packetQueue.empty()) {
         count += 1;
         drainManager = dm;
-    }
+        DPRINTF(Drain, "SimpleMemory Queue has requests, waiting to drain\n");
+     }
 
     if (count)
         setDrainState(Drainable::Draining);