From 7dfbed68feba4de8fe44aa205fc7437c92d5fd92 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Constantin=20F=C3=BCrst?= <c@fuersten.info>
Date: Wed, 10 Jan 2024 19:04:44 +0100
Subject: [PATCH] handle allocation slightly different, introduce a separate
 function for cleaner code that does on-node memory allocation, first querry
 the available size and do not rely on numa_alloc_onnode to report nullptr if
 the size is not really available

---
 offloading-cacher/cache.hpp | 56 +++++++++++++++++++++++++++++--------
 1 file changed, 44 insertions(+), 12 deletions(-)

diff --git a/offloading-cacher/cache.hpp b/offloading-cacher/cache.hpp
index 22b23f8..cce0439 100644
--- a/offloading-cacher/cache.hpp
+++ b/offloading-cacher/cache.hpp
@@ -69,6 +69,12 @@ namespace dsacache {
         // as this is set as the "optimal placement" node
         void GetCacheNode(uint8_t* src, const size_t size, int* OUT_DST_NODE, int* OUT_SRC_NODE) const;
 
+        // allocates memory of size "size" on the numa node "node"
+        // and returns nullptr if this is not possible, also may
+        // try to flush the cache of the requested node to
+        // alleviate encountered shortage
+        uint8_t* AllocOnNode(const size_t size, const int node);
+
         // checks whether the cache contains an entry for
         // the given data in the given memory node and
         // returns it, otherwise returns nullptr
@@ -164,32 +170,58 @@ inline std::unique_ptr<dsacache::CacheData> dsacache::Cache::Access(uint8_t* dat
     return std::move(task);
 }
 
-inline void dsacache::Cache::SubmitTask(CacheData* task, const int dst_node, const int src_node) {
-    std::cout << "[+] Allocating " << task->size_ << "B on node " << dst_node << " for " << std::hex << (uint64_t)task->src_ << std::dec << std::endl;
-
+inline uint8_t* dsacache::Cache::AllocOnNode(const size_t size, const int node) {
     // allocate data on this node and flush the unused parts of the
     // cache if the operation fails and retry once
     // TODO: smarter flush strategy could keep some stuff cached
 
-    uint8_t* dst = reinterpret_cast<uint8_t*>(numa_alloc_onnode(task->size_, dst_node));
+    // check currently free memory to see if the data fits
 
-    if (dst == nullptr) {
-        std::cout << "[!] First allocation try failed for " << task->size_ << "B on node " << dst_node << std::endl;
+    long long int free_space = 0;
+    numa_node_size64(node, &free_space);
+
+    if (free_space < size) {
+        std::cout << "[!] Memory shortage when allocating " << size << "B on node " << node << std::endl;
 
-        // allocation on dst_node failed so we flush the cache for this
+        // dst node lacks memory space so we flush the cache for this
         // node hoping to free enough currently unused entries to make
         // the second allocation attempt successful
 
-        Flush(dst_node);
+        Flush(node);
 
-        dst = reinterpret_cast<uint8_t*>(numa_alloc_onnode(task->size_, dst_node));
+        // re-test by getting the free space and checking again
 
-        if (dst == nullptr) {
-            std::cerr << "[x] Second allocation try failed for " << task->size_ << "B on node " << dst_node << std::endl;
-            return;
+        numa_node_size64(node, &free_space);
+
+        if (free_space < size) {
+            std::cout << "[x] Memory shortage after flush when allocating " << size << "B on node " << node << std::endl;
+
+            return nullptr;
         }
     }
 
+    uint8_t* dst = reinterpret_cast<uint8_t*>(numa_alloc_onnode(size, node));
+
+    if (dst == nullptr) {
+        std::cout << "[x] Allocation try failed for " << size << "B on node " << node << std::endl;
+
+        return nullptr;
+    }
+
+    return dst;
+}
+
+inline void dsacache::Cache::SubmitTask(CacheData* task, const int dst_node, const int src_node) {
+    std::cout << "[+] Allocating " << task->size_ << "B on node " << dst_node << " for " << std::hex << (uint64_t)task->src_ << std::dec << std::endl;
+
+
+    uint8_t* dst = AllocOnNode(task->size_, dst_node);
+
+    if (dst == nullptr) {
+        std::cout << "[x] Allocation failed so we can not cache" << std::endl;
+        return;
+    }
+
     task->incomplete_cache_ = dst;
 
     // querry copy policy function for the nodes to use for the copy