#731 Add @RequiresUnitOfWork annotation

extensions/persist/src/com/google/inject/persist/PersistModule.java

@@ -37,15 +37,21 @@ protected final void configure() {
     requireBinding(UnitOfWork.class);
     /*if[AOP]*/
     // wrapping in an if[AOP] just to allow this to compile in NO_AOP -- it won't be used
-
-    // class-level @Transacational
+    
+    // class-level @Transactional and @RequiresUnitOfWork
     bindInterceptor(annotatedWith(Transactional.class), any(), getTransactionInterceptor());
-    // method-level @Transacational
+    bindInterceptor(annotatedWith(RequiresUnitOfWork.class), any(),
+          getRequiresUnitOfWorkInterceptor());
+    // method-level @Transactional and @RequiresUnitOfWork
     bindInterceptor(any(), annotatedWith(Transactional.class), getTransactionInterceptor());
+    bindInterceptor(any(), annotatedWith(RequiresUnitOfWork.class),
+          getRequiresUnitOfWorkInterceptor());
     /*end[AOP]*/
   }
 
   protected abstract void configurePersistence();
 
   protected abstract MethodInterceptor getTransactionInterceptor();
+
+  protected abstract MethodInterceptor getRequiresUnitOfWorkInterceptor();
 }

extensions/persist/src/com/google/inject/persist/RequiresUnitOfWork.java

@@ -0,0 +1,19 @@
+package com.google.inject.persist;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Inherited;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * <p> Any method or class marked with this annotation will require the existence
+ * of a unit of work.
+ * <p>Marking a method {@code @RequiresUnitOfWork} will start a unit of work if none
+ * exists before the method executes and end it if it was started after the method returns.
+ */
+@Target({ ElementType.METHOD, ElementType.TYPE })
+@Retention(RetentionPolicy.RUNTIME)
+@Inherited
+public @interface RequiresUnitOfWork {
+}

extensions/persist/src/com/google/inject/persist/jpa/JpaLocalTxnInterceptor.java

@@ -18,7 +18,10 @@
 
 import com.google.inject.Inject;
 import com.google.inject.persist.Transactional;
-import com.google.inject.persist.UnitOfWork;
+
+import org.aopalliance.intercept.MethodInterceptor;
+import org.aopalliance.intercept.MethodInvocation;
+
 import java.lang.reflect.Method;
 import javax.persistence.EntityManager;
 import javax.persistence.EntityTransaction;
