Java: Add FUNCTION KILL command. (#1560)

* Java: Add `FUNCTION KILL` command. (#328) * Add `FUNCTION KILL` command. * Frantic and rampant code clean up. * More tests for the god of the tests! * Address PR comments. * Typo fix. * Fixing the tests s02ep06. * Fixing the tests s03ep01. * Address PR comments. Signed-off-by: Yury-Fridlyand <[email protected]> Co-authored-by: aaron-congo <[email protected]>
valkey-io · Jun 13, 2024 · 850fc10 · 850fc10
1 parent ed84725
commit 850fc10
Show file tree

Hide file tree

Showing 12 changed files with 610 additions and 26 deletions.
diff --git a/glide-core/src/protobuf/redis_request.proto b/glide-core/src/protobuf/redis_request.proto
@@ -200,6 +200,7 @@ enum RequestType {
     BLMPop = 158;
     XLen = 159;
     Sort = 160;
+    FunctionKill = 161;
     LSet = 165;
     XDel = 166;
     XRange = 167;

diff --git a/glide-core/src/request_type.rs b/glide-core/src/request_type.rs
@@ -170,6 +170,7 @@ pub enum RequestType {
     BLMPop = 158,
     XLen = 159,
     Sort = 160,
+    FunctionKill = 161,
     LSet = 165,
     XDel = 166,
     XRange = 167,
@@ -357,6 +358,7 @@ impl From<::protobuf::EnumOrUnknown<ProtobufRequestType>> for RequestType {
             ProtobufRequestType::ExpireTime => RequestType::ExpireTime,
             ProtobufRequestType::PExpireTime => RequestType::PExpireTime,
             ProtobufRequestType::XLen => RequestType::XLen,
+            ProtobufRequestType::FunctionKill => RequestType::FunctionKill,
             ProtobufRequestType::LSet => RequestType::LSet,
             ProtobufRequestType::XDel => RequestType::XDel,
             ProtobufRequestType::XRange => RequestType::XRange,
@@ -541,6 +543,7 @@ impl RequestType {
             RequestType::ExpireTime => Some(cmd("EXPIRETIME")),
             RequestType::PExpireTime => Some(cmd("PEXPIRETIME")),
             RequestType::XLen => Some(cmd("XLEN")),
+            RequestType::FunctionKill => Some(get_two_word_command("FUNCTION", "KILL")),
             RequestType::LSet => Some(cmd("LSET")),
             RequestType::XDel => Some(cmd("XDEL")),
             RequestType::XRange => Some(cmd("XRANGE")),

diff --git a/java/client/src/main/java/glide/api/RedisClient.java b/java/client/src/main/java/glide/api/RedisClient.java
@@ -20,6 +20,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.FlushAll;
 import static redis_request.RedisRequestOuterClass.RequestType.FunctionDelete;
 import static redis_request.RedisRequestOuterClass.RequestType.FunctionFlush;
+import static redis_request.RedisRequestOuterClass.RequestType.FunctionKill;
 import static redis_request.RedisRequestOuterClass.RequestType.FunctionList;
 import static redis_request.RedisRequestOuterClass.RequestType.FunctionLoad;
 import static redis_request.RedisRequestOuterClass.RequestType.Info;
@@ -277,4 +278,9 @@ public CompletableFuture<Boolean> copy(
         }
         return commandManager.submitNewCommand(Copy, arguments, this::handleBooleanResponse);
     }
+
+    @Override
+    public CompletableFuture<String> functionKill() {
+        return commandManager.submitNewCommand(FunctionKill, new String[0], this::handleStringResponse);
+    }
 }
diff --git a/java/client/src/main/java/glide/api/RedisClusterClient.java b/java/client/src/main/java/glide/api/RedisClusterClient.java
@@ -22,6 +22,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.FlushAll;
 import static redis_request.RedisRequestOuterClass.RequestType.FunctionDelete;
 import static redis_request.RedisRequestOuterClass.RequestType.FunctionFlush;
+import static redis_request.RedisRequestOuterClass.RequestType.FunctionKill;
 import static redis_request.RedisRequestOuterClass.RequestType.FunctionList;
 import static redis_request.RedisRequestOuterClass.RequestType.FunctionLoad;
 import static redis_request.RedisRequestOuterClass.RequestType.Info;
@@ -575,4 +576,15 @@ public CompletableFuture<ClusterValue<Object>> fcall(
                                 ? ClusterValue.ofSingleValue(handleObjectOrNullResponse(response))
                                 : ClusterValue.ofMultiValue(handleMapResponse(response)));
     }
+
+    @Override
+    public CompletableFuture<String> functionKill() {
+        return commandManager.submitNewCommand(FunctionKill, new String[0], this::handleStringResponse);
+    }
+
+    @Override
+    public CompletableFuture<String> functionKill(@NonNull Route route) {
+        return commandManager.submitNewCommand(
+                FunctionKill, new String[0], route, this::handleStringResponse);
+    }
 }
diff --git a/java/client/src/main/java/glide/api/commands/ScriptingAndFunctionsClusterCommands.java b/java/client/src/main/java/glide/api/commands/ScriptingAndFunctionsClusterCommands.java
@@ -339,4 +339,37 @@ CompletableFuture<ClusterValue<Map<String, Object>[]>> functionList(
      * }</pre>
      */
     CompletableFuture<ClusterValue<Object>> fcall(String function, String[] arguments, Route route);
+
+    /**
+     * Kills a function that is currently executing.<br>
+     * <code>FUNCTION KILL</code> terminates read-only functions only.<br>
+     * The command will be routed to all primary nodes.
+     *
+     * @since Redis 7.0 and above.
+     * @see <a href="https://redis.io/docs/latest/commands/function-kill/">redis.io</a> for details.
+     * @return <code>OK</code> if function is terminated. Otherwise, throws an error.
+     * @example
+     *     <pre>{@code
+     * String response = client.functionKill().get();
+     * assert response.equals("OK");
+     * }</pre>
+     */
+    CompletableFuture<String> functionKill();
+
+    /**
+     * Kills a function that is currently executing.<br>
+     * <code>FUNCTION KILL</code> terminates read-only functions only.
+     *
+     * @since Redis 7.0 and above.
+     * @see <a href="https://redis.io/docs/latest/commands/function-kill/">redis.io</a> for details.
+     * @param route Specifies the routing configuration for the command. The client will route the
+     *     command to the nodes defined by <code>route</code>.
+     * @return <code>OK</code> if function is terminated. Otherwise, throws an error.
+     * @example
+     *     <pre>{@code
+     * String response = client.functionKill(RANDOM).get();
+     * assert response.equals("OK");
+     * }</pre>
+     */
+    CompletableFuture<String> functionKill(Route route);
 }
diff --git a/java/client/src/main/java/glide/api/commands/ScriptingAndFunctionsCommands.java b/java/client/src/main/java/glide/api/commands/ScriptingAndFunctionsCommands.java
@@ -142,4 +142,19 @@ public interface ScriptingAndFunctionsCommands {
      * }</pre>
      */
     CompletableFuture<Object> fcall(String function);
+
+    /**
+     * Kills a function that is currently executing.<br>
+     * <code>FUNCTION KILL</code> terminates read-only functions only.
+     *
+     * @since Redis 7.0 and above.
+     * @see <a href="https://redis.io/docs/latest/commands/function-kill/">redis.io</a> for details.
+     * @return <code>OK</code> if function is terminated. Otherwise, throws an error.
+     * @example
+     *     <pre>{@code
+     * String response = client.functionKill().get();
+     * assert response.equals("OK");
+     * }</pre>
+     */
+    CompletableFuture<String> functionKill();
 }
