diff --git a/aosp/apex_handler_android.cc b/aosp/apex_handler_android.cc
index cdbc983..38ec410 100644
--- a/aosp/apex_handler_android.cc
+++ b/aosp/apex_handler_android.cc
@@ -23,32 +23,15 @@
 
 namespace chromeos_update_engine {
 
-// Don't change this path... apexd relies on it.
-constexpr const char* kApexReserveSpaceDir = "/data/apex/ota_reserved";
+namespace {
 
-uint64_t ApexHandlerAndroid::CalculateSize(
-    const std::vector<ApexInfo>& apex_infos) const {
-  return CalculateSize(apex_infos, GetApexService());
-}
-
-uint64_t ApexHandlerAndroid::CalculateSize(
-    const std::vector<ApexInfo>& apex_infos,
-    android::sp<android::apex::IApexService> apex_service) const {
-  // The safest option is to allocate space for every compressed APEX
-  uint64_t size_required_default = 0;
-
-  // We might not need to decompress every APEX. Communicate with apexd to get
-  // accurate requirement.
-  int64_t size_from_apexd;
+android::apex::CompressedApexInfoList CreateCompressedApexInfoList(
+    const std::vector<ApexInfo>& apex_infos) {
   android::apex::CompressedApexInfoList compressed_apex_info_list;
-
   for (const auto& apex_info : apex_infos) {
     if (!apex_info.is_compressed()) {
       continue;
     }
-
-    size_required_default += apex_info.decompressed_size();
-
     android::apex::CompressedApexInfo compressed_apex_info;
     compressed_apex_info.moduleName = apex_info.package_name();
     compressed_apex_info.versionCode = apex_info.version();
@@ -56,33 +39,41 @@
     compressed_apex_info_list.apexInfos.emplace_back(
         std::move(compressed_apex_info));
   }
-  if (size_required_default == 0 || apex_service == nullptr) {
-    return size_required_default;
+  return compressed_apex_info_list;
+}
+
+}  // namespace
+
+android::base::Result<uint64_t> ApexHandlerAndroid::CalculateSize(
+    const std::vector<ApexInfo>& apex_infos) const {
+  // We might not need to decompress every APEX. Communicate with apexd to get
+  // accurate requirement.
+  auto apex_service = GetApexService();
+  if (apex_service == nullptr) {
+    return android::base::Error() << "Failed to get hold of apexservice";
   }
 
+  auto compressed_apex_info_list = CreateCompressedApexInfoList(apex_infos);
+  int64_t size_from_apexd;
   auto result = apex_service->calculateSizeForCompressedApex(
       compressed_apex_info_list, &size_from_apexd);
   if (!result.isOk()) {
-    return size_required_default;
+    return android::base::Error()
+           << "Failed to get size required from apexservice";
   }
   return size_from_apexd;
 }
 
-bool ApexHandlerAndroid::AllocateSpace(const uint64_t size_required) const {
-  return AllocateSpace(size_required, kApexReserveSpaceDir);
-}
-
-bool ApexHandlerAndroid::AllocateSpace(const uint64_t size_required,
-                                       const std::string& dir_path) const {
-  if (size_required == 0) {
-    return true;
+bool ApexHandlerAndroid::AllocateSpace(
+    const std::vector<ApexInfo>& apex_infos) const {
+  auto apex_service = GetApexService();
+  if (apex_service == nullptr) {
+    return false;
   }
-  base::FilePath path{dir_path};
-  // The filename is not important, it just needs to be under
-  // kApexReserveSpaceDir. We call it "full.tmp" because the current space
-  // estimation is simply adding up all decompressed sizes.
-  path = path.Append("full.tmp");
-  return utils::ReserveStorageSpace(path.value().c_str(), size_required);
+  auto compressed_apex_info_list = CreateCompressedApexInfoList(apex_infos);
+  auto result =
+      apex_service->reserveSpaceForCompressedApex(compressed_apex_info_list);
+  return result.isOk();
 }
 
 android::sp<android::apex::IApexService> ApexHandlerAndroid::GetApexService()
diff --git a/aosp/apex_handler_android.h b/aosp/apex_handler_android.h
index aac1cd9..00f3a80 100644
--- a/aosp/apex_handler_android.h
+++ b/aosp/apex_handler_android.h
@@ -30,17 +30,12 @@
 
 class ApexHandlerAndroid : virtual public ApexHandlerInterface {
  public:
-  uint64_t CalculateSize(const std::vector<ApexInfo>& apex_infos) const;
-  bool AllocateSpace(const uint64_t size_required) const;
+  android::base::Result<uint64_t> CalculateSize(
+      const std::vector<ApexInfo>& apex_infos) const;
+  bool AllocateSpace(const std::vector<ApexInfo>& apex_infos) const;
 
  private:
-  friend class ApexHandlerAndroidTest;
   android::sp<android::apex::IApexService> GetApexService() const;
-  uint64_t CalculateSize(
-      const std::vector<ApexInfo>& apex_infos,
-      android::sp<android::apex::IApexService> apex_service) const;
-  bool AllocateSpace(const uint64_t size_required,
-                     const std::string& dir_path) const;
 };
 
 }  // namespace chromeos_update_engine
diff --git a/aosp/apex_handler_android_unittest.cc b/aosp/apex_handler_android_unittest.cc
index 3a99f79..981ae9d 100644
--- a/aosp/apex_handler_android_unittest.cc
+++ b/aosp/apex_handler_android_unittest.cc
@@ -29,81 +29,45 @@
 
 namespace fs = std::filesystem;
 
-class ApexHandlerAndroidTest : public ::testing::Test {
- protected:
-  ApexHandlerAndroidTest() = default;
-
-  android::sp<android::apex::IApexService> GetApexService() const {
-    return apex_handler_.GetApexService();
-  }
-
-  uint64_t CalculateSize(
-      const std::vector<ApexInfo>& apex_infos,
-      android::sp<android::apex::IApexService> apex_service) const {
-    return apex_handler_.CalculateSize(apex_infos, apex_service);
-  }
-
-  bool AllocateSpace(const uint64_t size_required,
-                     const std::string& dir_path) const {
-    return apex_handler_.AllocateSpace(size_required, dir_path);
-  }
-
-  ApexInfo CreateApexInfo(const std::string& package_name,
-                          int version,
-                          bool is_compressed,
-                          int decompressed_size) {
-    ApexInfo result;
-    result.set_package_name(package_name);
-    result.set_version(version);
-    result.set_is_compressed(is_compressed);
-    result.set_decompressed_size(decompressed_size);
-    return std::move(result);
-  }
-
-  ApexHandlerAndroid apex_handler_;
-};
-
-// TODO(b/172911822): Once apexd has more optimized response for CalculateSize,
-//  improve this test
-TEST_F(ApexHandlerAndroidTest, CalculateSize) {
-  std::vector<ApexInfo> apex_infos;
-  ApexInfo compressed_apex_1 = CreateApexInfo("sample1", 1, true, 10);
-  ApexInfo compressed_apex_2 = CreateApexInfo("sample2", 2, true, 20);
-  apex_infos.push_back(compressed_apex_1);
-  apex_infos.push_back(compressed_apex_2);
-  auto apex_service = GetApexService();
-  EXPECT_TRUE(apex_service != nullptr) << "Apexservice not found";
-  int required_size = CalculateSize(apex_infos, apex_service);
-  EXPECT_EQ(required_size, 30);
+ApexInfo CreateApexInfo(const std::string& package_name,
+                        int version,
+                        bool is_compressed,
+                        int decompressed_size) {
+  ApexInfo result;
+  result.set_package_name(package_name);
+  result.set_version(version);
+  result.set_is_compressed(is_compressed);
+  result.set_decompressed_size(decompressed_size);
+  return std::move(result);
 }
 
-TEST_F(ApexHandlerAndroidTest, AllocateSpace) {
-  // Allocating 0 space should be a no op
-  TemporaryDir td;
-  EXPECT_TRUE(AllocateSpace(0, td.path));
-  EXPECT_TRUE(fs::is_empty(td.path));
+TEST(ApexHandlerAndroidTest, CalculateSize) {
+  ApexHandlerAndroid apex_handler;
+  std::vector<ApexInfo> apex_infos;
+  ApexInfo compressed_apex_1 = CreateApexInfo("sample1", 1, true, 1);
+  ApexInfo compressed_apex_2 = CreateApexInfo("sample2", 2, true, 2);
+  ApexInfo uncompressed_apex = CreateApexInfo("uncompressed", 1, false, 4);
+  apex_infos.push_back(compressed_apex_1);
+  apex_infos.push_back(compressed_apex_2);
+  apex_infos.push_back(uncompressed_apex);
+  auto result = apex_handler.CalculateSize(apex_infos);
+  ASSERT_TRUE(result.ok());
+  EXPECT_EQ(*result, 3u);
+}
 
-  // Allocating non-zero space should create a file with tmp suffix
-  EXPECT_TRUE(AllocateSpace(2 << 20, td.path));
-  EXPECT_FALSE(fs::is_empty(td.path));
-  int num_of_file = 0;
-  for (const auto& entry : fs::directory_iterator(td.path)) {
-    num_of_file++;
-    EXPECT_TRUE(EndsWith(entry.path().string(), ".tmp"));
-    EXPECT_EQ(fs::file_size(entry.path()), 2u << 20);
-  }
-  EXPECT_EQ(num_of_file, 1);
+TEST(ApexHandlerAndroidTest, AllocateSpace) {
+  ApexHandlerAndroid apex_handler;
+  std::vector<ApexInfo> apex_infos;
+  ApexInfo compressed_apex_1 = CreateApexInfo("sample1", 1, true, 1);
+  ApexInfo compressed_apex_2 = CreateApexInfo("sample2", 2, true, 2);
+  ApexInfo uncompressed_apex = CreateApexInfo("uncompressed", 1, false, 4);
+  apex_infos.push_back(compressed_apex_1);
+  apex_infos.push_back(compressed_apex_2);
+  apex_infos.push_back(uncompressed_apex);
+  EXPECT_TRUE(apex_handler.AllocateSpace(apex_infos));
 
-  // AllocateSpace should be safe to call twice
-  EXPECT_TRUE(AllocateSpace(100, td.path));
-  EXPECT_FALSE(fs::is_empty(td.path));
-  num_of_file = 0;
-  for (const auto& entry : fs::directory_iterator(td.path)) {
-    num_of_file++;
-    EXPECT_TRUE(EndsWith(entry.path().string(), ".tmp"));
-    EXPECT_EQ(fs::file_size(entry.path()), 100u);
-  }
-  EXPECT_EQ(num_of_file, 1);
+  // Should be able to pass empty list
+  EXPECT_TRUE(apex_handler.AllocateSpace({}));
 }
 
 }  // namespace chromeos_update_engine
diff --git a/aosp/apex_handler_interface.h b/aosp/apex_handler_interface.h
index c3399b6..b9b6c96 100644
--- a/aosp/apex_handler_interface.h
+++ b/aosp/apex_handler_interface.h
@@ -19,6 +19,8 @@
 
 #include <vector>
 
+#include <android-base/result.h>
+
 #include "update_engine/update_metadata.pb.h"
 
 namespace chromeos_update_engine {
@@ -26,9 +28,9 @@
 class ApexHandlerInterface {
  public:
   virtual ~ApexHandlerInterface() = default;
-  virtual uint64_t CalculateSize(
+  virtual android::base::Result<uint64_t> CalculateSize(
       const std::vector<ApexInfo>& apex_infos) const = 0;
-  virtual bool AllocateSpace(const uint64_t size_required) const = 0;
+  virtual bool AllocateSpace(const std::vector<ApexInfo>& apex_infos) const = 0;
 };
 
 }  // namespace chromeos_update_engine
diff --git a/aosp/update_attempter_android.cc b/aosp/update_attempter_android.cc
index c685855..1080acb 100644
--- a/aosp/update_attempter_android.cc
+++ b/aosp/update_attempter_android.cc
@@ -321,7 +321,8 @@
     const vector<string>& key_value_pair_headers,
     brillo::ErrorPtr* error) {
   // update_engine state must be checked before modifying payload_fd_ otherwise
-  // already running update will be terminated (existing file descriptor will be closed)
+  // already running update will be terminated (existing file descriptor will be
+  // closed)
   if (status_ == UpdateStatus::UPDATED_NEED_REBOOT) {
     return LogAndSetError(
         error, FROM_HERE, "An update already applied, waiting for reboot");
@@ -983,7 +984,14 @@
                                    manifest.apex_info().end());
   uint64_t apex_size_required = 0;
   if (apex_handler_android_ != nullptr) {
-    apex_size_required = apex_handler_android_->CalculateSize(apex_infos);
+    auto result = apex_handler_android_->CalculateSize(apex_infos);
+    if (!result.ok()) {
+      LogAndSetError(error,
+                     FROM_HERE,
+                     "Failed to calculate size required for compressed APEX");
+      return 0;
+    }
+    apex_size_required = *result;
   }
 
   string payload_id = GetPayloadId(headers);
@@ -1006,7 +1014,7 @@
   }
 
   if (apex_size_required > 0 && apex_handler_android_ != nullptr &&
-      !apex_handler_android_->AllocateSpace(apex_size_required)) {
+      !apex_handler_android_->AllocateSpace(apex_infos)) {
     LOG(ERROR) << "Insufficient space for apex decompression: "
                << apex_size_required << " bytes";
     return apex_size_required;