@@ -28,32 +31,28 @@
 /** @author Dhanji R. Prasanna (dhanji@gmail.com) */
 class JpaLocalTxnInterceptor implements MethodInterceptor {
 
-  // TODO(gak): Move these args to the cxtor & make these final.
-  @Inject private JpaPersistService emProvider = null;
-
-  @Inject private UnitOfWork unitOfWork = null;
-
   @Transactional
   private static class Internal {}
-
-  // Tracks if the unit of work was begun implicitly by this transaction.
-  private final ThreadLocal<Boolean> didWeStartWork = new ThreadLocal<>();
-
+
+  // TODO(gak): Move this arg to the cxtor & make this final.
+  @Inject
+  private UnitOfWorkHandler unitOfWorkHandler;
+
   @Override
   public Object invoke(MethodInvocation methodInvocation) throws Throwable {
 
-    // Should we start a unit of work?
-    if (!emProvider.isWorking()) {
-      emProvider.begin();
-      didWeStartWork.set(true);
-    }
+    unitOfWorkHandler.requireUnitOfWork();
 
     Transactional transactional = readTransactionMetadata(methodInvocation);
-    EntityManager em = this.emProvider.get();
+    EntityManager em = unitOfWorkHandler.getEntityManager();
 
     // Allow 'joining' of transactions if there is an enclosing @Transactional method.
     if (em.getTransaction().isActive()) {
-      return methodInvocation.proceed();
+      try {
+        return methodInvocation.proceed();
+      } finally {
+        unitOfWorkHandler.endRequireUnitOfWork();
+      }
     }
 
     final EntityTransaction txn = em.getTransaction();
@@ -72,10 +71,8 @@ public Object invoke(MethodInvocation methodInvocation) throws Throwable {
       //propagate whatever exception is thrown anyway
       throw e;
     } finally {
-      // Close the em if necessary (guarded so this code doesn't run unless catch fired).
-      if (null != didWeStartWork.get() && !txn.isActive()) {
-        didWeStartWork.remove();
-        unitOfWork.end();
+      if (!txn.isActive()) {
+        unitOfWorkHandler.endRequireUnitOfWork();
       }
     }
 
@@ -84,11 +81,7 @@ public Object invoke(MethodInvocation methodInvocation) throws Throwable {
     try {
       txn.commit();
     } finally {
-      //close the em if necessary
-      if (null != didWeStartWork.get()) {
-        didWeStartWork.remove();
-        unitOfWork.end();
-      }
+      unitOfWorkHandler.endRequireUnitOfWork();
     }
 
     //or return result

extensions/persist/src/com/google/inject/persist/jpa/JpaPersistModule.java

@@ -26,6 +26,10 @@
 import com.google.inject.persist.UnitOfWork;
 import com.google.inject.persist.finder.DynamicFinder;
 import com.google.inject.persist.finder.Finder;
+
+import org.aopalliance.intercept.MethodInterceptor;
+import org.aopalliance.intercept.MethodInvocation;
+
 import java.lang.reflect.AccessibleObject;
 import java.lang.reflect.InvocationHandler;
 import java.lang.reflect.Method;
@@ -53,6 +57,7 @@ public JpaPersistModule(String jpaUnit) {
 
   private Map<?, ?> properties;
   private MethodInterceptor transactionInterceptor;
+  private MethodInterceptor requiresUnitOfWorkInterceptor;
 
   @Override
   protected void configurePersistence() {
@@ -65,9 +70,13 @@ protected void configurePersistence() {
     bind(EntityManager.class).toProvider(JpaPersistService.class);
     bind(EntityManagerFactory.class)
         .toProvider(JpaPersistService.EntityManagerFactoryProvider.class);
+    bind(UnitOfWorkHandler.class);
 
     transactionInterceptor = new JpaLocalTxnInterceptor();
     requestInjection(transactionInterceptor);
+
+    requiresUnitOfWorkInterceptor = new RequiresUnitOfWorkInterceptor();
+    requestInjection(requiresUnitOfWorkInterceptor);
 
     // Bind dynamic finders.
     for (Class<?> finder : dynamicFinders) {
@@ -80,9 +89,11 @@ protected MethodInterceptor getTransactionInterceptor() {
     return transactionInterceptor;
   }
 
-  @Provides
-  @Jpa
-  Map<?, ?> provideProperties() {
+  @Override protected MethodInterceptor getRequiresUnitOfWorkInterceptor() {
+    return requiresUnitOfWorkInterceptor;
+  }
+
+  @Provides @Jpa Map<?, ?> provideProperties() {
     return properties;
   }
 

extensions/persist/src/com/google/inject/persist/jpa/RequiresUnitOfWorkInterceptor.java

@@ -0,0 +1,43 @@
+package com.google.inject.persist.jpa;
+
+import java.lang.reflect.Method;
+
+import org.aopalliance.intercept.MethodInterceptor;
+import org.aopalliance.intercept.MethodInvocation;
+
+import com.google.inject.Inject;
+import com.google.inject.persist.RequiresUnitOfWork;
+
+public class RequiresUnitOfWorkInterceptor implements MethodInterceptor {
+  @Inject
+  private UnitOfWorkHandler unitOfWorkHandler;
+
+  public Object invoke(MethodInvocation methodInvocation) throws Throwable {
+    RequiresUnitOfWork annotation = readAnnotationMetadata(methodInvocation);
+    if (annotation == null) {
+      // Avoid creating the unit of work in Object class methods
+      return methodInvocation.proceed();
+
+    } else {
+      unitOfWorkHandler.requireUnitOfWork();
+      try {
+        return methodInvocation.proceed();
+      } finally {
+        unitOfWorkHandler.endRequireUnitOfWork();
+      }
+    }
+  }
+
+  private RequiresUnitOfWork readAnnotationMetadata(MethodInvocation methodInvocation) {
+    RequiresUnitOfWork annotation;
+    Method method = methodInvocation.getMethod();
+
+    // Annotation in method or class
+    annotation = method.getAnnotation(RequiresUnitOfWork.class);
+    if (annotation == null) {
+      annotation = method.getClass().getAnnotation(RequiresUnitOfWork.class);
+    }
+
+    return annotation;
+  }
+}

extensions/persist/src/com/google/inject/persist/jpa/UnitOfWorkHandler.java

@@ -0,0 +1,54 @@
+package com.google.inject.persist.jpa;
+
+import javax.persistence.EntityManager;
+
+import com.google.inject.Inject;
+import com.google.inject.persist.UnitOfWork;
+
+/**
+ * Utility class able to handle easily the start and end of units of work
+ * in a nested context. It starts the unit of work if there isn't any in
+ * the current thread, and ends it if the unit of work was started by this
+ * instance.
+ */
+public class UnitOfWorkHandler {
+  private JpaPersistService emProvider = null;
+  private UnitOfWork unitOfWork = null;
+
+  /** Tracks if the unit of work was begun implicitly by this handler. */
+  private final ThreadLocal<Integer> didWeStartWork = new ThreadLocal<Integer>();
+
+  @Inject
+  public UnitOfWorkHandler(JpaPersistService emProvider, UnitOfWork unitOfWork) {
+    this.emProvider = emProvider;
+    this.unitOfWork = unitOfWork;
+  }
+
+  public void requireUnitOfWork() {
+    if (!emProvider.isWorking()) {
+      unitOfWork.begin();
+      didWeStartWork.set(1);
+    } else {
+      Integer level = didWeStartWork.get();
+      if ((level != null) && (level.intValue() >= 0)) {
+        didWeStartWork.set(level + 1);
+      }
+    }
+  }
+
+  public void endRequireUnitOfWork() {
+    Integer level = didWeStartWork.get();
+    if ((level != null) && (level.intValue() > 0)) {
+      if (level.intValue() == 1) {
+        didWeStartWork.remove();
+        unitOfWork.end();
+      } else {
+        didWeStartWork.set(level - 1);
+      }
+    }
+  }
+
+  public EntityManager getEntityManager() {
+    return emProvider.get();
+  }
+}