From 438ae59d60259744d3185a202b6b22ccd22f7c20 Mon Sep 17 00:00:00 2001
From: Tyler Rockwood <rockwood@redpanda.com>
Date: Tue, 30 Jan 2024 21:19:44 -0600
Subject: [PATCH] storage/kvstore: remove oversized alloc

It's possible to have enough entries in the kvstore that we run into
oversized allocs. Prevent that by using a btree_map

Signed-off-by: Tyler Rockwood <rockwood@redpanda.com>
(cherry picked from commit a8c3a83d18366ce5ee054e554895e3b41f73ec00)
---
 src/v/bytes/bytes.h      | 6 ++++++
 src/v/bytes/iobuf.cc     | 8 ++++++++
 src/v/storage/kvstore.cc | 1 -
 src/v/storage/kvstore.h  | 4 ++--
 4 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/v/bytes/bytes.h b/src/v/bytes/bytes.h
index 7c62dcddbfff..074f1b1ddc7c 100644
--- a/src/v/bytes/bytes.h
+++ b/src/v/bytes/bytes.h
@@ -177,3 +177,9 @@ operator^(const std::array<char, Size>& a, const std::array<char, Size>& b) {
       a.begin(), a.end(), b.begin(), out.begin(), std::bit_xor<>());
     return out;
 }
+
+struct bytes_type_cmp {
+    using is_transparent = std::true_type;
+    bool operator()(const bytes& lhs, const bytes_view& rhs) const;
+    bool operator()(const bytes& lhs, const bytes& rhs) const;
+};
diff --git a/src/v/bytes/iobuf.cc b/src/v/bytes/iobuf.cc
index 265f4553096b..06fa16a2a668 100644
--- a/src/v/bytes/iobuf.cc
+++ b/src/v/bytes/iobuf.cc
@@ -338,3 +338,11 @@ std::ostream& operator<<(std::ostream& os, const bytes_view& b) {
     return os;
 }
 } // namespace std
+
+bool bytes_type_cmp::operator()(const bytes& lhs, const bytes_view& rhs) const {
+    return lhs < rhs;
+}
+
+bool bytes_type_cmp::operator()(const bytes& lhs, const bytes& rhs) const {
+    return lhs < rhs;
+}
diff --git a/src/v/storage/kvstore.cc b/src/v/storage/kvstore.cc
index 84b1e25018b2..3981c0fc575f 100644
--- a/src/v/storage/kvstore.cc
+++ b/src/v/storage/kvstore.cc
@@ -480,7 +480,6 @@ ss::future<> kvstore::load_snapshot_from_reader(snapshot_reader& reader) {
     }
 
     auto lock = co_await _db_mut.get_units();
-    _db.reserve(batch.header().record_count);
     co_await batch.for_each_record_async([this](model::record r) {
         auto key = iobuf_to_bytes(r.release_key());
         _probe.add_cached_bytes(key.size() + r.value().size_bytes());
diff --git a/src/v/storage/kvstore.h b/src/v/storage/kvstore.h
index 2788ea78569b..61ba3a9bfac6 100644
--- a/src/v/storage/kvstore.h
+++ b/src/v/storage/kvstore.h
@@ -25,7 +25,7 @@
 #include <seastar/core/gate.hh>
 #include <seastar/core/timer.hh>
 
-#include <absl/container/node_hash_map.h>
+#include <absl/container/btree_map.h>
 
 namespace storage {
 
@@ -162,7 +162,7 @@ class kvstore {
     // Protect _db and _next_offset across asynchronous mutations.
     mutex _db_mut;
     model::offset _next_offset;
-    absl::node_hash_map<bytes, iobuf, bytes_type_hash, bytes_type_eq> _db;
+    absl::btree_map<bytes, iobuf, bytes_type_cmp> _db;
     std::optional<ntp_sanitizer_config> _ntp_sanitizer_config;
 
     ss::future<> put(key_space ks, bytes key, std::optional<iobuf> value);