svn commit: r1568402 - /tomee/tomee/trunk/container/openejb-core/src/main/java/org/apache/openejb/core/stateful/StatefulContainer.java

[rmannibucau]

Author: rmannibucau
Date: Fri Feb 14 17:54:00 2014
New Revision: 1568402

URL: http://svn.apache.org/r1568402
Log:
TOMEE-1125 @Remove doesnt trigger stateful cleanup

Modified:
    
tomee/tomee/trunk/container/openejb-core/src/main/java/org/apache/openejb/core/stateful/StatefulContainer.java

Modified: 
tomee/tomee/trunk/container/openejb-core/src/main/java/org/apache/openejb/core/stateful/StatefulContainer.java
URL: 
http://svn.apache.org/viewvc/tomee/tomee/trunk/container/openejb-core/src/main/java/org/apache/openejb/core/stateful/StatefulContainer.java?rev=1568402&r1=1568401&r2=1568402&view=diff
==============================================================================
--- 
tomee/tomee/trunk/container/openejb-core/src/main/java/org/apache/openejb/core/stateful/StatefulContainer.java
 (original)
+++ 
tomee/tomee/trunk/container/openejb-core/src/main/java/org/apache/openejb/core/stateful/StatefulContainer.java
 Fri Feb 14 17:54:00 2014
@@ -589,7 +589,7 @@ public class StatefulContainer implement
                         callContext.setCurrentOperation(Operation.REMOVE);
                     }
 
-                    discardInstance(callContext);
+                    discardInstance(primKey, instance);
                 }
 
                 // un register EntityManager
@@ -799,19 +799,24 @@ public class StatefulContainer implement
         }
     }
 
-    private void discardInstance(final ThreadContext threadContext) {
-        final Object primaryKey = threadContext.getPrimaryKey();
+    private void discardInstance(final Object primaryKey, final Instance 
instance) {
         if (primaryKey == null) {
             return;
         }
 
-        final Instance instance = checkedOutInstances.remove(primaryKey);
-        if (instance != null && 
!containsExtendedPersistenceContext(instance.beanContext)) {
-            cache.remove(primaryKey);
+        final Instance i;
+        if (instance == null) {
+            i = checkedOutInstances.remove(primaryKey);
+        } else {
+            checkedOutInstances.remove(primaryKey);
+            i = instance;
         }
 
-        if (null != instance && null != instance.creationalContext) {
-            instance.creationalContext.release();
+        if (!containsExtendedPersistenceContext(instance.beanContext)) {
+            cache.remove(primaryKey);
+        }
+        if (i != null && null != i.creationalContext) {
+            i.creationalContext.release();
         }
     }
 
@@ -829,7 +834,7 @@ public class StatefulContainer implement
 
         final ExceptionType type = 
callContext.getBeanContext().getExceptionType(e);
         if (type == ExceptionType.SYSTEM) {
-            discardInstance(callContext);
+            discardInstance(callContext.getPrimaryKey(), null);
             EjbTransactionUtil.handleSystemException(txPolicy, e, callContext);
         } else {
             EjbTransactionUtil.handleApplicationException(txPolicy, e, type == 
ExceptionType.APPLICATION_ROLLBACK);
@@ -1074,7 +1079,7 @@ public class StatefulContainer implement
                     txPolicy.setRollbackOnly(e);
 
                     // [3] Discard the instance
-                    discardInstance(callContext);
+                    discardInstance(callContext.getPrimaryKey(), instance);
 
                     // [4] throw the java.rmi.RemoteException to the client
                     throw new OpenEJBRuntimeException(message, e);
@@ -1116,7 +1121,7 @@ public class StatefulContainer implement
                     // Transaction is complete so can not be rolled back
 
                     // [3] Discard the instance
-                    discardInstance(callContext);
+                    discardInstance(callContext.getPrimaryKey(), instance);
 
                     // [4] throw throw first exception to the client
                     if (firstException == null)
@@ -1148,7 +1153,7 @@ public class StatefulContainer implement
 
                 interceptorStack.invoke();
             } catch (Throwable callbackException) {
-                discardInstance(threadContext);
+                discardInstance(threadContext.getPrimaryKey(), instance);
                 
EjbTransactionUtil.handleSystemException(threadContext.getTransactionPolicy(), 
callbackException, threadContext);
             } finally {
                 ThreadContext.exit(oldContext);