From d6593c10f924d8f7877f1a8bcd6d0e8c1436d2c8 Mon Sep 17 00:00:00 2001 From: slhafzjw Date: Thu, 9 Apr 2026 12:19:20 +0800 Subject: [PATCH] refactor(agent): support register registry shutdown hooks while agent launching --- .../runtime/interaction/WebSocketGateway.java | 27 ----------------- .../slhaf/partner/framework/agent/Agent.java | 16 ++++++++++ .../agent/factory/context/AgentContext.kt | 30 +++++++++++++++++++ .../agent/interaction/AgentGatewayRegistry.kt | 8 ++++- 4 files changed, 53 insertions(+), 28 deletions(-) diff --git a/Partner-Core/src/main/java/work/slhaf/partner/runtime/interaction/WebSocketGateway.java b/Partner-Core/src/main/java/work/slhaf/partner/runtime/interaction/WebSocketGateway.java index c7e68ad9..be77649c 100644 --- a/Partner-Core/src/main/java/work/slhaf/partner/runtime/interaction/WebSocketGateway.java +++ b/Partner-Core/src/main/java/work/slhaf/partner/runtime/interaction/WebSocketGateway.java @@ -46,7 +46,6 @@ public class WebSocketGateway extends WebSocketServer implements AgentGateway { - try { - // 先断开所有客户端连接 - for (WebSocket webSocket : getConnections()) { - try { - webSocket.close(1001, "Server shutting down"); - } catch (Exception e) { - log.warn("关闭客户端连接时出错: ", e); - } - } - //关闭WebSocketServer，给10秒超时时间确保连接正确关闭 - this.stop(10000); - log.info("WebSocketServer 已关闭"); - } catch (IllegalStateException e) { - log.warn("无法添加关闭钩子，JVM可能已在关闭过程中: ", e); - } catch (Exception e) { - log.error("WebSocketServer关闭失败: ", e); - } - })); - } catch (IllegalStateException e) { - log.warn("无法添加关闭钩子，JVM可能已在关闭过程中: ", e); - } - } - @Override public void onOpen(WebSocket webSocket, ClientHandshake clientHandshake) { log.info("新连接: {}", webSocket.getRemoteSocketAddress()); diff --git a/Partner-Framework/src/main/java/work/slhaf/partner/framework/agent/Agent.java b/Partner-Framework/src/main/java/work/slhaf/partner/framework/agent/Agent.java index ef0ebe10..a3a6b6e1 100644 --- a/Partner-Framework/src/main/java/work/slhaf/partner/framework/agent/Agent.java +++ b/Partner-Framework/src/main/java/work/slhaf/partner/framework/agent/Agent.java @@ -4,6 +4,7 @@ import lombok.extern.slf4j.Slf4j; import work.slhaf.partner.framework.agent.config.ConfigCenter; import work.slhaf.partner.framework.agent.exception.AgentLaunchFailedException; import work.slhaf.partner.framework.agent.factory.AgentRegisterFactory; +import work.slhaf.partner.framework.agent.factory.context.AgentContext; import work.slhaf.partner.framework.agent.interaction.AgentGatewayRegistration; import work.slhaf.partner.framework.agent.interaction.AgentGatewayRegistry; import work.slhaf.partner.framework.agent.model.ModelRuntimeRegistry; @@ -58,6 +59,21 @@ public final class Agent { for (AgentGatewayRegistration registration : gatewayRegistrations) { registration.register(); } + AgentContext.INSTANCE.addPreShutdownHook( + "agent-gateway-registry-close", + 0, + AgentGatewayRegistry.INSTANCE::close + ); + AgentContext.INSTANCE.addPostShutdownHook( + "state-center-save", + 0, + StateCenter.INSTANCE::save + ); + AgentContext.INSTANCE.addPostShutdownHook( + "config-center-close", + 100, + ConfigCenter.INSTANCE::close + ); Path externalModuleDir = ConfigCenter.INSTANCE.getPaths().getResourcesDir().resolve("module"); AgentRegisterFactory.addScanDir(externalModuleDir.toString()); AgentRegisterFactory.launch(applicationClass.getPackageName()); diff --git a/Partner-Framework/src/main/java/work/slhaf/partner/framework/agent/factory/context/AgentContext.kt b/Partner-Framework/src/main/java/work/slhaf/partner/framework/agent/factory/context/AgentContext.kt index 31ebe92f..ed2f0ff8 100644 --- a/Partner-Framework/src/main/java/work/slhaf/partner/framework/agent/factory/context/AgentContext.kt +++ b/Partner-Framework/src/main/java/work/slhaf/partner/framework/agent/factory/context/AgentContext.kt @@ -36,6 +36,8 @@ object AgentContext { get() = _metadata private val shutdownHooks = mutableMapOf>() + private val preShutdownHooks = mutableListOf() + private val postShutdownHooks = mutableListOf() init { installShutdownHook() @@ -91,6 +93,14 @@ object AgentContext { return true } + fun addPreShutdownHook(name: String, order: Int = 0, action: Runnable) { + preShutdownHooks.add(LifecycleShutdownHookDesc(name, order, action)) + } + + fun addPostShutdownHook(name: String, order: Int = 0, action: Runnable) { + postShutdownHooks.add(LifecycleShutdownHookDesc(name, order, action)) + } + private fun installShutdownHook() { class Instances( @@ -157,13 +167,27 @@ object AgentContext { } } + fun triggerLifecycleHooks(hooks: List) { + val log = LoggerFactory.getLogger(AgentContext::class.java) + hooks.sortedBy { it.order } + .forEach { + try { + it.action.run() + } catch (e: Exception) { + log.error("Failed to invoke lifecycle shutdown hook {}", it.name, e) + } + } + } + Runtime.getRuntime().addShutdownHook(Thread { val instances = computeInstances() + triggerLifecycleHooks(preShutdownHooks) shutdownHooks[ShutdownHookDesc.Type.RUNNING]?.let { trigger(it, instances) } shutdownHooks[ShutdownHookDesc.Type.ADDITIONAL]?.let { trigger(it, instances) } shutdownHooks[ShutdownHookDesc.Type.STANDALONE]?.let { trigger(it, instances) } shutdownHooks[ShutdownHookDesc.Type.SUB]?.let { trigger(it, instances) } shutdownHooks[ShutdownHookDesc.Type.CAPABILITY]?.let { trigger(it, instances) } + triggerLifecycleHooks(postShutdownHooks) }) } @@ -248,3 +272,9 @@ data class ShutdownHookDesc( CAPABILITY } } + +data class LifecycleShutdownHookDesc( + val name: String, + val order: Int, + val action: Runnable +) diff --git a/Partner-Framework/src/main/java/work/slhaf/partner/framework/agent/interaction/AgentGatewayRegistry.kt b/Partner-Framework/src/main/java/work/slhaf/partner/framework/agent/interaction/AgentGatewayRegistry.kt index 497b3929..c6b870aa 100644 --- a/Partner-Framework/src/main/java/work/slhaf/partner/framework/agent/interaction/AgentGatewayRegistry.kt +++ b/Partner-Framework/src/main/java/work/slhaf/partner/framework/agent/interaction/AgentGatewayRegistry.kt @@ -11,7 +11,7 @@ import java.nio.file.Path import java.util.concurrent.locks.ReentrantLock import kotlin.concurrent.withLock -object AgentGatewayRegistry : Configurable, ConfigRegistration { +object AgentGatewayRegistry : Configurable, ConfigRegistration, AutoCloseable { private val log = LoggerFactory.getLogger(AgentGatewayRegistry::class.java) private val registryLock = ReentrantLock() @@ -56,6 +56,12 @@ object AgentGatewayRegistry : Configurable, ConfigRegistration