src/java_tools/junitrunner/java/com/google/testing/junit/runner/PersistentTestRunner.java - bazel - Git at Google

 // Copyright 2019 The Bazel Authors. All Rights Reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //    http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 package com.google.testing.junit.runner;

 import static java.nio.charset.StandardCharsets.UTF_8;

 import com.google.common.hash.HashCode;
 import com.google.common.hash.Hasher;
 import com.google.common.hash.Hashing;
 import com.google.common.hash.HashingInputStream;
 import com.google.common.io.ByteStreams;
 import com.google.common.io.Files;
 import com.google.devtools.build.lib.worker.WorkerProtocol.WorkRequest;
 import com.google.devtools.build.lib.worker.WorkerProtocol.WorkResponse;
 import java.io.ByteArrayOutputStream;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.PrintStream;
 import java.net.URL;
 import java.util.List;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;

 /** A utility class for running Java tests using persistent workers. */
 final class PersistentTestRunner {

   private PersistentTestRunner() {
     // utility class; should not be instantiated
   }

   /**
    * Runs new tests in the same process. Communicates with bazel using the worker's protocol. Reads
    * a {@link WorkRequest} sent by bazel and sends back a {@link WorkResponse}.
    *
    * <p>Uses three different classloaders:
    *
    * <ul>
    *   <li>A classloader for direct dependencies: loads the test's classes and the first two layers
    *       of its dependencies.
    *   <li>A classloader for transitive dependencies: loads the remaining classes from the
    *       transitive dependencies.
    *   <li>The system classloader: initialized at JVM startup time; loads the test runner's classes.
    * </ul>
    *
    * <p>The classloaders have a child-parent relationship: direct CL -> transitive CL -> system CL.
    *
    * <p>The direct dependencies classloader is the current thread's classloader.
    *
    * <p>The transitive dependencies classloader checks if a class was already loaded by the direct
    * dependencies classloader if it did not succeed in loading the class itself. This is required
    * for classes loaded by the transitive classloader that reference classes in the direct
    * dependencies classloader.
    *
    * <p>The default class loading logic in {@link ClassLoader} applies in all other cases.
    *
    * <p>The direct and transitive classloaders are rebuilt before every test run only if the
    * combined hash of the jars to be loaded has changed.
    */
   static int runPersistentTestRunner(
       String suitClassName, String workspacePrefix, SuiteTestRunner suiteTestRunner) {
     PrintStream originalStdOut = System.out;
     PrintStream originalStdErr = System.err;

     // TODO(elenairina): Remove this variable after cl/282553936 is released.
     String legacyTestRuntimeClasspathFile = System.getenv("TEST_RUNTIME_CLASSPATH_FILE");

     String directClasspathFile = System.getenv("TEST_DIRECT_CLASSPATH_FILE");
     String transitiveClasspathFile = System.getenv("TEST_TRANSITIVE_CLASSPATH_FILE");
     String absolutePathPrefix = System.getenv("JAVA_RUNFILES") + File.separator + workspacePrefix;

     PersistentTestRunnerClassLoader transitiveDepsClassLoader = null;
     PersistentTestRunnerClassLoader directDepsClassLoader = null;

     // Reading the work requests and solving them in sequence is not a problem because Bazel creates
     // up to --worker_max_instances (defaults to 4) instances per worker key.
     while (true) {
       try {
         WorkRequest request = WorkRequest.parseDelimitedFrom(System.in);

         if (request == null) {
           // null is only returned when the stream reaches EOF
           break;
         }

         if (legacyTestRuntimeClasspathFile != null) {
           // Re-use the same classloader variable in the legacy case for simplicity.
           directDepsClassLoader =
               maybeRecreateClassLoader(
                   getFilesWithAbsolutePathPrefixFromFile(
                       legacyTestRuntimeClasspathFile, absolutePathPrefix),
                   ClassLoader.getSystemClassLoader(),
                   null);
         } else {
           transitiveDepsClassLoader =
               maybeRecreateClassLoader(
                   getFilesWithAbsolutePathPrefixFromFile(
                       transitiveClasspathFile, absolutePathPrefix),
                   ClassLoader.getSystemClassLoader(),
                   transitiveDepsClassLoader);

           directDepsClassLoader =
               maybeRecreateClassLoader(
                   getFilesWithAbsolutePathPrefixFromFile(directClasspathFile, absolutePathPrefix),
                   transitiveDepsClassLoader,
                   directDepsClassLoader);
           transitiveDepsClassLoader.setChild(directDepsClassLoader);
         }

         Thread.currentThread().setContextClassLoader(directDepsClassLoader);

         ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
         PrintStream printStream = new PrintStream(outputStream, true);
         System.setOut(printStream);
         System.setErr(printStream);
         String[] arguments = request.getArgumentsList().toArray(new String[0]);
         int exitCode = -1;
         try {
           exitCode =
               suiteTestRunner.runTestsInSuite(
                   suitClassName, arguments, directDepsClassLoader, /* resolve= */ true);
         } finally {
           System.setOut(originalStdOut);
           System.setErr(originalStdErr);
         }

         WorkResponse response =
             WorkResponse.newBuilder()
                 .setOutput(outputStream.toString())
                 .setExitCode(exitCode)
                 .build();

         response.writeDelimitedTo(System.out);
         System.out.flush();
       } catch (IOException e) {
         e.printStackTrace();
         return 1;
       }
     }
     return 0;
   }

   /**
    * Returns a {@link Class} with the given name. Loads the class for the given classloader, or from
    * the system classloader if none is specified.
    */
   static Class<?> getTestClass(String name, ClassLoader classLoader, boolean resolve) {
     if (name == null) {
       return null;
     }

     try {
       if (classLoader == null) {
         return Class.forName(name);
       }
       return Class.forName(name, resolve, classLoader);
     } catch (ClassNotFoundException e) {
       return null;
     }
   }

   /**
    * Returns a classloader that loads the given jars, only if the combined hash of their content is
    * different than the one for the given previous classloader.
    *
    * <p>Returns previousClassLoader if the hashes are the same.
    *
    * <p>Needs to be called before every test run to avoid having stale classes. A class already
    * loaded in a classloader can not be unloaded. To overcome this a new classloader has to be
    * created at every test run.
    */
   private static PersistentTestRunnerClassLoader maybeRecreateClassLoader(
       List<File> runtimeJars,
       ClassLoader parent,
       @Nullable PersistentTestRunnerClassLoader previousClassLoader)
       throws IOException {
     HashCode combinedHash = getCombinedHashForFiles(runtimeJars);
     if (previousClassLoader != null && combinedHash.equals(previousClassLoader.getChecksum())) {
       return previousClassLoader;
     }

     return new PersistentTestRunnerClassLoader(
         convertFileListToURLArray(runtimeJars), parent, combinedHash);
   }

   private static HashCode getCombinedHashForFiles(List<File> files) throws IOException {
     Hasher hasher = Hashing.sha256().newHasher();
     for (File file : files) {
       hasher.putBytes(getFileHash(file).asBytes());
     }
     return hasher.hash();
   }

   private static HashCode getFileHash(File file) throws IOException {
     InputStream inputStream = new FileInputStream(file);
     HashingInputStream hashingStream = new HashingInputStream(Hashing.sha256(), inputStream);
     ByteStreams.copy(hashingStream, ByteStreams.nullOutputStream());
     return hashingStream.hash();
   }

   private static URL[] convertFileListToURLArray(List<File> jars) throws IOException {
     URL[] urls = new URL[jars.size()];
     for (int i = 0; i < jars.size(); i++) {
       urls[i] = jars.get(i).toURI().toURL();
     }
     return urls;
   }

   private static List<File> getFilesWithAbsolutePathPrefixFromFile(
       String runtimeClasspathFilename, String absolutePathPrefix) throws IOException {
     return Files.readLines(new File(runtimeClasspathFilename), UTF_8).stream()
         .map(entry -> new File(absolutePathPrefix + entry))
         .collect(Collectors.toList());
   }

   static boolean isPersistentTestRunner() {
     return Boolean.parseBoolean(System.getenv("PERSISTENT_TEST_RUNNER"));
   }
 }
	// Copyright 2019 The Bazel Authors. All Rights Reserved.
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	package com.google.testing.junit.runner;

	import static java.nio.charset.StandardCharsets.UTF_8;

	import com.google.common.hash.HashCode;
	import com.google.common.hash.Hasher;
	import com.google.common.hash.Hashing;
	import com.google.common.hash.HashingInputStream;
	import com.google.common.io.ByteStreams;
	import com.google.common.io.Files;
	import com.google.devtools.build.lib.worker.WorkerProtocol.WorkRequest;
	import com.google.devtools.build.lib.worker.WorkerProtocol.WorkResponse;
	import java.io.ByteArrayOutputStream;
	import java.io.File;
	import java.io.FileInputStream;
	import java.io.IOException;
	import java.io.InputStream;
	import java.io.PrintStream;
	import java.net.URL;
	import java.util.List;
	import java.util.stream.Collectors;
	import javax.annotation.Nullable;

	/** A utility class for running Java tests using persistent workers. */
	final class PersistentTestRunner {

	private PersistentTestRunner() {
	// utility class; should not be instantiated
	}

	/**
	* Runs new tests in the same process. Communicates with bazel using the worker's protocol. Reads
	* a {@link WorkRequest} sent by bazel and sends back a {@link WorkResponse}.
	*
	* <p>Uses three different classloaders:
	*
	* <ul>
	* <li>A classloader for direct dependencies: loads the test's classes and the first two layers
	* of its dependencies.
	* <li>A classloader for transitive dependencies: loads the remaining classes from the
	* transitive dependencies.
	* <li>The system classloader: initialized at JVM startup time; loads the test runner's classes.
	* </ul>
	*
	* <p>The classloaders have a child-parent relationship: direct CL -> transitive CL -> system CL.
	*
	* <p>The direct dependencies classloader is the current thread's classloader.
	*
	* <p>The transitive dependencies classloader checks if a class was already loaded by the direct
	* dependencies classloader if it did not succeed in loading the class itself. This is required
	* for classes loaded by the transitive classloader that reference classes in the direct
	* dependencies classloader.
	*
	* <p>The default class loading logic in {@link ClassLoader} applies in all other cases.
	*
	* <p>The direct and transitive classloaders are rebuilt before every test run only if the
	* combined hash of the jars to be loaded has changed.
	*/
	static int runPersistentTestRunner(
	String suitClassName, String workspacePrefix, SuiteTestRunner suiteTestRunner) {
	PrintStream originalStdOut = System.out;
	PrintStream originalStdErr = System.err;

	// TODO(elenairina): Remove this variable after cl/282553936 is released.
	String legacyTestRuntimeClasspathFile = System.getenv("TEST_RUNTIME_CLASSPATH_FILE");

	String directClasspathFile = System.getenv("TEST_DIRECT_CLASSPATH_FILE");
	String transitiveClasspathFile = System.getenv("TEST_TRANSITIVE_CLASSPATH_FILE");
	String absolutePathPrefix = System.getenv("JAVA_RUNFILES") + File.separator + workspacePrefix;

	PersistentTestRunnerClassLoader transitiveDepsClassLoader = null;
	PersistentTestRunnerClassLoader directDepsClassLoader = null;

	// Reading the work requests and solving them in sequence is not a problem because Bazel creates
	// up to --worker_max_instances (defaults to 4) instances per worker key.
	while (true) {
	try {
	WorkRequest request = WorkRequest.parseDelimitedFrom(System.in);

	if (request == null) {
	// null is only returned when the stream reaches EOF
	break;
	}

	if (legacyTestRuntimeClasspathFile != null) {
	// Re-use the same classloader variable in the legacy case for simplicity.
	directDepsClassLoader =
	maybeRecreateClassLoader(
	getFilesWithAbsolutePathPrefixFromFile(
	legacyTestRuntimeClasspathFile, absolutePathPrefix),
	ClassLoader.getSystemClassLoader(),
	null);
	} else {
	transitiveDepsClassLoader =
	maybeRecreateClassLoader(
	getFilesWithAbsolutePathPrefixFromFile(
	transitiveClasspathFile, absolutePathPrefix),
	ClassLoader.getSystemClassLoader(),
	transitiveDepsClassLoader);

	directDepsClassLoader =
	maybeRecreateClassLoader(
	getFilesWithAbsolutePathPrefixFromFile(directClasspathFile, absolutePathPrefix),
	transitiveDepsClassLoader,
	directDepsClassLoader);
	transitiveDepsClassLoader.setChild(directDepsClassLoader);
	}

	Thread.currentThread().setContextClassLoader(directDepsClassLoader);

	ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
	PrintStream printStream = new PrintStream(outputStream, true);
	System.setOut(printStream);
	System.setErr(printStream);
	String[] arguments = request.getArgumentsList().toArray(new String[0]);
	int exitCode = -1;
	try {
	exitCode =
	suiteTestRunner.runTestsInSuite(
	suitClassName, arguments, directDepsClassLoader, /* resolve= */ true);
	} finally {
	System.setOut(originalStdOut);
	System.setErr(originalStdErr);
	}

	WorkResponse response =
	WorkResponse.newBuilder()
	.setOutput(outputStream.toString())
	.setExitCode(exitCode)
	.build();

	response.writeDelimitedTo(System.out);
	System.out.flush();
	} catch (IOException e) {
	e.printStackTrace();
	return 1;
	}
	}
	return 0;
	}

	/**
	* Returns a {@link Class} with the given name. Loads the class for the given classloader, or from
	* the system classloader if none is specified.
	*/
	static Class<?> getTestClass(String name, ClassLoader classLoader, boolean resolve) {
	if (name == null) {
	return null;
	}

	try {
	if (classLoader == null) {
	return Class.forName(name);
	}
	return Class.forName(name, resolve, classLoader);
	} catch (ClassNotFoundException e) {
	return null;
	}
	}

	/**
	* Returns a classloader that loads the given jars, only if the combined hash of their content is
	* different than the one for the given previous classloader.
	*
	* <p>Returns previousClassLoader if the hashes are the same.
	*
	* <p>Needs to be called before every test run to avoid having stale classes. A class already
	* loaded in a classloader can not be unloaded. To overcome this a new classloader has to be
	* created at every test run.
	*/
	private static PersistentTestRunnerClassLoader maybeRecreateClassLoader(
	List<File> runtimeJars,
	ClassLoader parent,
	@Nullable PersistentTestRunnerClassLoader previousClassLoader)
	throws IOException {
	HashCode combinedHash = getCombinedHashForFiles(runtimeJars);
	if (previousClassLoader != null && combinedHash.equals(previousClassLoader.getChecksum())) {
	return previousClassLoader;
	}

	return new PersistentTestRunnerClassLoader(
	convertFileListToURLArray(runtimeJars), parent, combinedHash);
	}

	private static HashCode getCombinedHashForFiles(List<File> files) throws IOException {
	Hasher hasher = Hashing.sha256().newHasher();
	for (File file : files) {
	hasher.putBytes(getFileHash(file).asBytes());
	}
	return hasher.hash();
	}

	private static HashCode getFileHash(File file) throws IOException {
	InputStream inputStream = new FileInputStream(file);
	HashingInputStream hashingStream = new HashingInputStream(Hashing.sha256(), inputStream);
	ByteStreams.copy(hashingStream, ByteStreams.nullOutputStream());
	return hashingStream.hash();
	}

	private static URL[] convertFileListToURLArray(List<File> jars) throws IOException {
	URL[] urls = new URL[jars.size()];
	for (int i = 0; i < jars.size(); i++) {
	urls[i] = jars.get(i).toURI().toURL();
	}
	return urls;
	}

	private static List<File> getFilesWithAbsolutePathPrefixFromFile(
	String runtimeClasspathFilename, String absolutePathPrefix) throws IOException {
	return Files.readLines(new File(runtimeClasspathFilename), UTF_8).stream()
	.map(entry -> new File(absolutePathPrefix + entry))
	.collect(Collectors.toList());
	}

	static boolean isPersistentTestRunner() {
	return Boolean.parseBoolean(System.getenv("PERSISTENT_TEST_RUNNER"));
	}
	}