Refactoring of buffer handling system

src/Communicate/BufferHandler.h

@@ -0,0 +1,170 @@
+#ifndef IPPL_BUFFER_HANDLER_H
+#define IPPL_BUFFER_HANDLER_H
+
+#include <memory>
+#include <set>
+
+#include "Communicate/Archive.h"
+
+namespace ippl {
+
+    /**
+     * @brief Interface for memory buffer handling.
+     *
+     * Defines methods for acquiring, freeing, and managing memory buffers.
+     * Implementations are responsible for managing buffers efficiently,
+     * ensuring that allocated buffers are reused where possible.
+     *
+     * @tparam MemorySpace The memory space type used for buffer allocation.
+     */
+    template <typename MemorySpace>
+    class BufferHandler {
+    public:
+        using archive_type = ippl::detail::Archive<MemorySpace>;
+        using buffer_type  = std::shared_ptr<archive_type>;
+        using size_type    = ippl::detail::size_type;
+
+        virtual ~BufferHandler() {}
+
+        /**
+         * @brief Requests a memory buffer of a specified size.
+         *
+         * Provides a buffer of at least the specified size, with the option
+         * to allocate additional space based on an overallocation multiplier.
+         * This function attempts to reuse available buffers if possible.
+         *
+         * @param size The required size of the buffer, in bytes.
+         * @param overallocation A multiplier to allocate extra space, which may help
+         *                       avoid frequent reallocation in some use cases.
+         * @return A shared pointer to the allocated buffer.
+         */
+        virtual buffer_type getBuffer(size_type size, double overallocation) = 0;
+
+        /**
+         * @brief Frees a specified buffer.
+         *
+         * Moves the specified buffer to a free state, making it available
+         * for reuse in future buffer requests.
+         *
+         * @param buffer The buffer to be freed.
+         */
+        virtual void freeBuffer(buffer_type buffer) = 0;
+
+        /**
+         * @brief Frees all currently used buffers.
+         *
+         * Transfers all used buffers to the free state, making them available
+         * for reuse. This does not deallocate memory but resets buffer usage.
+         */
+        virtual void freeAllBuffers() = 0;
+
+        /**
+         * @brief Deletes all buffers.
+         *
+         * Releases all allocated memory buffers, both used and free.
+         * After this call, no buffers are available until new allocations.
+         */
+        virtual void deleteAllBuffers() = 0;
+
+        /**
+         * @brief Gets the size of all buffers in use.
+         *
+         * @return Total size of buffers that are in use in bytes.
+         */
+        virtual size_type getUsedSize() const = 0;
+
+        /**
+         * @brief Gets the size of all free buffers.
+         *
+         * @return Total size of free buffers in bytes.
+         */
+        virtual size_type getFreeSize() const = 0;
+    };
+
+    /**
+     * @class DefaultBufferHandler
+     * @brief Concrete implementation of BufferHandler for managing memory buffers.
+     *
+     * This class implements the BufferHandler interface, providing concrete behavior for
+     * buffer allocation, freeing, and memory management. It maintains two sorted sets of free and
+     * in-use buffers to allow for efficient queries.
+     *
+     * @tparam MemorySpace The memory space type for the buffer (e.g., `Kokkos::HostSpace`).
+     */
+    template <typename MemorySpace>
+    class DefaultBufferHandler : public BufferHandler<MemorySpace> {
+    public:
+        using typename BufferHandler<MemorySpace>::archive_type;
+        using typename BufferHandler<MemorySpace>::buffer_type;
+        using typename BufferHandler<MemorySpace>::size_type;
+
+        ~DefaultBufferHandler() override;
+
+        /**
+         * @brief Acquires a buffer of at least the specified size.
+         *
+         * Requests a memory buffer of the specified size, with the option
+         * to request a buffer larger than the base size by an overallocation
+         * multiplier. If a sufficiently large buffer is available, it is returned. If not, the
+         * largest free buffer is reallocated. If there are no free buffers available, only then a
+         * new buffer is allocated.
+         *
+         * @param size The required buffer size.
+         * @param overallocation A multiplier to allocate additional buffer space.
+         * @return A shared pointer to the allocated buffer.
+         */
+        buffer_type getBuffer(size_type size, double overallocation) override;
+
+        /**
+         * @copydoc BufferHandler::freeBuffer
+         */
+        void freeBuffer(buffer_type buffer) override;
+
+        /**
+         * @copydoc BufferHandler::freeBuffer
+         */
+        void freeAllBuffers() override;
+
+        /**
+         * @copydoc BufferHandler::freeBuffer
+         */
+        void deleteAllBuffers() override;
+
+        /**
+         * @copydoc BufferHandler::freeBuffer
+         */
+        size_type getUsedSize() const override;
+
+        /**
+         * @copydoc BufferHandler::freeBuffer
+         */
+        size_type getFreeSize() const override;
+
+    private:
+        using buffer_comparator_type = bool (*)(const buffer_type&, const buffer_type&);
+        using buffer_set_type        = std::set<buffer_type, buffer_comparator_type>;
+
+        static bool bufferSizeComparator(const buffer_type& lhs, const buffer_type& rhs);
+
+        bool isBufferUsed(buffer_type buffer) const;
+        void releaseUsedBuffer(buffer_type buffer);
+        buffer_type findFreeBuffer(size_type requiredSize);
+        buffer_set_type::iterator findSmallestSufficientBuffer(size_type requiredSize);
+        buffer_type getFreeBuffer(buffer_type buffer);
+        buffer_type reallocateLargestFreeBuffer(size_type requiredSize);
+        buffer_type allocateNewBuffer(size_type requiredSize);
+
+        size_type usedSize_m;  ///< Total size of all allocated buffers
+        size_type freeSize_m;  ///< Total size of all free buffers
+
+    protected:
+        buffer_set_type used_buffers{
+            &DefaultBufferHandler::bufferSizeComparator};  ///< Set of used buffers
+        buffer_set_type free_buffers{
+            &DefaultBufferHandler::bufferSizeComparator};  ///< Set of free buffers
+    };
+}  // namespace ippl
+
+#include "Communicate/BufferHandler.hpp"
+
+#endif

