Merge pull request redpanda-data#23946 from redpanda-data/stephan/mem…

…-sampler-off-by-1-admin

samplers: Fix off by one shard id check

src/v/redpanda/admin/debug.cc

@@ -31,6 +31,7 @@
 #include "serde/rw/rw.h"
 #include "storage/kvstore.h"
 
+#include <seastar/core/shard_id.hh>
 #include <seastar/core/sstring.hh>
 #include <seastar/http/httpd.hh>
 #include <seastar/json/json_elements.hh>
@@ -494,6 +495,14 @@ void admin_server::register_debug_routes() {
 
 using admin::apply_validator;
 
+void check_shard_id(seastar::shard_id id) {
+    auto max_shard_id = ss::smp::count - 1;
+    if (id > max_shard_id) {
+        throw ss::httpd::bad_param_exception(
+          fmt::format("Shard id too high, max shard id is {}", max_shard_id));
+    }
+}
+
 ss::future<ss::json::json_return_type>
 admin_server::cpu_profile_handler(std::unique_ptr<ss::http::request> req) {
     vlog(adminlog.info, "Request to sampled cpu profile");
@@ -509,12 +518,7 @@ admin_server::cpu_profile_handler(std::unique_ptr<ss::http::request> req) {
     }
 
     if (shard_id.has_value()) {
-        auto all_cpus = ss::smp::all_cpus();
-        auto max_shard_id = std::max_element(all_cpus.begin(), all_cpus.end());
-        if (*shard_id > *max_shard_id) {
-            throw ss::httpd::bad_param_exception(fmt::format(
-              "Shard id too high, max shard id is {}", *max_shard_id));
-        }
+        check_shard_id(*shard_id);
     }
 
     std::optional<std::chrono::milliseconds> wait_ms;
@@ -694,11 +698,7 @@ admin_server::sampled_memory_profile_handler(
     }
 
     if (shard_id.has_value()) {
-        auto max_shard_id = ss::smp::count;
-        if (*shard_id > max_shard_id) {
-            throw ss::httpd::bad_param_exception(fmt::format(
-              "Shard id too high, max shard id is {}", max_shard_id));
-        }
+        check_shard_id(*shard_id);
     }
 
     auto profiles = co_await _memory_sampling_service.local()

tests/rptest/tests/memory_sampling_test.py

@@ -11,6 +11,7 @@
 from rptest.services.cluster import cluster
 from rptest.tests.redpanda_test import RedpandaTest
 from rptest.utils.mode_checks import skip_debug_mode
+from requests.exceptions import HTTPError
 
 BOOTSTRAP_CONFIG = {
     'memory_enable_memory_sampling': True,
@@ -65,3 +66,14 @@ def test_get_per_shard_stacks(self):
         assert 'size' in profile[0]['allocation_sites'][0]
         assert 'count' in profile[0]['allocation_sites'][0]
         assert 'backtrace' in profile[0]['allocation_sites'][0]
+
+        # test out of bounds shard
+        failed = False
+        try:
+            max_shard_id_plus_one = self.redpanda.shards()[1] + 1
+            admin.get_sampled_memory_profile(shard=max_shard_id_plus_one)
+        except HTTPError as e:
+            if 'shard id too high' in e.response.text.lower():
+                failed = True
+
+        assert failed, "Asking for too high shard id didn't fail"