diff --git a/src/v/storage/spill_key_index.cc b/src/v/storage/spill_key_index.cc index 0406c22e83f0..02e9b6baa0f7 100644 --- a/src/v/storage/spill_key_index.cc +++ b/src/v/storage/spill_key_index.cc @@ -129,8 +129,7 @@ ss::future<> spill_key_index::add_key(compaction_key b, value_type v) { node.key(), node.mapped(), [this](const bytes& k, value_type o) { - _keys_mem_usage -= k.size(); - _mem_units.return_units(k.size()); + release_entry_memory(k); return spill(compacted_index::entry_type::key, k, o); }); }); @@ -241,9 +240,7 @@ ss::future<> spill_key_index::drain_all_keys() { }, [this] { auto node = _midx.extract(_midx.begin()); - auto mem_usage = entry_mem_usage(node.key()); - _keys_mem_usage -= mem_usage; - _mem_units.return_units(mem_usage); + release_entry_memory(node.key()); return ss::do_with( node.key(), node.mapped(), [this](const bytes& k, value_type o) { return spill(compacted_index::entry_type::key, k, o); diff --git a/src/v/storage/spill_key_index.h b/src/v/storage/spill_key_index.h index 9c3ca1514e1f..637724530921 100644 --- a/src/v/storage/spill_key_index.h +++ b/src/v/storage/spill_key_index.h @@ -92,7 +92,7 @@ class spill_key_index final : public compacted_index_writer::impl { return debug::AllocatedByteSize(_midx); } - size_t entry_mem_usage(const compaction_key& k) const { + size_t entry_mem_usage(const bytes& k) const { // One entry in a node hash map: key and value // are allocated together, and the key is a basic_sstring with // internal buffer that may be spilled if key was longer. @@ -100,6 +100,16 @@ class spill_key_index final : public compacted_index_writer::impl { return (is_external ? sizeof(k) + k.size() : sizeof(k)) + value_sz; } + void release_entry_memory(const bytes& k) { + auto entry_memory = entry_mem_usage(k); + _keys_mem_usage -= entry_memory; + + // Handle the case of a double-release, in case this comes up + // during retries/exception handling. + auto release_units = std::min(entry_memory, _mem_units.count()); + _mem_units.return_units(release_units); + } + ss::future<> drain_all_keys(); ss::future<> add_key(compaction_key, value_type); ss::future<> spill(compacted_index::entry_type, bytes_view, value_type);