Fix handling of affinity lock when specific CPU is requested. Fixes OpenHFT#43.

epickrram · epickrram · commit 4ada7ab44f37 · 2018-01-11T11:58:00.000Z
diff --git a/affinity/src/main/java/net/openhft/affinity/AffinityLock.java b/affinity/src/main/java/net/openhft/affinity/AffinityLock.java
@@ -46,6 +46,7 @@ public class AffinityLock implements Closeable {
     public static final BitSet RESERVED_AFFINITY;
     private static final Logger LOGGER = LoggerFactory.getLogger(AffinityLock.class);
     private static final LockInventory LOCK_INVENTORY;
+    static final int ANY_CPU = -1;
 
     static {
         int processors = Runtime.getRuntime().availableProcessors();
@@ -164,7 +165,7 @@ public static AffinityLock acquireCore() {
      * @return A handle for an affinity lock.
      */
     public static AffinityLock acquireLock(boolean bind) {
-        return acquireLock(bind, -1, AffinityStrategies.ANY);
+        return acquireLock(bind, ANY_CPU, AffinityStrategies.ANY);
     }
 
 
@@ -190,7 +191,7 @@ public static AffinityLock acquireLock(int cpuId) {
      * @return A handle for an affinity lock.
      */
     public static AffinityLock acquireCore(boolean bind) {
-        return acquireCore(bind, -1, AffinityStrategies.ANY);
+        return acquireCore(bind, ANY_CPU, AffinityStrategies.ANY);
     }
 
     private static AffinityLock acquireLock(boolean bind, int cpuId, @NotNull AffinityStrategy... strategies) {
diff --git a/affinity/src/main/java/net/openhft/affinity/LockInventory.java b/affinity/src/main/java/net/openhft/affinity/LockInventory.java
@@ -83,24 +83,33 @@ public final synchronized void set(CpuLayout cpuLayout) {
     }
 
     public final synchronized AffinityLock acquireLock(boolean bind, int cpuId, AffinityStrategy... strategies) {
+        final boolean specificCpuRequested = !isAnyCpu(cpuId);
+        if (specificCpuRequested && cpuId != 0) {
+            final AffinityLock required = logicalCoreLocks[cpuId];
+            if (required.canReserve() && anyStrategyMatches(cpuId, cpuId, strategies)) {
+                updateLockForCurrentThread(bind, required, false);
+                return required;
+            }
+        }
 
+        LOGGER.warn("Unable to acquire lock on CPU {} for thread {}, trying to find another CPU",
+                cpuId, Thread.currentThread());
 
         for (AffinityStrategy strategy : strategies) {
             // consider all processors except cpu 0 which is usually used by the OS.
             // if you have only one core, this library is not appropriate in any case.
             for (int i = logicalCoreLocks.length - 1; i > 0; i--) {
                 AffinityLock al = logicalCoreLocks[i];
-                if (al.canReserve() && (cpuId < 0 || strategy.matches(cpuId, al.cpuId()))) {
-                    al.assignCurrentThread(bind, false);
-                    LockCheck.updateCpu(al.cpuId());
+                if (al.canReserve() && (isAnyCpu(cpuId) || strategy.matches(cpuId, al.cpuId()))) {
+                    updateLockForCurrentThread(bind, al, false);
                     return al;
                 }
             }
         }
 
         LOGGER.warn("No reservable CPU for {}", Thread.currentThread());
 
-        return newLock(-1, false, false);
+        return newLock(AffinityLock.ANY_CPU, false, false);
     }
 
     public final synchronized AffinityLock acquireCore(boolean bind, int cpuId, AffinityStrategy... strategies) {
@@ -112,8 +121,7 @@ public final synchronized AffinityLock acquireCore(boolean bind, int cpuId, Affi
                         continue LOOP;
 
                 final AffinityLock al = als[0];
-                al.assignCurrentThread(bind, true);
-                LockCheck.updateCpu(al.cpuId());
+                updateLockForCurrentThread(bind, al, true);
                 return al;
             }
         }
@@ -196,4 +204,22 @@ private void releaseAffinityLock(final Thread t, final AffinityLock al, final St
             LOGGER.warn("Failed to delete lock file at " + lockFilePath);
         }
     }
+
+    private static boolean anyStrategyMatches(final int cpuOne, final int cpuTwo, final AffinityStrategy[] strategies) {
+        for (AffinityStrategy strategy : strategies) {
+            if (strategy.matches(cpuOne, cpuTwo)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    private static boolean isAnyCpu(final int cpuId) {
+        return cpuId == AffinityLock.ANY_CPU;
+    }
+
+    private static void updateLockForCurrentThread(final boolean bind, final AffinityLock al, final boolean b) {
+        al.assignCurrentThread(bind, b);
+        LockCheck.updateCpu(al.cpuId());
+    }
 }
diff --git a/affinity/src/main/java/net/openhft/affinity/NonForkingAffinityLock.java b/affinity/src/main/java/net/openhft/affinity/NonForkingAffinityLock.java
@@ -68,7 +68,7 @@ public static AffinityLock acquireCore() {
      * @return A handle for an affinity lock.
      */
     public static AffinityLock acquireLock(boolean bind) {
-        return acquireLock(bind, -1, AffinityStrategies.ANY);
+        return acquireLock(bind, AffinityLock.ANY_CPU, AffinityStrategies.ANY);
     }
 
     /**
@@ -80,7 +80,7 @@ public static AffinityLock acquireLock(boolean bind) {
      * @return A handle for an affinity lock.
      */
     public static AffinityLock acquireCore(boolean bind) {
-        return acquireCore(bind, -1, AffinityStrategies.ANY);
+        return acquireCore(bind, AffinityLock.ANY_CPU, AffinityStrategies.ANY);
     }
 
     private static AffinityLock acquireLock(boolean bind, int cpuId, @NotNull AffinityStrategy... strategies) {
diff --git a/affinity/src/test/java/net/openhft/affinity/AffinityLockTest.java b/affinity/src/test/java/net/openhft/affinity/AffinityLockTest.java
@@ -33,6 +33,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotSame;
 import static org.junit.Assert.assertThat;
+import static org.junit.Assume.assumeTrue;
 
 /**
  * @author peter.lawrey
@@ -222,6 +223,15 @@ public void run() {
         displayStatus();
     }
 
+    @Test
+    public void shouldReturnLockForSpecifiedCpu() {
+        assumeTrue(Runtime.getRuntime().availableProcessors() > 3);
+
+        try (final AffinityLock affinityLock = AffinityLock.acquireLock(2)) {
+            assertThat(affinityLock.cpuId(), is(2));
+        }
+    }
+
     @Test
     public void lockFilesShouldBeRemovedOnRelease() {
         final AffinityLock lock = AffinityLock.acquireLock();
