DiskPersistence.cpp

Go to the documentation of this file.

#include "RobotAPI/libraries/armem/server/ltm/persistence/disk/DiskPersistence.h"
 
#include <fcntl.h>
#include <unistd.h>
 
#include <set>
 
#include "RobotAPI/libraries/armem/server/ltm/detail/mixins/util/filesystem.h"
 
namespace armarx::armem::server::ltm::persistence
{
 
    std::vector<std::string>
    DiskPersistence::getContainerKeys(const armarx::armem::MemoryID& id)
    {
        std::vector<std::string> containers;
 
        if (!enabled_)
        {
            return containers;
        }
 
        // If it has a nice form
        if (!id.hasEntityName() || id.hasTimestamp())
        {
            std::vector<std::filesystem::path> dirs = getAllDirectories(id);
 
 
            for (auto& path : dirs)
            {
                std::string container = path.filename().string();
 
                if (!container.empty())
                {
                    containers.emplace_back(container);
                }
            }
        }
        // If it has not a nice form
        else
        {
            std::vector<std::filesystem::path> dayDirs = getAllDirectories(id);
 
            for (std::filesystem::path& dayDir : dayDirs)
            {
                if (!util::fs::detail::isDateString(dayDir.filename()))
                {
                    ARMARX_WARNING << "Found a non-date folder inside an entity '" << id.str()
                                   << "' with name '" << dayDir.filename() << "'. "
                                   << "Ignoring this folder, however this is a bad situation.";
                    continue;
                }
 
                std::vector<std::filesystem::path> secondDirs = util::fs::getAllDirectories(dayDir);
 
                for (std::filesystem::path& secondDir : secondDirs)
                {
                    if (!util::fs::detail::isNumberString(secondDir.filename()))
                    {
                        ARMARX_WARNING << "Found a non-timestamp folder inside an entity '"
                                       << id.str() << "' hours folder with name '"
                                       << secondDir.filename() << "'. "
                                       << "Ignoring this folder, however this is a bad situation.";
                        continue;
                    }
 
                    std::vector<std::filesystem::path> timestampDirs =
                        util::fs::getAllDirectories(secondDir);
 
                    for (std::filesystem::path& timestampDir : timestampDirs)
                    {
                        if (!util::fs::detail::isNumberString(timestampDir.filename()))
                        {
                            ARMARX_WARNING
                                << "Found a non-timestamp folder inside an entity '" << id.str()
                                << "' seconds folder with name '" << timestampDir.filename()
                                << "'. "
                                << "Ignoring this folder, however this is a bad situation.";
                            continue;
                        }
 
                        std::string container = timestampDir.filename().string();
 
                        if (!container.empty())
                        {
                            containers.emplace_back(container);
                        }
                    }
                }
            }
        }
 
        return containers;
    }
 
    std::vector<std::string>
    DiskPersistence::getItemKeys(const armarx::armem::MemoryID& id)
    {
        if (!enabled_)
        {
            return std::vector<std::string>();
        }
 
        // THREAD-SAFETY: Acquire per-directory mutex to get consistent file listing
        std::string directoryPath = getFullPath(id).string();
 
        std::mutex* dirMutex = nullptr;
        {
            std::lock_guard mapLock(directoryMutexMapLock_);
            auto& mutexPtr = directoryMutexes_[directoryPath];
            if (!mutexPtr)
            {
                mutexPtr = std::make_unique<std::mutex>();
            }
            dirMutex = mutexPtr.get();
        }
 
        std::lock_guard dirLock(*dirMutex);
 
        std::vector<std::filesystem::path> files = getAllFiles(id);
        std::vector<std::string> filesStr;
 
        for (auto& path : files)
        {
            std::string item = path.filename().string();
 
            if (!item.empty())
            {
                filesStr.emplace_back(item);
            }
        }
 
        return filesStr;
    }
 
    bool
    DiskPersistence::containsContainer(const armarx::armem::MemoryID& id, std::string key)
    {
        if (!enabled_)
        {
            return false;
        }
 
        //from id get the directory for this memory item and append the key to it, then check if the directory exists
        auto path_to_id = getFullPath(id);
        auto correct_container_path = path_to_id / key;
        bool contains_container = util::fs::directoryExists(correct_container_path);
 
        return contains_container;
    }
 
    bool
    DiskPersistence::containsItem(const armarx::armem::MemoryID& id, std::string key)
    {
        if (!enabled_)
        {
            return false;
        }
 
        return fileExists(id, key);
    }
 
