3 files changed, 36 insertions, 19 deletions

diff --git a/core/java/android/service/dataloader/DataLoaderService.java b/core/java/android/service/dataloader/DataLoaderService.java
index 60373ace2d8a..e35b8b731f14 100644
--- a/core/java/android/service/dataloader/DataLoaderService.java
+++ b/core/java/android/service/dataloader/DataLoaderService.java
@@ -109,9 +109,7 @@ public abstract class DataLoaderService extends Service {
                 @NonNull IDataLoaderStatusListener listener)
                 throws RuntimeException {
             try {
-                if (!nativeCreateDataLoader(id, control, params, listener)) {
-                    Slog.e(TAG, "Failed to create native loader for " + id);
-                }
+                nativeCreateDataLoader(id, control, params, listener);
             } catch (Exception ex) {
                 Slog.e(TAG, "Failed to create native loader for " + id, ex);
                 destroy(id);
diff --git a/services/incremental/IncrementalService.cpp b/services/incremental/IncrementalService.cpp
index 68b7ac0001e5..79699bb77a92 100644
--- a/services/incremental/IncrementalService.cpp
+++ b/services/incremental/IncrementalService.cpp
@@ -160,8 +160,8 @@ const bool IncrementalService::sEnablePerfLogging =
 
 IncrementalService::IncFsMount::~IncFsMount() {
     if (dataLoaderStub) {
-        dataLoaderStub->requestDestroy();
-        dataLoaderStub->waitForDestroy();
+        dataLoaderStub->cleanupResources();
+        dataLoaderStub = {};
     }
     LOG(INFO) << "Unmounting and cleaning up mount " << mountId << " with root '" << root << '\'';
     for (auto&& [target, _] : bindPoints) {
@@ -999,7 +999,8 @@ bool IncrementalService::startLoading(StorageId storage) const {
             return false;
         }
     }
-    return dataLoaderStub->requestStart();
+    dataLoaderStub->requestStart();
+    return true;
 }
 
 void IncrementalService::mountExistingImages() {
@@ -1466,11 +1467,17 @@ IncrementalService::DataLoaderStub::DataLoaderStub(IncrementalService& service,
         mParams(std::move(params)),
         mControl(std::move(control)),
         mListener(externalListener ? *externalListener : DataLoaderStatusListener()) {
-    //
 }
 
-IncrementalService::DataLoaderStub::~DataLoaderStub() {
-    waitForDestroy();
+IncrementalService::DataLoaderStub::~DataLoaderStub() = default;
+
+void IncrementalService::DataLoaderStub::cleanupResources() {
+    requestDestroy();
+    mParams = {};
+    mControl = {};
+    waitForStatus(IDataLoaderStatusListener::DATA_LOADER_DESTROYED, std::chrono::seconds(60));
+    mListener = {};
+    mId = kInvalidStorageId;
 }
 
 bool IncrementalService::DataLoaderStub::requestCreate() {
@@ -1485,10 +1492,6 @@ bool IncrementalService::DataLoaderStub::requestDestroy() {
     return setTargetStatus(IDataLoaderStatusListener::DATA_LOADER_DESTROYED);
 }
 
-bool IncrementalService::DataLoaderStub::waitForDestroy(Clock::duration duration) {
-    return waitForStatus(IDataLoaderStatusListener::DATA_LOADER_DESTROYED, duration);
-}
-
 bool IncrementalService::DataLoaderStub::setTargetStatus(int status) {
     {
         std::unique_lock lock(mStatusMutex);
@@ -1541,6 +1544,10 @@ bool IncrementalService::DataLoaderStub::destroy() {
 }
 
 bool IncrementalService::DataLoaderStub::fsmStep() {
+    if (!isValid()) {
+        return false;
+    }
+
     int currentStatus;
     int targetStatus;
     {
@@ -1580,6 +1587,15 @@ bool IncrementalService::DataLoaderStub::fsmStep() {
 }
 
 binder::Status IncrementalService::DataLoaderStub::onStatusChanged(MountId mountId, int newStatus) {
+    if (!isValid()) {
+        return binder::Status::
+                fromServiceSpecificError(-EINVAL, "onStatusChange came to invalid DataLoaderStub");
+    }
+    if (mId != mountId) {
+        LOG(ERROR) << "Mount ID mismatch: expected " << mId << ", but got: " << mountId;
+        return binder::Status::fromServiceSpecificError(-EPERM, "Mount ID mismatch.");
+    }
+
     {
         std::unique_lock lock(mStatusMutex);
         if (mCurrentStatus == newStatus) {
diff --git a/services/incremental/IncrementalService.h b/services/incremental/IncrementalService.h
index 8b28bac26f9a..bd01d7760a01 100644
--- a/services/incremental/IncrementalService.h
+++ b/services/incremental/IncrementalService.h
@@ -173,13 +173,14 @@ private:
                        FileSystemControlParcel&& control,
                        const DataLoaderStatusListener* externalListener);
         ~DataLoaderStub();
+        // Cleans up the internal state and invalidates DataLoaderStub. Any subsequent calls will
+        // result in an error.
+        void cleanupResources();
 
         bool requestCreate();
         bool requestStart();
         bool requestDestroy();
 
-        bool waitForDestroy(Clock::duration duration = std::chrono::seconds(60));
-
         void onDump(int fd);
 
         MountId id() const { return mId; }
@@ -188,6 +189,8 @@ private:
     private:
         binder::Status onStatusChanged(MountId mount, int newStatus) final;
 
+        bool isValid() const { return mId != kInvalidStorageId; }
+
         bool create();
         bool start();
         bool destroy();
@@ -198,10 +201,10 @@ private:
         bool fsmStep();
 
         IncrementalService& mService;
-        MountId const mId;
-        DataLoaderParamsParcel const mParams;
-        FileSystemControlParcel const mControl;
-        DataLoaderStatusListener const mListener;
+        MountId mId = kInvalidStorageId;
+        DataLoaderParamsParcel mParams;
+        FileSystemControlParcel mControl;
+        DataLoaderStatusListener mListener;
 
         std::mutex mStatusMutex;
         std::condition_variable mStatusCondition;