基于java agent实现服务热更新

日常服务更新时，通常需要进行代码修改、编译构建、生成镜像、上传镜像、停止旧容器启动新容器、健康检查等待服务启动完成等步骤，开发者每次修改完代码后一般需要完成以上步骤才能完成功能的更新，耗时且费力，通过java agent实现服务可实现代码秒级更新

核心是通过ASM修改字节码，流程如下：

可以围绕该技术进行服务配置、文件上传、埋点数据上报、数据统计分析等功能

核心代码：
agent端定时扫描指定文件夹下上传的class文件，通过Instrumentation重新定义加载class文件：

public class AgentMain {

    protected static final Logger LOG = LoggerFactory.getLogger(AgentMain.class);

    private static final String FOLDER_PATH = "/xxx/upload/class"; // 目标文件夹路径

    private static final Set<String>  failedClassSet = new HashSet<>();
    private static final Object lockObject = new Object();

    private static final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);
    public static void agentmain(String args, Instrumentation inst) {

        synchronized(lockObject) {
            Class[] allLoadedClasses = inst.getAllLoadedClasses();
            for (Class loadedClass : allLoadedClasses) {
                // 找到启动类的加载器，通过启动类的加载器加载要替换的class文件
                if (loadedClass.getSimpleName().equals("Application")) {
                    try {
                        ClassLoader classLoader = loadedClass.getClassLoader();
                        // 调度每10秒执行一次的任务
                        premain(classLoader,inst);
                        break;
                    } catch (Exception var13) {
                        throw new RuntimeException(var13);
                    }
                }
            }
        }
    }

    public static void premain(ClassLoader classLoader, Instrumentation instrumentation) {
        ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);
        Random random = new Random();
        int randomIntInRange = random.nextInt(5);
        scheduler.scheduleAtFixedRate(() -> {
            try {
                List<File> classFiles = getClassFilesFromFolder(FOLDER_PATH);
                for (File classFile : classFiles) {
                    byte[] classBytes = readClassFile(classFile.getAbsolutePath());
                    String className = extractClassName(classFile.getName());
                    if (!failedClassSet.contains(className)) {
                        try {
                            Class<?> clazz = classLoader.loadClass(className);
                            ClassDefinition classDefinition = new ClassDefinition(clazz, classBytes);
                            instrumentation.redefineClasses(classDefinition);
                            LOG.info("Class {} has been redefined.",className);
                            // 更新完删除class文件
                            if (classFile.delete()) {
                                LOG.info("Class file {} deleted successfully.",classFile.getName());
                            } else {
                                LOG.error("Failed to delete class file {}" , classFile.getName());
                            }
                        } catch (ClassNotFoundException e) {
                            LOG.error("Class {} not found: ",className);
                            failedClassSet.add(className);
                        } catch (Exception e) {
                            LOG.error("Failed to redefine class " + className + ": " + e.getMessage());
                            failedClassSet.add(className);
                        }
                    }
                }
            } catch (Exception e) {
                LOG.error("Failed to redefine class");
            }
        }, randomIntInRange, 5+randomIntInRange, TimeUnit.SECONDS); // 延迟0秒后开始，每5秒执行一次
    }

    public static String extractClassName(String fileName) {
        fileName = fileName.replace('_', '.');;
        if (fileName.endsWith(".class")) {
            String className = fileName.substring(0, fileName.length() - 6);
            return className;
        } else {
            throw new IllegalArgumentException("The file name does not end with .class");
        }
    }

    // 获取指定文件夹下所有的 .class 文件
    public static List<File> getClassFilesFromFolder(String folderPath) throws IOException {
        try (var paths = Files.walk(Paths.get(folderPath))) {
            return paths.filter(path -> Files.isRegularFile(path) && path.toString().endsWith(".class"))
                    .map(Path::toFile)
                    .collect(Collectors.toList());
        }
    }

    // 读取 .class 文件字节码
    private static byte[] readClassFile(String classFilePath) throws IOException {
        Path path = Paths.get(classFilePath);
        return Files.readAllBytes(path);
    }
}

server端 扫描指定配置的服务，将agent 加载到服务对应的JVM进程：

public class AttachApplication {

    private static final long INITIAL_DELAY = 0;
    private static final long PERIOD = 10; // 时间间隔，单位是秒
    protected static final Logger LOG = LoggerFactory.getLogger(AttachApplication.class);

    // 存储已附加的虚拟机 ID
    private static final Set<String> attachedVMs = new HashSet<>();

    public static void main(String[] args) {
        SpringApplication.run(AttachApplication.class, args);
       startScheduledTask();
    }

    @Bean
    public ScheduledExecutorService scheduledExecutorService() {
        return Executors.newScheduledThreadPool(1);
    }

    private static void startScheduledTask() {
        ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);
        Runnable task = () -> {
            try {
                LOG.info("AttachApplication start");
                List<VirtualMachineDescriptor> vmDescriptors = VirtualMachine.list();
                LOG.info("vmDescriptors size :{}",vmDescriptors.size());
                for (VirtualMachineDescriptor vmd : vmDescriptors) {
                    String vmId = vmd.id();
                    if (attachedVMs.add(vmId)) {
                        LOG.info("vmId add successful：{}",vmId);
                        // 指定的jvm进程 
                        if (vmd.displayName().contains("XXXXApplication")) {
                            LOG.info("attach JVM process vmId：{}",vmId);
                            VirtualMachine virtualMachine = VirtualMachine.attach(vmId);
                            LOG.info("attach JVM process successful vmId：{}",vmId);
                            try {
                                LOG.info("load agent");
                                // 加载 agent
                                virtualMachine.loadAgent(getAgentPath());
                                LOG.info("virtualMachine loadAgent successful vmd.displayName: "+vmd.displayName());
                            } catch (Throwable e) {
                                LOG.error("attach jvm process fail",e);
                            } finally {
                                virtualMachine.detach();
                            }
                        }
                    }
                }

            } catch (Exception e) {
                LOG.error("attach application error:",e);
            }
            LOG.info("AttachApplication execute end.");
        };

        scheduler.scheduleAtFixedRate(task, INITIAL_DELAY, PERIOD, TimeUnit.SECONDS);
    }

    private static String getAgentPath() {
        try {

            // 从 resources 目录中获取 agent 文件的路径
            InputStream inputStream = AttachApplication.class.getClassLoader().getResourceAsStream("java-agent.jar");
            if (inputStream == null) {
                throw new IOException("can not find agent JAR file");
            }

            // 将资源文件保存到临时目录中
            Path tempFile = Files.createTempFile("java-agent", ".jar");
            Files.copy(inputStream, tempFile, java.nio.file.StandardCopyOption.REPLACE_EXISTING);
            tempFile.toFile().deleteOnExit(); // 确保程序退出时删除临时文件
            return tempFile.toAbsolutePath().toString();

        } catch (IOException e) {
            throw new RuntimeException("can not find agent JAR file path", e);
        }
    }
}