--- C:\dev\hibernate\Branch_3_2\Hibernate3\src\org\hibernate\cache\TreeCache.java	Mon Feb 19 17:05:11 2007
+++ C:\dev\jboss\jboss-4.2\ejb3\src\main\org\jboss\ejb3\entity\JBCCache.java	Mon Feb 19 17:12:35 2007
@@ -1,5 +1,25 @@
-//$Id: TreeCache.java 9965 2006-05-30 18:00:28Z steve.ebersole@jboss.com $
-package org.hibernate.cache;
+/*
+ * JBoss, Home of Professional Open Source.
+ * Copyright 2006, Red Hat Middleware LLC, and individual contributors
+ * as indicated by the @author tags. See the copyright.txt file in the
+ * distribution for a full listing of individual contributors.
+ *
+ * This is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU Lesser General Public License as
+ * published by the Free Software Foundation; either version 2.1 of
+ * the License, or (at your option) any later version.
+ *
+ * This software is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with this software; if not, write to the Free
+ * Software Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA
+ * 02110-1301 USA, or see the FSF site: http://www.fsf.org.
+ */
+package org.jboss.ejb3.entity;
 
 import java.util.HashMap;
 import java.util.Iterator;
@@ -12,17 +32,24 @@
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+import org.hibernate.cache.Cache;
+import org.hibernate.cache.CacheException;
+import org.hibernate.cache.StandardQueryCache;
+import org.hibernate.cache.UpdateTimestampsCache;
 import org.jboss.cache.Fqn;
+import org.jboss.cache.config.Option;
 import org.jboss.cache.lock.TimeoutException;
 
 /**
- * Represents a particular region within the given JBossCache TreeCache.
+ * Subclass of the standard <code>org.hibernate.cache.TreeCache</code> used as
+ * a workaround until issues related to JBCLUSTER-150 are resolved in Hibernate.
  *
  * @author Gavin King
+ * @author Brian Stansberry
  */
