diff --git a/fs_mgr/libsnapshot/include_test/libsnapshot/test_helpers.h b/fs_mgr/libsnapshot/include_test/libsnapshot/test_helpers.h
index 7aef086..ff0047e 100644
--- a/fs_mgr/libsnapshot/include_test/libsnapshot/test_helpers.h
+++ b/fs_mgr/libsnapshot/include_test/libsnapshot/test_helpers.h
@@ -51,8 +51,8 @@
 extern std::unique_ptr<SnapshotManager> sm;
 extern class TestDeviceInfo* test_device;
 extern std::string fake_super;
-static constexpr uint64_t kSuperSize = 16_MiB + 4_KiB;
-static constexpr uint64_t kGroupSize = 16_MiB;
+static constexpr uint64_t kSuperSize = 32_MiB + 4_KiB;
+static constexpr uint64_t kGroupSize = 32_MiB;
 
 // Redirect requests for "super" to our fake super partition.
 class TestPartitionOpener final : public android::fs_mgr::PartitionOpener {
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index f0646fc..21b720e 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -676,6 +676,8 @@
         }
     }
 
+    bool compression_enabled = false;
+
     uint64_t total_cow_file_size = 0;
     DmTargetSnapshot::Status initial_target_values = {};
     for (const auto& snapshot : snapshots) {
@@ -692,6 +694,8 @@
             return false;
         }
         total_cow_file_size += snapshot_status.cow_file_size();
+
+        compression_enabled |= snapshot_status.compression_enabled();
     }
 
     if (cow_file_size) {
@@ -703,6 +707,7 @@
     initial_status.set_sectors_allocated(initial_target_values.sectors_allocated);
     initial_status.set_total_sectors(initial_target_values.total_sectors);
     initial_status.set_metadata_sectors(initial_target_values.metadata_sectors);
+    initial_status.set_compression_enabled(compression_enabled);
 
     // Point of no return - mark that we're starting a merge. From now on every
     // snapshot must be a merge target.
@@ -1405,7 +1410,7 @@
             return false;
         }
 
-        CHECK(base_sectors == target.spec.length);
+        CHECK(base_sectors <= target.spec.length);
 
         if (!snapuserd_client_->AttachDmUser(misc_name)) {
             // This error is unrecoverable. We cannot proceed because reads to
@@ -2285,9 +2290,15 @@
 }
 
 bool SnapshotManager::WriteUpdateState(LockedFile* lock, UpdateState state) {
-    SnapshotUpdateStatus status = {};
+    SnapshotUpdateStatus status;
     status.set_state(state);
-    status.set_compression_enabled(IsCompressionEnabled());
+
+    // If we're transitioning between two valid states (eg, we're not beginning
+    // or ending an OTA), then make sure to propagate the compression bit.
+    if (!(state == UpdateState::Initiated || state == UpdateState::None)) {
+        SnapshotUpdateStatus old_status = ReadSnapshotUpdateStatus(lock);
+        status.set_compression_enabled(old_status.compression_enabled());
+    }
     return WriteSnapshotUpdateStatus(lock, status);
 }
 
@@ -2477,6 +2488,12 @@
     auto lock = LockExclusive();
     if (!lock) return Return::Error();
 
+    auto update_state = ReadUpdateState(lock.get());
+    if (update_state != UpdateState::Initiated) {
+        LOG(ERROR) << "Cannot create update snapshots in state " << update_state;
+        return Return::Error();
+    }
+
     // TODO(b/134949511): remove this check. Right now, with overlayfs mounted, the scratch
     // partition takes up a big chunk of space in super, causing COW images to be created on
     // retrofit Virtual A/B devices.
@@ -2571,6 +2588,14 @@
         return Return::Error();
     }
 
+    SnapshotUpdateStatus status = {};
+    status.set_state(update_state);
+    status.set_compression_enabled(cow_creator.compression_enabled);
+    if (!WriteSnapshotUpdateStatus(lock.get(), status)) {
+        LOG(ERROR) << "Unable to write new update state";
+        return Return::Error();
+    }
+
     created_devices.Release();
     LOG(INFO) << "Successfully created all snapshots for target slot " << target_suffix;
 
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index bb44425..80a6074 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -854,15 +854,15 @@
         group_->add_partition_names("prd");
         sys_ = manifest_.add_partitions();
         sys_->set_partition_name("sys");
-        sys_->set_estimate_cow_size(6_MiB);
+        sys_->set_estimate_cow_size(2_MiB);
         SetSize(sys_, 3_MiB);
         vnd_ = manifest_.add_partitions();
         vnd_->set_partition_name("vnd");
-        vnd_->set_estimate_cow_size(6_MiB);
+        vnd_->set_estimate_cow_size(2_MiB);
         SetSize(vnd_, 3_MiB);
         prd_ = manifest_.add_partitions();
         prd_->set_partition_name("prd");
-        prd_->set_estimate_cow_size(6_MiB);
+        prd_->set_estimate_cow_size(2_MiB);
         SetSize(prd_, 3_MiB);
 
         // Initialize source partition metadata using |manifest_|.
@@ -1050,11 +1050,17 @@
         ASSERT_TRUE(sm->UnmapUpdateSnapshot(name));
     }
 
-    // Grow all partitions.
+    // Grow all partitions. Set |prd| large enough that |sys| and |vnd|'s COWs
+    // fit in super, but not |prd|.
     constexpr uint64_t partition_size = 3788_KiB;
     SetSize(sys_, partition_size);
     SetSize(vnd_, partition_size);
-    SetSize(prd_, partition_size);
+    SetSize(prd_, 18_MiB);
+
+    // Make sure |prd| does not fit in super at all. On VABC, this means we
+    // fake an extra large COW for |vnd| to fill up super.
+    vnd_->set_estimate_cow_size(30_MiB);
+    prd_->set_estimate_cow_size(30_MiB);
 
     AddOperationForPartitions();
 