src/Communicate/BufferHandler.hpp

@@ -0,0 +1,147 @@
+#ifndef IPPL_BUFFER_HANDLER_HPP
+#define IPPL_BUFFER_HANDLER_HPP
+
+namespace ippl {
+
+    template <typename MemorySpace>
+    DefaultBufferHandler<MemorySpace>::~DefaultBufferHandler() {}
+
+    template <typename MemorySpace>
+    typename DefaultBufferHandler<MemorySpace>::buffer_type DefaultBufferHandler<MemorySpace>::getBuffer(
+        size_type size, double overallocation) {
+        size_type requiredSize = static_cast<size_type>(size * overallocation);
+
+        auto freeBuffer = findFreeBuffer(requiredSize);
+        if (freeBuffer != nullptr) {
+            return getFreeBuffer(freeBuffer);
+        }
+
+        if (!free_buffers.empty()) {
+            return reallocateLargestFreeBuffer(requiredSize);
+        }
+
+        return allocateNewBuffer(requiredSize);
+    }
+
+    template <typename MemorySpace>
+    void DefaultBufferHandler<MemorySpace>::freeBuffer(buffer_type buffer) {
+        if (isBufferUsed(buffer)) {
+            releaseUsedBuffer(buffer);
+        }
+    }
+
+    template <typename MemorySpace>
+    void DefaultBufferHandler<MemorySpace>::freeAllBuffers() {
+        free_buffers.insert(used_buffers.begin(), used_buffers.end());
+        used_buffers.clear();
+
+        freeSize_m += usedSize_m;
+        usedSize_m = 0;
+    }
+
+    template <typename MemorySpace>
+    void DefaultBufferHandler<MemorySpace>::deleteAllBuffers() {
+        freeSize_m      = 0;
+        usedSize_m = 0;
+
+        used_buffers.clear();
+        free_buffers.clear();
+    }
+
+    template <typename MemorySpace>
+    typename DefaultBufferHandler<MemorySpace>::size_type DefaultBufferHandler<MemorySpace>::getUsedSize()
+        const {
+        return usedSize_m;
+    }
+
+    template <typename MemorySpace>
+    typename DefaultBufferHandler<MemorySpace>::size_type DefaultBufferHandler<MemorySpace>::getFreeSize() const {
+        return freeSize_m;
+    }
+
+    template <typename MemorySpace>
+    bool DefaultBufferHandler<MemorySpace>::bufferSizeComparator(const buffer_type& lhs,
+                                                          const buffer_type& rhs) {
+        if (lhs->getBufferSize() != rhs->getBufferSize()) {
+            return lhs->getBufferSize() < rhs->getBufferSize();
+        }
+
+        // Use memory address as a tie-breaker to enforce total ordering of buffers.
+        return lhs < rhs;
+    }
+
+    template <typename MemorySpace>
+    bool DefaultBufferHandler<MemorySpace>::isBufferUsed(buffer_type buffer) const {
+        return used_buffers.find(buffer) != used_buffers.end();
+    }
+
+    template <typename MemorySpace>
+    void DefaultBufferHandler<MemorySpace>::releaseUsedBuffer(buffer_type buffer) {
+        auto it = used_buffers.find(buffer);
+
+        usedSize_m -= buffer->getBufferSize();
+        freeSize_m += buffer->getBufferSize();
+
+        used_buffers.erase(it);
+        free_buffers.insert(buffer);
+    }
+
+    template <typename MemorySpace>
+    typename DefaultBufferHandler<MemorySpace>::buffer_type DefaultBufferHandler<MemorySpace>::findFreeBuffer(
+        size_type requiredSize) {
+        auto it = findSmallestSufficientBuffer(requiredSize);
+        if (it != free_buffers.end()) {
+            return *it;
+        }
+        return nullptr;
+    }
+
+    template <typename MemorySpace>
+    typename DefaultBufferHandler<MemorySpace>::buffer_set_type::iterator
+    DefaultBufferHandler<MemorySpace>::findSmallestSufficientBuffer(size_type requiredSize) {
+        return std::find_if(free_buffers.begin(), free_buffers.end(),
+                            [requiredSize](const buffer_type& buffer) {
+                                return buffer->getBufferSize() >= requiredSize;
+                            });
+    }
+
+    template <typename MemorySpace>
+    typename DefaultBufferHandler<MemorySpace>::buffer_type
+    DefaultBufferHandler<MemorySpace>::getFreeBuffer(buffer_type buffer) {
+        freeSize_m -= buffer->getBufferSize();
+        usedSize_m += buffer->getBufferSize();
+
+        free_buffers.erase(buffer);
+        used_buffers.insert(buffer);
+        return buffer;
+    }
+
+    template <typename MemorySpace>
+    typename DefaultBufferHandler<MemorySpace>::buffer_type
+    DefaultBufferHandler<MemorySpace>::reallocateLargestFreeBuffer(size_type requiredSize) {
+        auto largest_it    = std::prev(free_buffers.end());
+        buffer_type buffer = *largest_it;
+
+        freeSize_m -= buffer->getBufferSize();
+        usedSize_m += requiredSize;
+
+        free_buffers.erase(buffer);
+        buffer->reallocBuffer(requiredSize);
+
+        used_buffers.insert(buffer);
+        return buffer;
+    }
+
+    template <typename MemorySpace>
+    typename DefaultBufferHandler<MemorySpace>::buffer_type DefaultBufferHandler<MemorySpace>::allocateNewBuffer(
+        size_type requiredSize) {
+        buffer_type newBuffer = std::make_shared<archive_type>(requiredSize);
+
+        usedSize_m += newBuffer->getBufferSize();
+        used_buffers.insert(newBuffer);
+        return newBuffer;
+    }
+
+}  // namespace ippl
+
+#endif

