[fbuffers] Future file buffer implementation

A file buffer that can be fulfilled using string keys. The extract method waits until the file is provided.

Author: jesusmb1995

src/llama-mmap.cpp

@@ -338,6 +338,112 @@ template <> void llama_file_buffer<true>::write_u32(uint32_t val) const {
 template struct llama_file_buffer<false>;
 template struct llama_file_buffer<true>;
 
+// llama_future_file_buffer implementation
+
+namespace {
+std::string final_key(const std::string & promise_key, const std::string & context) {
+    return promise_key + ":" + context;
+}
+
+std::mutex promise_registry_mutex;
+
+std::map<std::string, std::promise<std::unique_ptr<llama_file_buffer<false>>>> promise_registry_ro;
+std::map<std::string, std::promise<std::unique_ptr<llama_file_buffer<true>>>>  promise_registry_rw;
+
+template <bool Writable>
+std::map<std::string, std::promise<std::unique_ptr<llama_file_buffer<Writable>>>> & promise_registry() {
+    if constexpr (Writable) {
+        return promise_registry_rw;
+    } else {
+        return promise_registry_ro;
+    }
+}
+
+/// @brief Ensures a promise exists in the registry for the given key.
+/// If it doesn't exist, creates it. Returns an iterator to the promise.
+/// Thread-safe.
+template <bool Writable>
+typename std::map<std::string, std::promise<std::unique_ptr<llama_file_buffer<Writable>>>>::iterator
+ensure_promise_registry(const std::string & key) {
+    std::lock_guard<std::mutex> lock(promise_registry_mutex);
+    auto                        it = promise_registry<Writable>().find(key);
+    if (it != promise_registry<Writable>().end()) {
+        return it;
+    }
+    auto result =
+        promise_registry<Writable>().emplace(key, std::promise<std::unique_ptr<llama_file_buffer<Writable>>>());
+    LLAMA_LOG_CMAKE_DEBUG("%s: created future file buffer %p for %s\n", __func__, (void *) &(*it), key.c_str());
+    return result.first;
+}
+}  // namespace
+
+template <bool Writable>
+llama_future_file_buffer<Writable>::llama_future_file_buffer(const std::string & promise_key,
+                                                             const std::string & context) :
+    file_buffer_future(),
+    file_buffer() {
+    std::string key              = final_key(promise_key, context);
+    file_buffer_promise_iterator = ensure_promise_registry<Writable>(key);
+    file_buffer_future           = file_buffer_promise_iterator->second.get_future();
+}
+
+template <bool Writable>
+llama_future_file_buffer<Writable>::llama_future_file_buffer(llama_future_file_buffer && other) noexcept :
+    file_buffer_promise_iterator(std::move(other.file_buffer_promise_iterator)),
+    file_buffer_future(std::move(other.file_buffer_future)),
+    file_buffer(std::move(other.file_buffer)) {
+    // Set the other object's iterator to end() to mark it as moved from
+    // to avoid early erasure at destruction of the moved other object
+    other.file_buffer_promise_iterator = promise_registry<Writable>().end();
+}
+
+template <bool Writable>
+llama_future_file_buffer<Writable> & llama_future_file_buffer<Writable>::operator=(
+    llama_future_file_buffer && other) noexcept {
+    if (this != &other) {
+        file_buffer_promise_iterator       = std::move(other.file_buffer_promise_iterator);
+        file_buffer_future                 = std::move(other.file_buffer_future);
+        file_buffer                        = std::move(other.file_buffer);
+        other.file_buffer_promise_iterator = promise_registry<Writable>().end();
+    }
+    return *this;
+}
+
+template <bool Writable> llama_future_file_buffer<Writable>::~llama_future_file_buffer() {
+    std::lock_guard<std::mutex> lock(promise_registry_mutex);
+    if (file_buffer_promise_iterator != promise_registry<Writable>().end()) {
+        promise_registry<Writable>().erase(file_buffer_promise_iterator);
+    }
+}
+
+template <bool Writable>
+bool llama_future_file_buffer<Writable>::fulfill_promise(const std::string & promise_key, const std::string & context,
+                                                         std::unique_ptr<llama_file_buffer<Writable>> && value) {
+    std::string key = final_key(promise_key, context);
+    auto        it  = ensure_promise_registry<Writable>(key);
+    if (it != promise_registry<Writable>().end()) {
+        LLAMA_LOG_CMAKE_DEBUG("fulfilling future file buffer %p for %s\n", (void *) &(*it), key.c_str());
+        it->second.set_value(std::move(value));
+        return true;
+    }
+    return false;
+}
+
+template <bool Writable>
+std::unique_ptr<llama_file_buffer<Writable>> llama_future_file_buffer<Writable>::extract() const {
+    if (file_buffer) {
+        return std::move(file_buffer);
+    }
+
+    auto future_result = file_buffer_future.get();
+    file_buffer        = std::move(future_result);
+    return std::move(file_buffer);
+}
+
+// Explicit instantiations for llama_future_file_buffer
+template struct llama_future_file_buffer<false>;
+template struct llama_future_file_buffer<true>;
+
 // llama_mmap
 
 struct llama_mmap::impl {

src/llama-mmap.h

@@ -3,6 +3,10 @@
 #include <cstdint>
 #include <memory>
 #include <vector>
+#include "uint8-buff-stream.h"
+#include <future>
+#include <string>
+#include <map>
 
 struct llama_file;
 struct llama_mmap;
@@ -74,9 +78,43 @@ template <bool Writable> struct llama_file_buffer : public llama_file {
     std::unique_ptr<std::basic_streambuf<uint8_t>> streambuf;
 };
 
+template <bool Writable> struct llama_future_file_buffer {
+    /// @brief A file buffer object whose operations will block
+    /// until the given promise key is set with a file buffer.
+    /// @param promise_key The key to use for the promise (e.g. a file path).
+    /// @param context The context to use for the promise, used to distinguish same promise key (e.g. for a same file opened twice).
+    llama_future_file_buffer(const std::string & promise_key, const std::string & context);
+
+    // Delete copy constructor and copy assignment operator
+    llama_future_file_buffer(const llama_future_file_buffer &)             = delete;
+    llama_future_file_buffer & operator=(const llama_future_file_buffer &) = delete;
+
+    llama_future_file_buffer(llama_future_file_buffer && other) noexcept;
+    llama_future_file_buffer & operator=(llama_future_file_buffer && other) noexcept;
+
+    ~llama_future_file_buffer();
+
+    /// @brief Sets the given key and context with a file buffer so that
+    /// operations can resume/start.
+    static bool fulfill_promise(const std::string & promise_key, const std::string & context,
+                                std::unique_ptr<llama_file_buffer<Writable>> && value);
+
+    /// @brief Waits for future buffer or obtains current if already
+    /// fulfilled and moves the future contents outside the registry.
+    std::unique_ptr<llama_file_buffer<Writable>> extract() const;
+
+  private:
+    typename std::map<std::string, std::promise<std::unique_ptr<llama_file_buffer<Writable>>>>::iterator
+                                                                      file_buffer_promise_iterator;
+    mutable std::future<std::unique_ptr<llama_file_buffer<Writable>>> file_buffer_future;
+    mutable std::unique_ptr<llama_file_buffer<Writable>>              file_buffer;
+};
+
 // Type aliases for convenience
 using llama_file_buffer_ro = llama_file_buffer<false>;
 using llama_file_buffer_rw = llama_file_buffer<true>;
+using llama_future_file_buffer_ro = llama_future_file_buffer<false>;
+using llama_future_file_buffer_rw = llama_future_file_buffer<true>;
 
 struct llama_mmap {
     llama_mmap(const llama_mmap &) = delete;