Close and reopen security agent logs

newrelic-security-agent/build.gradle

@@ -78,6 +78,7 @@ dependencies {
     shadowIntoJar 'com.github.oshi:oshi-core:6.4.1'
     shadowIntoJar 'com.google.code.gson:gson:2.14.0'
     shadowIntoJar 'org.apache.httpcomponents:httpclient:4.5.14'
+    shadowIntoJar 'org.crac:crac:1.5.0'
     implementation "com.newrelic.agent.java:newrelic-api:${nrAPIVersion}"
 }
 

newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/filelogging/CracResource.java

@@ -0,0 +1,54 @@
+package com.newrelic.agent.security.intcodeagent.filelogging;
+
+import org.crac.CheckpointException;
+import org.crac.Context;
+import org.crac.Resource;
+import org.crac.RestoreException;
+
+import java.io.Closeable;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.Phaser;
+
+class CracResource implements Resource, Runnable {
+    private final Phaser phaser = new Phaser(2);
+    private final List<Closeable> toClose = new ArrayList<>();
+    private final List<Runnable> toReopen = new ArrayList<>();
+
+    public synchronized void addAction(Closeable close, Runnable reopen) {
+        toClose.add(close);
+        toReopen.add(reopen);
+    }
+
+    public void beforeCheckpoint(Context<? extends Resource> ctx) throws CheckpointException {
+        FileLoggerThreadPool threadPool = FileLoggerThreadPool.getInstance();
+        assert !threadPool.isLoggingToStdOut;
+        threadPool.getExecutor().submit(this);
+        if (phaser.arriveAndAwaitAdvance() < 0) {
+            throw new CheckpointException("Failed to close the writer");
+        }
+    }
+
+    public void afterRestore(Context<? extends Resource> ctx) throws RestoreException {
+        if (phaser.arriveAndAwaitAdvance() < 0) {
+            throw new RestoreException("Cannot restore writer");
+        }
+    }
+
+    @Override
+    public synchronized void run() {
+        try {
+            for (Closeable c : toClose) {
+                c.close();
+            }
+            phaser.arriveAndAwaitAdvance();
+            // checkpoint/restore happens here
+            phaser.arriveAndAwaitAdvance();
+            for (Runnable r : toReopen) {
+                r.run();
+            }
+        } catch (Exception e) {
+            phaser.forceTermination();
+        }
+    }
+}

newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/filelogging/FileLoggerThreadPool.java

@@ -10,6 +10,7 @@
 import com.newrelic.agent.security.intcodeagent.websocket.JsonConverter;
 import com.newrelic.api.agent.security.utils.logging.LogLevel;
 import org.apache.commons.lang3.StringUtils;
+import org.crac.Core;
 
 import java.io.IOException;
 import java.util.concurrent.*;
@@ -31,6 +32,8 @@ public class FileLoggerThreadPool {
 
     private static OSVariables osVariables;
 
+    private static CracResource cracResource;
+
     private FileLoggerThreadPool() throws IOException {
         maxfiles = LogFileHelper.logFileCount();
         maxfilesize = LogFileHelper.logFileLimit()* 1024L;
@@ -87,8 +90,19 @@ public Thread newThread(Runnable r) {
                 return t;
             }
         });
+
+        if (!isLoggingToStdOut) {
+            // We create the resource centrally here because both LogWriter and InitLogWriter
+            // are tightly coupled to this single-threaded executor
+            cracResource = new CracResource();
+            Core.getGlobalContext().register(cracResource);
+        }
     }
 
+    CracResource getCracResource() {
+        assert !isLoggingToStdOut;
+        return cracResource;
+    }
 
     public void shutDownThreadPoolExecutor() {
 

newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/filelogging/InitLogWriter.java

@@ -93,6 +93,8 @@ private static Boolean createLogFile() {
             writer.flush();
             maxFileSize = FileLoggerThreadPool.getInstance().maxfilesize;
 
+            FileLoggerThreadPool.getInstance().getCracResource().addAction(writer, InitLogWriter::createLogFile);
+
             // k2.log.handler.maxfilesize=10
             // k2.log.handler.maxfilesize.unit=MB
             if (!osVariables.getWindows()) {

newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/filelogging/LogWriter.java

@@ -77,6 +77,8 @@ private static boolean createLogFile() {
 
             maxFileSize = FileLoggerThreadPool.getInstance().maxfilesize;
 
+            FileLoggerThreadPool.getInstance().getCracResource().addAction(writer, LogWriter::createLogFile);
+
             if (!osVariables.getWindows()) {
                 Files.setPosixFilePermissions(currentLogFile.toPath(), PosixFilePermissions.fromString(IUtilConstants.FILE_PERMISSIONS));
             }