DiskPersistence.h

Go to the documentation of this file.

#pragma once
 
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <filesystem>
#include <memory>
#include <mutex>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>
 
#include "RobotAPI/libraries/armem/core/MemoryID.h"
#include "RobotAPI/libraries/armem/server/ltm/persistence/MemoryPersistenceStrategy.h"
 
namespace armarx::armem::server::ltm::persistence
{
    /**
     * @brief Statistics for batch write operations.
     */
    struct BatchWriteStatistics
    {
        std::atomic<uint64_t> totalBatchesWritten{0};
        std::atomic<uint64_t> totalItemsBatched{0};
        std::atomic<uint64_t> totalBytesWritten{0};
        std::atomic<uint64_t> totalFlushTimeNs{0};
        std::atomic<uint64_t> maxBatchSize{0};
        std::atomic<uint64_t> flushBySize{0};    // Flushes triggered by size threshold
        std::atomic<uint64_t> flushByTime{0};    // Flushes triggered by time threshold
        std::atomic<uint64_t> flushByExplicit{0}; // Explicit flush calls
 
        void reset()
        {
            totalBatchesWritten = 0;
            totalItemsBatched = 0;
            totalBytesWritten = 0;
            totalFlushTimeNs = 0;
            maxBatchSize = 0;
            flushBySize = 0;
            flushByTime = 0;
            flushByExplicit = 0;
        }
 
        double getAvgFlushTimeMs() const
        {
            uint64_t count = totalBatchesWritten.load(std::memory_order_relaxed);
            if (count == 0) return 0.0;
            return totalFlushTimeNs.load(std::memory_order_relaxed) / (count * 1e6);
        }
 
        double getAvgBatchSize() const
        {
            uint64_t count = totalBatchesWritten.load(std::memory_order_relaxed);
            if (count == 0) return 0.0;
            return static_cast<double>(totalItemsBatched.load(std::memory_order_relaxed)) / count;
        }
    };
 
    /**
     * @brief A single item pending batch write.
     */
    struct BatchWriteItem
    {
        armarx::armem::MemoryID id;
        std::string key;
        std::vector<unsigned char> data;
        std::chrono::steady_clock::time_point enqueuedTime;
    };
 
    class FileIdentifier : public ItemIdentifier
    {
    public:
        FileIdentifier(std::string& filename, std::string& fileType) :
            filename_(filename),
            fileType_(fileType){
 
            };
 
        virtual ~FileIdentifier()
        {
        }
 
        std::string
        getKey() override
        {
            return filename_ + fileType_;
        }
 
    private:
        std::string filename_;
        std::string fileType_;
    };
 
    /**
   * @brief Persistence strategy that writes items (e.g. json files) to a specific container (a directory)
   * Use it to write the data of a WM to disk.
   *
   * Where are the items written (=:location)?
   *  /memoryParentPath/exportName/path(id)/
   *
   * How is the file of an item named?
   * -> Just they key (e.g. filename = key = "data.aron.json")
   *    If you might want a more sophisticated solution fell free to implement your custom ItemIdentifier today!
   */
    class DiskPersistence : public MemoryPersistenceStrategy
    {
    public:
        DiskPersistence() : DiskPersistence(std::filesystem::path("."))
        {
        }
 
        DiskPersistence(const std::filesystem::path& memoryParentPath) :
            DiskPersistence("Disk", "DefaultExport", memoryParentPath)
        {
        }
 
        /**
       * @param identifier basically a unique name for the strategy (important if you use different strategies @see RedundantPersistenceStrategy)
       * @param exportName identifier for the exported memory. A new directory with name is created beneath the memoryParentPath. Everything is stored inside it.
       * @param memoryParentPath path where the memory should be exported to
       */
        DiskPersistence(const std::string& identifier,
                        const std::string& exportName,
                        const std::filesystem::path& memoryParentPath) :
            MemoryPersistenceStrategy(identifier, exportName), memoryParentPath_(memoryParentPath)
        {
        }
 
        /**
         * @brief Destructor - flushes pending batch and stops batch thread.
         */
        virtual ~DiskPersistence()
        {
            stopBatchWriter();
        }
 
        /**
       * Returns all containers for the current id.
       * @return containers <=> directories at current location (=/memoryParentPath/exportName/path(id)/)
       */
        std::vector<std::string> getContainerKeys(const armarx::armem::MemoryID& id) override;
 
        /**
       * Returns all items for the current id.
       * @return items <=> files at the current location (=/memoryParentPath/exportName/path(id)/)
       */
        std::vector<std::string> getItemKeys(const armarx::armem::MemoryID& id) override;
 
        /**
       * Checks if the container is available for the current memory id.
       * @return true if the current location contains the directory with name 'key'
       */
        bool containsContainer(const armarx::armem::MemoryID& id, std::string key) override;
 
