From e5d4a62ae75178e110c73d4cf37ad06e0a68c715 Mon Sep 17 00:00:00 2001
From: Yoganand Rajasekaran <yrajas@microsoft.com>
Date: Tue, 15 Oct 2024 18:43:36 -0700
Subject: [PATCH] Use generics for CustomProcedure.

---
 libs/server/Custom/CustomProcedureWrapper.cs     | 3 ++-
 main/GarnetServer/Extensions/SetStringAndList.cs | 2 +-
 main/GarnetServer/Extensions/Sum.cs              | 2 +-
 test/Garnet.test/RespCustomCommandTests.cs       | 6 +++---
 4 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/libs/server/Custom/CustomProcedureWrapper.cs b/libs/server/Custom/CustomProcedureWrapper.cs
index 70a7dc5ddf..b0e5852a01 100644
--- a/libs/server/Custom/CustomProcedureWrapper.cs
+++ b/libs/server/Custom/CustomProcedureWrapper.cs
@@ -15,7 +15,8 @@ public abstract class CustomProcedure : CustomFunctions
         /// Custom command implementation
         /// </summary>
         /// <param name="garnetApi"></param>
-        public abstract bool Execute(IGarnetApi garnetApi, ArgSlice input, ref MemoryResult<byte> output);
+        public abstract bool Execute<TGarnetApi>(TGarnetApi garnetApi, ArgSlice input, ref MemoryResult<byte> output)
+            where TGarnetApi : IGarnetApi;
     }
 
     class CustomProcedureWrapper
diff --git a/main/GarnetServer/Extensions/SetStringAndList.cs b/main/GarnetServer/Extensions/SetStringAndList.cs
index c8ddfba24b..0b7ccd49a4 100644
--- a/main/GarnetServer/Extensions/SetStringAndList.cs
+++ b/main/GarnetServer/Extensions/SetStringAndList.cs
@@ -8,7 +8,7 @@ namespace Garnet
 {
     class SetStringAndList : CustomProcedure
     {
-        public override bool Execute(IGarnetApi garnetApi, ArgSlice input, ref MemoryResult<byte> output)
+        public override bool Execute<TGarnetApi>(TGarnetApi garnetApi, ArgSlice input, ref MemoryResult<byte> output)
         {
             var offset = 0;
             var key = GetNextArg(input, ref offset);
diff --git a/main/GarnetServer/Extensions/Sum.cs b/main/GarnetServer/Extensions/Sum.cs
index 3e227cfe10..37ce9607e1 100644
--- a/main/GarnetServer/Extensions/Sum.cs
+++ b/main/GarnetServer/Extensions/Sum.cs
@@ -8,7 +8,7 @@ namespace Garnet
 {
     class Sum : CustomProcedure
     {
-        public override bool Execute(IGarnetApi garnetApi, ArgSlice input, ref MemoryResult<byte> output)
+        public override bool Execute<TGarnetApi>(TGarnetApi garnetApi, ArgSlice input, ref MemoryResult<byte> output)
         {
             var offset = 0;
             var sum = 0;
diff --git a/test/Garnet.test/RespCustomCommandTests.cs b/test/Garnet.test/RespCustomCommandTests.cs
index ad59e7e7d9..8b211d5559 100644
--- a/test/Garnet.test/RespCustomCommandTests.cs
+++ b/test/Garnet.test/RespCustomCommandTests.cs
@@ -21,9 +21,9 @@ namespace Garnet.test
 {
     public class LargeGet : CustomProcedure
     {
-        public override bool Execute(IGarnetApi garnetApi, ArgSlice input, ref MemoryResult<byte> output)
+        public override bool Execute<TGarnetApi>(TGarnetApi garnetApi, ArgSlice input, ref MemoryResult<byte> output)
         {
-            static bool ResetBuffer(IGarnetApi garnetApi, ref MemoryResult<byte> output, int buffOffset)
+            static bool ResetBuffer(TGarnetApi garnetApi, ref MemoryResult<byte> output, int buffOffset)
             {
                 bool status = garnetApi.ResetScratchBuffer(buffOffset);
                 if (!status)
@@ -92,7 +92,7 @@ public override void Main<TGarnetApi>(TGarnetApi garnetApi, ArgSlice input, ref
 
     public class OutOfOrderFreeBuffer : CustomProcedure
     {
-        public override bool Execute(IGarnetApi garnetApi, ArgSlice input, ref MemoryResult<byte> output)
+        public override bool Execute<TGarnetApi>(TGarnetApi garnetApi, ArgSlice input, ref MemoryResult<byte> output)
         {
             var offset = 0;
             var key = GetNextArg(input, ref offset);