diff --git a/java/client/src/test/java/glide/api/RedisClientTest.java b/java/client/src/test/java/glide/api/RedisClientTest.java
@@ -86,6 +86,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.FlushAll;
 import static redis_request.RedisRequestOuterClass.RequestType.FunctionDelete;
 import static redis_request.RedisRequestOuterClass.RequestType.FunctionFlush;
+import static redis_request.RedisRequestOuterClass.RequestType.FunctionKill;
 import static redis_request.RedisRequestOuterClass.RequestType.FunctionList;
 import static redis_request.RedisRequestOuterClass.RequestType.FunctionLoad;
 import static redis_request.RedisRequestOuterClass.RequestType.GeoAdd;
@@ -5228,6 +5229,27 @@ public void fcall_returns_success() {
         assertEquals(value, payload);
     }
 
+    @SneakyThrows
+    @Test
+    public void functionKill_returns_success() {
+        // setup
+        String[] args = new String[0];
+        CompletableFuture<String> testResponse = new CompletableFuture<>();
+        testResponse.complete(OK);
+
+        // match on protobuf request
+        when(commandManager.<String>submitNewCommand(eq(FunctionKill), eq(args), any()))
+                .thenReturn(testResponse);
+
+        // exercise
+        CompletableFuture<String> response = service.functionKill();
+        String payload = response.get();
+
+        // verify
+        assertEquals(testResponse, response);
+        assertEquals(OK, payload);
+    }
+
     @SneakyThrows
     @Test
     public void bitcount_returns_success() {

diff --git a/java/client/src/test/java/glide/api/RedisClusterClientTest.java b/java/client/src/test/java/glide/api/RedisClusterClientTest.java
@@ -30,6 +30,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.FlushAll;
 import static redis_request.RedisRequestOuterClass.RequestType.FunctionDelete;
 import static redis_request.RedisRequestOuterClass.RequestType.FunctionFlush;
+import static redis_request.RedisRequestOuterClass.RequestType.FunctionKill;
 import static redis_request.RedisRequestOuterClass.RequestType.FunctionList;
 import static redis_request.RedisRequestOuterClass.RequestType.FunctionLoad;
 import static redis_request.RedisRequestOuterClass.RequestType.Info;
@@ -1543,4 +1544,46 @@ public void fcall_without_keys_and_with_route_returns_success() {
         assertEquals(testResponse, response);
         assertEquals(value, payload);
     }
+
+    @SneakyThrows
+    @Test
+    public void functionKill_returns_success() {
+        // setup
+        String[] args = new String[0];
+        CompletableFuture<String> testResponse = new CompletableFuture<>();
+        testResponse.complete(OK);
+
+        // match on protobuf request
+        when(commandManager.<String>submitNewCommand(eq(FunctionKill), eq(args), any()))
+                .thenReturn(testResponse);
+
+        // exercise
+        CompletableFuture<String> response = service.functionKill();
+        String payload = response.get();
+
+        // verify
+        assertEquals(testResponse, response);
+        assertEquals(OK, payload);
+    }
+
+    @SneakyThrows
+    @Test
+    public void functionKill_with_route_returns_success() {
+        // setup
+        String[] args = new String[0];
+        CompletableFuture<String> testResponse = new CompletableFuture<>();
+        testResponse.complete(OK);
+
+        // match on protobuf request
+        when(commandManager.<String>submitNewCommand(eq(FunctionKill), eq(args), eq(RANDOM), any()))
+                .thenReturn(testResponse);
+
+        // exercise
+        CompletableFuture<String> response = service.functionKill(RANDOM);
+        String payload = response.get();
+
+        // verify
+        assertEquals(testResponse, response);
+        assertEquals(OK, payload);
+    }
 }