        /**
       * Checks if current container contains the item defined by its key.
       * @return true if the current location contains a file with the name 'key'
       */
        bool containsItem(const armarx::armem::MemoryID& id, std::string key) override;
 
        /**
       * Create a new file with name 'key' and stores the data inside it.
       */
        void storeItem(const armarx::armem::MemoryID& id,
                       std::string key,
                       std::vector<unsigned char>& data) override;
 
        /**
       * Reads the data of the file with name 'key' at the current location.
       * @return data if a file was found, an empty vector if the file is empty or was not found
       */
        std::vector<unsigned char> retrieveItem(const armarx::armem::MemoryID& id,
                                                std::string key) override;
 
        void
        createPropertyDefinitions(PropertyDefinitionsPtr& defs, const std::string& prefix) override
        {
            // Nothing to do
        }
 
        void
        setMinAvailableDiskSpace(const int minDiskSpace)
        {
            this->minDiskSpace = minDiskSpace;
            if (!enoughDiskSpaceLeft())
            {
                ARMARX_WARNING << "Not enough available disk space for DiskPersistance Strategy. "
                                  "You need at least "
                               << this->minDiskSpace
                               << " GB available disk space to record into LTM using this strategy";
            }
        }
 
        int minDiskSpace = 50; // in GB
 
        // ==================== Compression Configuration ====================
 
        /**
         * @brief Enable transparent gzip compression of exported JSON files.
         *
         * When enabled, items whose key ends in ".json" and whose payload is at least
         * getCompressionMinBytes() large are gzipped and stored under "<key>.gz". Reading is
         * unaffected: containsItem()/retrieveItem() accept the plain (logical) key and transparently
         * fall back to the ".gz" variant, and getItemKeys() reports logical keys. Exports written
         * before this feature existed therefore keep working, and an export may freely mix both.
         *
         * The compressed file deliberately gets a ".gz" suffix rather than being compressed in
         * place under the original name: a reader that does not know about compression then simply
         * does not find the file, instead of choking on unparsable content.
         */
        void setCompressionEnabled(bool enable) { compressionEnabled_ = enable; }
 
        bool isCompressionEnabled() const { return compressionEnabled_; }
 
        /// zlib level, 1 (fastest) to 9 (smallest). Default 1 - see gzipCompress().
        void setCompressionLevel(int level) { compressionLevel_ = level; }
 
        int getCompressionLevel() const { return compressionLevel_; }
 
        /**
         * @brief Only compress payloads of at least this many bytes.
         *
         * Defaults to one filesystem block (4096). Compressing anything smaller cannot free a
         * single block, so it would cost compatibility (older readers no longer find the file) for
         * no disk saving at all. In practice this leaves metadata.aron.json and type.aron.json as
         * plain JSON and only compresses the instance data, which is where the size actually is.
         */
        void setCompressionMinBytes(std::size_t bytes) { compressionMinBytes_ = bytes; }
 
        std::size_t getCompressionMinBytes() const { return compressionMinBytes_; }
 
        /// Suffix appended to the logical key when an item is stored compressed.
        static constexpr const char* GZIP_SUFFIX = ".gz";
 
        // ==================== Batch Write Configuration ====================
 
        /**
         * @brief Enable or disable batch writing.
         * When enabled, writes are accumulated and flushed in batches for better I/O performance.
         * @param enable True to enable batching, false for immediate writes (default behavior)
         */
        void setBatchWriteEnabled(bool enable);
 
        /**
         * @brief Check if batch writing is enabled.
         */
        bool isBatchWriteEnabled() const { return batchWriteEnabled_.load(std::memory_order_acquire); }
 
        /**
         * @brief Set the maximum number of items to accumulate before auto-flushing.
         * @param size Maximum batch size (default: 100)
         */
        void setBatchSizeThreshold(size_t size) { batchSizeThreshold_ = size; }
 
        /**
         * @brief Set the maximum time to hold items before auto-flushing.
         * @param ms Maximum time in milliseconds (default: 100ms)
         */
        void setBatchTimeThresholdMs(size_t ms) { batchTimeThresholdMs_ = ms; }
 
        /**
         * @brief Explicitly flush all pending batch writes to disk.
         * This is automatically called when batch thresholds are reached.
         */
        void flushBatch();
 
        /**
         * @brief Get the current number of items pending in the batch buffer.
         */
        size_t getBatchPendingCount() const;
 
        /**
         * @brief Get batch write statistics.
         */
        const BatchWriteStatistics& getBatchStatistics() const { return batchStats_; }
 
        /**
         * @brief Reset batch write statistics.
         */
        void resetBatchStatistics() { batchStats_.reset(); }
 
 
    public:
        size_t getStorageErrorCount() const
        {
            return storageErrorCount_.load(std::memory_order_acquire);
        }
 
        void resetStorageErrorCount()
        {
            storageErrorCount_.store(0, std::memory_order_release);
        }
 