    void
    DiskPersistence::storeItem(const armarx::armem::MemoryID& id,
                               std::string key,
                               std::vector<unsigned char>& data)
    {
        if (!enabled_)
        {
            return;
        }
 
        // If batch writing is enabled, add to batch buffer instead of writing immediately
        if (batchWriteEnabled_.load(std::memory_order_acquire))
        {
            enqueueBatchItem(id, key, std::move(data));
            return;
        }
 
        // Original immediate write path
        // THREAD-SAFETY: Get per-directory mutex for fine-grained locking
        // This allows parallel writes to different directories (different entities)
        // while preventing filesystem corruption within the same directory
        std::string directoryPath = getFullPath(id).string();
 
        // Get or create mutex for this directory
        std::mutex* dirMutex = nullptr;
        {
            std::lock_guard mapLock(directoryMutexMapLock_);
            auto& mutexPtr = directoryMutexes_[directoryPath];
            if (!mutexPtr)
            {
                mutexPtr = std::make_unique<std::mutex>();
            }
            dirMutex = mutexPtr.get();
        }
 
        // Lock this specific directory for the duration of the write
        // Other threads can write to different directories in parallel
        std::lock_guard dirLock(*dirMutex);
 
        try
        {
            auto dir = getFullPath(id);
            auto parentDir = dir.parent_path();
 
            // Check permissions before attempting write
            if (util::fs::directoryExists(parentDir) && !util::fs::hasWritePermission(parentDir))
            {
                ARMARX_ERROR << "No write permission for directory: " << parentDir
                             << ". Cannot store " << id.str() << "/" << key;
                storageErrorCount_.fetch_add(1, std::memory_order_relaxed);
                return;
            }
 
            ensureFullPathExists(id, true);
 
            if (enoughDiskSpaceLeft())
            {
                writeDataToFile(id, key, data);
            }
            else
            {
                ARMARX_ERROR << "Not enough disk space available for DiskPersistence. "
                             << "Skipping storage of " << id.str() << "/" << key;
                storageErrorCount_.fetch_add(1, std::memory_order_relaxed);
            }
        }
        catch (const armarx::LocalException& e)
        {
            ARMARX_ERROR << "ArmarX exception while storing " << id.str() << "/" << key
                         << ": " << e.what();
            storageErrorCount_.fetch_add(1, std::memory_order_relaxed);
        }
        catch (const std::filesystem::filesystem_error& e)
        {
            ARMARX_ERROR << "Filesystem error while storing " << id.str() << "/" << key
                         << ": " << e.what() << " (error code: " << e.code() << ")";
            storageErrorCount_.fetch_add(1, std::memory_order_relaxed);
        }
        catch (const std::system_error& e)
        {
            ARMARX_ERROR << "System error while storing " << id.str() << "/" << key
                         << ": " << e.what() << " (error code: " << e.code() << ")";
            storageErrorCount_.fetch_add(1, std::memory_order_relaxed);
        }
        catch (const std::exception& e)
        {
            ARMARX_ERROR << "Unexpected exception while storing " << id.str() << "/" << key
                         << ": " << e.what();
            storageErrorCount_.fetch_add(1, std::memory_order_relaxed);
        }
        catch (...)
        {
            ARMARX_ERROR << "Unknown exception while storing " << id.str() << "/" << key;
            storageErrorCount_.fetch_add(1, std::memory_order_relaxed);
        }
    }
 
    std::vector<unsigned char>
    DiskPersistence::retrieveItem(const armarx::armem::MemoryID& id, std::string key)
    {
        if (!enabled_)
        {
            return std::vector<unsigned char>();
        }
 
        // THREAD-SAFETY: Acquire per-directory mutex to prevent reading during writes
        // This ensures we don't read partial/empty files while another thread is writing
        std::string directoryPath = getFullPath(id).string();
 
        std::mutex* dirMutex = nullptr;
        {
            std::lock_guard mapLock(directoryMutexMapLock_);
            auto& mutexPtr = directoryMutexes_[directoryPath];
            if (!mutexPtr)
            {
                mutexPtr = std::make_unique<std::mutex>();
            }
            dirMutex = mutexPtr.get();
        }
 
        std::lock_guard dirLock(*dirMutex);
 
        if (fileExists(id, key))
        {
            return readDataFromFile(id, key);
        }
 
        return std::vector<unsigned char>();
    }
 
    std::filesystem::path
    DiskPersistence::getMemoryParentPath()
    {
        std::string p = memoryParentPath_.string();
 
        ArmarXDataPath::ReplaceEnvVars(p);
 
        return p;
    }
 
