SED-4443 ClassCastException in FunctionMessageHandler

step-functions/step-functions-handler/src/main/java/step/functions/handler/FileApplicationContextFactory.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright (C) 2025, exense GmbH
+ *
+ * This file is part of Step
+ *
+ * Step is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * Step is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with Step.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package step.functions.handler;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import step.grid.contextbuilder.ApplicationContextFactory;
+import step.grid.contextbuilder.ClassPathHelper;
+
+import java.io.File;
+import java.net.URL;
+import java.net.URLClassLoader;
+import java.util.List;
+
+public class FileApplicationContextFactory extends ApplicationContextFactory {
+
+    private static final Logger logger = LoggerFactory.getLogger(FileApplicationContextFactory.class);
+
+    private final File jar;
+
+    public FileApplicationContextFactory(File jar) {
+        this.jar = jar;
+    }
+
+    public String getId() {
+        return this.jar.getAbsolutePath();
+    }
+
+    public boolean requiresReload() {
+        return false;
+    }
+
+    public ClassLoader buildClassLoader(ClassLoader parentClassLoader) {
+        if (logger.isDebugEnabled()) {
+            logger.debug("Creating URLClassLoader from extracted local file {}", this.jar.getAbsolutePath());
+        }
+        List<URL> urls = ClassPathHelper.forSingleFile(this.jar);
+        URL[] urlArray = urls.toArray(new URL[urls.size()]);
+        return new URLClassLoader(urlArray, parentClassLoader);
+    }
+
+    public void onClassLoaderClosed() {
+    }
+}

step-functions/step-functions-handler/src/main/java/step/functions/handler/FunctionMessageHandler.java

@@ -30,9 +30,9 @@
 import step.grid.agent.handler.context.OutputMessageBuilder;
 import step.grid.agent.tokenpool.AgentTokenWrapper;
 import step.grid.agent.tokenpool.TokenReservationSession;
+import step.grid.bootstrap.ResourceExtractor;
 import step.grid.contextbuilder.ApplicationContextBuilder;
 import step.grid.contextbuilder.ApplicationContextControl;
-import step.grid.contextbuilder.LocalResourceApplicationContextFactory;
 import step.grid.contextbuilder.RemoteApplicationContextFactory;
 import step.grid.filemanager.FileVersionId;
 import step.grid.io.InputMessage;
@@ -41,8 +41,10 @@
 import step.livereporting.client.LiveReportingClient;
 import step.reporting.LiveReporting;
 
+import java.io.File;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Proxy;
+import java.nio.file.Files;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -71,6 +73,7 @@ public class FunctionMessageHandler extends AbstractMessageHandler {
 	private ApplicationContextBuilder applicationContextBuilder;
 
 	public FunctionHandlerFactory functionHandlerFactory;
+	private File functionHandlerInitializerJar;
 
 	public FunctionMessageHandler() {
 		super();
@@ -115,6 +118,7 @@ public void init(AgentTokenServices agentTokenServices) {
 		applicationContextBuilder.forkCurrentContext(BRANCH_HANDLER_INITIALIZER);
 
 		functionHandlerFactory = new FunctionHandlerFactory(applicationContextBuilder, agentTokenServices.getFileManagerClient());
+		functionHandlerInitializerJar = ResourceExtractor.extractResource(this.getClass().getClassLoader(), "step-functions-handler-initializer.jar");
 	}
 
 	@Override
@@ -179,8 +183,8 @@ public OutputMessage handle(AgentTokenWrapper token, InputMessage inputMessage)
 	}
 
 	private LiveReporting initializeLiveReporting(Map<String, String> properties, TokenReservationSession tokenReservationSession) throws Exception {
-		ApplicationContextControl applicationContextControl = applicationContextBuilder.pushContext(BRANCH_HANDLER_INITIALIZER, new LocalResourceApplicationContextFactory(this.getClass().getClassLoader(), "step-functions-handler-initializer.jar"), true);
-		// The usage of this application context will be released when the session is closed, underlying registered file won't be cleanable before this release happens
+		ApplicationContextControl applicationContextControl = applicationContextBuilder.pushContext(BRANCH_HANDLER_INITIALIZER, new FileApplicationContextFactory(functionHandlerInitializerJar), false);
+		// The usage of this application context will be released when the session is closed
 		tokenReservationSession.registerObjectToBeClosedWithSession(applicationContextControl);
 		return applicationContextBuilder.runInContext(BRANCH_HANDLER_INITIALIZER, () -> {
 			// There's no easy way to do this in the AbstractFunctionHandler itself, because
@@ -273,5 +277,8 @@ public void close() throws Exception {
 		if (liveReportingExecutor != null) {
 			liveReportingExecutor.shutdownNow();
 		}
+		if (functionHandlerInitializerJar != null) {
+			Files.deleteIfExists(functionHandlerInitializerJar.toPath());
+		}
 	}
 }