diff --git a/java/integTest/src/test/java/glide/TestUtilities.java b/java/integTest/src/test/java/glide/TestUtilities.java
@@ -181,4 +181,36 @@ public static String generateLuaLibCode(
         }
         return code.toString();
     }
+
+    /**
+     * Create a lua lib with a RO function which runs an endless loop up to timeout sec.<br>
+     * Execution takes at least 5 sec regardless of the timeout configured.<br>
+     * If <code>readOnly</code> is <code>false</code>, function sets a dummy value to the first key
+     * given.
+     */
+    public static String createLuaLibWithLongRunningFunction(
+            String libName, String funcName, int timeout, boolean readOnly) {
+        String code =
+                "#!lua name=$libName\n"
+                        + "local function $libName_$funcName(keys, args)\n"
+                        + "  local started = tonumber(redis.pcall('time')[1])\n"
+                        // fun fact - redis does no writes if 'no-writes' flag is set
+                        + "  redis.pcall('set', keys[1], 42)\n"
+                        + "  while (true) do\n"
+                        + "    local now = tonumber(redis.pcall('time')[1])\n"
+                        + "    if now > started + $timeout then\n"
+                        + "      return 'Timed out $timeout sec'\n"
+                        + "    end\n"
+                        + "  end\n"
+                        + "  return 'OK'\n"
+                        + "end\n"
+                        + "redis.register_function{\n"
+                        + "function_name='$funcName',\n"
+                        + "callback=$libName_$funcName,\n"
+                        + (readOnly ? "flags={ 'no-writes' }\n" : "")
+                        + "}";
+        return code.replace("$timeout", Integer.toString(timeout))
+                .replace("$funcName", funcName)
+                .replace("$libName", libName);
+    }
 }