fix(txn): discard partial writes from failed EXEC commands

src/commands/cmd_function.cc

@@ -121,8 +121,9 @@ uint64_t GenerateFunctionFlags(uint64_t flags, const std::vector<std::string> &a
 }
 
 uint64_t GenerateFCallFlags(uint64_t flags, const std::vector<std::string> &, const Config &config) {
+  flags |= kCmdSkipTxnSavepoint;
   if (!config.lua_strict_key_accessing) {
-    return flags | kCmdExclusive;
+    flags |= kCmdExclusive;
   }
 
   return flags;

src/commands/cmd_script.cc

@@ -126,8 +126,9 @@ uint64_t GenerateScriptFlags(uint64_t flags, const std::vector<std::string> &arg
 }
 
 uint64_t GenerateEvalFlags(uint64_t flags, const std::vector<std::string> &, const Config &config) {
+  flags |= kCmdSkipTxnSavepoint;
   if (!config.lua_strict_key_accessing) {
-    return flags | kCmdExclusive;
+    flags |= kCmdExclusive;
   }
 
   return flags;

src/commands/cmd_txn.cc

@@ -78,16 +78,25 @@ class CommandExec : public Commander {
     }
 
     auto storage = srv->storage;
+    // If EXEC stops early, some commands may not have replies yet.
+    // Use the queued command count to keep the error array length correct.
+    auto command_count = conn->GetMultiExecCommands()->size();
     // Execute multi-exec commands
     conn->SetInExec();
     auto s = storage->BeginTxn();
     if (s.IsOK()) {
-      conn->ExecuteCommands(conn->GetMultiExecCommands());
+      s = conn->ExecuteCommands(conn->GetMultiExecCommands());
       // In Redis, errors happening after EXEC instead are not handled in a special way:
       // all the other commands will be executed even if some command fails during
       // the transaction.
       // So, if conn->IsMultiError(), the transaction should still be committed.
-      s = storage->CommitTxn();
+      // Only fatal internal errors make ExecuteCommands return a non-OK status
+      // (e.g. savepoint failures) and abort the shared batch.
+      if (s.IsOK()) {
+        s = storage->CommitTxn();
+      } else if (auto abort_s = storage->AbortTxn(); !abort_s.IsOK()) {
+        s = abort_s;
+      }
     }
 
     conn->ResetMultiExec();
@@ -96,7 +105,7 @@ class CommandExec : public Commander {
     if (s) {
       conn->Reply(Array(conn->GetQueuedReplies()));
     } else {
-      conn->Reply(Array(std::vector<std::string>(conn->GetQueuedReplies().size(), redis::Error(s))));
+      conn->Reply(Array(std::vector<std::string>(command_count, redis::Error(s))));
     }
 
     conn->ClearQueuedReplies();

src/commands/commander.h

@@ -89,6 +89,8 @@ enum CommandFlags : uint64_t {
   kCmdAdmin = 1ULL << 11,
   // "skip-monitor" flag, for commands that should skip monitor feed
   kCmdSkipMonitor = 1ULL << 12,
+  // "skip-txn-savepoint" flag, for commands that should skip savepoint
+  kCmdSkipTxnSavepoint = 1ULL << 13,
 };
 
 enum class CommandCategory : uint8_t {

src/server/redis_connection.cc

@@ -113,7 +113,11 @@ void Connection::OnRead([[maybe_unused]] struct bufferevent *bev) {
     return;
   }
 
-  ExecuteCommands(req_.GetCommands());
+  s = ExecuteCommands(req_.GetCommands());
+  if (!s.IsOK()) {
+    Reply(redis::Error(s));
+    return;
+  }
   if (IsFlagEnabled(kCloseAsync)) {
     Close();
   }
@@ -423,7 +427,7 @@ static bool IsCmdAllowedInStaleData(const std::string &cmd_name) {
   return cmd_name == "info" || cmd_name == "slaveof" || cmd_name == "config";
 }
 
-void Connection::ExecuteCommands(std::deque<CommandTokens> *to_process_cmds) {
+Status Connection::ExecuteCommands(std::deque<CommandTokens> *to_process_cmds) {
   const Config *config = srv_->GetConfig();
   std::string reply;
   const std::string &password = config->requirepass;
@@ -447,7 +451,7 @@ void Connection::ExecuteCommands(std::deque<CommandTokens> *to_process_cmds) {
             "[connection] A likely HTTP request is detected in the RESP connection, indicating a potential "
             "Cross-Protocol Scripting attack. Connection aborted.");
         EnableFlag(kCloseAsync);
-        return;
+        return Status::OK();
       }
       Reply(redis::Error(
           {Status::NotOK,
@@ -473,7 +477,7 @@ void Connection::ExecuteCommands(std::deque<CommandTokens> *to_process_cmds) {
     if (srv_->PauseConnIfNeeded(this, cmd_name, cmd_flags)) {
       multi_error_exit.Disable();  // Don't mark transaction as failed - we're deferring, not erroring
       to_process_cmds->push_front(std::move(cmd_tokens));
-      return;
+      return Status::OK();
     }
 
     if (GetNamespace().empty()) {
@@ -589,6 +593,12 @@ void Connection::ExecuteCommands(std::deque<CommandTokens> *to_process_cmds) {
       }
       engine::Context ctx(srv_->storage);
 
+      bool use_txn_savepoint = in_exec_ && (cmd_flags & kCmdWrite) && !(cmd_flags & kCmdSkipTxnSavepoint);
+      if (use_txn_savepoint) {
+        s = srv_->storage->SetTxnSavePoint();
+        if (!s.IsOK()) return s.Prefixed("failed to set transaction command savepoint");
+      }
+
       std::vector<GlobalIndexer::RecordResult> index_records;
       if (!srv_->index_mgr.index_map.empty() && IsCmdForIndexing(cmd_flags, attributes->category) &&
           !config->cluster_enabled) {
@@ -609,12 +619,38 @@ void Connection::ExecuteCommands(std::deque<CommandTokens> *to_process_cmds) {
       }
 
       s = ExecuteCommand(ctx, cmd_name, cmd_tokens, current_cmd.get(), &reply);
-      for (const auto &record : index_records) {
-        auto s = GlobalIndexer::Update(ctx, record);
-        if (!s.IsOK() && !s.Is<Status::TypeMismatched>()) {
-          WARN("[connection] index updating failed for key: {}", record.key);
+
+      if (use_txn_savepoint) {
+        if (s.IsOK()) {
+          for (const auto &record : index_records) {
+            auto index_s = GlobalIndexer::Update(ctx, record);
+            if (index_s.IsOK() || index_s.Is<Status::TypeMismatched>()) continue;
+
+            WARN("[connection] index updating failed for key: {}", record.key);
+            s = index_s;
+            break;
+          }
+        }
+
+        if (s.IsOK()) {
+          auto pop_s = srv_->storage->PopTxnSavePoint();
+          if (!pop_s.IsOK()) return pop_s.Prefixed("failed to pop transaction command savepoint");
+        } else {
+          auto rollback_s = srv_->storage->RollbackTxnToSavePoint();
+          if (!rollback_s.IsOK()) {
+            return rollback_s.Prefixed("failed to rollback transaction command savepoint");
+          }
+        }
+      } else {
+        for (const auto &record : index_records) {
+          auto index_s = GlobalIndexer::Update(ctx, record);
+          if (!index_s.IsOK() && !index_s.Is<Status::TypeMismatched>()) {
+            WARN("[connection] index updating failed for key: {}", record.key);
+          }
         }
       }
+
+      if (!exec_error_.IsOK()) return exec_error_;
     }
 
     if (!(cmd_flags & redis::kCmdSkipMonitor)) {
@@ -643,11 +679,13 @@ void Connection::ExecuteCommands(std::deque<CommandTokens> *to_process_cmds) {
     if (!reply.empty()) Reply(reply);
     reply.clear();
   }
+  return Status::OK();
 }
 
 void Connection::ResetMultiExec() {
   in_exec_ = false;
   multi_error_ = false;
+  exec_error_ = Status::OK();
   multi_cmds_.clear();
   DisableFlag(Connection::kMultiExec);
 }

src/server/redis_connection.h

@@ -175,7 +175,7 @@ class Connection : public EvbufCallbackBase<Connection> {
   evbuffer *Input() { return bufferevent_get_input(bev_); }
   evbuffer *Output() { return bufferevent_get_output(bev_); }
   bufferevent *GetBufferEvent() { return bev_; }
-  void ExecuteCommands(std::deque<CommandTokens> *to_process_cmds);
+  Status ExecuteCommands(std::deque<CommandTokens> *to_process_cmds);
   Status ExecuteCommand(engine::Context &ctx, const std::string &cmd_name, const std::vector<std::string> &cmd_tokens,
                         Commander *current_cmd, std::string *reply);
   bool IsProfilingEnabled(const std::string &cmd);
@@ -194,6 +194,7 @@ class Connection : public EvbufCallbackBase<Connection> {
   bool IsInExec() const { return in_exec_; }
   bool IsInScript() const { return in_script_; }
   bool IsMultiError() const { return multi_error_; }
+  void SetExecError(Status status) { exec_error_ = std::move(status); }
   void ResetMultiExec();
   std::deque<redis::CommandTokens> *GetMultiExecCommands() { return &multi_cmds_; }
 
@@ -235,6 +236,8 @@ class Connection : public EvbufCallbackBase<Connection> {
   Server *srv_;
   bool in_exec_ = false;
   bool multi_error_ = false;
+  // Fatal EXEC error from nested paths (e.g. script savepoint failure); aborts the whole EXEC.
+  Status exec_error_;
   std::atomic<bool> is_running_ = false;
   std::deque<redis::CommandTokens> multi_cmds_;
   bool in_script_ = false;

src/storage/scripting.cc

@@ -887,8 +887,29 @@ int RedisGenericCommand(lua_State *lua, int raise_error) {
     return raise_error ? RaiseError(lua) : 1;
   }
 
+  bool use_txn_savepoint = conn->IsInExec() && (cmd_flags & redis::kCmdWrite);
+  if (use_txn_savepoint) {
+    auto savepoint_s = srv->storage->SetTxnSavePoint();
+    if (!savepoint_s.IsOK()) {
+      auto err = savepoint_s.Prefixed("failed to set script command savepoint");
+      PushError(lua, err.Msg().c_str());
+      conn->SetExecError(std::move(err));
+      return RaiseError(lua);
+    }
+  }
+
   std::string output;
   s = conn->ExecuteCommand(*script_run_ctx->ctx, cmd_name, args, cmd.get(), &output);
+  if (use_txn_savepoint) {
+    auto savepoint_s = s.IsOK() ? srv->storage->PopTxnSavePoint() : srv->storage->RollbackTxnToSavePoint();
+    if (!savepoint_s.IsOK()) {
+      auto err = savepoint_s.Prefixed(s.IsOK() ? "failed to pop script command savepoint"
+                                               : "failed to rollback script command savepoint");
+      PushError(lua, err.Msg().c_str());
+      conn->SetExecError(std::move(err));
+      return RaiseError(lua);
+    }
+  }
   if (!s) {
     PushError(lua, s.Msg().data());
     return raise_error ? RaiseError(lua) : 1;

src/storage/storage.cc

@@ -48,6 +48,7 @@
 #include "rocksdb/write_batch.h"
 #include "rocksdb_crc32c.h"
 #include "server/server.h"
+#include "status.h"
 #include "storage/batch_indexer.h"
 #include "string_util.h"
 #include "table_properties_collector.h"
@@ -999,6 +1000,45 @@ Status Storage::CommitTxn() {
   return {Status::NotOK, s.ToString()};
 }
 
+Status Storage::AbortTxn() {
+  if (!is_txn_mode_) {
+    return Status{Status::NotOK, "cannot abort while not in transaction mode"};
+  }
+  is_txn_mode_ = false;
+  txn_write_batch_.reset();
+  return Status::OK();
+}
+
+Status Storage::SetTxnSavePoint() {
+  if (!is_txn_mode_) {
+    return Status{Status::NotOK, "cannot set savepoint while not in transaction mode"};
+  }
+  txn_write_batch_->SetSavePoint();
+  return Status::OK();
+}
+
+Status Storage::PopTxnSavePoint() {
+  if (!is_txn_mode_) {
+    return Status{Status::NotOK, "cannot pop savepoint while not in transaction mode"};
+  }
+  auto s = txn_write_batch_->PopSavePoint();
+  if (!s.ok()) {
+    return Status{Status::NotOK, s.ToString()};
+  }
+  return Status::OK();
+}
+
+Status Storage::RollbackTxnToSavePoint() {
+  if (!is_txn_mode_) {
+    return Status{Status::NotOK, "cannot rollback savepoint while not in transaction mode"};
+  }
+  auto s = txn_write_batch_->RollbackToSavePoint();
+  if (!s.ok()) {
+    return Status{Status::NotOK, s.ToString()};
+  }
+  return Status::OK();
+}
+
 ObserverOrUniquePtr<rocksdb::WriteBatchBase> Storage::GetWriteBatchBase() {
   if (is_txn_mode_) {
     return ObserverOrUniquePtr<rocksdb::WriteBatchBase>(txn_write_batch_.get(), ObserverOrUnique::Observer);

src/storage/storage.h

@@ -311,6 +311,10 @@ class Storage {
 
   Status BeginTxn();
   Status CommitTxn();
+  Status AbortTxn();
+  Status SetTxnSavePoint();
+  Status PopTxnSavePoint();
+  Status RollbackTxnToSavePoint();
   ObserverOrUniquePtr<rocksdb::WriteBatchBase> GetWriteBatchBase();
 
   Storage(const Storage &) = delete;

tests/gocase/unit/multi/multi_test.go

@@ -22,6 +22,7 @@ package multi
 import (
 	"context"
 	"fmt"
+	"strings"
 	"testing"
 
 	"github.com/apache/kvrocks/tests/gocase/util"
@@ -98,6 +99,79 @@ func TestMulti(t *testing.T) {
 		require.Equal(t, map[string]string{"f0": "v0", "f1": "v1", "f2": "v2"}, hgetall.Val())
 	})
 
+	t.Run("Runtime command errors don't leave partial writes", func(t *testing.T) {
+		largeA := strings.Repeat("A", 64)
+		largeB := strings.Repeat("B", 64)
+		largeC := strings.Repeat("C", 64)
+
+		require.NoError(t, rdb.ConfigSet(ctx, "rocksdb.write_options.write_batch_max_bytes", "0").Err())
+		require.NoError(t, rdb.Del(ctx, "txhash", "after").Err())
+		require.NoError(t, rdb.HSet(ctx, "txhash", "f1", "old1", "f2", "old2", "f3", "old3").Err())
+		require.NoError(t, rdb.ConfigSet(ctx, "rocksdb.write_options.write_batch_max_bytes", "180").Err())
+		defer func() {
+			require.NoError(t, rdb.ConfigSet(ctx, "rocksdb.write_options.write_batch_max_bytes", "0").Err())
+		}()
+
+		require.NoError(t, rdb.Do(ctx, "MULTI").Err())
+		require.NoError(t, rdb.Do(ctx, "HSET", "txhash", "f1", largeA, "f2", largeB, "f3", largeC).Err())
+		require.NoError(t, rdb.Do(ctx, "SET", "after", "ok").Err())
+		require.NoError(t, rdb.Do(ctx, "HMGET", "txhash", "f1", "f2", "f3").Err())
+
+		replies := rdb.Do(ctx, "EXEC").Val().([]interface{})
+		require.Len(t, replies, 3)
+		require.ErrorContains(t, replies[0].(error), "Memory limit reached")
+		require.Equal(t, "OK", replies[1])
+		require.Equal(t, []interface{}{"old1", "old2", "old3"}, replies[2])
+		require.Equal(t, []interface{}{"old1", "old2", "old3"}, rdb.HMGet(ctx, "txhash", "f1", "f2", "f3").Val())
+		require.Equal(t, "ok", rdb.Get(ctx, "after").Val())
+	})
+
+	t.Run("EVAL runtime errors preserve successful nested commands", func(t *testing.T) {
+		require.NoError(t, rdb.ConfigSet(ctx, "rocksdb.write_options.write_batch_max_bytes", "0").Err())
+		require.NoError(t, rdb.Del(ctx, "scriptkey", "after-script").Err())
+
+		require.NoError(t, rdb.Do(ctx, "MULTI").Err())
+		require.NoError(t, rdb.Do(ctx, "EVAL",
+			"redis.call('SET', KEYS[1], 'v'); return redis.call('HGET', KEYS[1], 'field')", "1", "scriptkey").Err())
+		require.NoError(t, rdb.Do(ctx, "SET", "after-script", "ok").Err())
+
+		replies := rdb.Do(ctx, "EXEC").Val().([]interface{})
+		require.Len(t, replies, 2)
+		require.ErrorContains(t, replies[0].(error), "WRONGTYPE")
+		require.Equal(t, "OK", replies[1])
+		require.Equal(t, "v", rdb.Get(ctx, "scriptkey").Val())
+		require.Equal(t, "ok", rdb.Get(ctx, "after-script").Val())
+	})
+
+	t.Run("redis.pcall rolls back partial writes from its failed command", func(t *testing.T) {
+		largeA := strings.Repeat("A", 64)
+		largeB := strings.Repeat("B", 64)
+		largeC := strings.Repeat("C", 64)
+
+		require.NoError(t, rdb.ConfigSet(ctx, "rocksdb.write_options.write_batch_max_bytes", "0").Err())
+		require.NoError(t, rdb.Del(ctx, "pcall-hash", "pcall-after").Err())
+		require.NoError(t, rdb.HSet(ctx, "pcall-hash", "f1", "old1", "f2", "old2", "f3", "old3").Err())
+		require.NoError(t, rdb.ConfigSet(ctx, "rocksdb.write_options.write_batch_max_bytes", "180").Err())
+		defer func() {
+			require.NoError(t, rdb.ConfigSet(ctx, "rocksdb.write_options.write_batch_max_bytes", "0").Err())
+		}()
+
+		script := `
+			local result = redis.pcall('HSET', KEYS[1], 'f1', ARGV[1], 'f2', ARGV[2], 'f3', ARGV[3])
+			redis.call('SET', KEYS[2], 'ok')
+			return result
+		`
+		require.NoError(t, rdb.Do(ctx, "MULTI").Err())
+		require.NoError(t, rdb.Do(ctx, "EVAL", script, "2", "pcall-hash", "pcall-after", largeA, largeB, largeC).Err())
+
+		replies := rdb.Do(ctx, "EXEC").Val().([]interface{})
+		require.Len(t, replies, 1)
+		require.ErrorContains(t, replies[0].(error), "Memory limit reached")
+		require.Equal(t, []interface{}{"old1", "old2", "old3"},
+			rdb.HMGet(ctx, "pcall-hash", "f1", "f2", "f3").Val())
+		require.Equal(t, "ok", rdb.Get(ctx, "pcall-after").Val())
+	})
+
 	t.Run("EXEC fails if there are errors while queueing commands #1", func(t *testing.T) {
 		require.NoError(t, rdb.Del(ctx, "foo1", "foo2").Err())
 		require.NoError(t, rdb.Do(ctx, "MULTI").Err())