Merge pull request #2250 from newrelic/crac_testing

jbedell-newrelic · web-flow · commit 66d1f1327185 · 2025-03-12T07:59:24.000-04:00
Implement support for CRaC
diff --git a/newrelic-agent/build.gradle b/newrelic-agent/build.gradle
@@ -101,6 +101,8 @@ dependencies {
 
     shadowIntoJar("org.yaml:snakeyaml:1.33")
 
+    shadowIntoJar("org.crac:crac:1.5.0")
+
     implementation("javax.management.j2ee:management-api:1.1-rev-1") {
         transitive = false
     }
diff --git a/newrelic-agent/src/main/java/com/newrelic/agent/MetricNames.java b/newrelic-agent/src/main/java/com/newrelic/agent/MetricNames.java
@@ -416,6 +416,9 @@ public class MetricNames {
 
     public static final String SUPPORTABILITY_AGENT_CONNECT_BACKOFF_ATTEMPTS = "Supportability/Agent/Collector/Connect/BackoffAttempts";
 
+    public static final String SUPPORTABILITY_AGENT_CRAC_CHECKPOINT = "Supportability/Agent/CracCheckpoint";
+    public static final String SUPPORTABILITY_AGENT_CRAC_RESTORE = "Supportability/Agent/CracRestore";
+
     // expected errors
     public static final String SUPPORTABILITY_API_EXPECTED_ERROR_API_MESSAGE = "ExpectedError/Api/Message";
     public static final String SUPPORTABILITY_API_EXPECTED_ERROR_API_THROWABLE = "ExpectedError/Api/Throwable";
diff --git a/newrelic-agent/src/main/java/com/newrelic/agent/instrumentation/context/FinalClassTransformer.java b/newrelic-agent/src/main/java/com/newrelic/agent/instrumentation/context/FinalClassTransformer.java
@@ -145,14 +145,17 @@ public MethodVisitor visitMethod(int access, String name, String desc, String si
     }
 
     private void writeClassFiles(String className, InstrumentationContext context, byte[] classBytes) {
+        PrintWriter oldFileWriter = null;
+        PrintWriter newFileWriter = null;
         try {
             File old = File.createTempFile(className.replace('/', '_'), ".old", BootstrapLoader.getTempDir());
-            Utils.print(context.bytes, new PrintWriter(old));
-
+            oldFileWriter = new PrintWriter(old);
+            Utils.print(context.bytes, oldFileWriter);
             Agent.LOG.debug("Wrote " + old.getAbsolutePath());
-            File newFile = File.createTempFile(className.replace('/', '_'), ".new", BootstrapLoader.getTempDir());
 
-            Utils.print(classBytes, new PrintWriter(newFile));
+            File newFile = File.createTempFile(className.replace('/', '_'), ".new", BootstrapLoader.getTempDir());
+            newFileWriter = new PrintWriter(newFile);
+            Utils.print(classBytes, newFileWriter);
             Agent.LOG.debug("Wrote " + newFile.getAbsolutePath());
 
             File newClassFile = File.createTempFile(className.replace('/', '_'), ".new.class", BootstrapLoader.getTempDir());
@@ -162,6 +165,9 @@ private void writeClassFiles(String className, InstrumentationContext context, b
             Agent.LOG.debug("Wrote " + newClassFile.getAbsolutePath());
         } catch (Throwable t) {
             Agent.LOG.log(Level.FINEST, t, "Error writing debug bytecode for {0}", className);
+        } finally {
+            if (oldFileWriter != null) oldFileWriter.close();
+            if (newFileWriter != null) newFileWriter.close();
         }
     }
 
diff --git a/newrelic-agent/src/main/java/com/newrelic/agent/logging/Log4jLogManager.java b/newrelic-agent/src/main/java/com/newrelic/agent/logging/Log4jLogManager.java
@@ -11,9 +11,15 @@
 import com.newrelic.agent.config.AgentConfig;
 import com.newrelic.agent.config.AgentConfigImpl;
 import com.newrelic.agent.config.AgentJarHelper;
+import com.newrelic.api.agent.NewRelic;
 import org.apache.logging.log4j.Level;
+import org.apache.logging.log4j.core.config.Configurator;
+import org.apache.logging.log4j.core.config.builder.api.ConfigurationBuilder;
+import org.apache.logging.log4j.core.config.builder.api.ConfigurationBuilderFactory;
+import org.apache.logging.log4j.core.config.builder.impl.BuiltConfiguration;
 
 import java.io.File;
+import java.net.MalformedURLException;
 import java.net.URL;
 import java.text.MessageFormat;
 import java.util.HashMap;
@@ -78,27 +84,6 @@ private Log4jLogger initializeRootLogger(String name) {
                 // and we want to clear any other log4j system properties
                 clearAllLog4jSystemProperties(systemProps);
 
-                URL log4jConfigXmlUrl = null;
-                if (jarFileName.endsWith(".jar")) {
-                    // it isn't enough to specify the jar, we have to specify the path within the jar
-                    log4jConfigXmlUrl = new URL(new StringBuilder("jar:file:")
-                            .append(jarFileName)
-                            .append("!")
-                            .append(AGENT_JAR_LOG4J_CONFIG_FILE)
-                            .toString());
-                } else {
-                    // we likely have received a path to a set of class files (this happens when running tests)
-                    try {
-                        // guava's Resources class is usually smart enough to figure out where to find the log4j2.xml file
-                        log4jConfigXmlUrl = Resources.getResource(this.getClass(), AGENT_JAR_LOG4J_CONFIG_FILE);
-                    } catch (IllegalArgumentException iae) {
-                        // fallback on path
-                        log4jConfigXmlUrl = new File(jarFileName).toURI().toURL();
-                    }
-                }
-
-                System.setProperty(CONFIG_FILE_PROP, log4jConfigXmlUrl.toString());
-                System.setProperty(LEGACY_CONFIG_FILE_PROP, log4jConfigXmlUrl.toString());
                 // Log4j won't be able to find log4j-provider.properties because it isn't on the classpath (it's in our agent) so this sets it manually
                 System.setProperty(CONTEXT_FACTORY_PROP, "org.apache.logging.log4j.core.impl.Log4jContextFactory");
 
@@ -111,6 +96,16 @@ private Log4jLogger initializeRootLogger(String name) {
                 System.setProperty(LEGACY_CLASSLOADER_PROP, "true");
                 System.setProperty(CLASSLOADER_PROP, "true");
 
+                // This is in place in case we need to revert the programmatic initialization of our logger to
+                // use the log4j XML config file. This has to use a system property or env var because the config
+                // service is not properly bootstrapped yet.
+                if (Boolean.getBoolean("newrelic.config.agent_root_logger_init_with_file") ||
+                        Boolean.parseBoolean(System.getenv("NEW_RELIC_AGENT_ROOT_LOGGER_INIT_WITH_FILE"))) {
+                    initLog4jViaFile(jarFileName);
+                } else {
+                    initLog4jViaConfigurationBuilder();
+                }
+
                 try {
                     logger = createRootLogger(name);
                 } finally {
@@ -275,6 +270,39 @@ private boolean canWriteLogFile(String logFileName) {
         }
     }
 
+    private void initLog4jViaConfigurationBuilder() {
+        ConfigurationBuilder<BuiltConfiguration> builder = ConfigurationBuilderFactory.newConfigurationBuilder();
+        builder.setConfigurationName("Default");
+        builder.add(builder.newRootLogger(Level.INFO));
+        builder.setShutdownHook("disable");
+        builder.add(builder.newLogger("com.newrelic.agent.deps.org.reflections", Level.OFF));
+        Configurator.initialize(builder.build());
+    }
+
+    private void initLog4jViaFile(String jarFileName) throws MalformedURLException {
+        URL log4jConfigXmlUrl = null;
+        if (jarFileName.endsWith(".jar")) {
+            // it isn't enough to specify the jar, we have to specify the path within the jar
+            log4jConfigXmlUrl = new URL(new StringBuilder("jar:file:")
+                    .append(jarFileName)
+                    .append("!")
+                    .append(AGENT_JAR_LOG4J_CONFIG_FILE)
+                    .toString());
+        } else {
+            // we likely have received a path to a set of class files (this happens when running tests)
+            try {
+                // guava's Resources class is usually smart enough to figure out where to find the log4j2.xml file
+                log4jConfigXmlUrl = Resources.getResource(this.getClass(), AGENT_JAR_LOG4J_CONFIG_FILE);
+            } catch (IllegalArgumentException iae) {
+                // fallback on path
+                log4jConfigXmlUrl = new File(jarFileName).toURI().toURL();
+            }
+        }
+
+        System.setProperty(CONFIG_FILE_PROP, log4jConfigXmlUrl.toString());
+        System.setProperty(LEGACY_CONFIG_FILE_PROP, log4jConfigXmlUrl.toString());
+    }
+
     @Override
     public void addConsoleHandler() {
         rootLogger.addConsoleAppender();
diff --git a/newrelic-agent/src/main/java/com/newrelic/agent/logging/Log4jLogger.java b/newrelic-agent/src/main/java/com/newrelic/agent/logging/Log4jLogger.java
@@ -20,6 +20,9 @@
 import org.apache.logging.log4j.core.layout.PatternLayout;
 import org.apache.logging.log4j.message.Message;
 import org.apache.logging.log4j.message.ParameterizedMessage;
+import org.crac.Context;
+import org.crac.Core;
+import org.crac.Resource;
 
 import java.security.AccessController;
 import java.security.PrivilegedAction;
@@ -30,7 +33,7 @@
 
 import static com.newrelic.agent.logging.FileAppenderFactory.FILE_APPENDER_NAME;
 
-class Log4jLogger implements IAgentLogger {
+class Log4jLogger implements IAgentLogger, Resource {
 
     /**
      * The name of the console appender.
@@ -48,6 +51,12 @@ class Log4jLogger implements IAgentLogger {
     private final Logger logger;
     private final Map<String, IAgentLogger> childLoggers = new ConcurrentHashMap<>();
 
+    private String fileName;
+    private long logLimitBytes;
+    private int fileCount;
+    private boolean isDaily;
+    private String path;
+
     /**
      * Creates this Log4jLogger.
      *
@@ -57,6 +66,7 @@ private Log4jLogger(final String name, boolean isAgentRoot) {
         logger = LogManager.getLogger(name);
 
         if (isAgentRoot) {
+            Core.getGlobalContext().register(this);
             LoggerContext ctx = (LoggerContext) LogManager.getContext(false);
             Configuration config = ctx.getConfiguration();
             LoggerConfig loggerConfig = config.getLoggerConfig(LogManager.ROOT_LOGGER_NAME);
@@ -251,6 +261,41 @@ public void addConsoleAppender() {
         ctx.updateLoggers();
     }
 
+    /**
+     * Stops the file appender for the taking of a checkpoint
+     */
+    public void stopFileAppender() {
+        LoggerContext ctx = (LoggerContext) LogManager.getContext(false);
+        Configuration config = ctx.getConfiguration();
+        LoggerConfig loggerConfig = config.getLoggerConfig(logger.getName());
+
+        Appender checkpointFileAppender = loggerConfig.getAppenders().get(FILE_APPENDER_NAME);
+        if (checkpointFileAppender != null) {
+            // remove it from the list, so we don't try to write to it any longer
+            // this could cause missed messages in the log file, but would still
+            // go to the console
+            loggerConfig.removeAppender(checkpointFileAppender.getName());
+            // stop to close the open file
+            checkpointFileAppender.stop();
+            ctx.updateLoggers();
+        }
+    }
+
+    /**
+     * Starts the file appender after a restore
+     */
+    public void startFileAppender() {
+        LoggerContext ctx = (LoggerContext) LogManager.getContext(false);
+        Configuration config = ctx.getConfiguration();
+        LoggerConfig loggerConfig = config.getLoggerConfig(logger.getName());
+
+        Appender checkpointFileAppender = loggerConfig.getAppenders().get(FILE_APPENDER_NAME);
+        if (checkpointFileAppender == null && fileName != null) { // don't add it if it's already  there
+            addFileAppender(fileName, logLimitBytes, fileCount, isDaily, path);
+            ctx.updateLoggers();
+        }
+    }
+
     /**
      * Adds a file appender.
      *
@@ -259,6 +304,12 @@ public void addConsoleAppender() {
      * @param fileCount The number of files.
      */
     public void addFileAppender(String fileName, long logLimitBytes, int fileCount, boolean isDaily, String path) {
+        this.fileName = fileName;
+        this.logLimitBytes = logLimitBytes;
+        this.fileCount = fileCount;
+        this.isDaily = isDaily;
+        this.path = path;
+
         LoggerContext ctx = (LoggerContext) LogManager.getContext(false);
         Configuration config = ctx.getConfiguration();
         LoggerConfig loggerConfig = config.getLoggerConfig(logger.getName());
@@ -479,4 +530,16 @@ public void logToChild(String childName, Level level, String pattern, Object par
         }
     }
 
+    @Override
+    public void beforeCheckpoint(Context<? extends Resource> context) throws Exception {
+        Agent.LOG.info("Stopping Log4jLogger for CRaC checkpoint, log messages may be missing from the log file between here and restore, but should still appear in the console");
+        stopFileAppender();
+    }
+
+    @Override
+    public void afterRestore(Context<? extends Resource> context) throws Exception {
+        Agent.LOG.info("Restarting Log4jLogger for CRaC restore");
+        startFileAppender();
+    }
+
 }
diff --git a/newrelic-agent/src/main/java/com/newrelic/agent/transport/apache/ApacheHttpClientWrapper.java b/newrelic-agent/src/main/java/com/newrelic/agent/transport/apache/ApacheHttpClientWrapper.java
@@ -14,6 +14,7 @@
 import com.newrelic.agent.transport.HostConnectException;
 import com.newrelic.agent.transport.HttpClientWrapper;
 import com.newrelic.agent.transport.ReadResult;
+import com.newrelic.api.agent.NewRelic;
 import org.apache.http.Header;
 import org.apache.http.HttpEntity;
 import org.apache.http.HttpResponse;
@@ -37,6 +38,9 @@
 import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
 import org.apache.http.message.BasicHeader;
 import org.apache.http.protocol.HttpContext;
+import org.crac.Context;
+import org.crac.Core;
+import org.crac.Resource;
 
 import javax.net.ssl.SSLContext;
 import java.io.BufferedReader;
@@ -54,15 +58,22 @@
 
 import static com.newrelic.agent.transport.DataSenderImpl.GZIP_ENCODING;
 
-public class ApacheHttpClientWrapper implements HttpClientWrapper {
+public class ApacheHttpClientWrapper implements HttpClientWrapper, Resource {
     private final ApacheProxyManager proxyManager;
-    private final PoolingHttpClientConnectionManager connectionManager;
-    private final CloseableHttpClient httpClient;
+    private PoolingHttpClientConnectionManager connectionManager;
+    private CloseableHttpClient httpClient;
+    private SSLContext sslContext;
+    private final int defaultTimeoutInMillis;
 
     public ApacheHttpClientWrapper(ApacheProxyManager proxyManager, SSLContext sslContext, int defaultTimeoutInMillis) {
         this.proxyManager = proxyManager;
         this.connectionManager = createHttpClientConnectionManager(sslContext);
         this.httpClient = createHttpClient(defaultTimeoutInMillis);
+
+        this.sslContext = sslContext;
+        this.defaultTimeoutInMillis = defaultTimeoutInMillis;
+
+        Core.getGlobalContext().register(this);
     }
 
     private static final String USER_AGENT_HEADER_VALUE = initUserHeaderValue();
@@ -259,4 +270,21 @@ private BufferedReader getBufferedReader(HttpResponse response, InputStream is)
         }
         return new BufferedReader(new InputStreamReader(is, StandardCharsets.UTF_8));
     }
+
+    @Override
+    public void beforeCheckpoint(Context<? extends Resource> context) throws Exception {
+        Agent.LOG.info("Stopping collector connection for CRaC checkpoint");
+        NewRelic.getAgent().getMetricAggregator().incrementCounter(MetricNames.SUPPORTABILITY_AGENT_CRAC_CHECKPOINT);
+        connectionManager.close();
+        httpClient.close();
+    }
+
+    @Override
+    public void afterRestore(Context<? extends Resource> context) throws Exception {
+        Agent.LOG.info("Restarting collector connection for CRaC restore");
+        NewRelic.getAgent().getMetricAggregator().incrementCounter(MetricNames.SUPPORTABILITY_AGENT_CRAC_RESTORE);
+        connectionManager = createHttpClientConnectionManager(sslContext);
+        httpClient = createHttpClient(defaultTimeoutInMillis);
+    }
+
 }

Original file line number	Diff line number	Diff line change
`@@ -101,6 +101,8 @@ dependencies {`
`101`	`101`
`102`	`102`	`shadowIntoJar("org.yaml:snakeyaml:1.33")`
`103`	`103`
	`104`	`+ shadowIntoJar("org.crac:crac:1.5.0")`
	`105`	`+`
`104`	`106`	`implementation("javax.management.j2ee:management-api:1.1-rev-1") {`
`105`	`107`	`transitive = false`
`106`	`108`	`}`