Support ZSTD in userspace COW

Perf:
Lz4:
Update took 429 seconds
merge time 35.4s
cow size: 3.18GB

ZSTD:
Update took 676 seconds
merge time 49.4s
cow size: 2.62GB

Gz:
Update took 1057 seconds
merge time: 50.0s
cow size: 2.55GB

In summary, ZSTD is a middle point between LZ4 and GZ.

Speed: LZ4 > ZSTD > GZ
Space: LZ4 > ZSTD > GZ

Bug: 274129758

Change-Id: I203bf088b7c2a9ce429f75478799da0e7126febf

fastboot/Android.bp

@@ -196,6 +196,7 @@ cc_binary {
         "libfastbootshim",
         "libsnapshot_cow",
         "liblz4",
+        "libzstd",
         "libsnapshot_nobinder",
         "update_metadata-protos",
         "liburing",

fs_mgr/libsnapshot/Android.bp

@@ -163,6 +163,7 @@ cc_defaults {
         "libbrotli",
         "libz",
         "liblz4",
+        "libzstd",
     ],
     export_include_dirs: ["include"],
 }

fs_mgr/libsnapshot/include/libsnapshot/cow_format.h

@@ -157,7 +157,8 @@ enum CowCompressionAlgorithm : uint8_t {
     kCowCompressNone = 0,
     kCowCompressGz = 1,
     kCowCompressBrotli = 2,
-    kCowCompressLz4 = 3
+    kCowCompressLz4 = 3,
+    kCowCompressZstd = 4,
 };
 
 static constexpr uint8_t kCowReadAheadNotStarted = 0;

fs_mgr/libsnapshot/libsnapshot_cow/cow_compress.cpp

@@ -29,6 +29,7 @@
 #include <libsnapshot/cow_writer.h>
 #include <lz4.h>
 #include <zlib.h>
