2 files changed, 54 insertions, 1 deletions

diff --git a/src/mem/cache/base_cache.cc b/src/mem/cache/base_cache.cc
index 1c519fb86..c26d7782b 100644
--- a/src/mem/cache/base_cache.cc
+++ b/src/mem/cache/base_cache.cc
@@ -140,6 +140,9 @@ BaseCache::CachePort::recvRetry()
             }
             waitingOnRetry = false;
         }
+        // Check if we're done draining once this list is empty
+        if (drainList.empty())
+            cache->checkDrain();
     }
     else if (!isCpuSide)
     {
@@ -338,6 +341,10 @@ BaseCache::CacheEvent::process()
         cachePort->drainList.push_back(pkt);
         cachePort->waitingOnRetry = true;
     }
+
+    // Check if we're done draining once this list is empty
+    if (cachePort->drainList.empty())
+        cachePort->cache->checkDrain();
 }
 
 const char *
@@ -599,3 +606,18 @@ BaseCache::regStats()
         ;
 
 }
+
+unsigned int
+BaseCache::drain(Event *de)
+{
+    // Set status
+    if (!canDrain()) {
+        drainEvent = de;
+
+        changeState(SimObject::Draining);
+        return 1;
+    }
+
+    changeState(SimObject::Drained);
+    return 0;
+}
diff --git a/src/mem/cache/base_cache.hh b/src/mem/cache/base_cache.hh
index 565280aef..ea7544fbb 100644
--- a/src/mem/cache/base_cache.hh
+++ b/src/mem/cache/base_cache.hh
@@ -105,6 +105,8 @@ class BaseCache : public MemObject
 
         void clearBlocked();
 
+        bool canDrain() { return drainList.empty(); }
+
         bool blocked;
 
         bool mustSendRetry;
@@ -227,6 +229,9 @@ class BaseCache : public MemObject
     /** The number of misses to trigger an exit event. */
     Counter missCount;
 
+    /** The drain event. */
+    Event *drainEvent;
+
   public:
     // Statistics
     /**
@@ -340,7 +345,7 @@ class BaseCache : public MemObject
     BaseCache(const std::string &name, Params &params)
         : MemObject(name), blocked(0), blockedSnoop(0), masterRequests(0),
           slaveRequests(0), blkSize(params.blkSize),
-          missCount(params.maxMisses)
+          missCount(params.maxMisses), drainEvent(NULL)
     {
         //Start ports at null if more than one is created we should panic
         cpuSidePort = NULL;
@@ -477,6 +482,7 @@ class BaseCache : public MemObject
     {
         uint8_t flag = 1<<cause;
         masterRequests &= ~flag;
+        checkDrain();
     }
 
     /**
@@ -512,6 +518,7 @@ class BaseCache : public MemObject
     {
         uint8_t flag = 1<<cause;
         slaveRequests &= ~flag;
+        checkDrain();
     }
 
     /**
@@ -589,6 +596,30 @@ class BaseCache : public MemObject
             return;
         }
     }
+
+    virtual unsigned int drain(Event *de);
+
+    void checkDrain()
+    {
+        if (drainEvent && canDrain()) {
+            drainEvent->process();
+            changeState(SimObject::Drained);
+            // Clear the drain event
+            drainEvent = NULL;
+        }
+    }
+
+    bool canDrain()
+    {
+        if (doMasterRequest() || doSlaveRequest()) {
+            return false;
+        } else if (memSidePort && !memSidePort->canDrain()) {
+            return false;
+        } else if (cpuSidePort && !cpuSidePort->canDrain()) {
+            return false;
+        }
+        return true;
+    }
 };
 
 #endif //__BASE_CACHE_HH__