sim: Add an option to forward work items to Python

[gem5.git] / src / sim / drain.cc

diff --git a/src/sim/drain.cc b/src/sim/drain.cc
index 384374099522d5963e2556891c2a5830a1ba8e52..c6073487607a745607ac83b8cb3693b4a61e44a3 100644 (file)
--- a/src/sim/drain.cc
+++ b/src/sim/drain.cc
@@ -68,7 +68,7 @@ DrainManager::tryDrain()
     DPRINTF(Drain, "Trying to drain %u objects.\n", drainableCount());
     _state = DrainState::Draining;
     for (auto *obj : _allDrainable)
-        _count += obj->drain(&_instance);
+        _count += obj->dmDrain() == DrainState::Drained ? 0 : 1;
 
     if (_count == 0) {
         DPRINTF(Drain, "Drain done.\n");
@@ -98,7 +98,7 @@ DrainManager::resume()
     DPRINTF(Drain, "Resuming %u objects.\n", drainableCount());
     _state = DrainState::Running;
     for (auto *obj : _allDrainable)
-        obj->drainResume();
+        obj->dmDrainResume();
 }
 
 void
@@ -150,7 +150,7 @@ DrainManager::drainableCount() const
 
 Drainable::Drainable()
     : _drainManager(DrainManager::instance()),
-      _drainState(DrainState::Running)
+      _drainState(_drainManager.state())
 {
     _drainManager.registerDrainable(this);
 }
@@ -160,8 +160,23 @@ Drainable::~Drainable()
     _drainManager.unregisterDrainable(this);
 }
 
+DrainState
+Drainable::dmDrain()
+{
+    _drainState = DrainState::Draining;
+    _drainState = drain();
+    assert(_drainState == DrainState::Draining ||
+           _drainState == DrainState::Drained);
+
+    return _drainState;
+}
+
 void
-Drainable::drainResume()
+Drainable::dmDrainResume()
 {
+    panic_if(_drainState != DrainState::Drained,
+             "Trying to resume an object that hasn't been drained\n");
+
     _drainState = DrainState::Running;
+    drainResume();
 }