src/Communicate/Buffers.cpp

@@ -32,9 +32,16 @@ namespace ippl {
         }
 
         void Communicator::deleteAllBuffers() {
-            buffers_m.forAll([]<typename Map>(Map&& m) {
-                m.clear();
+            buffer_handlers_m.forAll([]<typename BufferHandler>(BufferHandler&& bh) {
+                bh.deleteAllBuffers();
             });
         }
+
+        void Communicator::freeAllBuffers() {
+            buffer_handlers_m.forAll([]<typename BufferHandler>(BufferHandler&& bh) {
+                bh.freeAllBuffers();
+            });
+        }
+
     }  // namespace mpi
 }  // namespace ippl

src/Communicate/Buffers.hpp

@@ -24,20 +24,21 @@ namespace ippl {
     namespace mpi {
 
         template <typename MemorySpace, typename T>
-        Communicator::buffer_type<MemorySpace> Communicator::getBuffer(int id, size_type size,
-                                                                       double overallocation) {
-            auto& buffers = buffers_m.get<MemorySpace>();
-            size *= sizeof(T);
-            if (buffers.contains(id)) {
-                if (buffers[id]->getBufferSize() < size) {
-                    buffers[id]->reallocBuffer(size);
-                }
-                return buffers[id];
-            }
-            buffers[id] = std::make_shared<archive_type<MemorySpace>>(
-                (size_type)(size * std::max(overallocation, defaultOveralloc_m)));
-            return buffers[id];
+        Communicator::buffer_type<MemorySpace> Communicator::getBuffer(size_type size,
+                                                                        double overallocation) {
+            auto& buffer_handler = buffer_handlers_m.get<MemorySpace>();
+
+            return buffer_handler.getBuffer(size * sizeof(T), std::max(overallocation, defaultOveralloc_m));
+        }
+
+
+        template <typename MemorySpace>
+        void Communicator::freeBuffer(Communicator::buffer_type<MemorySpace> buffer) {
+            auto& buffer_handler = buffer_handlers_m.get<MemorySpace>();
+
+            buffer_handler.freeBuffer(buffer);
         }
+
     }  // namespace mpi
 
 }  // namespace ippl

src/Communicate/CMakeLists.txt

@@ -1,11 +1,16 @@
 set (_SRCS
     Communicator.cpp
+    CommunicatorLogging.cpp
     Environment.cpp
     Buffers.cpp
     Request.cpp
+    LogEntry.cpp
     )
 
 set (_HDRS
+    LogEntry.h
+    BufferHandler.h
+    LoggingBufferHandler.h
     Archive.h
     Archive.hpp
     Buffers.hpp
@@ -23,6 +28,7 @@ set (_HDRS
     Window.h
     Window.hpp
     PointToPoint.hpp
+    CommunicatorLogging.hpp
     )
 
 include_directories (