org.datanucleus.enhancer.EnhancerProcessor Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of datanucleus-enhancer
DataNucleus Enhancer is a Java byte-code enhancer for use with DataNucleus.
There is a newer version: 3.1.1
/**********************************************************************
Copyright (c) 2009 Erik Bengtson and others. 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.

Contributors:
    ...
 **********************************************************************/
package org.datanucleus.enhancer;

import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;

import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.ProcessingEnvironment;
import javax.annotation.processing.RoundEnvironment;
import javax.annotation.processing.SupportedAnnotationTypes;
import javax.annotation.processing.SupportedOptions;
import javax.annotation.processing.SupportedSourceVersion;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.TypeElement;
import javax.lang.model.util.ElementFilter;
import javax.tools.FileObject;
import javax.tools.StandardLocation;

/**
 * Enhance classes after javac.
 * PS. This only works for classes that are annotated. If you have a class that is defined in XML yet has
 * no annotations then it will not be discovered by this process.
 */
@SupportedAnnotationTypes("*")
@SupportedSourceVersion(SourceVersion.RELEASE_5)
@SupportedOptions({EnhancerProcessor.OPTION_API})
public class EnhancerProcessor extends AbstractProcessor
{
    // use "javac -AenhanceAPI=JPA" to process JPA annotations
    public final static String OPTION_API = "enhanceAPI";

    DataNucleusEnhancer enhancer = null;

    Set classNames = new HashSet();
    Set classpathURLs = new HashSet();
    final Set locations = new HashSet();

    /** enhancement runnable **/
    EnhanceRunnable enhanceRunnable = new EnhanceRunnable();

    /** enhancement thread **/
    Thread enhancementThread = new Thread(enhanceRunnable);

    /** a jvm shutdown hook **/
    Thread hook = new Thread(new ShutDownHook());

    /** whether an error has been raised **/
    boolean errorRaised = false;

    /**
     * Default public constructor
     */
    public EnhancerProcessor()
    {
    }

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv)
    {
        super.init(processingEnv);
        
        // Get the enhance API option and initialise the Enhancer
        String enhanceAPI = processingEnv.getOptions().get(OPTION_API);
        if (enhanceAPI != null && enhanceAPI.equalsIgnoreCase("JPA"))
        {
            enhancer = new DataNucleusEnhancer(enhanceAPI, "ASM");
        }
        else
        {
            enhancer = new DataNucleusEnhancer();
        }
    }


    public boolean process(Set arg0, RoundEnvironment roundEnv)
    {
        if (roundEnv.errorRaised())
        {
            errorRaised = true;
        }
        //create classpath roots (classoutput, sourceoutput, source path, or current dir
        String classpath = "";
        try
        {
            FileObject f = processingEnv.getFiler().getResource(StandardLocation.CLASS_OUTPUT, "", "d");
            classpath = new File(new URL(f.toUri().toString()).toURI()).getParentFile().getAbsolutePath() + File.separatorChar;
            classpathURLs.add(new URL("file:/"+classpath));
        }
        catch (Exception e1)
        {
            try
            {
                FileObject f = processingEnv.getFiler().getResource(StandardLocation.SOURCE_OUTPUT, "", "d");
                classpath = new File(new URL(f.toUri().toString()).toURI()).getParentFile().getAbsolutePath() + File.separatorChar;
                classpathURLs.add(new URL("file:/"+classpath));
            }
            catch (Exception e2)
            {
                try
                {
                    FileObject f = processingEnv.getFiler().getResource(StandardLocation.SOURCE_PATH, "", "d");
                    classpath = new File(new URL(f.toUri().toString()).toURI()).getParentFile().getAbsolutePath() + File.separatorChar;
                    classpathURLs.add(new URL("file:/"+classpath));
                }
                catch (Exception e3)
                {
                }
                try
                {
                    classpath = System.getProperty("user.dir") + File.separatorChar;
                    classpathURLs.add(new File(classpath).toURI().toURL());
                }
                catch (Exception e3)
                {
                }
            }
        }
        Set classes = new HashSet();
        classes.addAll(ElementFilter.typesIn(roundEnv.getRootElements()));
        Iterator elems = classes.iterator();
        while (elems.hasNext())
        {
            final TypeElement el = elems.next();
            try
            {
                File location = new File(classpath + el.getQualifiedName().toString().replace('.', File.separatorChar) + ".class");
                locations.add(location);
                classNames.add(el.getQualifiedName().toString());
                classpathURLs.add(location.toURI().toURL());
            }
            catch (IOException e)
            {
                e.printStackTrace();
            }
        }
        if (roundEnv.processingOver())
        {
            Runtime.getRuntime().addShutdownHook(hook);
            enhancementThread.start();
        }
        return false;
    }

    /**
     * Performs the enhancement
     */
    public class EnhanceRunnable implements Runnable
    {
        public boolean running;
        public boolean finished;
        public void run()
        {
            running = true;
            try
            {
                // why running 1000 times? to ensure all classes were compiled and saved to disk
                for (int i = 0; i < 1000 && !existsAll() && !errorRaised; i++)
                {
                    try
                    {
                        Thread.sleep(100);
                    }
                    catch (InterruptedException e)
                    {
                    }
                }
                // this must be invoked from this running thread
                enhancer.setClassLoader(new URLClassLoader(classpathURLs.toArray(new URL[classpathURLs.size()]), EnhancerProcessor.class.getClassLoader()));
                enhancer.addClasses(classNames.toArray(new String[classNames.size()]));
                enhancer.enhance();
            }
            finally
            {
                running = false;
                finished = true;
                try
                {
                    //remove shutdown hook
                    Runtime.getRuntime().removeShutdownHook(hook);
                }
                catch(IllegalStateException ex)
                {
                    //will happen if shutdown has started
                }
            }
        }

        /**
         * Ensure the files exist before we start enhancing
         */
        private boolean existsAll()
        {
            Iterator elems = new HashSet(locations).iterator();
            while (elems.hasNext())
            {
                if (!elems.next().exists())
                {
                    return false;
                }
            }
            return true;
        }
    }
    
    /**
     * Ensure that JVM is not shutdown before the enhancer has finished
     */
    private class ShutDownHook implements Runnable
    {
        public void run()
        {
            while(!enhanceRunnable.finished)
            {
                try
                {
                    Thread.sleep(500);
                }
                catch (InterruptedException e)
                {
                    e.printStackTrace();
                }
            }                    
        }  
    }
}