Merge pull request #104 from spodila/Issue87

Fix for Issue #87

Author: spodila

fenzo-core/src/main/java/com/netflix/fenzo/AssignableVirtualMachine.java

@@ -353,21 +353,21 @@ void removeExpiredLeases(boolean all) {
     int expireLimitedLeases(AssignableVMs.VMRejectLimiter vmRejectLimiter) {
         if(singleLeaseMode)
             return 0;
-        int rejected=0;
-        Iterator<Map.Entry<String,VirtualMachineLease>> iterator = leasesMap.entrySet().iterator();
         long now = System.currentTimeMillis();
-        while(iterator.hasNext()) {
-            VirtualMachineLease l = iterator.next().getValue();
-            if(l.getOfferedTime() < (now - leaseOfferExpirySecs*1000) && vmRejectLimiter.reject()) {
-                leaseIdToHostnameMap.remove(l.getId());
-                if(logger.isDebugEnabled())
-                    logger.debug(getHostname() + ": expiring lease offer id " + l.getId());
-                leaseRejectAction.call(l);
-                iterator.remove();
-                rejected++;
+        for (VirtualMachineLease l: leasesMap.values()) {
+            if (l.getOfferedTime() < (now - leaseOfferExpirySecs * 1000) && vmRejectLimiter.reject()) {
+                for (VirtualMachineLease vml: leasesMap.values()) {
+                    leaseIdToHostnameMap.remove(vml.getId());
+                    if(logger.isDebugEnabled())
+                        logger.debug(getHostname() + ": expiring lease offer id " + l.getId());
+                    leaseRejectAction.call(vml);
+                }
+                int size = leasesMap.values().size();
+                leasesMap.clear();
+                return size;
             }
         }
-        return rejected;
+        return 0;
     }
 
     String getCurrVMId() {

fenzo-core/src/test/java/com/netflix/fenzo/OfferRejectionsTest.java

@@ -21,10 +21,7 @@
 import junit.framework.Assert;
 import org.junit.Test;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
@@ -198,7 +195,7 @@ public void call(VirtualMachineLease virtualMachineLease) {
     @Test
     public void testTimedOfferRejects() throws Exception {
         final AtomicInteger expireCount = new AtomicInteger();
-        final int leaseExpirySecs=3;
+        final int leaseExpirySecs=2;
         final TaskScheduler scheduler = new TaskScheduler.Builder()
                 .withLeaseRejectAction(new Action1<VirtualMachineLease>() {
                     @Override
@@ -216,7 +213,7 @@ public void call(VirtualMachineLease virtualMachineLease) {
             leases.clear();
             Thread.sleep(1000);
         }
-        leases = LeaseProvider.getLeases(nLeases, 4, 4000, 1, 10);
+        leases = LeaseProvider.getLeases(nLeases, nLeases, 4, 4000, 1, 10);
         for(int i=0; i<leaseExpirySecs-1; i++) {
             scheduler.scheduleOnce(Collections.<TaskRequest>emptyList(), leases);
             leases.clear();
@@ -230,4 +227,35 @@ public void call(VirtualMachineLease virtualMachineLease) {
         }
         Assert.assertEquals(nLeases*2, expireCount.get());
     }
+
+    // test that all offers of a VM are rejected when one of them expires
+    @Test
+    public void testRejectAllOffersOfVm() throws Exception {
+        final AtomicInteger expireCount = new AtomicInteger();
+        final int leaseExpirySecs=2;
+        final Set<String> hostsRejectedFrom = new HashSet<>();
+        final TaskScheduler scheduler = new TaskScheduler.Builder()
+                .withLeaseRejectAction(new Action1<VirtualMachineLease>() {
+                    @Override
+                    public void call(VirtualMachineLease virtualMachineLease) {
+                        expireCount.incrementAndGet();
+                        hostsRejectedFrom.add(virtualMachineLease.hostname());
+                    }
+                })
+                .withLeaseOfferExpirySecs(leaseExpirySecs)
+                .withMaxOffersToReject(1)
+                .build();
+        final int nhosts = 2;
+        List<VirtualMachineLease> leases = LeaseProvider.getLeases(nhosts, 4, 4000, 1, 10);
+        // add the same leases with same hostnames twice again, so there are 3 offers for each of the nHosts.
+        leases.addAll(LeaseProvider.getLeases(nhosts, 4, 4000, 1, 10));
+        leases.addAll(LeaseProvider.getLeases(nhosts, 4, 4000, 1, 10));
+        for (int i=0; i<leaseExpirySecs+1; i++) {
+            scheduler.scheduleOnce(Collections.<TaskRequest>emptyList(), leases);
+            leases.clear();
+            Thread.sleep(1000);
+        }
+        Assert.assertEquals(3, expireCount.get());
+        Assert.assertEquals(1, hostsRejectedFrom.size());
+    }
 }