From 5ffb9767aff7199a65d5b1ec42a89b2377d20e62 Mon Sep 17 00:00:00 2001
From: John Spray <jcs@redpanda.com>
Date: Tue, 2 Aug 2022 15:55:36 +0100
Subject: [PATCH] storage: make memory units release safer

While in general the accounting is robust, this is a little
fragile in error paths (or when the code is changed in future),
as it is an exception-generating condition to release more
units from a sem_units than you took: a double release
could perhaps occur if we released before spill(), then
there was an exception in spill() that caused us to iterate.
---
 src/v/storage/spill_key_index.cc |  7 ++-----
 src/v/storage/spill_key_index.h  | 12 +++++++++++-
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/v/storage/spill_key_index.cc b/src/v/storage/spill_key_index.cc
index 0406c22e83f0b..02e9b6baa0f77 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 60a6577ec7a7b..6aec5c0ee48df 100644
--- a/src/v/storage/spill_key_index.h
+++ b/src/v/storage/spill_key_index.h
@@ -91,7 +91,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.
@@ -99,6 +99,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);