From 132fd1e3f2867e1cbe80ada7ea79c9b2195e47ab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20Paszkowski?= <lukasz.paszkowski@scylladb.com>
Date: Thu, 6 Mar 2025 08:43:53 +0100
Subject: [PATCH] replica/exceptions: Add a new custom replica exception

The new exception `critical_disk_utilization_exception` is thrown
when the user table mutation writes are being blocked due to e.g.
reaching a critical disk utilization level.

This new exception, is then correctly handled on the coordinator
side when transforming into `mutation_write_failure_exception` with
a meaningful error message: "Write rejected due to critical disk
utilization".
---
 idl/replica_exception.idl.hh |  7 ++++++-
 replica/exceptions.cc        |  2 ++
 replica/exceptions.hh        | 20 +++++++++++++++++++-
 service/storage_proxy.cc     |  5 +++++
 4 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/idl/replica_exception.idl.hh b/idl/replica_exception.idl.hh
index b5a434ffc8..fe9810981b 100644
--- a/idl/replica_exception.idl.hh
+++ b/idl/replica_exception.idl.hh
@@ -23,12 +23,17 @@ class stale_topology_exception {
 class abort_requested_exception {
 };
 
+class critical_disk_utilization_exception {
+    sstring failed_action();
+};
+
 struct exception_variant {
     std::variant<replica::unknown_exception,
             replica::no_exception,
             replica::rate_limit_exception,
             replica::stale_topology_exception,
-            replica::abort_requested_exception
+            replica::abort_requested_exception,
+            replica::critical_disk_utilization_exception
     > reason;
 };
 
diff --git a/replica/exceptions.cc b/replica/exceptions.cc
index dfb3eb06a9..7ffb266108 100644
--- a/replica/exceptions.cc
+++ b/replica/exceptions.cc
@@ -23,6 +23,8 @@ exception_variant try_encode_replica_exception(std::exception_ptr eptr) {
         return e;
     } catch (abort_requested_exception&) {
         return abort_requested_exception();
+    } catch (const critical_disk_utilization_exception& e) {
+        return e;
     } catch (...) {
         return no_exception{};
     }
diff --git a/replica/exceptions.hh b/replica/exceptions.hh
index 1f4d2810a3..dea2afbb88 100644
--- a/replica/exceptions.hh
+++ b/replica/exceptions.hh
@@ -63,6 +63,23 @@ public:
     virtual const char* what() const noexcept override { return _message.c_str(); }
 };
 
+class critical_disk_utilization_exception final: public replica_exception {
+    seastar::sstring _failed_action;
+    seastar::sstring _message;
+public:
+    critical_disk_utilization_exception(std::string_view failed_action) noexcept
+        : replica_exception()
+        , _failed_action(failed_action)
+        , _message(seastar::format("Critical disk utilization: {}", failed_action))
+    { }
+
+    const seastar::sstring& failed_action() const {
+        return _failed_action;
+    }
+
+    virtual const char* what() const noexcept override { return _message.c_str(); }
+};
+
 using abort_requested_exception = seastar::abort_requested_exception;
 
 struct exception_variant {
@@ -70,7 +87,8 @@ struct exception_variant {
             no_exception,
             rate_limit_exception,
             stale_topology_exception,
-            abort_requested_exception
+            abort_requested_exception,
+            critical_disk_utilization_exception
     > reason;
 
     exception_variant()
diff --git a/service/storage_proxy.cc b/service/storage_proxy.cc
index d649ca012e..36eb863246 100644
--- a/service/storage_proxy.cc
+++ b/service/storage_proxy.cc
@@ -774,6 +774,9 @@ private:
                     } else if constexpr (std::is_same_v<Ex, replica::abort_requested_exception>) {
                         msg = e.what();
                         return error::FAILURE;
+                    } else if constexpr (std::is_same_v<Ex, replica::critical_disk_utilization_exception>) {
+                        msg = e.what();
+                        return error::FAILURE;
                     }
                 }, exception->reason);
             }
@@ -4491,6 +4494,8 @@ void storage_proxy::send_to_live_endpoints(storage_proxy::response_id_type respo
                 err = error::TIMEOUT;
             } else if (auto* e = try_catch<db::virtual_table_update_exception>(eptr)) {
                 msg = e->grab_cause();
+            } else if (auto* e = try_catch<replica::critical_disk_utilization_exception>(eptr)) {
+                msg = e->what();
             } else {
                 slogger.error("exception during mutation write to {}: {}", coordinator, eptr);
             }