+#include <zstd.h>
 
 namespace android {
 namespace snapshot {
@@ -40,6 +41,8 @@ std::optional<CowCompressionAlgorithm> CompressionAlgorithmFromString(std::strin
         return {kCowCompressBrotli};
     } else if (name == "lz4") {
         return {kCowCompressLz4};
+    } else if (name == "zstd") {
+        return {kCowCompressZstd};
     } else if (name == "none" || name.empty()) {
         return {kCowCompressNone};
     } else {
@@ -112,6 +115,23 @@ std::basic_string<uint8_t> CompressWorker::Compress(CowCompressionAlgorithm comp
             }
             return buffer;
         }
+        case kCowCompressZstd: {
+            std::basic_string<uint8_t> buffer(ZSTD_compressBound(length), '\0');
+            const auto compressed_size =
+                    ZSTD_compress(buffer.data(), buffer.size(), data, length, 0);
+            if (compressed_size <= 0) {
+                LOG(ERROR) << "ZSTD compression failed " << compressed_size;
+                return {};
+            }
+            // Don't run compression if the compressed output is larger
+            if (compressed_size >= length) {
+                buffer.resize(length);
+                memcpy(buffer.data(), data, length);
+            } else {
+                buffer.resize(compressed_size);
+            }
+            return buffer;
+        }
         default:
             LOG(ERROR) << "unhandled compression type: " << compression;
             break;

fs_mgr/libsnapshot/libsnapshot_cow/cow_decompress.cpp

@@ -17,12 +17,15 @@
 #include "cow_decompress.h"
 
 #include <array>
+#include <cstring>
 #include <utility>
+#include <vector>
 
 #include <android-base/logging.h>
 #include <brotli/decode.h>
 #include <lz4.h>
 #include <zlib.h>
+#include <zstd.h>
 
 namespace android {
 namespace snapshot {
@@ -336,9 +339,56 @@ class Lz4Decompressor final : public IDecompressor {
     }
 };
 
+class ZstdDecompressor final : public IDecompressor {
+  public:
+    ssize_t Decompress(void* buffer, size_t buffer_size, size_t decompressed_size,
+                       size_t ignore_bytes = 0) override {
+        if (buffer_size < decompressed_size - ignore_bytes) {
+            LOG(INFO) << "buffer size " << buffer_size
+                      << " is not large enough to hold decompressed data. Decompressed size "
+                      << decompressed_size << ", ignore_bytes " << ignore_bytes;
+            return -1;
+        }
+        if (ignore_bytes == 0) {
+            if (!Decompress(buffer, decompressed_size)) {
+                return -1;
+            }
+            return decompressed_size;
+        }
+        std::vector<unsigned char> ignore_buf(decompressed_size);
+        if (!Decompress(buffer, decompressed_size)) {
+            return -1;
+        }
+        memcpy(buffer, ignore_buf.data() + ignore_bytes, buffer_size);
+        return decompressed_size;
+    }
+    bool Decompress(void* output_buffer, const size_t output_size) {
+        std::string input_buffer;
+        input_buffer.resize(stream_->Size());
+        size_t bytes_read = stream_->Read(input_buffer.data(), input_buffer.size());
+        if (bytes_read != input_buffer.size()) {
+            LOG(ERROR) << "Failed to read all input at once. Expected: " << input_buffer.size()
+                       << " actual: " << bytes_read;
+            return false;
+        }
+        const auto bytes_decompressed = ZSTD_decompress(output_buffer, output_size,
+                                                        input_buffer.data(), input_buffer.size());
+        if (bytes_decompressed != output_size) {
+            LOG(ERROR) << "Failed to decompress ZSTD block, expected output size: " << output_size
+                       << ", actual: " << bytes_decompressed;
+            return false;
+        }
+        return true;
+    }
+};
+
 std::unique_ptr<IDecompressor> IDecompressor::Lz4() {
     return std::make_unique<Lz4Decompressor>();
 }
 
+std::unique_ptr<IDecompressor> IDecompressor::Zstd() {
+    return std::make_unique<ZstdDecompressor>();
+}
+
 }  // namespace snapshot
 }  // namespace android

fs_mgr/libsnapshot/libsnapshot_cow/cow_decompress.h

@@ -47,6 +47,7 @@ class IDecompressor {
     static std::unique_ptr<IDecompressor> Gz();
     static std::unique_ptr<IDecompressor> Brotli();
     static std::unique_ptr<IDecompressor> Lz4();
+    static std::unique_ptr<IDecompressor> Zstd();
 
     static std::unique_ptr<IDecompressor> FromString(std::string_view compressor);
 

fs_mgr/libsnapshot/libsnapshot_cow/cow_reader.cpp

@@ -30,6 +30,7 @@
 #include <zlib.h>
 
 #include "cow_decompress.h"
+#include "libsnapshot/cow_format.h"
 
 namespace android {
 namespace snapshot {
@@ -777,6 +778,11 @@ ssize_t CowReader::ReadData(const CowOperation* op, void* buffer, size_t buffer_
         case kCowCompressBrotli:
             decompressor = IDecompressor::Brotli();
             break;
+        case kCowCompressZstd:
+            if (header_.block_size != op->data_length) {
+                decompressor = IDecompressor::Zstd();
+            }
+            break;
         case kCowCompressLz4:
             if (header_.block_size != op->data_length) {
                 decompressor = IDecompressor::Lz4();

fs_mgr/libsnapshot/snapuserd/Android.bp

@@ -112,6 +112,7 @@ cc_defaults {
         "liblz4",
         "libext4_utils",
         "liburing",
+        "libzstd",
     ],
 
     header_libs: [

init/Android.bp

@@ -169,6 +169,7 @@ libinit_cc_defaults {
         "libfsverity_init",
         "liblmkd_utils",
         "liblz4",
+        "libzstd",
         "libmini_keyctl_static",
         "libmodprobe",
         "libprocinfo",
@@ -370,6 +371,7 @@ cc_binary {
         "libprotobuf-cpp-lite",
         "libsnapshot_cow",
         "liblz4",
+        "libzstd",
         "libsnapshot_init",
         "update_metadata-protos",
         "libprocinfo",