    std::filesystem::path
    DiskPersistence::getFullPath(const armarx::armem::MemoryID& id)
    {
        auto p = getMemoryParentPath() / getExportName();
 
        auto cleanID =
            id.cleanID(); //somehow, the iDs are jumbled when loading the LTM from disk, this solves it for now
 
        auto fullPath = util::fs::toPath(p, cleanID);
 
        return fullPath;
    }
 
    bool
    DiskPersistence::fullPathExists(const armarx::armem::MemoryID& id)
    {
        auto p = getFullPath(id);
        return util::fs::directoryExists(p);
    }
 
    bool
    DiskPersistence::fileExists(const armarx::armem::MemoryID& id, const std::string& filename)
    {
        auto p = getFullPath(id) / filename;
        return util::fs::fileExists(p);
    }
 
    void
    DiskPersistence::ensureFullPathExists(const armarx::armem::MemoryID& id,
                                          bool createIfNotExistent)
    {
        auto p = getFullPath(id);
        util::fs::ensureDirectoryExists(p, createIfNotExistent);
    }
 
    void
    DiskPersistence::ensureFileExists(const armarx::armem::MemoryID& id,
                                      const std::string& filename,
                                      bool createIfNotExistent)
    {
        auto p = getFullPath(id) / filename;
        util::fs::ensureFileExists(p, createIfNotExistent);
    }
 
    void
    DiskPersistence::writeDataToFile(const armarx::armem::MemoryID& id,
                                     const std::string& filename,
                                     const std::vector<unsigned char>& data)
    {
        auto p = getFullPath(id) / filename;
        util::fs::writeDataToFile(p, data);
    }
 
    std::vector<unsigned char>
    DiskPersistence::readDataFromFile(const armarx::armem::MemoryID& id,
                                      const std::string& filename)
    {
        auto p = getFullPath(id) / filename;
        return util::fs::readDataFromFile(p);
    }
 
    std::vector<std::filesystem::path>
    DiskPersistence::getAllFiles(const armarx::armem::MemoryID& id)
    {
        if (fullPathExists(id))
        {
            auto p = getFullPath(id);
            return util::fs::getAllFiles(p);
        }
 
        return std::vector<std::filesystem::path>();
    }
 
    std::vector<std::filesystem::path>
    DiskPersistence::getAllDirectories(const armarx::armem::MemoryID& id)
    {
        if (fullPathExists(id))
        {
            auto p = getFullPath(id);
            return util::fs::getAllDirectories(p);
        }
 
        return std::vector<std::filesystem::path>();
    }
 
    bool
    DiskPersistence::enoughDiskSpaceLeft()
    {
        std::string path_to_disk = this->getMemoryParentPath();
        bool debug_info_output_enabled = false;
 
        if (std::filesystem::exists(path_to_disk))
        {
            try
            {
                auto space_info = std::filesystem::space(path_to_disk);
                int const conversion_factor = 1024;
 
                auto available_space = space_info.available /
                                       (conversion_factor * conversion_factor * conversion_factor);
 
                if (debug_info_output_enabled)
                {
                    ARMARX_DEBUG << "Capacity: "
                                 << space_info.capacity /
                                        (conversion_factor * conversion_factor * conversion_factor)
                                 << " GB\n";
                    ARMARX_DEBUG << "Free space: "
                                 << space_info.free /
                                        (conversion_factor * conversion_factor * conversion_factor)
                                 << " GB\n";
                    ARMARX_DEBUG << "Available space: "
                                 << space_info.available /
                                        (conversion_factor * conversion_factor * conversion_factor)
                                 << " GB\n";
 
                    ARMARX_DEBUG << "Min disk space: " << this->minDiskSpace << " GB\n";
                }
                return static_cast<bool>(available_space >= this->minDiskSpace);
            }
            catch (const std::filesystem::filesystem_error& e)
            {
                ARMARX_WARNING << "Error: " << e.what() << '\n';
                return false;
            }
            catch (...)
            {
                ARMARX_DEBUG << "Error while trying to get info on available disk space";
                return false;
            }
        }
        else
        {
            ARMARX_DEBUG << "Cannot find path to disk and thus cannot check if enough space is "
                            "still available";
            return true;
        }
    }
 
    void
    DiskPersistence::validateStoragePermissions()
    {
        auto basePath = getMemoryParentPath();
 
        if (!util::fs::canCreateFiles(basePath))
        {
            ARMARX_ERROR << "LTM storage path '" << basePath
                         << "' is not writable! Data will not be persisted.";
            ARMARX_ERROR << "Please check:\n"
                         << "  1. Directory exists\n"
                         << "  2. Current user has write permissions\n"
                         << "  3. Filesystem is not read-only\n"
                         << "  4. No SELinux/AppArmor restrictions";
        }
        else
        {
            ARMARX_INFO << "LTM storage path validated: " << basePath;
        }
    }
 
