Fixed FULL_LAZY to pick proper chain-of-states in rubah.runtime.states.States

src/main/java/rubah/bytecode/transformers/AddTraverseMethod.java

@@ -20,6 +20,7 @@
  *******************************************************************************/
 package rubah.bytecode.transformers;
 
+import java.lang.ThreadLocal;
 import java.lang.ref.Reference;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
@@ -63,7 +64,6 @@ public class AddTraverseMethod extends RubahTransformer {
 		// Throwable's backtrace is something handled by native code, don't touch
 		Throwable.class.getName(),
 		Rubah.class.getPackage().getName() + ".",
-		ThreadLocal.class.getName(),
 		Object.class.getPackage().getName(),
 		"sun.security",
 		"sun.nio.ch.SocketChannelImpl",
@@ -73,11 +73,14 @@ public class AddTraverseMethod extends RubahTransformer {
 
 	public static boolean isAllowed(String fqn) {
 
-		if (fqn.startsWith(Object.class.getName()) || fqn.startsWith(Class.class.getName()) || fqn.startsWith(Reference.class.getPackage().getName()))
+		if (fqn.startsWith(Object.class.getName())
+			|| fqn.startsWith(Class.class.getName())
+			|| fqn.startsWith(Reference.class.getPackage().getName())
+			|| fqn.startsWith(ThreadLocal.class.getName()))
 			return true;
 
-		for (String allowedPak : blackListedPackages) {
-			if (fqn.startsWith(allowedPak)) {
+		for (String disallowedPak : blackListedPackages) {
+			if (fqn.startsWith(disallowedPak)) {
 				return false;
 			}
 		}

src/main/java/rubah/runtime/state/MigratingProgramState.java

@@ -24,6 +24,7 @@
 import java.io.File;
 import java.io.FileWriter;
 import java.io.IOException;
+import java.lang.reflect.Field;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
@@ -131,11 +132,13 @@ public void doStart() {
 
 			// Threads
 			for (StoppedThread stoppedThread : this.state.getStopped()) {
+                                migrateThreadLocals(stoppedThread.rubahThread, stoppedThread.rubahThread);
 				Runnable r = (Runnable)this.strategy.migrate(stoppedThread.rubahThread.getTarget());
 				stoppedThread.rubahThread.setTarget(r);
 			}
 
 			for (Entry<RubahThread, RubahThread> entry : this.redirectedThreads.entrySet()) {
+                                migrateThreadLocals(entry.getValue(), entry.getKey());
 				entry.getKey().setTarget(entry.getValue());
 			}
 		} catch (IllegalArgumentException e) {
@@ -211,4 +214,17 @@ public void ensureStaticFieldsMigrated(Class<?> c) {
 		this.migrateStaticFields(Arrays.asList(new Class<?>[]{ c }));
 	}
 
+        private void migrateThreadLocals(RubahThread source, RubahThread target) {
+                try {
+                    Field f = Thread.class.getDeclaredField("threadLocals");
+                    f.setAccessible(true);
+                    Object o = f.get(source);  
+                    if (o != null) { // Check if Thread.threadLocals is set
+                        o = this.strategy.migrate(o);
+                        f.set(target, o);
+                    }
+                } catch (Exception ex) {
+                    throw new Error("ThreadLocal migration failed, should not happen...");
+                }
+        }
 }

src/main/java/rubah/runtime/state/migrator/ReferenceMigratorFactory.java

@@ -36,7 +36,9 @@ public ReferenceMigratorFactory(MigrationStrategy strategy) {
 
 	@Override
 	public boolean canMigrate(Class<?> preClass) {
-		return Reference.class.isAssignableFrom(preClass) || AtomicReference.class.isAssignableFrom(preClass);
+		return Reference.class.isAssignableFrom(preClass) 
+                        || AtomicReference.class.isAssignableFrom(preClass)
+                        || (Reference.class.isAssignableFrom(preClass) && preClass.getName().startsWith(ThreadLocal.class.getName()));
 	}
 
 	@Override
@@ -66,11 +68,19 @@ public Object doMigrate(Object obj) {
 					AtomicReference ref = (AtomicReference) obj;
 
 					ret = new AtomicReference(strategy.migrate(ref.get()));
-				}
-				else
+				} else if (obj.getClass().getName().startsWith(ThreadLocal.class.getName())) {
+                                        ret = obj;
+
+                                        for (long offset : UnsafeUtils.getInstance().getOffsets(ret.getClass()).getOffsets())
+                                                strategy.migrate(ret, offset, ret, offset);
+                                }
+                                else {
+                                        // FIXME: The following cannot properly migrate subclassed references
+
 					// TODO find a more general way to migrated references
 					// Taking a chance here...
 					return ret;
+                                }
 			}
 
 			if (ret == obj) {

src/main/java/rubah/runtime/state/strategy/FullyLazyMonolithic.java

@@ -407,8 +407,17 @@ public Object migrate(Object obj) {
 				AtomicReference ref = (AtomicReference) obj;
 
 				ret = new AtomicReference(this.migrate(ref.get()));
-			}
-			else {
+			} else if (obj.getClass().getName().startsWith(ThreadLocal.class.getName())) {
+				// Migrate java.lang.ThreadLocal$ThreadLocalMap$Entry (which is a subclass of WeakReference)
+
+				// ...$Entry is a key-value-tuple that uses a WeakReference to its key.
+				// The key always points to a ThreadLocal which we assume never to be updated itself,
+				// Therefore, we traverse only the "value" field of Entry and do not rebuilt the WeakReference
+				for (long offset : info.fieldOffsets)
+					this.traverse(ret, offset, ret, offset);
+			} else {
+                                // FIXME: The following cannot properly migrate subclassed references
+
 				// TODO find a more general way to migrated references
 				// Taking a chance here...
 				// Follow reference here
@@ -534,6 +543,23 @@ private ClassConversionInfo buildConversionInfo(Class<?> c) {
 				ret.migrateAction = MigrateAction.MIGRATE_CLASS;
 				ret.mappingAction = MappingAction.MAP;
 				ret.traverseAction = TraverseAction.MIGRATE;
+			} else if (Reference.class.isAssignableFrom(c) && c.getName().startsWith(ThreadLocal.class.getName())) {
+				// Migrate java.lang.ThreadLocal$ThreadLocalMap$Entry (which is a subclass of WeakReference)
+
+				// Migrate reference as regular
+				ret.migrateAction = MigrateAction.MIGRATE_REFERENCE;
+				ret.traverseAction = TraverseAction.MIGRATE;
+				ret.proxyClass = Class.forName(ProxyGenerator.generateProxyName(c.getName()), false, Rubah.getLoader());
+				ret.forwardFieldOffset = -1;
+				ret.mappingAction = MappingAction.MAP;
+
+				// Addidionally collect fieldOffsets to later on migrate the "value" field               
+				LinkedList<Long> offsets = UnsafeUtils.getInstance().getOffsets(c).getOffsets();
+				ret.fieldOffsets = new long[offsets.size()];
+				int i = 0;
+				for (Long offset : offsets) {
+					ret.fieldOffsets[i++] = offset;
+				}
 			} else if (Reference.class.isAssignableFrom(c) || AtomicReference.class.isAssignableFrom(c)) {
 				ret.migrateAction = MigrateAction.MIGRATE_REFERENCE;
 				ret.traverseAction = TraverseAction.MIGRATE;

src/main/java/rubah/tools/updater/ParsingArguments.java

@@ -171,6 +171,7 @@ public final void setState(UpdateState state) {
 					break;
 				case FULL_LAZY:
 					migrationStrategy = new FullyLazyMonolithic();
+					state.setLazy(true);
 					break;
 				case EAGER_LAZY:
 					migrationStrategy = new EagerLazy(mappingStrategy, this.nThreads);