1 files changed, 10 insertions, 3 deletions
diff --git a/src/sim/system.cc b/src/sim/system.cc
index 5ec7f4b30..259ed3e88 100644
--- a/src/sim/system.cc
+++ b/src/sim/system.cc
@@ -181,7 +181,7 @@ System::getMasterPort(const std::string &if_name, PortID idx)
 void
 System::setMemoryMode(Enums::MemoryMode mode)
 {
-    assert(getState() == Drained);
+    assert(getDrainState() == Drainable::Drained);
     memoryMode = mode;
 }
 
@@ -328,10 +328,17 @@ System::isMemAddr(Addr addr) const
     return physmem.isMemAddr(addr);
 }
 
+unsigned int
+System::drain(DrainManager *dm)
+{
+    setDrainState(Drainable::Drained);
+    return 0;
+}
+
 void
-System::resume()
+System::drainResume()
 {
-    SimObject::resume();
+    Drainable::drainResume();
     totalNumInsts = 0;
 }