-public class TreeCache implements Cache {
+public class JBCCache implements Cache {
 	
-	private static final Log log = LogFactory.getLog(TreeCache.class);
+	private static final Log log = LogFactory.getLog(JBCCache.class);
 
 	private static final String ITEM = "item";
 
@@ -30,13 +57,40 @@
 	private final String regionName;
 	private final Fqn regionFqn;
 	private final TransactionManager transactionManager;
+    private boolean localWritesOnly;
 
-	public TreeCache(org.jboss.cache.TreeCache cache, String regionName, TransactionManager transactionManager) 
+	public JBCCache(org.jboss.cache.TreeCache cache, String regionName, TransactionManager transactionManager) 
 	throws CacheException {
 		this.cache = cache;
 		this.regionName = regionName;
 		this.regionFqn = Fqn.fromString( regionName.replace( '.', '/' ) );
 		this.transactionManager = transactionManager;
+        if (cache.getUseRegionBasedMarshalling())
+        {           
+           localWritesOnly = StandardQueryCache.class.getName().equals(regionName);
+           
+           boolean fetchState = cache.getFetchInMemoryState();
+           try
+           {
+              // We don't want a state transfer for the StandardQueryCache,
+              // as it can include classes from multiple scoped classloaders
+              if (localWritesOnly)
+                 cache.setFetchInMemoryState(false);
+              
+              // We always activate
+              activateCacheRegion(regionFqn.toString());
+           }
+           finally
+           {
+              // Restore the normal state transfer setting
+              if (localWritesOnly)
+                 cache.setFetchInMemoryState(fetchState);              
+           }
+        }
+        else
+        {
+           log.debug("TreeCache is not configured for region based marshalling");
+        }
 	}
 
 	public Object get(Object key) throws CacheException {
@@ -60,7 +114,14 @@
 
 	public void update(Object key, Object value) throws CacheException {
 		try {
-			cache.put( new Fqn( regionFqn, key ), ITEM, value );
+            if (localWritesOnly) {
+               Option option = new Option();
+               option.setCacheModeLocal(true);
+               cache.put( new Fqn( regionFqn, key ), ITEM, value, option );
+            }
+            else {               
+                cache.put( new Fqn( regionFqn, key ), ITEM, value );
+            }
 		}
 		catch (Exception e) {
 			throw new CacheException(e);
@@ -70,8 +131,23 @@
 	public void put(Object key, Object value) throws CacheException {
 		Transaction tx = suspend();
 		try {
-			//do the failfast put outside the scope of the JTA txn
-			cache.putFailFast( new Fqn( regionFqn, key ), ITEM, value, 0 );
+           if (localWritesOnly) {
+              Option option = new Option();
+              option.setCacheModeLocal(true);
+              // Overloaded method isn't available, so have to use InvocationContext
+              cache.getInvocationContext().setOptionOverrides(option);
+              try {
+                  // do the failfast put outside the scope of the JTA txn
+                  cache.putFailFast( new Fqn( regionFqn, key ), ITEM, value, 0 );
+              }
+              finally {
+                 cache.getInvocationContext().setOptionOverrides(null);
+              }
+           }
+           else {               
+               //do the failfast put outside the scope of the JTA txn
+			   cache.putFailFast( new Fqn( regionFqn, key ), ITEM, value, 0 );
+           }
 		}
 		catch (TimeoutException te) {
 			//ignore!
@@ -109,7 +185,14 @@
 
 	public void remove(Object key) throws CacheException {
 		try {
-			cache.remove( new Fqn( regionFqn, key ) );
+           if (localWritesOnly) {
+              Option option = new Option();
+              option.setCacheModeLocal(true);
+              cache.remove( new Fqn( regionFqn, key ), option );
+           }
+           else {               
+               cache.remove( new Fqn( regionFqn, key ) );
+           }
 		}
 		catch (Exception e) {
 			throw new CacheException(e);
@@ -127,13 +210,21 @@
 
 	public void destroy() throws CacheException {
 		try {
-			// NOTE : evict() operates locally only (i.e., does not propogate
-			// to any other nodes in the potential cluster).  This is
-			// exactly what is needed when we destroy() here; destroy() is used
-			// as part of the process of shutting down a SessionFactory; thus
-			// these removals should not be propogated
-			cache.evict( regionFqn );
-		}
+			// NOTE : Hibernate's class uses evict() but that isn't recursive!
+			//cache.evict( regionFqn );
+            Option opt = new Option();
+            opt.setCacheModeLocal(true);
+            cache.remove(regionFqn, opt);
+            
+            if (cache.getUseRegionBasedMarshalling() && !isSharedClassLoaderRegion(regionName))
+            {
+               inactivateCacheRegion();
+            }
+		}
+        catch (CacheException e)
+        {
+           throw e;
+        }
 		catch( Exception e ) {
 			throw new CacheException( e );
 		}
@@ -201,5 +292,59 @@
 	public String toString() {
 		return "TreeCache(" + regionName + ')';
 	}
-	
+       
+    private boolean isSharedClassLoaderRegion(String regionName)
+    {
+       return (StandardQueryCache.class.getName().equals(regionName) 
+                || UpdateTimestampsCache.class.getName().equals(regionName));
+    }
+    
+    private void activateCacheRegion(String regionName) throws CacheException
+    {
+       String fqnString = regionFqn.toString();
+       // FIXME -- find a way that doesn't involve this API
+       if (cache.getMarshaller().isInactive(fqnString))
+       {
+          try
+          {
+             // Only register the classloader if it's not a shared region.  
+             // If it's shared, no single classloader is valid
+             if (!isSharedClassLoaderRegion(regionName))
+             {
+                cache.registerClassLoader(fqnString, Thread.currentThread().getContextClassLoader());
+             }
+             cache.activateRegion(fqnString);
+          }
+          catch (Exception e)
+          {
+             throw new CacheException("Problem activating region " + regionName, e);
+          }
+       }
+       else
+       {
+          log.debug("activateCacheRegion(): Region " + fqnString + " is already active");
+       }
+    }
+    
+    private void inactivateCacheRegion() throws CacheException
+    {
+       String fqnString = regionFqn.toString();
+       // FIXME -- find a way that doesn't involve this API
+       if (!cache.getMarshaller().isInactive(fqnString))
+       {
+          try
+          {
+             cache.inactivateRegion(fqnString);
+             cache.unregisterClassLoader(fqnString);
+          }
+          catch (Exception e)
+          {
+             throw new CacheException("Problem activating region " + fqnString, e);
+          }
+       }     
+       else
+       {
+          log.debug("inactivateCacheRegion(): Region " + fqnString + " is already inactive");
+       }
+    }	
 }