    // ==================== Batch Write Implementation ====================
 
    void
    DiskPersistence::setBatchWriteEnabled(bool enable)
    {
        bool wasEnabled = batchWriteEnabled_.exchange(enable, std::memory_order_acq_rel);
 
        if (enable && !wasEnabled)
        {
            // Starting batch mode - start the background flush thread
            startBatchWriter();
            ARMARX_INFO << "Batch write mode ENABLED (threshold: " << batchSizeThreshold_
                        << " items or " << batchTimeThresholdMs_ << "ms)";
        }
        else if (!enable && wasEnabled)
        {
            // Stopping batch mode - flush remaining items and stop thread
            flushBatch();
            stopBatchWriter();
            ARMARX_INFO << "Batch write mode DISABLED";
        }
    }
 
    void
    DiskPersistence::enqueueBatchItem(const armarx::armem::MemoryID& id,
                                      const std::string& key,
                                      std::vector<unsigned char> data)
    {
        bool shouldFlush = false;
        {
            std::lock_guard<std::mutex> lock(batchMutex_);
 
            // Set batch start time on first item
            if (batchBuffer_.empty())
            {
                batchStartTime_ = std::chrono::steady_clock::now();
            }
 
            // Add item to batch
            batchBuffer_.push_back({id, key, std::move(data), std::chrono::steady_clock::now()});
 
            // Check if we should flush due to size threshold
            if (batchBuffer_.size() >= batchSizeThreshold_)
            {
                shouldFlush = true;
            }
        }
 
        // Flush outside the lock if needed
        if (shouldFlush)
        {
            flushBatchInternal(0); // 0 = flush by size
        }
        else
        {
            // Notify the background thread that there's work
            batchCondition_.notify_one();
        }
    }
 
    void
    DiskPersistence::flushBatch()
    {
        flushBatchInternal(2); // 2 = explicit flush
    }
 
    size_t
    DiskPersistence::getBatchPendingCount() const
    {
        std::lock_guard<std::mutex> lock(batchMutex_);
        return batchBuffer_.size();
    }
 
    void
    DiskPersistence::flushBatchInternal(int reason)
    {
        std::vector<BatchWriteItem> itemsToWrite;
 
        {
            std::lock_guard<std::mutex> lock(batchMutex_);
            if (batchBuffer_.empty())
            {
                return;
            }
            itemsToWrite = std::move(batchBuffer_);
            batchBuffer_.clear();
        }
 
        // Update statistics for flush reason
        switch (reason)
        {
            case 0: batchStats_.flushBySize.fetch_add(1, std::memory_order_relaxed); break;
            case 1: batchStats_.flushByTime.fetch_add(1, std::memory_order_relaxed); break;
            case 2: batchStats_.flushByExplicit.fetch_add(1, std::memory_order_relaxed); break;
        }
 
        // Write the batch
        auto startTime = std::chrono::steady_clock::now();
        writeBatch(itemsToWrite);
        auto endTime = std::chrono::steady_clock::now();
 
        // Update statistics
        uint64_t durationNs = std::chrono::duration_cast<std::chrono::nanoseconds>(endTime - startTime).count();
        batchStats_.totalBatchesWritten.fetch_add(1, std::memory_order_relaxed);
        batchStats_.totalItemsBatched.fetch_add(itemsToWrite.size(), std::memory_order_relaxed);
        batchStats_.totalFlushTimeNs.fetch_add(durationNs, std::memory_order_relaxed);
 
        // Update max batch size
        uint64_t currentMax = batchStats_.maxBatchSize.load(std::memory_order_relaxed);
        while (itemsToWrite.size() > currentMax)
        {
            if (batchStats_.maxBatchSize.compare_exchange_weak(currentMax, itemsToWrite.size(), std::memory_order_relaxed))
            {
                break;
            }
        }
 
        ARMARX_DEBUG << "Batch flush: " << itemsToWrite.size() << " items in "
                     << (durationNs / 1e6) << "ms (reason=" << reason << ")";
    }
 
