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
diff --git a/fastboot/Android.bp b/fastboot/Android.bp
index 7794c4b..56cac88 100644
--- a/fastboot/Android.bp
+++ b/fastboot/Android.bp
@@ -196,6 +196,7 @@
         "libfastbootshim",
         "libsnapshot_cow",
         "liblz4",
+        "libzstd",
         "libsnapshot_nobinder",
         "update_metadata-protos",
         "liburing",
diff --git a/fs_mgr/libsnapshot/Android.bp b/fs_mgr/libsnapshot/Android.bp
index 3dd1f1a..d3bd904 100644
--- a/fs_mgr/libsnapshot/Android.bp
+++ b/fs_mgr/libsnapshot/Android.bp
@@ -163,6 +163,7 @@
         "libbrotli",
         "libz",
         "liblz4",
+        "libzstd",
     ],
     export_include_dirs: ["include"],
 }
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/cow_format.h b/fs_mgr/libsnapshot/include/libsnapshot/cow_format.h
index 6baf4be..c3ca00a 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/cow_format.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/cow_format.h
@@ -157,7 +157,8 @@
     kCowCompressNone = 0,
     kCowCompressGz = 1,
     kCowCompressBrotli = 2,
-    kCowCompressLz4 = 3
+    kCowCompressLz4 = 3,
+    kCowCompressZstd = 4,
 };
 
 static constexpr uint8_t kCowReadAheadNotStarted = 0;
diff --git a/fs_mgr/libsnapshot/libsnapshot_cow/cow_compress.cpp b/fs_mgr/libsnapshot/libsnapshot_cow/cow_compress.cpp
index d06c904..a4a0ad6 100644
--- a/fs_mgr/libsnapshot/libsnapshot_cow/cow_compress.cpp
+++ b/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 @@
         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 @@
             }
             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;
diff --git a/fs_mgr/libsnapshot/libsnapshot_cow/cow_decompress.cpp b/fs_mgr/libsnapshot/libsnapshot_cow/cow_decompress.cpp
index 3d34413..da90cc0 100644
--- a/fs_mgr/libsnapshot/libsnapshot_cow/cow_decompress.cpp
+++ b/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 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
diff --git a/fs_mgr/libsnapshot/libsnapshot_cow/cow_decompress.h b/fs_mgr/libsnapshot/libsnapshot_cow/cow_decompress.h
index 9e83f3c..52b3d76 100644
--- a/fs_mgr/libsnapshot/libsnapshot_cow/cow_decompress.h
+++ b/fs_mgr/libsnapshot/libsnapshot_cow/cow_decompress.h
@@ -47,6 +47,7 @@
     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);
 
diff --git a/fs_mgr/libsnapshot/libsnapshot_cow/cow_reader.cpp b/fs_mgr/libsnapshot/libsnapshot_cow/cow_reader.cpp
index 4c5a8bf..6749cf1 100644
--- a/fs_mgr/libsnapshot/libsnapshot_cow/cow_reader.cpp
+++ b/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 @@
         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();
diff --git a/fs_mgr/libsnapshot/snapuserd/Android.bp b/fs_mgr/libsnapshot/snapuserd/Android.bp
index 1e03683..9fe567a 100644
--- a/fs_mgr/libsnapshot/snapuserd/Android.bp
+++ b/fs_mgr/libsnapshot/snapuserd/Android.bp
@@ -112,6 +112,7 @@
         "liblz4",
         "libext4_utils",
         "liburing",
+        "libzstd",
     ],
 
     header_libs: [
diff --git a/init/Android.bp b/init/Android.bp
index 7b52903..41c7a95 100644
--- a/init/Android.bp
+++ b/init/Android.bp
@@ -169,6 +169,7 @@
         "libfsverity_init",
         "liblmkd_utils",
         "liblz4",
+        "libzstd",
         "libmini_keyctl_static",
         "libmodprobe",
         "libprocinfo",
@@ -370,6 +371,7 @@
         "libprotobuf-cpp-lite",
         "libsnapshot_cow",
         "liblz4",
+        "libzstd",
         "libsnapshot_init",
         "update_metadata-protos",
         "libprocinfo",