Allow an AffinityLock to be release without resetting the current thread. #65

Author: peter-lawrey

affinity/src/main/java/net/openhft/affinity/AffinityLock.java

@@ -85,6 +85,7 @@ public class AffinityLock implements Closeable {
     @Nullable
     Thread assignedThread;
     Throwable boundHere;
+    private boolean resetAffinity = true;
 
     AffinityLock(int cpuId, boolean base, boolean reservable, LockInventory lockInventory) {
         this.lockInventory = lockInventory;
@@ -274,6 +275,22 @@ private static boolean areAssertionsEnabled() {
         return debug;
     }
 
+    /**
+     * @return Whether to reset the affinity, false indicates the thread is about to die anyway.
+     */
+    public boolean resetAffinity() {
+        return resetAffinity;
+    }
+
+    /**
+     * @param resetAffinity Whether to reset the affinity, false indicates the thread is about to die anyway.
+     * @return this
+     */
+    public AffinityLock resetAffinity(boolean resetAffinity) {
+        this.resetAffinity = resetAffinity;
+        return this;
+    }
+
     /**
      * Assigning the current thread has a side effect of preventing the lock being used again until
      * it is released.
@@ -358,7 +375,9 @@ public void release() {
         if (cpuId == ANY_CPU)
             return;
         // expensive if not actually used.
-        lockInventory.release();
+        boolean resetAffinity = this.resetAffinity;
+        this.resetAffinity = true;
+        lockInventory.release(resetAffinity);
     }
 
     @Override

affinity/src/main/java/net/openhft/affinity/LockInventory.java

@@ -179,7 +179,7 @@ public final synchronized void bindWholeCore(int logicalCoreID) {
         }
     }
 
-    public final synchronized void release() {
+    public final synchronized void release(boolean resetAffinity) {
         Thread t = Thread.currentThread();
         for (AffinityLock al : logicalCoreLocks) {
             Thread at = al.assignedThread;
@@ -189,7 +189,8 @@ public final synchronized void release() {
                 releaseAffinityLock(t, al, "Releasing cpu {} from {} as it is not alive.");
             }
         }
-        Affinity.resetToBaseAffinity();
+        if (resetAffinity)
+            Affinity.resetToBaseAffinity();
     }
 
     public final synchronized String dumpLocks() {

affinity/src/test/java/net/openhft/affinity/AffinityLockTest.java

@@ -149,6 +149,20 @@ public void assignReleaseThread() throws IOException {
         assertEquals(AffinityLock.BASE_AFFINITY, Affinity.getAffinity());
     }
 
+    @Test
+    public void resetAffinity() {
+        assertTrue(Affinity.getAffinity().cardinality() > 1);
+        try (AffinityLock lock = AffinityLock.acquireLock()) {
+            assertEquals(1, Affinity.getAffinity().cardinality());
+            assertTrue(lock.resetAffinity());
+            lock.resetAffinity(false);
+        }
+        assertEquals(1, Affinity.getAffinity().cardinality());
+        try (AffinityLock lock = AffinityLock.acquireLock()) {
+        }
+        assertTrue(Affinity.getAffinity().cardinality() > 1);
+    }
+
     @Test
     public void testIssue21() throws IOException {
         if (!new File("/proc/cpuinfo").exists()) {