From 23c73eeec69643a50f10aa9533de825bc4827197 Mon Sep 17 00:00:00 2001
From: Eugene Zhuravlev <jeka@jetbrains.com>
Date: Mon, 2 Aug 2021 20:28:41 +0200
Subject: [PATCH] ensure thread context classloader for annotation processor is
 the same one that was used to load the processor (IDEA-274284)

GitOrigin-RevId: f2be3d7feaf3395d429d90d98e9f6fc391bb8f77
---
 .../org/jetbrains/jps/javac/APIWrappers.java  | 36 +++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/jps/jps-builders-6/src/org/jetbrains/jps/javac/APIWrappers.java b/jps/jps-builders-6/src/org/jetbrains/jps/javac/APIWrappers.java
index 8af4042c1940..afd2e72528d7 100644
--- a/jps/jps-builders-6/src/org/jetbrains/jps/javac/APIWrappers.java
+++ b/jps/jps-builders-6/src/org/jetbrains/jps/javac/APIWrappers.java
@@ -2,6 +2,7 @@
 package org.jetbrains.jps.javac;
 
 import com.intellij.openapi.util.Pair;
+import com.intellij.openapi.util.Ref;
 import com.intellij.util.Function;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -132,6 +133,7 @@ public class APIWrappers {
 
     public void init(ProcessingEnvironment processingEnv) {
       myProcessingEnv = processingEnv;
+      final Ref<ClassLoader> oldCtxLoader = setupContextClassLoader();
       try {
         getWrapperDelegate().init(wrap(ProcessingEnvironment.class, new ProcessingEnvironmentWrapper(processingEnv, myFileManager)));
       }
@@ -139,9 +141,13 @@ public class APIWrappers {
         sendDiagnosticWarning(processingEnv, e);
         throw e;
       }
+      finally {
+        restoreContextClassLoader(oldCtxLoader);
+      }
     }
 
     public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
+      final Ref<ClassLoader> oldCtxLoader = setupContextClassLoader();
       try {
         return getWrapperDelegate().process(annotations, roundEnv);
       }
@@ -149,6 +155,36 @@ public class APIWrappers {
         sendDiagnosticWarning(myProcessingEnv, e);
         throw e;
       }
+      finally {
+        restoreContextClassLoader(oldCtxLoader);
+      }
+    }
+
+    /*
+      Some processors may use libraries/frameworks requiring sophisticated setup via thread context class loader.
+      This ensures that the context loader is the same as the one used to load processor itself:
+      it might help to avoid possible conflicts with JPS core classes
+     */
+    @Nullable
+    private Ref<ClassLoader> setupContextClassLoader() {
+      final Processor delegate = getWrapperDelegate();
+      if (delegate != null) {
+        final Thread currentThread = Thread.currentThread();
+        final ClassLoader processorLoader = delegate.getClass().getClassLoader();
+        final ClassLoader currentCtxLoader = currentThread.getContextClassLoader();
+        if (processorLoader != currentCtxLoader) {
+          currentThread.setContextClassLoader(processorLoader);
+          return Ref.create(currentCtxLoader);
+        }
+      }
+      return null;
+    }
+
+    @SuppressWarnings("MethodMayBeStatic")
+    private void restoreContextClassLoader(@Nullable Ref<ClassLoader> loaderRef) {
+      if (loaderRef != null) {
+        Thread.currentThread().setContextClassLoader(loaderRef.get());
+      }
     }
 
     private void sendDiagnosticWarning(ProcessingEnvironment processingEnv, Throwable e) {