From b8c64426e014a776da61e37c9c8819e19a1f414e Mon Sep 17 00:00:00 2001
From: Vladimir Buhtoyarov <v.bukhtoyarov@vk.team>
Date: Mon, 3 Jun 2024 14:50:23 +0300
Subject: [PATCH] Example of RemoteTaskExecution with embedded HotRod server

---
 pom.xml                                       |  5 +++
 .../embeddedhotrod/ServerTaskExample.java     | 32 +++++++++++++++++++
 .../SimpleEmbeddedHotRodServer.java           | 18 +++++++++++
 .../services/org.infinispan.tasks.ServerTask  |  1 +
 4 files changed, 56 insertions(+)
 create mode 100644 src/main/java/org/infinispan/playground/embeddedhotrod/ServerTaskExample.java
 create mode 100644 src/main/resources/META-INF/services/org.infinispan.tasks.ServerTask

diff --git a/pom.xml b/pom.xml
index d990ff8..b99e521 100644
--- a/pom.xml
+++ b/pom.xml
@@ -16,6 +16,11 @@
          <artifactId>infinispan-core</artifactId>
          <version>${version.infinispan}</version>
       </dependency>
+      <dependency>
+         <groupId>org.infinispan</groupId>
+         <artifactId>infinispan-server-runtime</artifactId>
+         <version>${version.infinispan}</version>
+      </dependency>
       <dependency>
          <groupId>org.infinispan</groupId>
          <artifactId>infinispan-server-hotrod</artifactId>
diff --git a/src/main/java/org/infinispan/playground/embeddedhotrod/ServerTaskExample.java b/src/main/java/org/infinispan/playground/embeddedhotrod/ServerTaskExample.java
new file mode 100644
index 0000000..2286ae2
--- /dev/null
+++ b/src/main/java/org/infinispan/playground/embeddedhotrod/ServerTaskExample.java
@@ -0,0 +1,32 @@
+package org.infinispan.playground.embeddedhotrod;
+
+import java.util.Map;
+
+import org.infinispan.tasks.ServerTask;
+import org.infinispan.tasks.TaskContext;
+
+public class ServerTaskExample implements ServerTask {
+
+    public static final String TASK_NAME = "just-example";
+    private static final ThreadLocal<TaskContext> treadLocalContext = new ThreadLocal();
+
+    @Override
+    public void setTaskContext(TaskContext taskContext) {
+        treadLocalContext.set(taskContext);
+    }
+
+    @Override
+    public Object call() throws Exception {
+        TaskContext taskContext = treadLocalContext.get();
+        Map<String, Object> params = (Map<String, Object>) taskContext.getParameters().get();
+        int first = (Integer) params.get("first");
+        int second = (Integer) params.get("second");
+        return first + second;
+    }
+
+    @Override
+    public String getName() {
+        return TASK_NAME;
+    }
+
+}
diff --git a/src/main/java/org/infinispan/playground/embeddedhotrod/SimpleEmbeddedHotRodServer.java b/src/main/java/org/infinispan/playground/embeddedhotrod/SimpleEmbeddedHotRodServer.java
index 1725db9..4e0690c 100644
--- a/src/main/java/org/infinispan/playground/embeddedhotrod/SimpleEmbeddedHotRodServer.java
+++ b/src/main/java/org/infinispan/playground/embeddedhotrod/SimpleEmbeddedHotRodServer.java
@@ -1,15 +1,20 @@
 package org.infinispan.playground.embeddedhotrod;
 
 import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
 
 import org.infinispan.Cache;
 import org.infinispan.client.hotrod.RemoteCache;
 import org.infinispan.client.hotrod.RemoteCacheManager;
 import org.infinispan.client.hotrod.configuration.ConfigurationBuilder;
+import org.infinispan.factories.GlobalComponentRegistry;
 import org.infinispan.manager.DefaultCacheManager;
+import org.infinispan.server.Extensions;
 import org.infinispan.server.hotrod.HotRodServer;
 import org.infinispan.server.hotrod.configuration.HotRodServerConfiguration;
 import org.infinispan.server.hotrod.configuration.HotRodServerConfigurationBuilder;
+import org.infinispan.tasks.TaskManager;
 
 public class SimpleEmbeddedHotRodServer {
 
@@ -18,6 +23,13 @@ public static void main(String[] args) throws IOException {
         DefaultCacheManager defaultCacheManager = new DefaultCacheManager("infinispan.xml");
         Cache<String, String> embeddedCache = defaultCacheManager.getCache("default");
 
+        // Load remote tasks into task manager
+        GlobalComponentRegistry gcr = defaultCacheManager.getGlobalComponentRegistry();
+        Extensions extensions = new Extensions();
+        extensions.load(SimpleEmbeddedHotRodServer.class.getClassLoader());
+        TaskManager taskManager = gcr.getComponent(TaskManager.class);
+        taskManager.registerTaskEngine(extensions.getServerTaskEngine(defaultCacheManager));
+
         // Create a Hot Rod server which exposes the cache manager
         HotRodServerConfiguration build = new HotRodServerConfigurationBuilder().build();
         HotRodServer server = new HotRodServer();
@@ -71,6 +83,12 @@ public static void main(String[] args) throws IOException {
            System.out.printf("%s...", s);
         }
 
+        System.out.print("\nVerifying remote task execution");
+        Map<String, Object> params = new HashMap<>();
+        params.put("first", 40);
+        params.put("second", 2);
+        assert Integer.valueOf(42).equals(remoteCache.execute(ServerTaskExample.TASK_NAME, params));
+
         System.out.println("\nDone !");
         remoteCacheManager.stop();
         server.stop();
diff --git a/src/main/resources/META-INF/services/org.infinispan.tasks.ServerTask b/src/main/resources/META-INF/services/org.infinispan.tasks.ServerTask
new file mode 100644
index 0000000..d4ed40a
--- /dev/null
+++ b/src/main/resources/META-INF/services/org.infinispan.tasks.ServerTask
@@ -0,0 +1 @@
+org.infinispan.playground.embeddedhotrod.ServerTaskExample
\ No newline at end of file