1 files changed, 7 insertions, 20 deletions

diff --git a/src/dev/arm/flash_device.cc b/src/dev/arm/flash_device.cc
index e4bfcf5c9..b651a1eeb 100644
--- a/src/dev/arm/flash_device.cc
+++ b/src/dev/arm/flash_device.cc
@@ -86,7 +86,6 @@ FlashDevice::FlashDevice(const FlashDeviceParams* p):
     pagesPerDisk(0),
     blocksPerDisk(0),
     planeMask(numPlanes - 1),
-    drainManager(NULL),
     planeEventQueue(numPlanes),
     planeEvent(this)
 {
@@ -587,26 +586,16 @@ FlashDevice::unserialize(CheckpointIn &cp)
  * Drain; needed to enable checkpoints
  */
 
-unsigned int
-FlashDevice::drain(DrainManager *dm)
+DrainState
+FlashDevice::drain()
 {
-    unsigned int count = 0;
-
     if (planeEvent.scheduled()) {
-        count = 1;
-        drainManager = dm;
-    } else {
-        DPRINTF(Drain, "Flash device in drained state\n");
-    }
-
-    if (count) {
         DPRINTF(Drain, "Flash device is draining...\n");
-        setDrainState(DrainState::Draining);
+        return DrainState::Draining;
     } else {
-        DPRINTF(Drain, "Flash device drained\n");
-        setDrainState(DrainState::Drained);
+        DPRINTF(Drain, "Flash device in drained state\n");
+        return DrainState::Drained;
     }
-    return count;
 }
 
 /**
@@ -616,15 +605,13 @@ FlashDevice::drain(DrainManager *dm)
 void
 FlashDevice::checkDrain()
 {
-    if (drainManager == NULL) {
+    if (drainState() == DrainState::Draining)
         return;
-    }
 
     if (planeEvent.when() > curTick()) {
         DPRINTF(Drain, "Flash device is still draining\n");
     } else {
         DPRINTF(Drain, "Flash device is done draining\n");
-        drainManager->signalDrainDone();
-        drainManager = NULL;
+        signalDrainDone();
     }
 }