3 files changed, 15 insertions, 12 deletions

diff --git a/src/dev/copy_engine.cc b/src/dev/copy_engine.cc
index 809649718..bb15abab6 100644
--- a/src/dev/copy_engine.cc
+++ b/src/dev/copy_engine.cc
@@ -49,6 +49,7 @@
 #include "base/cp_annotate.hh"
 #include "base/trace.hh"
 #include "debug/DMACopyEngine.hh"
+#include "debug/Drain.hh"
 #include "dev/copy_engine.hh"
 #include "mem/packet.hh"
 #include "mem/packet_access.hh"
@@ -638,7 +639,7 @@ bool
 CopyEngine::CopyEngineChannel::inDrain()
 {
     if (ce->getState() == SimObject::Draining) {
-        DPRINTF(DMACopyEngine, "processing drain\n");
+        DPRINTF(Drain, "CopyEngine done draining, processing drain event\n");
         assert(drainEvent);
         drainEvent->process();
         drainEvent = NULL;
@@ -655,7 +656,7 @@ CopyEngine::CopyEngineChannel::drain(Event *de)
     unsigned int count = 1;
     count += cePort.drain(de);
 
-    DPRINTF(DMACopyEngine, "unable to drain, returning %d\n", count);
+    DPRINTF(Drain, "CopyEngineChannel not drained\n");
     drainEvent = de;
     return count;
 }
@@ -673,7 +674,7 @@ CopyEngine::drain(Event *de)
     else
         changeState(Drained);
 
-    DPRINTF(DMACopyEngine, "call to CopyEngine::drain() returning %d\n", count);
+    DPRINTF(Drain, "CopyEngine not drained\n");
     return count;
 }
 
diff --git a/src/dev/dma_device.cc b/src/dev/dma_device.cc
index ddc8e7386..5a2f52df1 100644
--- a/src/dev/dma_device.cc
+++ b/src/dev/dma_device.cc
@@ -43,6 +43,7 @@
 
 #include "base/chunk_generator.hh"
 #include "debug/DMA.hh"
+#include "debug/Drain.hh"
 #include "dev/dma_device.hh"
 #include "sim/system.hh"
 
@@ -103,7 +104,7 @@ DmaPort::recvTimingResp(PacketPtr pkt)
         delete pkt->req;
         delete pkt;
 
-        if (pendingCount == 0 && drainEvent) {
+        if (pendingCount == 0 && transmitList.empty() && drainEvent) {
             drainEvent->process();
             drainEvent = NULL;
         }
@@ -142,9 +143,10 @@ DmaDevice::drain(Event *de)
 unsigned int
 DmaPort::drain(Event *de)
 {
-    if (pendingCount == 0)
+    if (transmitList.empty() && pendingCount == 0)
         return 0;
     drainEvent = de;
+    DPRINTF(Drain, "DmaPort not drained\n");
     return 1;
 }
 
@@ -182,8 +184,6 @@ void
 DmaPort::dmaAction(Packet::Command cmd, Addr addr, int size, Event *event,
                    uint8_t *data, Tick delay, Request::Flags flag)
 {
-    assert(device->getState() == SimObject::Running);
-
     DmaReqState *reqState = new DmaReqState(event, size, delay);
 
 
@@ -287,7 +287,8 @@ DmaPort::sendDma()
         assert(pendingCount >= 0);
         delete pkt;
 
-        if (pendingCount == 0 && drainEvent) {
+        if (pendingCount == 0 && transmitList.empty() && drainEvent) {
+            DPRINTF(Drain, "DmaPort done draining, processing drain event\n");
             drainEvent->process();
             drainEvent = NULL;
         }
diff --git a/src/dev/i8254xGBe.cc b/src/dev/i8254xGBe.cc
index 1f2c92425..14d767288 100644
--- a/src/dev/i8254xGBe.cc
+++ b/src/dev/i8254xGBe.cc
@@ -44,6 +44,7 @@
 
 #include "base/inet.hh"
 #include "base/trace.hh"
+#include "debug/Drain.hh"
 #include "debug/EthernetAll.hh"
 #include "dev/i8254xGBe.hh"
 #include "mem/packet.hh"
@@ -2072,12 +2073,12 @@ IGbE::drain(Event *de)
     if (tickEvent.scheduled())
         deschedule(tickEvent);
 
-    if (count)
+    if (count) {
+        DPRINTF(Drain, "IGbE not drained\n");
         changeState(Draining);
-    else
+    } else
         changeState(Drained);
 
-    DPRINTF(EthernetSM, "got drain() returning %d", count);
     return count;
 }
 
@@ -2100,12 +2101,12 @@ IGbE::checkDrain()
     if (!drainEvent)
         return;
 
-    DPRINTF(EthernetSM, "checkDrain() in drain\n");
     txFifoTick = false;
     txTick = false;
     rxTick = false;
     if (!rxDescCache.hasOutstandingEvents() &&
         !txDescCache.hasOutstandingEvents()) {
+        DPRINTF(Drain, "IGbE done draining, processing drain event\n");
         drainEvent->process();
         drainEvent = NULL;
     }