Fixes OpenHFT#30 - added synchronization for lastAffinityLock

dpisklov · dpisklov · commit d6f855245c3b · 2018-01-04T16:55:26.000Z
diff --git a/affinity/src/main/java/net/openhft/affinity/AffinityThreadFactory.java b/affinity/src/main/java/net/openhft/affinity/AffinityThreadFactory.java
@@ -21,6 +21,7 @@
 import org.jetbrains.annotations.Nullable;
 
 import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * This is a ThreadFactory which assigns threads based the strategies provided.
@@ -56,17 +57,19 @@ public synchronized Thread newThread(@NotNull final Runnable r) {
         Thread t = new Thread(new Runnable() {
             @Override
             public void run() {
-                AffinityLock al = lastAffinityLock == null ? AffinityLock.acquireLock() : lastAffinityLock.acquireLock(strategies);
-                try {
-                    if (al.cpuId() >= 0)
-                        lastAffinityLock = al;
+                try (AffinityLock ignored = acquireLockBasedOnLast()) {
                     r.run();
-                } finally {
-                    al.release();
                 }
             }
         }, name2);
         t.setDaemon(daemon);
         return t;
     }
+
+    private synchronized AffinityLock acquireLockBasedOnLast() {
+        AffinityLock al = lastAffinityLock == null ? AffinityLock.acquireLock() : lastAffinityLock.acquireLock(strategies);
+        if (al.cpuId() >= 0)
+            lastAffinityLock = al;
+        return al;
+    }
 }
diff --git a/affinity/src/test/java/net/openhft/affinity/impl/AbstractAffinityImplTest.java b/affinity/src/test/java/net/openhft/affinity/impl/AbstractAffinityImplTest.java
@@ -70,11 +70,9 @@ public void setAffinityCompletesGracefully() {
     }
 
     @Test
-    @Ignore("TODO FIX")
     public void getAffinityReturnsValuePreviouslySet() {
         final IAffinity impl = getImpl();
-        final int cores = CORES;
-        for (int core = 0; core < cores; core++) {
+        for (int core = 0; core < CORES; core++) {
             final BitSet mask = new BitSet();
             mask.set(core, true);
             getAffinityReturnsValuePreviouslySet(impl, mask);
