4 files changed, 15 insertions, 48 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();
     }
 }
diff --git a/src/dev/arm/flash_device.hh b/src/dev/arm/flash_device.hh
index 330299451..891217cbf 100644
--- a/src/dev/arm/flash_device.hh
+++ b/src/dev/arm/flash_device.hh
@@ -62,7 +62,7 @@ class FlashDevice : public AbstractNVM
     ~FlashDevice();
 
     /** Checkpoint functions*/
-    unsigned int drain(DrainManager *dm);
+    DrainState drain() M5_ATTR_OVERRIDE;
     void checkDrain();
 
     void serialize(CheckpointOut &cp) const M5_ATTR_OVERRIDE;
@@ -176,13 +176,6 @@ class FlashDevice : public AbstractNVM
     uint32_t planeMask;
 
     /**
-     * drain manager
-     * Needed to be able to implement checkpoint functionality
-     */
-
-    DrainManager *drainManager;
-
-    /**
      * when the disk is first started we are unsure of the number of
      * used pages, this variable will help determining what we do know.
      */
diff --git a/src/dev/arm/ufs_device.cc b/src/dev/arm/ufs_device.cc
index 696aeba6f..61b125ef5 100644
--- a/src/dev/arm/ufs_device.cc
+++ b/src/dev/arm/ufs_device.cc
@@ -733,7 +733,6 @@ UFSHostDevice::UFSHostDevice(const UFSHostDeviceParams* p) :
     transferTrack(0),
     taskCommandTrack(0),
     idlePhaseStart(0),
-    drainManager(NULL),
     SCSIResumeEvent(this),
     UTPEvent(this)
 {
@@ -2316,18 +2315,15 @@ UFSHostDevice::unserialize(CheckpointIn &cp)
  * Drain; needed to enable checkpoints
  */
 
-unsigned int
-UFSHostDevice::drain(DrainManager *dm)
+DrainState
+UFSHostDevice::drain()
 {
     if (UFSHCIMem.TRUTRLDBR) {
-        drainManager = dm;
         DPRINTF(UFSHostDevice, "UFSDevice is draining...\n");
-        setDrainState(DrainState::Draining);
-        return 1;
+        return DrainState::Draining;
     } else {
         DPRINTF(UFSHostDevice, "UFSDevice drained\n");
-        setDrainState(DrainState::Drained);
-        return 0;
+        return DrainState::Drained;
     }
 }
 
@@ -2338,16 +2334,14 @@ UFSHostDevice::drain(DrainManager *dm)
 void
 UFSHostDevice::checkDrain()
 {
-    if (drainManager == NULL) {
+    if (drainState() != DrainState::Draining)
         return;
-    }
 
     if (UFSHCIMem.TRUTRLDBR) {
         DPRINTF(UFSHostDevice, "UFSDevice is still draining; with %d active"
             " doorbells\n", activeDoorbells);
     } else {
         DPRINTF(UFSHostDevice, "UFSDevice is done draining\n");
-        drainManager->signalDrainDone();
-        drainManager = NULL;
+        signalDrainDone();
     }
 }
diff --git a/src/dev/arm/ufs_device.hh b/src/dev/arm/ufs_device.hh
index 15e983ad8..716b1bdcb 100644
--- a/src/dev/arm/ufs_device.hh
+++ b/src/dev/arm/ufs_device.hh
@@ -173,7 +173,7 @@ class UFSHostDevice : public DmaDevice
 
     UFSHostDevice(const UFSHostDeviceParams* p);
 
-    unsigned int drain(DrainManager *dm);
+    DrainState drain() M5_ATTR_OVERRIDE;
     void checkDrain();
     void serialize(CheckpointOut &cp) const M5_ATTR_OVERRIDE;
     void unserialize(CheckpointIn &cp) M5_ATTR_OVERRIDE;
@@ -1053,13 +1053,6 @@ class UFSHostDevice : public DmaDevice
     Tick idlePhaseStart;
 
     /**
-     * drain manager
-     * Needed to be able to implement checkpoint functionality
-     */
-
-    DrainManager *drainManager;
-
-    /**
      * logic units connected to the UFS Host device
      * Note again that the "device" as such is represented by one or multiple
      * logic units.