[jboss-cvs] JBossAS SVN: r60643 - branches/Branch_4_2/ejb3/src/main/org/jboss/ejb3/cache.

[jboss-cvs-commits at lists.jboss.org]

Author: bstansberry at jboss.com
Date: 2007-02-19 03:51:12 -0500 (Mon, 19 Feb 2007)
New Revision: 60643

Modified:
   branches/Branch_4_2/ejb3/src/main/org/jboss/ejb3/cache/StatefulReplicationInterceptor.java
Log:
[EJBTHREE-871] Invocation on nested SFSB does not trigger replication of parent context

Modified: branches/Branch_4_2/ejb3/src/main/org/jboss/ejb3/cache/StatefulReplicationInterceptor.java
===================================================================
--- branches/Branch_4_2/ejb3/src/main/org/jboss/ejb3/cache/StatefulReplicationInterceptor.java	2007-02-19 08:50:44 UTC (rev 60642)
+++ branches/Branch_4_2/ejb3/src/main/org/jboss/ejb3/cache/StatefulReplicationInterceptor.java	2007-02-19 08:51:12 UTC (rev 60643)
@@ -21,6 +21,10 @@
  */
 package org.jboss.ejb3.cache;
 
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Stack;
+
 import org.jboss.aop.advice.Interceptor;
 import org.jboss.aop.joinpoint.Invocation;
 import org.jboss.ejb3.stateful.StatefulBeanContext;
@@ -28,13 +32,18 @@
 import org.jboss.ejb3.stateful.StatefulContainerInvocation;
 
 /**
- * Replicate SFSB if it is modified
+ * Replicate SFSB if it is modified.
  *
  * @author <a href="mailto:bill at jboss.org">Bill Burke</a>
+ * @author Brian Stansberry
+ * 
  * @version $Revision$
  */
 public class StatefulReplicationInterceptor implements Interceptor
 {
+   private static final ThreadLocal<Map<StatefulBeanContext, Stack<Boolean>>> replicationContext = 
+            new ThreadLocal<Map<StatefulBeanContext, Stack<Boolean>>>();
+   
    public String getName()
    {
       return this.getClass().getName();
@@ -42,27 +51,139 @@
 
    public Object invoke(Invocation invocation) throws Throwable
    {
-      Object rtn = invocation.invokeNext();
+      // Find the ultimate parent context for the tree of SFSBs the target
+      // bean is part of.  This "tree" could just be the bean itself, or
+      // a multi-layer tree of nested SFSBs.
+      StatefulContainerInvocation ejbInv = (StatefulContainerInvocation) invocation;
+      StatefulBeanContext ctx = (StatefulBeanContext) ejbInv.getBeanContext();
+      StatefulBeanContext root = ctx.getUltimateContainedIn();
+      
+      // Find out if the ultimate parent is clustered
+      boolean clustered = false;
+      StatefulContainer container = (StatefulContainer) root.getContainer();
+      ClusteredStatefulCache clusteredCache = null;
+      if (container.getCache() instanceof ClusteredStatefulCache)
+      {
+         clustered = true;
+         clusteredCache = (ClusteredStatefulCache) container.getCache();
+      }
+      
+      // Track nested calls to this tree so we know when the outer call
+      // returns -- that's when we replicate
+      if (clustered)         
+         pushCallStack(root);
+      
+      boolean stackUnwound = false;
+      Object rtn = null;
+      try
+      {
+         rtn = invocation.invokeNext();
+      }
+      finally
+      {
+         stackUnwound = (clustered && isCallStackUnwound(root));
+      }
 
+
+      // We only replicate if the ultimate parent is clustered
+      // TODO should we fail somehow during bean creation otherwise??
+      boolean mustReplicate = clustered;
+      
+      // If the bean implements Optimized, we call isModified() even
+      // if we know we won't replicate, as the bean might be expecting 
+      // us to call the method
       Object obj = invocation.getTargetObject();
-
       if (obj instanceof Optimized)
       {
-         if (!((Optimized) obj).isModified())
+         if (((Optimized) obj).isModified() == false)
          {
-            return rtn;
+            mustReplicate = false;
          }
       }
-      StatefulContainerInvocation ejb = (StatefulContainerInvocation) invocation;
-      StatefulBeanContext ctx = (StatefulBeanContext) ejb.getBeanContext();
-      StatefulContainer container = (StatefulContainer) ejb.getAdvisor();
-      // has not been overriden
-      if (container.getCache() instanceof ClusteredStatefulCache)
+      
+      if (mustReplicate)
       {
-         ClusteredStatefulCache cache = (ClusteredStatefulCache) container.getCache();
-         cache.replicate(ctx);
+         // Mark the bean for replication. If the call stack is not
+         // unwound yet this will tell the outer caller the tree is 
+         // dirty even if the outer bean's isModified() returns false
+         root.markedForReplication = true;
       }
       
+      if (stackUnwound && root.markedForReplication)
+      {
+         clusteredCache.replicate(root);
+      }
+      
+      if (ctx != root && ctx.markedForReplication)
+      {
+         // ctx is a ProxiedStatefulBeanContext that may have failed over
+         // and needs to invalidate any remote nodes that hold stale refs
+         // to their delegate. So we replicate it.
+         container = (StatefulContainer) ctx.getContainer();
+         StatefulCache cache = container.getCache();
+         if (cache instanceof ClusteredStatefulCache)
+         {
+            clusteredCache = (ClusteredStatefulCache) cache;
+            clusteredCache.replicate(ctx);
+         }
+         else
+         {
+            // not replicable
+            ctx.markedForReplication = false;
+         }
+      }
+      
       return rtn;
    }
+   
+   private static void pushCallStack(StatefulBeanContext ctx)
+   {
+      Stack<Boolean> callStack = null;
+      Map<StatefulBeanContext, Stack<Boolean>> map = replicationContext.get();
+      if (map == null)
+      {
+         map = new HashMap<StatefulBeanContext, Stack<Boolean>>();
+         replicationContext.set(map);
+      }
+      else
+      {
+         callStack = map.get(ctx);
+      }
+
+      if (callStack == null)
+      {
+         callStack = new Stack<Boolean>();
+         map.put(ctx, callStack);
+      }
+      
+      callStack.push(Boolean.TRUE);
+   }
+   
+   private static boolean isCallStackUnwound(StatefulBeanContext ctx)
+   {
+      Map<StatefulBeanContext, Stack<Boolean>> map = replicationContext.get();
+      if (map == null)
+      {
+         throw new IllegalStateException("replicationContext contains no Map");
+      }
+      Stack<Boolean> callStack = map.get(ctx);
+      if (callStack == null)
+      {
+         throw new IllegalStateException("replicationContext contains no call stack");
+      }
+      
+      callStack.pop();
+      boolean unwound = (callStack.size() == 0);
+      
+      if (unwound)
+      {
+         map.remove(ctx);
+         if (map.size() == 0)
+         {
+            replicationContext.set(null);
+         }
+      }
+      
+      return unwound;
+   }
 }