Add repodata.zst support (#2113)

.gitignore

@@ -80,3 +80,4 @@ __cache__/
 .rendered*
 installed.json
 tmp/
+test_7.json.state

CMakeLists.txt

@@ -32,7 +32,9 @@ if (MSVC)
     # /W4 : warnings level 4 (default in visual studio projects and recommended minimum level).
     # /external:I $ENV{CONDA_PREFIX}: consider the conda env prefix libraries headers as "external" to this project.
     # /external:W0 : set the warning level to 1 for external libraries headers (severe warnings, default when unspecified) to only see important warnings coming from dependencies headers.
-    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /DNOMINMAX /EHsc /Zc:__cplusplus /MP /W4 /external:I $ENV{CONDA_PREFIX} /external:W1")
+    string(REGEX REPLACE "/W[0-9]" "/W4" CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}")
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /D_CRT_SECURE_NO_WARNINGS /DNOMINMAX /EHsc /Zc:__cplusplus /MP /experimental:external /external:I $ENV{CONDA_PREFIX} /external:W1")
+    # Force release mode to avoid debug libraries to be linked
     set(CMAKE_BUILD_TYPE Release)
     # add_definitions("-DUNICODE -D_UNICODE")
 else()

libmamba/CMakeLists.txt

@@ -430,6 +430,7 @@ macro(libmamba_create_target target_name linkage deps_linkage output_name)
         find_package(CURL REQUIRED)
         find_package(LibArchive REQUIRED)
         find_package(zstd REQUIRED)
+        find_package(BZip2 REQUIRED)
         find_package(OpenSSL REQUIRED)
         find_package(yaml-cpp CONFIG REQUIRED)
         find_package(reproc++ CONFIG REQUIRED)
@@ -444,6 +445,8 @@ macro(libmamba_create_target target_name linkage deps_linkage output_name)
             zstd::libzstd_shared
             ${CURL_LIBRARIES}
             ${OPENSSL_LIBRARIES}
+            zstd::libzstd_shared
+            BZip2::BZip2
             yaml-cpp
             reproc++
             reproc

libmamba/include/mamba/core/context.hpp

@@ -236,11 +236,12 @@ namespace mamba
 
         bool override_channels_enabled = true;
 
-
         std::vector<std::string> pinned_packages = {};
 
         bool use_only_tar_bz2 = false;
 
+        bool repodata_use_zst = false;
+        std::vector<std::string> repodata_has_zst = { "https://conda.anaconda.org/conda-forge" };
 
         // usernames on anaconda.org can have a underscore, which influences the
         // first two characters

libmamba/include/mamba/core/error_handling.hpp

@@ -30,6 +30,7 @@ namespace mamba
         lockfile_failure,
         selfupdate_failure,
         satisfiablitity_error,
+        user_interrupted,
     };
 
     class mamba_error : public std::runtime_error

libmamba/include/mamba/core/fetch.hpp

@@ -15,8 +15,8 @@ extern "C"
 
 #include <string>
 #include <vector>
-
-#include "nlohmann/json.hpp"
+#include <zstd.h>
+#include <bzlib.h>
 
 #include "progress_bar.hpp"
 #include "validate.hpp"
