From b7e6a4d4149037740038779f1b10b7bce63b61f6 Mon Sep 17 00:00:00 2001
From: Travis Downs <travis.downs@redpanda.com>
Date: Fri, 3 May 2024 00:09:00 -0400
Subject: [PATCH] memory.cc: fix cross-shard shrinking realloc

Cross-shard shrinking realloc crashes because we assert that the
incoming pointer is shard local inside the shrink method but there is
no particular reason to assume this is the case with a realloc: the
initial allocation may have been made on another shard.

Fix this by falling through to the free/malloc/memcpy path. This also
means that realloc using the same size is a way to "rehome" a possibly
foreign pointer: this does nothing if the pointer is already local but
it will allocate a local pointer and copy the allocation contents if
not.

Fixes #2202.
---
 src/core/memory.cc       | 10 +++++---
 tests/unit/alloc_test.cc | 55 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/src/core/memory.cc b/src/core/memory.cc
index a52d9cebb01..c0a9c05a8fa 100644
--- a/src/core/memory.cc
+++ b/src/core/memory.cc
@@ -2286,8 +2286,6 @@ void* realloc(void* ptr, size_t size) {
         abort();
     }
     // if we're here, it's a non-null seastar memory ptr
-    // or original functions aren't available.
-    // at any rate, using the seastar allocator is OK now.
     auto old_size = ptr ? object_size(ptr) : 0;
     if (size == old_size) {
         return ptr;
@@ -2296,10 +2294,16 @@ void* realloc(void* ptr, size_t size) {
         ::free(ptr);
         return nullptr;
     }
-    if (size < old_size) {
+    if (size < old_size && cpu_pages::is_local_pointer(ptr)) {
+        // local pointers can sometimes be shrunk by returning freed
+        // pages to the buddy allocator
         seastar::memory::shrink(ptr, size);
         return ptr;
     }
+
+    // either a request to realloc larger than the existing allocation size
+    // or a cross-shard pointer: in either case we allocate a new local
+    // pointer and copy the contents
     auto nptr = malloc(size);
     if (!nptr) {
         return nptr;
diff --git a/tests/unit/alloc_test.cc b/tests/unit/alloc_test.cc
index f9284a49096..be587cb6a70 100644
--- a/tests/unit/alloc_test.cc
+++ b/tests/unit/alloc_test.cc
@@ -20,10 +20,12 @@
  */
 
 #include <seastar/core/memory.hh>
+#include <seastar/core/shard_id.hh>
 #include <seastar/core/smp.hh>
 #include <seastar/core/temporary_buffer.hh>
 #include <seastar/testing/perf_tests.hh>
 #include <seastar/testing/test_case.hh>
+#include <seastar/testing/thread_test_case.hh>
 #include <seastar/util/log.hh>
 #include <seastar/util/memory_diagnostics.hh>
 
@@ -155,6 +157,59 @@ SEASTAR_TEST_CASE(test_memory_diagnostics) {
     return make_ready_future<>();
 }
 
+SEASTAR_THREAD_TEST_CASE(test_cross_thread_realloc) {
+    // Tests that realloc seems to do the right thing with various sizes of
+    // buffer, including cases where the initial allocation is on another
+    // shard.
+    // Needs at least 2 shards to usefully test the cross-shard aspect but
+    // still passes when only 1 shard is used.
+    auto do_xshard_realloc = [](bool cross_shard, size_t initial_size, size_t realloc_size) { 
+        BOOST_TEST_CONTEXT("cross_shard=" << cross_shard << ", initial="
+                << initial_size << ", realloc_size=" << realloc_size) {
+                    
+            auto other_shard = (this_shard_id() + cross_shard) % smp::count;
+
+            char *p = static_cast<char *>(malloc(initial_size));
+
+            // write some sentinels and check them after realloc
+            // x start of region
+            // y end of realloc'd region (if it falls within the initial size)
+            // z end of initial region
+            if (initial_size > 0) {
+                p[0] = 'x';
+                p[initial_size - 1] = 'z';
+                if (realloc_size > 0 && realloc_size <= initial_size) {
+                    p[realloc_size - 1] = 'y';
+                }
+            }
+            smp::submit_to(other_shard, [=] {
+                char* p2 = static_cast<char *>(realloc(p, realloc_size));
+                if (initial_size > 0 && realloc_size > 0) {
+                    BOOST_REQUIRE_EQUAL(p2[0], 'x');
+                    if (realloc_size <= initial_size) {
+                        BOOST_REQUIRE_EQUAL(p2[realloc_size - 1], 'y');
+                    }
+                    if (realloc_size > initial_size) {
+                        BOOST_REQUIRE_EQUAL(p2[initial_size - 1], 'z');
+                    }
+                }
+                free(p2);
+            }).get();
+        }
+    };
+
+    for (auto& cross_shard : {false, true}) {
+        do_xshard_realloc(cross_shard, 0, 0);
+        do_xshard_realloc(cross_shard, 0, 1);
+        do_xshard_realloc(cross_shard, 1, 0);
+        do_xshard_realloc(cross_shard, 50, 100);
+        do_xshard_realloc(cross_shard, 100, 50);
+        do_xshard_realloc(cross_shard, 100000, 500000);
+        do_xshard_realloc(cross_shard, 500000, 100000);
+    }
+}
+
+
 #ifndef SEASTAR_DEFAULT_ALLOCATOR
 
 struct thread_alloc_info {