Add MyJavaFileManager override inside CachedCompiler

Author: yevgenp

src/main/java/net/openhft/compiler/CachedCompiler.java

@@ -34,6 +34,7 @@
 import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.function.Function;
 
 import static net.openhft.compiler.CompilerUtils.*;
 
@@ -45,6 +46,7 @@ public class CachedCompiler implements Closeable {
 
     private final Map<ClassLoader, Map<String, Class<?>>> loadedClassesMap = Collections.synchronizedMap(new WeakHashMap<>());
     private final Map<ClassLoader, MyJavaFileManager> fileManagerMap = Collections.synchronizedMap(new WeakHashMap<>());
+    public Function<StandardJavaFileManager, MyJavaFileManager> fileManagerOverride;
 
     @Nullable
     private final File sourceDir;
@@ -125,8 +127,7 @@ public void report(Diagnostic<? extends JavaFileObject> diagnostic) {
 
             // nothing to return due to compiler error
             return Collections.emptyMap();
-        }
-        else {
+        } else {
             Map<String, byte[]> result = fileManager.getAllBuffers();
 
             return result;
@@ -153,7 +154,8 @@ public Class loadFromJava(@NotNull ClassLoader classLoader,
         MyJavaFileManager fileManager = fileManagerMap.get(classLoader);
         if (fileManager == null) {
             StandardJavaFileManager standardJavaFileManager = s_compiler.getStandardFileManager(null, null, null);
-            fileManagerMap.put(classLoader, fileManager = new MyJavaFileManager(standardJavaFileManager));
+            fileManager = getFileManager(standardJavaFileManager);
+            fileManagerMap.put(classLoader, fileManager);
         }
         final Map<String, byte[]> compiled = compileFromJava(className, javaCode, printWriter, fileManager);
         for (Map.Entry<String, byte[]> entry : compiled.entrySet()) {
@@ -188,4 +190,10 @@ public Class loadFromJava(@NotNull ClassLoader classLoader,
         }
         return clazz;
     }
+
+    private @NotNull MyJavaFileManager getFileManager(StandardJavaFileManager fm) {
+        return fileManagerOverride != null
+                ? fileManagerOverride.apply(fm)
+                : new MyJavaFileManager(fm);
+    }
 }

src/main/java/net/openhft/compiler/MyJavaFileManager.java

@@ -39,7 +39,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
-class MyJavaFileManager implements JavaFileManager {
+public class MyJavaFileManager implements JavaFileManager {
     private static final Logger LOG = LoggerFactory.getLogger(MyJavaFileManager.class);
     private final static Unsafe unsafe;
     private static final long OVERRIDE_OFFSET;
@@ -67,7 +67,7 @@ class MyJavaFileManager implements JavaFileManager {
     // synchronizing due to ConcurrentModificationException
     private final Map<String, CloseableByteArrayOutputStream> buffers = Collections.synchronizedMap(new LinkedHashMap<>());
 
-    MyJavaFileManager(StandardJavaFileManager fileManager) {
+    public MyJavaFileManager(StandardJavaFileManager fileManager) {
         this.fileManager = fileManager;
     }