From 0c8bd6ffffa26b3b0cfd3bceaa5afa0890fae79c Mon Sep 17 00:00:00 2001
From: Vasileios Zois <96085550+vazois@users.noreply.github.com>
Date: Fri, 28 Feb 2025 15:52:21 -0800
Subject: [PATCH] Fix Connection Dispose when Object Store Disabled Exception
 is thrown (#1057)

* ensure connection does not close when issuing obj command and getting an uninitialized ERR

* add test for object store disabled bug

* fix formatting

* addressing comments

---------

Co-authored-by: Tal Zaccai <talzacc@microsoft.com>
---
 libs/common/GarnetException.cs                | 32 +++++++++++++++----
 libs/server/Resp/RespServerSession.cs         |  7 ++--
 .../Storage/Session/ObjectStore/Common.cs     |  2 +-
 test/Garnet.test/RespTests.cs                 | 30 +++++++++++++++++
 4 files changed, 61 insertions(+), 10 deletions(-)

diff --git a/libs/common/GarnetException.cs b/libs/common/GarnetException.cs
index 0e335604d8..379523a71d 100644
--- a/libs/common/GarnetException.cs
+++ b/libs/common/GarnetException.cs
@@ -18,44 +18,62 @@ public class GarnetException : Exception
         public LogLevel LogLevel { get; } = LogLevel.Trace;
         public bool ClientResponse { get; } = true;
         public bool Panic { get; } = false;
+        public bool DisposeSession { get; } = true;
 
         /// <summary>
-        /// Throw Garnet exception
+        /// Throw Garnet exception.
         /// </summary>
-        public GarnetException(LogLevel logLevel = LogLevel.Trace, bool clientResponse = true, bool panic = false)
+        /// <param name="logLevel"></param>
+        /// <param name="clientResponse"></param>
+        /// <param name="panic"></param>
+        /// <param name="disposeSession"></param>
+        public GarnetException(LogLevel logLevel = LogLevel.Trace, bool clientResponse = true, bool panic = false, bool disposeSession = true)
         {
             LogLevel = logLevel;
             ClientResponse = clientResponse;
             Panic = panic;
+            DisposeSession = disposeSession;
         }
 
         /// <summary>
-        /// Throw Garnet exception with message
+        /// Throw Garnet exception with message.
         /// </summary>
         /// <param name="message"></param>
-        public GarnetException(string message, LogLevel logLevel = LogLevel.Trace, bool clientResponse = true, bool panic = false) : base(message)
+        /// <param name="logLevel"></param>
+        /// <param name="clientResponse"></param>
+        /// <param name="panic"></param>
+        /// <param name="disposeSession"></param>
+        public GarnetException(string message, LogLevel logLevel = LogLevel.Trace, bool clientResponse = true, bool panic = false, bool disposeSession = true) : base(message)
         {
             LogLevel = logLevel;
             ClientResponse = clientResponse;
             Panic = panic;
+            DisposeSession = disposeSession;
         }
 
         /// <summary>
-        /// Throw Garnet exception with message and inner exception
+        /// Throw Garnet exception with message and inner exception.
         /// </summary>
         /// <param name="message"></param>
         /// <param name="innerException"></param>
-        public GarnetException(string message, Exception innerException, LogLevel logLevel = LogLevel.Trace, bool clientResponse = true, bool panic = false) : base(message, innerException)
+        /// <param name="logLevel"></param>
+        /// <param name="clientResponse"></param>
+        /// <param name="panic"></param>
+        /// <param name="disposeSession"></param>
+        public GarnetException(string message, Exception innerException, LogLevel logLevel = LogLevel.Trace, bool clientResponse = true, bool panic = false, bool disposeSession = true) : base(message, innerException)
         {
             LogLevel = logLevel;
             ClientResponse = clientResponse;
             Panic = panic;
+            DisposeSession = disposeSession;
         }
 
         /// <summary>
         /// Throw helper that throws a GarnetException.
         /// </summary>
-        /// <param name="message">Exception message.</param>
+        /// <param name="message"></param>
+        /// <param name="logLevel"></param>
+        /// <exception cref="GarnetException"></exception>
         [DoesNotReturn]
         public static void Throw(string message, LogLevel logLevel = LogLevel.Trace) =>
             throw new GarnetException(message, logLevel);
diff --git a/libs/server/Resp/RespServerSession.cs b/libs/server/Resp/RespServerSession.cs
index 8e42a4c51b..71b8e32b49 100644
--- a/libs/server/Resp/RespServerSession.cs
+++ b/libs/server/Resp/RespServerSession.cs
@@ -371,8 +371,11 @@ public override int TryConsumeMessages(byte* reqBuffer, int bytesReceived)
                     Environment.Exit(-1);
                 }
 
-                // The session is no longer usable, dispose it
-                networkSender.DisposeNetworkSender(true);
+                if (ex.DisposeSession)
+                {
+                    // The session is no longer usable, dispose it
+                    networkSender.DisposeNetworkSender(true);
+                }
             }
             catch (Exception ex)
             {
diff --git a/libs/server/Storage/Session/ObjectStore/Common.cs b/libs/server/Storage/Session/ObjectStore/Common.cs
index e5df24178a..ff993dabf3 100644
--- a/libs/server/Storage/Session/ObjectStore/Common.cs
+++ b/libs/server/Storage/Session/ObjectStore/Common.cs
@@ -558,7 +558,7 @@ public GarnetStatus ObjectScan<TObjectContext>(byte[] key, ref ObjectInput input
 
         [MethodImpl(MethodImplOptions.NoInlining)]
         static void ThrowObjectStoreUninitializedException()
-            => throw new GarnetException("Object store is disabled");
+            => throw new GarnetException("Object store is disabled", disposeSession: false);
 
         #endregion
 
diff --git a/test/Garnet.test/RespTests.cs b/test/Garnet.test/RespTests.cs
index 477923356d..bde2c6911b 100644
--- a/test/Garnet.test/RespTests.cs
+++ b/test/Garnet.test/RespTests.cs
@@ -1504,6 +1504,36 @@ public void SingleDeleteWithObjectStoreDisable_LTM()
             }
         }
 
+        [Test]
+        public void GarnetObjectStoreDisabledError()
+        {
+            TearDown();
+            TestUtils.DeleteDirectory(TestUtils.MethodTestDir, wait: true);
+            server = TestUtils.CreateGarnetServer(TestUtils.MethodTestDir, disableObjects: true);
+            server.Start();
+
+            using var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig());
+            var db = redis.GetDatabase(0);
+
+            var iter = 100;
+            var mykey = "mykey";
+            for (var i = 0; i < iter; i++)
+            {
+                var exception = Assert.Throws<StackExchange.Redis.RedisServerException>(() => _ = db.ListLength(mykey));
+                ClassicAssert.AreEqual("ERR Garnet Exception: Object store is disabled", exception.Message);
+            }
+
+            // Ensure connection is still healthy
+            for (var i = 0; i < iter; i++)
+            {
+                var myvalue = "myvalue" + i;
+                var result = db.StringSet(mykey, myvalue);
+                ClassicAssert.IsTrue(result);
+                var returned = (string)db.StringGet(mykey);
+                ClassicAssert.AreEqual(myvalue, returned);
+            }
+        }
+
         [Test]
         public void MultiKeyDelete([Values] bool withoutObjectStore)
         {