        void validateStoragePermissions();
 
    private:
        std::filesystem::path memoryParentPath_;
 
        // THREAD-SAFETY: Per-directory mutexes for parallel writes
        // Allows concurrent writes to different directories (different entities)
        // while preventing corruption within the same directory
        mutable std::mutex directoryMutexMapLock_;
        mutable std::unordered_map<std::string, std::unique_ptr<std::mutex>> directoryMutexes_;
 
        // Error tracking
        std::atomic<size_t> storageErrorCount_{0};
 
        // Compression state (see the setters above)
        std::atomic<bool> compressionEnabled_{false};
        std::atomic<int> compressionLevel_{1};
        std::atomic<std::size_t> compressionMinBytes_{4096};
 
        /* Internal compression helpers */
 
        /// Whether an item with this key is worth compressing (JSON payloads only; .png/.exr are
        /// already compressed and would only grow).
        static bool isCompressibleKey(const std::string& key);
 
        /// Drop a trailing ".gz" if present, yielding the logical key callers use.
        static std::string toLogicalKey(const std::string& key);
 
        /**
         * @brief Delete the other on-disk spelling of the item just written.
         *
         * One logical key must map to exactly one file: retrieveItem() and getItemKeys() both
         * assume the plain and the ".gz" form are mutually exclusive. Toggling compression (or
         * crossing the size threshold) on an export that already holds the other form would
         * otherwise leave a sibling behind that shadows the fresh write.
         *
         * Called after the payload is safely on disk, never before: a crash in the window leaves
         * both files, which is only the state this repairs, whereas deleting first could lose the
         * item outright.
         */
        void removeStaleVariant(const armarx::armem::MemoryID& id, const std::string& writtenKey);
 
        /// Read and gunzip a ".gz" item; returns empty and logs if the stream is unusable.
        std::vector<unsigned char> readCompressedFile(const armarx::armem::MemoryID& id,
                                                      const std::string& compressedKey);
 
        /* Internal disk logic */
 
        bool fullPathExists(const armarx::armem::MemoryID& id);
 
        std::vector<std::filesystem::path> getAllFiles(const armarx::armem::MemoryID& id);
 
        std::vector<std::filesystem::path> getAllDirectories(const armarx::armem::MemoryID& id);
 
        bool fileExists(const armarx::armem::MemoryID& id, const std::string& filename);
 
        std::filesystem::path getFullPath(const armarx::armem::MemoryID& id);
 
        void ensureFullPathExists(const armarx::armem::MemoryID& id,
                                  bool createIfNotExistent = false);
 
        void ensureFileExists(const armarx::armem::MemoryID& id,
                              const std::string& filename,
                              bool createIfNotExistent = false);
 
        void writeDataToFile(const armarx::armem::MemoryID& id,
                             const std::string& filename,
                             const std::vector<unsigned char>& data);
 
        std::vector<unsigned char> readDataFromFile(const armarx::armem::MemoryID& id,
                                                    const std::string& filename);
 
        std::filesystem::path getMemoryParentPath();
 
        bool enoughDiskSpaceLeft();
 
        // ==================== Batch Write Implementation ====================
 
        /**
         * @brief Add an item to the batch buffer (called by storeItem when batching is enabled).
         */
        void enqueueBatchItem(const armarx::armem::MemoryID& id,
                              const std::string& key,
                              std::vector<unsigned char> data);
 
        /**
         * @brief Internal method to flush the batch buffer.
         * @param reason Statistics tracking reason (size/time/explicit)
         */
        void flushBatchInternal(int reason);
 
        /**
         * @brief Write a batch of items to disk efficiently.
         * Groups items by directory, creates directories once, writes files, syncs once.
         */
        void writeBatch(std::vector<BatchWriteItem>& items);
 
        /**
         * @brief Start the background batch flush thread.
         */
        void startBatchWriter();
 
        /**
         * @brief Stop the background batch flush thread.
         */
        void stopBatchWriter();
 
        /**
         * @brief Background thread function for time-based batch flushing.
         */
        void batchWriterThread();
 
        // Batch write state
        std::atomic<bool> batchWriteEnabled_{false};
        size_t batchSizeThreshold_ = 100;      // Flush when batch reaches this size
        size_t batchTimeThresholdMs_ = 100;    // Flush after this many ms
 
        // Batch buffer
        mutable std::mutex batchMutex_;
        std::vector<BatchWriteItem> batchBuffer_;
        std::chrono::steady_clock::time_point batchStartTime_;
 
        // Background flush thread
        std::thread batchWriterThread_;
        std::atomic<bool> stopBatchWriter_{false};
        std::condition_variable batchCondition_;
 
        // Statistics
        mutable BatchWriteStatistics batchStats_;
    };
} // namespace armarx::armem::server::ltm::persistence