@@ -25,6 +25,78 @@ namespace mamba
 {
     void init_curl_ssl();
 
+    struct ZstdStream
+    {
+        constexpr static size_t BUFFER_SIZE = 256000;
+        ZstdStream(curl_write_callback write_callback, void* write_callback_data)
+            : stream(ZSTD_createDCtx())
+            , m_write_callback(write_callback)
+            , m_write_callback_data(write_callback_data)
+        {
+            ZSTD_initDStream(stream);
+        }
+
+        ~ZstdStream()
+        {
+            ZSTD_freeDCtx(stream);
+        }
+
+        size_t write(char* in, size_t size);
+
+        static size_t write_callback(char* ptr, size_t size, size_t nmemb, void* self)
+        {
+            return static_cast<ZstdStream*>(self)->write(ptr, size * nmemb);
+        }
+
+        ZSTD_DCtx* stream;
+        char buffer[BUFFER_SIZE];
+
+        // original curl callback
+        curl_write_callback m_write_callback;
+        void* m_write_callback_data;
+    };
+
+    struct Bzip2Stream
+    {
+        constexpr static size_t BUFFER_SIZE = 256000;
+
+        Bzip2Stream(curl_write_callback write_callback, void* write_callback_data)
+            : m_write_callback(write_callback)
+            , m_write_callback_data(write_callback_data)
+        {
+            stream.bzalloc = nullptr;
+            stream.bzfree = nullptr;
+            stream.opaque = nullptr;
+
+            error = BZ2_bzDecompressInit(&stream, 0, false);
+            if (error != BZ_OK)
+            {
+                throw std::runtime_error("BZ2_bzDecompressInit failed");
+            }
+        }
+
+        size_t write(char* in, size_t size);
+
+        static size_t write_callback(char* ptr, size_t size, size_t nmemb, void* self)
+        {
+            return static_cast<Bzip2Stream*>(self)->write(ptr, size * nmemb);
+        }
+
+        ~Bzip2Stream()
+        {
+            BZ2_bzDecompressEnd(&stream);
+        }
+
+        int error;
+        bz_stream stream;
+        char buffer[BUFFER_SIZE];
+
+        // original curl callback
+        curl_write_callback m_write_callback;
+        void* m_write_callback_data;
+    };
+
+
     class DownloadTarget
     {
     public:
@@ -45,11 +117,16 @@ namespace mamba
 
         static int progress_callback(
             void*, curl_off_t total_to_download, curl_off_t now_downloaded, curl_off_t, curl_off_t);
-        void set_mod_etag_headers(const nlohmann::json& mod_etag);
+        void set_mod_etag_headers(const std::string& mod, const std::string& etag);
         void set_progress_bar(ProgressProxy progress_proxy);
         void set_expected_size(std::size_t size);
+        void set_head_only(bool yes)
+        {
+            curl_easy_setopt(m_handle, CURLOPT_NOBODY, yes);
+        }
 
         const std::string& name() const;
+        const std::string& url() const;
         std::size_t expected_size() const;
 
         void init_curl_target(const std::string& url);
@@ -61,9 +138,9 @@ namespace mamba
         curl_off_t get_speed();
 
         template <class C>
-        inline void set_finalize_callback(bool (C::*cb)(), C* data)
+        inline void set_finalize_callback(bool (C::*cb)(const DownloadTarget&), C* data)
         {
-            m_finalize_callback = std::bind(cb, data);
+            m_finalize_callback = std::bind(cb, data, std::placeholders::_1);
         }
 
         inline void set_ignore_failure(bool yes)
@@ -97,7 +174,9 @@ namespace mamba
         std::string etag, mod, cache_control;
 
     private:
-        std::function<bool()> m_finalize_callback;
+        std::unique_ptr<ZstdStream> m_zstd_stream;
+        std::unique_ptr<Bzip2Stream> m_bzip2_stream;
+        std::function<bool(const DownloadTarget&)> m_finalize_callback;
 
         std::string m_name, m_filename, m_url;
 
@@ -145,6 +224,7 @@ namespace mamba
 
     const int MAMBA_DOWNLOAD_FAILFAST = 1 << 0;
     const int MAMBA_DOWNLOAD_SORT = 1 << 1;
+    const int MAMBA_NO_CLEAR_PROGRESS_BARS = 1 << 2;
 }  // namespace mamba
 
 #endif  // MAMBA_FETCH_HPP

libmamba/include/mamba/core/mamba_fs.hpp

@@ -4,6 +4,7 @@
 #include <filesystem>
 #include <fstream>
 #include <string>
+#include <fmt/format.h>
 
 #include "mamba/core/util_string.hpp"
 
@@ -1170,7 +1171,7 @@ namespace fs
     }
 
     // void last_write_time(const path& p, now _, error_code& ec) noexcept;
-    inline void last_write_time(const u8path& path, now _, std::error_code& ec) noexcept
+    inline void last_write_time(const u8path& path, now, std::error_code& ec) noexcept
     {
 #if defined(USE_UTIMENSAT)
         if (utimensat(AT_FDCWD, path.string().c_str(), NULL, 0) == -1)
@@ -1380,5 +1381,21 @@ struct std::hash<::fs::u8path>
     }
 };
 
+template <>
+struct fmt::formatter<::fs::u8path>
+{
+    constexpr auto parse(format_parse_context& ctx) -> decltype(ctx.begin())
+    {
+        // make sure that range is empty
+        if (ctx.begin() != ctx.end() && *ctx.begin() != '}')
+            throw format_error("invalid format");
+        return ctx.begin();
+    }
 
+    template <class FormatContext>
+    auto format(const ::fs::u8path& path, FormatContext& ctx)
+    {
+        return fmt::format_to(ctx.out(), "'{}'", path.string());
+    }
+};
 #endif

libmamba/include/mamba/core/package_handling.hpp