    void
    DiskPersistence::writeBatch(std::vector<BatchWriteItem>& items)
    {
        if (items.empty())
        {
            return;
        }
 
        // Check disk space once for the whole batch
        if (!enoughDiskSpaceLeft())
        {
            ARMARX_ERROR << "Not enough disk space for batch write of " << items.size() << " items. Dropping batch!";
            storageErrorCount_.fetch_add(items.size(), std::memory_order_relaxed);
            return;
        }
 
        // Group items by directory path for efficient directory creation
        std::unordered_map<std::string, std::vector<BatchWriteItem*>> itemsByDirectory;
        for (auto& item : items)
        {
            std::string dirPath = getFullPath(item.id).string();
            itemsByDirectory[dirPath].push_back(&item);
        }
 
        // Process each directory group
        std::set<std::string> createdDirectories;  // Track directories we've created
        std::vector<std::string> directoriesToSync; // Directories that need fsync
 
        for (auto& [dirPath, dirItems] : itemsByDirectory)
        {
            // Get or create directory mutex
            std::mutex* dirMutex = nullptr;
            {
                std::lock_guard mapLock(directoryMutexMapLock_);
                auto& mutexPtr = directoryMutexes_[dirPath];
                if (!mutexPtr)
                {
                    mutexPtr = std::make_unique<std::mutex>();
                }
                dirMutex = mutexPtr.get();
            }
 
            // Lock this directory for the batch write
            std::lock_guard dirLock(*dirMutex);
 
            try
            {
                // Create directory once for all items in this group
                if (createdDirectories.find(dirPath) == createdDirectories.end())
                {
                    // Use the first item's ID to create the directory
                    ensureFullPathExists(dirItems[0]->id, true);
                    createdDirectories.insert(dirPath);
                }
 
                // Write all files in this directory without individual fsync
                for (auto* item : dirItems)
                {
                    auto filePath = getFullPath(item->id) / item->key;
 
                    // Use non-atomic write for batch mode (we'll sync at the end)
                    util::fs::writeDataToFile(filePath, item->data, false);
 
                    batchStats_.totalBytesWritten.fetch_add(item->data.size(), std::memory_order_relaxed);
                }
 
                // Track this directory for final sync
                directoriesToSync.push_back(dirPath);
            }
            catch (const std::exception& e)
            {
                ARMARX_ERROR << "Error writing batch to directory " << dirPath << ": " << e.what();
                storageErrorCount_.fetch_add(dirItems.size(), std::memory_order_relaxed);
            }
        }
 
        // Single fsync for all directories at the end (deferred sync)
        for (const auto& dirPath : directoriesToSync)
        {
            int dfd = ::open(dirPath.c_str(), O_DIRECTORY | O_RDONLY);
            if (dfd >= 0)
            {
                ::fsync(dfd);
                ::close(dfd);
            }
        }
    }
 
    void
    DiskPersistence::startBatchWriter()
    {
        if (batchWriterThread_.joinable())
        {
            return; // Already running
        }
 
        stopBatchWriter_.store(false, std::memory_order_release);
        batchWriterThread_ = std::thread(&DiskPersistence::batchWriterThread, this);
    }
 
    void
    DiskPersistence::stopBatchWriter()
    {
        stopBatchWriter_.store(true, std::memory_order_release);
        batchCondition_.notify_all();
 
        if (batchWriterThread_.joinable())
        {
            batchWriterThread_.join();
        }
 
        // Flush any remaining items
        flushBatchInternal(2);
    }
 
    void
    DiskPersistence::batchWriterThread()
    {
        ARMARX_DEBUG << "Batch writer thread started";
 
        while (!stopBatchWriter_.load(std::memory_order_acquire))
        {
            bool shouldFlush = false;
 
            {
                std::unique_lock<std::mutex> lock(batchMutex_);
 
                // Wait for items or timeout
                auto timeout = std::chrono::milliseconds(batchTimeThresholdMs_);
                batchCondition_.wait_for(lock, timeout, [this]() {
                    return stopBatchWriter_.load(std::memory_order_acquire) ||
                           batchBuffer_.size() >= batchSizeThreshold_;
                });
 
                // Check if we should flush due to time threshold
                if (!batchBuffer_.empty())
                {
                    auto now = std::chrono::steady_clock::now();
                    auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
                        now - batchStartTime_).count();
 
                    if (elapsed >= static_cast<long long>(batchTimeThresholdMs_) ||
                        batchBuffer_.size() >= batchSizeThreshold_)
                    {
                        shouldFlush = true;
                    }
                }
            }
 
            if (shouldFlush)
            {
                flushBatchInternal(1); // 1 = flush by time
            }
        }
 
        ARMARX_DEBUG << "Batch writer thread stopped";
    }
 
} // namespace armarx::armem::server::ltm::persistence