Free ISnapshotManager once update is done

When update-engine finishes serving an update,
DyanmicPartitionControlAndroid::Cleanup() is called, giving us a chance
to free ISnapshotManager. To ensure that the instance is allocated again
when needed, replace all access to Snapshot manager with a getter
function

Test: th
Bug: 401952955

Change-Id: Id5c3c9906ed50dd2694ae2c398ac69953d575702
diff --git a/aosp/dynamic_partition_control_android.cc b/aosp/dynamic_partition_control_android.cc
index d1c3bf2..78a9a13 100644
--- a/aosp/dynamic_partition_control_android.cc
+++ b/aosp/dynamic_partition_control_android.cc
@@ -226,7 +226,7 @@
     // One exception is when /metadata is not mounted. Fallback to
     // CreateLogicalPartition as snapshots are not created in the first place.
     params.timeout_ms = kMapSnapshotTimeout;
-    success = snapshot_->MapUpdateSnapshot(params, path);
+    success = GetSnapshotManager()->MapUpdateSnapshot(params, path);
   } else {
     params.timeout_ms = kMapTimeout;
     success = CreateLogicalPartition(params, path);
@@ -311,7 +311,7 @@
     // a paused update. Clean up any underlying devices.
     if (ExpectMetadataMounted() &&
         !device_name.ends_with(kRWSourcePartitionSuffix)) {
-      success &= snapshot_->UnmapUpdateSnapshot(device_name);
+      success &= GetSnapshotManager()->UnmapUpdateSnapshot(device_name);
     } else {
       LOG(INFO) << "Skip UnmapUpdateSnapshot(" << device_name << ")";
     }
@@ -328,7 +328,7 @@
 }
 
 bool DynamicPartitionControlAndroid::UnmapAllPartitions() {
-  snapshot_->UnmapAllSnapshots();
+  GetSnapshotManager()->UnmapAllSnapshots();
   if (mapped_devices_.empty()) {
     return false;
   }
@@ -358,12 +358,9 @@
   LOG(INFO) << "UnmapAllPartitions done";
   metadata_device_.reset();
   if (GetVirtualAbFeatureFlag().IsEnabled()) {
-    snapshot_ = SnapshotManager::New();
-  } else {
-    snapshot_ = SnapshotManagerStub::New();
+    // Release ISnapshotManager instance so GSID can be gracefully shutdown
+    snapshot_ = nullptr;
   }
-  CHECK(snapshot_ != nullptr) << "Cannot initialize SnapshotManager.";
-  LOG(INFO) << "SnapshotManager initialized.";
 }
 
 bool DynamicPartitionControlAndroid::DeviceExists(const std::string& path) {
@@ -585,7 +582,7 @@
     // should not proceed because during next boot, snapshots will overlay on
     // the devices incorrectly.
     if (ExpectMetadataMounted()) {
-      TEST_AND_RETURN_FALSE(snapshot_->CancelUpdate());
+      TEST_AND_RETURN_FALSE(GetSnapshotManager()->CancelUpdate());
     } else {
       LOG(INFO) << "Skip canceling previous update because metadata is not "
                 << "mounted";
@@ -996,6 +993,20 @@
   return true;
 }
 
+android::snapshot::ISnapshotManager*
+DynamicPartitionControlAndroid::GetSnapshotManager() {
+  if (snapshot_ == nullptr) {
+    if (GetVirtualAbFeatureFlag().IsEnabled()) {
+      snapshot_ = SnapshotManager::New();
+    } else {
+      snapshot_ = SnapshotManagerStub::New();
+    }
+  }
+  CHECK(snapshot_ != nullptr) << "Cannot initialize SnapshotManager.";
+  LOG(INFO) << "SnapshotManager initialized.";
+  return snapshot_.get();
+}
+
 bool DynamicPartitionControlAndroid::PrepareSnapshotPartitionsForUpdate(
     uint32_t source_slot,
     uint32_t target_slot,
@@ -1018,11 +1029,11 @@
   TEST_AND_RETURN_FALSE(
       CheckSuperPartitionAllocatableSpace(builder.get(), manifest, true));
 
-  if (!snapshot_->BeginUpdate()) {
+  if (!GetSnapshotManager()->BeginUpdate()) {
     LOG(ERROR) << "Cannot begin new update.";
     return false;
   }
-  auto ret = snapshot_->CreateUpdateSnapshots(manifest);
+  auto ret = GetSnapshotManager()->CreateUpdateSnapshots(manifest);
   if (!ret) {
     LOG(ERROR) << "Cannot create update snapshots: " << ret.string();
     if (required_size != nullptr &&
@@ -1124,9 +1135,9 @@
 
 bool DynamicPartitionControlAndroid::FinishUpdate(bool powerwash_required) {
   if (ExpectMetadataMounted()) {
-    if (snapshot_->GetUpdateState() == UpdateState::Initiated) {
+    if (GetSnapshotManager()->GetUpdateState() == UpdateState::Initiated) {
       LOG(INFO) << "Snapshot writes are done.";
-      return snapshot_->FinishedSnapshotWrites(powerwash_required);
+      return GetSnapshotManager()->FinishedSnapshotWrites(powerwash_required);
     }
   } else {
     LOG(INFO) << "Skip FinishedSnapshotWrites() because /metadata is not "
@@ -1350,7 +1361,7 @@
     return std::make_unique<NoOpAction>();
   }
   return std::make_unique<CleanupPreviousUpdateAction>(
-      prefs, boot_control, snapshot_.get(), delegate);
+      prefs, boot_control, GetSnapshotManager(), delegate);
 }
 
 bool DynamicPartitionControlAndroid::ResetUpdate(PrefsInterface* prefs) {
@@ -1372,7 +1383,7 @@
       prefs, false /* quick */, false /* skip dynamic partitions metadata */));
 
   if (ExpectMetadataMounted()) {
-    TEST_AND_RETURN_FALSE(snapshot_->CancelUpdate());
+    TEST_AND_RETURN_FALSE(GetSnapshotManager()->CancelUpdate());
   } else {
     LOG(INFO) << "Skip cancelling update in ResetUpdate because /metadata is "
               << "not mounted";
@@ -1473,7 +1484,7 @@
   }
 
   if (metadata_device_ == nullptr) {
-    metadata_device_ = snapshot_->EnsureMetadataMounted();
+    metadata_device_ = GetSnapshotManager()->EnsureMetadataMounted();
   }
   return metadata_device_ != nullptr;
 }
@@ -1497,7 +1508,7 @@
       .timeout_ms = kMapSnapshotTimeout};
   // TODO(zhangkelvin) Open an APPEND mode CowWriter once there's an API to do
   // it.
-  return snapshot_->OpenSnapshotWriter(params, label);
+  return GetSnapshotManager()->OpenSnapshotWriter(params, label);
 }  // namespace chromeos_update_engine
 
 std::unique_ptr<FileDescriptor> DynamicPartitionControlAndroid::OpenCowFd(
@@ -1531,7 +1542,7 @@
 }
 
 bool DynamicPartitionControlAndroid::MapAllPartitions() {
-  return snapshot_->MapAllSnapshots(kMapSnapshotTimeout);
+  return GetSnapshotManager()->MapAllSnapshots(kMapSnapshotTimeout);
 }
 
 bool DynamicPartitionControlAndroid::IsDynamicPartition(
@@ -1555,7 +1566,7 @@
 
 bool DynamicPartitionControlAndroid::UpdateUsesSnapshotCompression() {
   return GetVirtualAbFeatureFlag().IsEnabled() &&
-         snapshot_->UpdateUsesCompression();
+         GetSnapshotManager()->UpdateUsesCompression();
 }
 
 FeatureFlag
diff --git a/aosp/dynamic_partition_control_android.h b/aosp/dynamic_partition_control_android.h
index 1f70184..e5ba86a 100644
--- a/aosp/dynamic_partition_control_android.h
+++ b/aosp/dynamic_partition_control_android.h
@@ -339,6 +339,7 @@
   bool SetTargetBuildVars(const DeltaArchiveManifest& manifest);
 
   std::string GetDeviceName(std::string partition_name, uint32_t slot) const;
+  android::snapshot::ISnapshotManager* GetSnapshotManager();
 
   std::set<std::string> mapped_devices_;
   const FeatureFlag dynamic_partitions_;
diff --git a/aosp/update_attempter_android.cc b/aosp/update_attempter_android.cc
index b46f584..461cece 100644
--- a/aosp/update_attempter_android.cc
+++ b/aosp/update_attempter_android.cc
@@ -875,6 +875,8 @@
     return;
   }
 
+  boot_control_->GetDynamicPartitionControl()->Cleanup();
+
   if (status_ == UpdateStatus::CLEANUP_PREVIOUS_UPDATE) {
     ClearUpdateCompletedMarker();
     LOG(INFO) << "Terminating cleanup previous update.";
@@ -884,8 +886,6 @@
     return;
   }
 
-  boot_control_->GetDynamicPartitionControl()->Cleanup();
-
   for (auto observer : daemon_state_->service_observers())
     observer->SendPayloadApplicationComplete(error_code);