@@ -17,6 +17,7 @@ namespace mamba
 {
     enum compression_algorithm
     {
+        none,
         bzip2,
         zip,
         zstd

libmamba/include/mamba/core/subdirdata.hpp

@@ -11,8 +11,6 @@
 #include <regex>
 #include <string>
 
-#include "nlohmann/json.hpp"
-
 #include "mamba/core/channel.hpp"
 #include "mamba/core/context.hpp"
 #include "mamba/core/error_handling.hpp"
@@ -22,15 +20,54 @@
 #include "mamba/core/pool.hpp"
 #include "mamba/core/repo.hpp"
 #include "mamba/core/util.hpp"
-
+#include "package_handling.hpp"
 
 namespace decompress
 {
-    bool raw(const std::string& in, const std::string& out);
+    bool raw(mamba::compression_algorithm ca, const std::string& in, const std::string& out);
 }
 
 namespace mamba
 {
+    struct subdir_metadata
+    {
+        struct checked_at
+        {
+            bool value;
+            std::time_t last_checked;
+
+            bool has_expired() const
+            {
+                // difference in seconds, check every 14 days
+                return std::difftime(std::time(nullptr), last_checked) > 60 * 60 * 24 * 14;
+            }
+        };
+
+        static tl::expected<subdir_metadata, mamba_error> from_stream(std::istream& in);
+
+        std::string url;
+        std::string etag;
+        std::string mod;
+        std::string cache_control;
+#ifdef _WIN32
+        std::chrono::system_clock::time_point stored_mtime;
+#else
+        fs::file_time_type stored_mtime;
+#endif
+        std::size_t stored_file_size;
+        std::optional<checked_at> has_zst;
+        std::optional<checked_at> has_bz2;
+        std::optional<checked_at> has_jlap;
+
+        void store_file_metadata(const fs::u8path& path);
+        bool check_valid_metadata(const fs::u8path& path);
+
+        void serialize_to_stream(std::ostream& out) const;
+        void serialize_to_stream_tiny(std::ostream& out) const;
+
+        bool check_zst(const Channel* channel);
+    };
+
 
     /**
      * Represents a channel subdirectory (i.e. a platform)
@@ -55,8 +92,8 @@ namespace mamba
         MSubdirData& operator=(MSubdirData&&);
 
         // TODO return seconds as double
-        fs::file_time_type::duration check_cache(const fs::u8path& cache_file,
-                                                 const fs::file_time_type::clock::time_point& ref);
+        fs::file_time_type::duration check_cache(
+            const fs::u8path& cache_file, const fs::file_time_type::clock::time_point& ref) const;
         bool loaded() const;
 
         bool forbid_cache();
@@ -65,9 +102,12 @@ namespace mamba
         expected_t<std::string> cache_path() const;
         const std::string& name() const;
 
+        std::vector<std::unique_ptr<DownloadTarget>>& check_targets();
         DownloadTarget* target();
-        bool finalize_transfer();
 
+        bool finalize_check(const DownloadTarget& target);
+        bool finalize_transfer(const DownloadTarget& target);
+        void finalize_checks();
         expected_t<MRepo&> create_repo(MPool& pool);
 
     private:
@@ -78,11 +118,13 @@ namespace mamba
                     const std::string& repodata_fn = "repodata.json");
 
         bool load(MultiPackageCache& caches);
-        bool decompress();
-        void create_target(nlohmann::json& mod_etag);
+        void check_repodata_existence();
+        void create_target();
         std::size_t get_cache_control_max_age(const std::string& val);
+        void refresh_last_write_time(const fs::u8path& json_file, const fs::u8path& solv_file);
 
         std::unique_ptr<DownloadTarget> m_target = nullptr;
+        std::vector<std::unique_ptr<DownloadTarget>> m_check_targets;
 
         bool m_json_cache_valid = false;
         bool m_solv_cache_valid = false;
@@ -92,6 +134,7 @@ namespace mamba
         fs::u8path m_writable_pkgs_dir;
 
         ProgressProxy m_progress_bar;
+        ProgressProxy m_progress_bar_check;
 
         bool m_loaded;
         bool m_download_complete;
@@ -100,7 +143,7 @@ namespace mamba
         std::string m_json_fn;
         std::string m_solv_fn;
         bool m_is_noarch;
-        nlohmann::json m_mod_etag;
+        subdir_metadata m_metadata;
         std::unique_ptr<TemporaryFile> m_temp_file;
         const Channel* p_channel = nullptr;
     };

libmamba/include/mamba/core/transaction.hpp

@@ -50,7 +50,7 @@ namespace mamba
 
         void write_repodata_record(const fs::u8path& base_path);
         void add_url();
-        bool finalize_callback();
+        bool finalize_callback(const DownloadTarget& target);
         bool finished();
         void validate();
         bool extract();