exec/java-exec/src/main/java/org/apache/drill/exec/compile/ClassTransformer.java - drill - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you 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 org.apache.drill.exec.compile;

 import java.io.IOException;
 import java.util.LinkedList;
 import java.util.Map;
 import java.util.Set;

 import org.apache.commons.lang3.tuple.Pair;
 import org.apache.drill.common.config.DrillConfig;
 import org.apache.drill.common.util.DrillFileUtils;
 import org.apache.drill.common.util.DrillStringUtils;
 import org.apache.drill.exec.ExecConstants;
 import org.apache.drill.exec.compile.MergeAdapter.MergedClassResult;
 import org.apache.drill.exec.exception.ClassTransformationException;
 import org.apache.drill.exec.expr.CodeGenerator;
 import org.apache.drill.exec.server.options.OptionSet;
 import org.codehaus.commons.compiler.CompileException;
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.tree.ClassNode;

 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;

 /**
  * Compiles generated code, merges the resulting class with the
  * template class, and performs byte-code cleanup on the resulting
  * byte codes. The most important transform is scalar replacement
  * which replaces occurrences of non-escaping objects with a
  * collection of member variables.
  */

 public class ClassTransformer {
   private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(ClassTransformer.class);

   private static final int MAX_SCALAR_REPLACE_CODE_SIZE = 2*1024*1024; // 2meg

   private final ByteCodeLoader byteCodeLoader = new ByteCodeLoader();
   private final DrillConfig config;
   private final OptionSet optionManager;

   @VisibleForTesting // although we need it even if it weren't used in testing
   public enum ScalarReplacementOption {
     OFF, // scalar replacement will not ever be used
     TRY, // scalar replacement will be attempted, and if there is an error, we fall back to not using it
     ON; // scalar replacement will always be used, and any errors cause user visible errors

     /**
      * Convert a string to an enum value.
      *
      * @param s the string
      * @return an enum value
      * @throws IllegalArgumentException if the string doesn't match any of the enum values
      */
     public static ScalarReplacementOption fromString(final String s) {
       switch(s) {
       case "off":
         return OFF;
       case "try":
         return TRY;
       case "on":
         return ON;
       default:
         throw new IllegalArgumentException("Invalid ScalarReplacementOption \"" + s + "\"");
       }
     }
   }

   public ClassTransformer(final DrillConfig config, final OptionSet optionManager) {
     this.config = config;
     this.optionManager = optionManager;
   }

   public static class ClassSet {
     public final ClassSet parent;
     public final ClassNames precompiled;
     public final ClassNames generated;

     public ClassSet(ClassSet parent, String precompiled, String generated) {
       Preconditions.checkArgument(!generated.startsWith(precompiled),
           String.format("The new name of a class cannot start with the old name of a class, otherwise class renaming will cause problems. Precompiled class name %s. Generated class name %s",
               precompiled, generated));
       this.parent = parent;
       this.precompiled = new ClassNames(precompiled);
       this.generated = new ClassNames(generated);
     }

     public ClassSet getChild(String precompiled, String generated) {
       return new ClassSet(this, precompiled, generated);
     }

     public ClassSet getChild(String precompiled) {
       return new ClassSet(this, precompiled, precompiled.replace(this.precompiled.dot, this.generated.dot));
     }

     @Override
     public int hashCode() {
       final int prime = 31;
       int result = 1;
       result = prime * result + ((generated == null) ? 0 : generated.hashCode());
       result = prime * result + ((parent == null) ? 0 : parent.hashCode());
       result = prime * result + ((precompiled == null) ? 0 : precompiled.hashCode());
       return result;
     }

     @Override
     public boolean equals(Object obj) {
       if (this == obj) {
         return true;
       }
       if (obj == null) {
         return false;
       }
       if (getClass() != obj.getClass()) {
         return false;
       }
       ClassSet other = (ClassSet) obj;
       if (generated == null) {
         if (other.generated != null) {
           return false;
         }
       } else if (!generated.equals(other.generated)) {
         return false;
       }
       if (parent == null) {
         if (other.parent != null) {
           return false;
         }
       } else if (!parent.equals(other.parent)) {
         return false;
       }
       if (precompiled == null) {
         if (other.precompiled != null) {
           return false;
         }
       } else if (!precompiled.equals(other.precompiled)) {
         return false;
       }
       return true;
     }
   }

   public static class ClassNames {
     public final String dot;
     public final String slash;
     public final String clazz;

     public ClassNames(String className) {
       dot = className;
       slash = className.replace('.', DrillFileUtils.SEPARATOR_CHAR);
       clazz = DrillFileUtils.SEPARATOR_CHAR + slash + ".class";
     }

     @Override
     public int hashCode() {
       final int prime = 31;
       int result = 1;
       result = prime * result + ((clazz == null) ? 0 : clazz.hashCode());
       result = prime * result + ((dot == null) ? 0 : dot.hashCode());
       result = prime * result + ((slash == null) ? 0 : slash.hashCode());
       return result;
     }

     @Override
     public boolean equals(Object obj) {
       if (this == obj) {
         return true;
       }
       if (obj == null) {
         return false;
       }
       if (getClass() != obj.getClass()) {
         return false;
       }
       ClassNames other = (ClassNames) obj;
       if (clazz == null) {
         if (other.clazz != null) {
           return false;
         }
       } else if (!clazz.equals(other.clazz)) {
         return false;
       }
       if (dot == null) {
         if (other.dot != null) {
           return false;
         }
       } else if (!dot.equals(other.dot)) {
         return false;
       }
       if (slash == null) {
         if (other.slash != null) {
           return false;
         }
       } else if (!slash.equals(other.slash)) {
         return false;
       }
       return true;
     }
   }

   public Class<?> getImplementationClass(CodeGenerator<?> cg) throws ClassTransformationException {
     final QueryClassLoader loader = new QueryClassLoader(config, optionManager);
     return getImplementationClass(loader, cg.getDefinition(),
         cg.getGeneratedCode(), cg.getMaterializedClassName());
   }

   public Class<?> getImplementationClass(
       final QueryClassLoader classLoader,
       final TemplateClassDefinition<?> templateDefinition,
       final String entireClass,
       final String materializedClassName) throws ClassTransformationException {
     // unfortunately, this hasn't been set up at construction time, so we have to do it here
     final ScalarReplacementOption scalarReplacementOption = ScalarReplacementOption.fromString(optionManager.getOption(ExecConstants.SCALAR_REPLACEMENT_VALIDATOR));

     try {
       final long t1 = System.nanoTime();
       final ClassSet set = new ClassSet(null, templateDefinition.getTemplateClassName(), materializedClassName);
       final byte[][] implementationClasses = classLoader.getClassByteCode(set.generated, entireClass);

       long totalBytecodeSize = 0;
       Map<String, Pair<byte[], ClassNode>> classesToMerge = Maps.newHashMap();
       for (byte[] clazz : implementationClasses) {
         totalBytecodeSize += clazz.length;
         ClassNode node = AsmUtil.classFromBytes(clazz, ClassReader.SKIP_FRAMES);
         if (!AsmUtil.isClassOk(logger, "implementationClasses", node)) {
           throw new IllegalStateException("Problem found with implementationClasses");
         }
         classesToMerge.put(node.name, Pair.of(clazz, node));
       }

       final LinkedList<ClassSet> names = Lists.newLinkedList();
       final Set<ClassSet> namesCompleted = Sets.newHashSet();
       names.add(set);

       while ( !names.isEmpty() ) {
         final ClassSet nextSet = names.removeFirst();
         if (namesCompleted.contains(nextSet)) {
           continue;
         }
         final ClassNames nextPrecompiled = nextSet.precompiled;
         final byte[] precompiledBytes = byteCodeLoader.getClassByteCodeFromPath(nextPrecompiled.clazz);
         final ClassNames nextGenerated = nextSet.generated;
         // keeps only classes that have not be merged
         Pair<byte[], ClassNode> classNodePair = classesToMerge.remove(nextGenerated.slash);
         final ClassNode generatedNode;
         if (classNodePair != null) {
           generatedNode = classNodePair.getValue();
         } else {
           generatedNode = null;
         }

         /*
          * TODO
          * We're having a problem with some cases of scalar replacement, but we want to get
          * the code in so it doesn't rot anymore.
          *
          *  Here, we use the specified replacement option. The loop will allow us to retry if
          *  we're using TRY.
          */
         MergedClassResult result = null;
         boolean scalarReplace = scalarReplacementOption != ScalarReplacementOption.OFF && entireClass.length() < MAX_SCALAR_REPLACE_CODE_SIZE;
         while(true) {
           try {
             result = MergeAdapter.getMergedClass(nextSet, precompiledBytes, generatedNode, scalarReplace);
             break;
           } catch(RuntimeException e) {
             // if we had a problem without using scalar replacement, then rethrow
             if (!scalarReplace) {
               throw e;
             }

             // if we did try to use scalar replacement, decide if we need to retry or not
             if (scalarReplacementOption == ScalarReplacementOption.ON) {
               // option is forced on, so this is a hard error
               throw e;
             }

             /*
              * We tried to use scalar replacement, with the option to fall back to not using it.
              * Log this failure before trying again without scalar replacement.
              */
             logger.info("scalar replacement failure (retrying)\n", e);
             scalarReplace = false;
           }
         }

         for (String s : result.innerClasses) {
           s = s.replace(DrillFileUtils.SEPARATOR_CHAR, '.');
           names.add(nextSet.getChild(s));
         }
         classLoader.injectByteCode(nextGenerated.dot, result.bytes);
         namesCompleted.add(nextSet);
       }

       // adds byte code of the classes that have not been merged to make them accessible for outer class
       for (Map.Entry<String, Pair<byte[], ClassNode>> clazz : classesToMerge.entrySet()) {
         classLoader.injectByteCode(clazz.getKey().replace(DrillFileUtils.SEPARATOR_CHAR, '.'), clazz.getValue().getKey());
       }
       Class<?> c = classLoader.findClass(set.generated.dot);
       if (templateDefinition.getExternalInterface().isAssignableFrom(c)) {
         logger.debug("Compiled and merged {}: bytecode size = {}, time = {} ms.",
              c.getSimpleName(),
              DrillStringUtils.readable(totalBytecodeSize),
              (System.nanoTime() - t1 + 500_000) / 1_000_000);
         return c;
       }

       throw new ClassTransformationException("The requested class did not implement the expected interface.");
     } catch (CompileException | IOException | ClassNotFoundException e) {
       throw new ClassTransformationException(String.format("Failure generating transformation classes for value: \n %s", entireClass), e);
     }
   }
 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you 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 org.apache.drill.exec.compile;

	import java.io.IOException;
	import java.util.LinkedList;
	import java.util.Map;
	import java.util.Set;

	import org.apache.commons.lang3.tuple.Pair;
	import org.apache.drill.common.config.DrillConfig;
	import org.apache.drill.common.util.DrillFileUtils;
	import org.apache.drill.common.util.DrillStringUtils;
	import org.apache.drill.exec.ExecConstants;
	import org.apache.drill.exec.compile.MergeAdapter.MergedClassResult;
	import org.apache.drill.exec.exception.ClassTransformationException;
	import org.apache.drill.exec.expr.CodeGenerator;
	import org.apache.drill.exec.server.options.OptionSet;
	import org.codehaus.commons.compiler.CompileException;
	import org.objectweb.asm.ClassReader;
	import org.objectweb.asm.tree.ClassNode;

	import com.google.common.annotations.VisibleForTesting;
	import com.google.common.base.Preconditions;
	import com.google.common.collect.Lists;
	import com.google.common.collect.Maps;
	import com.google.common.collect.Sets;

	/**
	* Compiles generated code, merges the resulting class with the
	* template class, and performs byte-code cleanup on the resulting
	* byte codes. The most important transform is scalar replacement
	* which replaces occurrences of non-escaping objects with a
	* collection of member variables.
	*/

	public class ClassTransformer {
	private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(ClassTransformer.class);

	private static final int MAX_SCALAR_REPLACE_CODE_SIZE = 210241024; // 2meg

	private final ByteCodeLoader byteCodeLoader = new ByteCodeLoader();
	private final DrillConfig config;
	private final OptionSet optionManager;

	@VisibleForTesting // although we need it even if it weren't used in testing
	public enum ScalarReplacementOption {
	OFF, // scalar replacement will not ever be used
	TRY, // scalar replacement will be attempted, and if there is an error, we fall back to not using it
	ON; // scalar replacement will always be used, and any errors cause user visible errors

	/**
	* Convert a string to an enum value.
	*
	* @param s the string
	* @return an enum value
	* @throws IllegalArgumentException if the string doesn't match any of the enum values
	*/
	public static ScalarReplacementOption fromString(final String s) {
	switch(s) {
	case "off":
	return OFF;
	case "try":
	return TRY;
	case "on":
	return ON;
	default:
	throw new IllegalArgumentException("Invalid ScalarReplacementOption \"" + s + "\"");
	}
	}
	}

	public ClassTransformer(final DrillConfig config, final OptionSet optionManager) {
	this.config = config;
	this.optionManager = optionManager;
	}

	public static class ClassSet {
	public final ClassSet parent;
	public final ClassNames precompiled;
	public final ClassNames generated;

	public ClassSet(ClassSet parent, String precompiled, String generated) {
	Preconditions.checkArgument(!generated.startsWith(precompiled),
	String.format("The new name of a class cannot start with the old name of a class, otherwise class renaming will cause problems. Precompiled class name %s. Generated class name %s",
	precompiled, generated));
	this.parent = parent;
	this.precompiled = new ClassNames(precompiled);
	this.generated = new ClassNames(generated);
	}

	public ClassSet getChild(String precompiled, String generated) {
	return new ClassSet(this, precompiled, generated);
	}

	public ClassSet getChild(String precompiled) {
	return new ClassSet(this, precompiled, precompiled.replace(this.precompiled.dot, this.generated.dot));
	}

	@Override
	public int hashCode() {
	final int prime = 31;
	int result = 1;
	result = prime * result + ((generated == null) ? 0 : generated.hashCode());
	result = prime * result + ((parent == null) ? 0 : parent.hashCode());
	result = prime * result + ((precompiled == null) ? 0 : precompiled.hashCode());
	return result;
	}

	@Override
	public boolean equals(Object obj) {
	if (this == obj) {
	return true;
	}
	if (obj == null) {
	return false;
	}
	if (getClass() != obj.getClass()) {
	return false;
	}
	ClassSet other = (ClassSet) obj;
	if (generated == null) {
	if (other.generated != null) {
	return false;
	}
	} else if (!generated.equals(other.generated)) {
	return false;
	}
	if (parent == null) {
	if (other.parent != null) {
	return false;
	}
	} else if (!parent.equals(other.parent)) {
	return false;
	}
	if (precompiled == null) {
	if (other.precompiled != null) {
	return false;
	}
	} else if (!precompiled.equals(other.precompiled)) {
	return false;
	}
	return true;
	}
	}

	public static class ClassNames {
	public final String dot;
	public final String slash;
	public final String clazz;

	public ClassNames(String className) {
	dot = className;
	slash = className.replace('.', DrillFileUtils.SEPARATOR_CHAR);
	clazz = DrillFileUtils.SEPARATOR_CHAR + slash + ".class";
	}

	@Override
	public int hashCode() {
	final int prime = 31;
	int result = 1;
	result = prime * result + ((clazz == null) ? 0 : clazz.hashCode());
	result = prime * result + ((dot == null) ? 0 : dot.hashCode());
	result = prime * result + ((slash == null) ? 0 : slash.hashCode());
	return result;
	}

	@Override
	public boolean equals(Object obj) {
	if (this == obj) {
	return true;
	}
	if (obj == null) {
	return false;
	}
	if (getClass() != obj.getClass()) {
	return false;
	}
	ClassNames other = (ClassNames) obj;
	if (clazz == null) {
	if (other.clazz != null) {
	return false;
	}
	} else if (!clazz.equals(other.clazz)) {
	return false;
	}
	if (dot == null) {
	if (other.dot != null) {
	return false;
	}
	} else if (!dot.equals(other.dot)) {
	return false;
	}
	if (slash == null) {
	if (other.slash != null) {
	return false;
	}
	} else if (!slash.equals(other.slash)) {
	return false;
	}
	return true;
	}
	}

	public Class<?> getImplementationClass(CodeGenerator<?> cg) throws ClassTransformationException {
	final QueryClassLoader loader = new QueryClassLoader(config, optionManager);
	return getImplementationClass(loader, cg.getDefinition(),
	cg.getGeneratedCode(), cg.getMaterializedClassName());
	}

	public Class<?> getImplementationClass(
	final QueryClassLoader classLoader,
	final TemplateClassDefinition<?> templateDefinition,
	final String entireClass,
	final String materializedClassName) throws ClassTransformationException {
	// unfortunately, this hasn't been set up at construction time, so we have to do it here
	final ScalarReplacementOption scalarReplacementOption = ScalarReplacementOption.fromString(optionManager.getOption(ExecConstants.SCALAR_REPLACEMENT_VALIDATOR));

	try {
	final long t1 = System.nanoTime();
	final ClassSet set = new ClassSet(null, templateDefinition.getTemplateClassName(), materializedClassName);
	final byte[][] implementationClasses = classLoader.getClassByteCode(set.generated, entireClass);

	long totalBytecodeSize = 0;
	Map<String, Pair<byte[], ClassNode>> classesToMerge = Maps.newHashMap();
	for (byte[] clazz : implementationClasses) {
	totalBytecodeSize += clazz.length;
	ClassNode node = AsmUtil.classFromBytes(clazz, ClassReader.SKIP_FRAMES);
	if (!AsmUtil.isClassOk(logger, "implementationClasses", node)) {
	throw new IllegalStateException("Problem found with implementationClasses");
	}
	classesToMerge.put(node.name, Pair.of(clazz, node));
	}

	final LinkedList<ClassSet> names = Lists.newLinkedList();
	final Set<ClassSet> namesCompleted = Sets.newHashSet();
	names.add(set);

	while ( !names.isEmpty() ) {
	final ClassSet nextSet = names.removeFirst();
	if (namesCompleted.contains(nextSet)) {
	continue;
	}
	final ClassNames nextPrecompiled = nextSet.precompiled;
	final byte[] precompiledBytes = byteCodeLoader.getClassByteCodeFromPath(nextPrecompiled.clazz);
	final ClassNames nextGenerated = nextSet.generated;
	// keeps only classes that have not be merged
	Pair<byte[], ClassNode> classNodePair = classesToMerge.remove(nextGenerated.slash);
	final ClassNode generatedNode;
	if (classNodePair != null) {
	generatedNode = classNodePair.getValue();
	} else {
	generatedNode = null;
	}

	/*
	* TODO
	* We're having a problem with some cases of scalar replacement, but we want to get
	* the code in so it doesn't rot anymore.
	*
	* Here, we use the specified replacement option. The loop will allow us to retry if
	* we're using TRY.
	*/
	MergedClassResult result = null;
	boolean scalarReplace = scalarReplacementOption != ScalarReplacementOption.OFF && entireClass.length() < MAX_SCALAR_REPLACE_CODE_SIZE;
	while(true) {
	try {
	result = MergeAdapter.getMergedClass(nextSet, precompiledBytes, generatedNode, scalarReplace);
	break;
	} catch(RuntimeException e) {
	// if we had a problem without using scalar replacement, then rethrow
	if (!scalarReplace) {
	throw e;
	}

	// if we did try to use scalar replacement, decide if we need to retry or not
	if (scalarReplacementOption == ScalarReplacementOption.ON) {
	// option is forced on, so this is a hard error
	throw e;
	}

	/*
	* We tried to use scalar replacement, with the option to fall back to not using it.
	* Log this failure before trying again without scalar replacement.
	*/
	logger.info("scalar replacement failure (retrying)\n", e);
	scalarReplace = false;
	}
	}

	for (String s : result.innerClasses) {
	s = s.replace(DrillFileUtils.SEPARATOR_CHAR, '.');
	names.add(nextSet.getChild(s));
	}
	classLoader.injectByteCode(nextGenerated.dot, result.bytes);
	namesCompleted.add(nextSet);
	}

	// adds byte code of the classes that have not been merged to make them accessible for outer class
	for (Map.Entry<String, Pair<byte[], ClassNode>> clazz : classesToMerge.entrySet()) {
	classLoader.injectByteCode(clazz.getKey().replace(DrillFileUtils.SEPARATOR_CHAR, '.'), clazz.getValue().getKey());
	}
	Class<?> c = classLoader.findClass(set.generated.dot);
	if (templateDefinition.getExternalInterface().isAssignableFrom(c)) {
	logger.debug("Compiled and merged {}: bytecode size = {}, time = {} ms.",
	c.getSimpleName(),
	DrillStringUtils.readable(totalBytecodeSize),
	(System.nanoTime() - t1 + 500_000) / 1_000_000);
	return c;
	}

	throw new ClassTransformationException("The requested class did not implement the expected interface.");
	} catch (CompileException \| IOException \| ClassNotFoundException e) {
	throw new ClassTransformationException(String.format("Failure generating transformation classes for value: \n %s", entireClass), e);
	}
	}
	}