@@ -1066,11 +1072,7 @@
     auto tgt = MetadataBuilder::New(*opener_, "super", 1);
     ASSERT_NE(tgt, nullptr);
     ASSERT_NE(nullptr, tgt->FindPartition("sys_b-cow"));
-    if (IsCompressionEnabled()) {
-        ASSERT_EQ(nullptr, tgt->FindPartition("vnd_b-cow"));
-    } else {
-        ASSERT_NE(nullptr, tgt->FindPartition("vnd_b-cow"));
-    }
+    ASSERT_NE(nullptr, tgt->FindPartition("vnd_b-cow"));
     ASSERT_EQ(nullptr, tgt->FindPartition("prd_b-cow"));
 
     // Write some data to target partitions.
@@ -1104,6 +1106,7 @@
 
     // Initiate the merge and wait for it to be completed.
     ASSERT_TRUE(init->InitiateMerge());
+    ASSERT_EQ(init->IsSnapuserdRequired(), IsCompressionEnabled());
     ASSERT_EQ(UpdateState::MergeCompleted, init->ProcessUpdateState());
 
     // Check that the target partitions have the same content after the merge.
@@ -1128,6 +1131,7 @@
     SetSize(sys_, 4_MiB);  // grows
     SetSize(vnd_, 2_MiB);  // shrinks
     // prd_b is unchanged
+    ASSERT_TRUE(sm->BeginUpdate());
     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
     ASSERT_EQ(4_MiB, GetSnapshotSize("sys_b").value_or(0));
 }
@@ -1137,6 +1141,7 @@
 TEST_F(SnapshotUpdateTest, CowPartitionDoNotTakeOldPartitions) {
     SetSize(sys_, 2_MiB);  // shrinks
     // vnd_b and prd_b are unchanged.
+    ASSERT_TRUE(sm->BeginUpdate());
     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
 
     auto tgt = MetadataBuilder::New(*opener_, "super", 1);
@@ -1260,6 +1265,11 @@
 
 // Test that at the second update, old COW partition spaces are reclaimed.
 TEST_F(SnapshotUpdateTest, ReclaimCow) {
+    // Make sure VABC cows are small enough that they fit in fake_super.
+    sys_->set_estimate_cow_size(64_KiB);
+    vnd_->set_estimate_cow_size(64_KiB);
+    prd_->set_estimate_cow_size(64_KiB);
+
     // Execute the first update.
     ASSERT_TRUE(sm->BeginUpdate());
     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
@@ -1376,9 +1386,13 @@
 
 TEST_F(SnapshotUpdateTest, MergeCannotRemoveCow) {
     // Make source partitions as big as possible to force COW image to be created.
-    SetSize(sys_, 5_MiB);
-    SetSize(vnd_, 5_MiB);
-    SetSize(prd_, 5_MiB);
+    SetSize(sys_, 10_MiB);
+    SetSize(vnd_, 10_MiB);
+    SetSize(prd_, 10_MiB);
+    sys_->set_estimate_cow_size(12_MiB);
+    vnd_->set_estimate_cow_size(12_MiB);
+    prd_->set_estimate_cow_size(12_MiB);
+
     src_ = MetadataBuilder::New(*opener_, "super", 0);
     ASSERT_NE(src_, nullptr);
     src_->RemoveGroupAndPartitions(group_->name() + "_a");
@@ -1676,6 +1690,8 @@
     SetSize(sys_, partition_size);
     AddOperation(sys_, data_size);
 
+    sys_->set_estimate_cow_size(partition_size + data_size);
+
     // Set hastree extents.
     sys_->mutable_hash_tree_data_extent()->set_start_block(0);
     sys_->mutable_hash_tree_data_extent()->set_num_blocks(data_size / block_size);
@@ -1716,6 +1732,10 @@
 
 // Test for overflow bit after update
 TEST_F(SnapshotUpdateTest, Overflow) {
+    if (IsCompressionEnabled()) {
+        GTEST_SKIP() << "No overflow bit set for userspace COWs";
+    }
+
     const auto actual_write_size = GetSize(sys_);
     const auto declared_write_size = actual_write_size - 1_MiB;
 
@@ -1743,12 +1763,15 @@
     auto userdata = std::make_unique<LowSpaceUserdata>();
     ASSERT_TRUE(userdata->Init(kMaxFree));
 
-    // Grow all partitions to 5_MiB, total 15_MiB. This requires 15 MiB of CoW space. After
-    // using the empty space in super (< 1 MiB), it uses at least 14 MiB of /userdata space.
-    constexpr uint64_t partition_size = 5_MiB;
+    // Grow all partitions to 10_MiB, total 30_MiB. This requires 30 MiB of CoW space. After
+    // using the empty space in super (< 1 MiB), it uses 30 MiB of /userdata space.
+    constexpr uint64_t partition_size = 10_MiB;
     SetSize(sys_, partition_size);
     SetSize(vnd_, partition_size);
     SetSize(prd_, partition_size);
+    sys_->set_estimate_cow_size(partition_size);
+    vnd_->set_estimate_cow_size(partition_size);
+    prd_->set_estimate_cow_size(partition_size);
 
     AddOperationForPartitions();
 
@@ -1758,7 +1781,7 @@
     ASSERT_FALSE(res);
     ASSERT_EQ(Return::ErrorCode::NO_SPACE, res.error_code());
     ASSERT_GE(res.required_size(), 14_MiB);
-    ASSERT_LT(res.required_size(), 15_MiB);
+    ASSERT_LT(res.required_size(), 40_MiB);
 }
 
 class AutoKill final {