Merge pull request #325 from mwl/refactor/292-offers-method

mwl · mwl · commit 580ec16a1585 · 2015-09-23T10:28:33.000+01:00
Reviewed and approved by @philwinder.
diff --git a/scheduler/src/main/java/org/apache/mesos/elasticsearch/scheduler/ElasticsearchScheduler.java b/scheduler/src/main/java/org/apache/mesos/elasticsearch/scheduler/ElasticsearchScheduler.java
@@ -30,6 +30,7 @@ public class ElasticsearchScheduler implements Scheduler {
     private Observable statusUpdateWatchers = new StatusUpdateObservable();
     private Boolean registered = false;
     private ClusterState clusterState;
+    OfferStrategy offerStrategy;
 
     public ElasticsearchScheduler(Configuration configuration, TaskInfoFactory taskInfoFactory) {
         this.configuration = configuration;
@@ -67,6 +68,7 @@ public void registered(SchedulerDriver driver, Protos.FrameworkID frameworkId, P
         LOGGER.info("Framework registered as " + frameworkId.getValue());
 
         clusterState = new ClusterState(configuration.getState(), frameworkState); // Must use new framework state. This is when we are allocated our FrameworkID.
+        offerStrategy = new OfferStrategy(configuration, clusterState);
         clusterMonitor = new ClusterMonitor(configuration, this, driver, new StatePath(configuration.getState()));
         clusterState.getTaskList().forEach(clusterMonitor::startMonitoringTask); // Get all previous executors and start monitoring them.
         statusUpdateWatchers.addObserver(clusterState);
@@ -88,34 +90,19 @@ public void reregistered(SchedulerDriver driver, Protos.MasterInfo masterInfo) {
         LOGGER.info("Framework re-registered");
     }
 
-    // Todo, this massive if statement needs to be performed better.
     @Override
     public void resourceOffers(SchedulerDriver driver, List<Protos.Offer> offers) {
         if (!registered) {
             LOGGER.debug("Not registered, can't accept resource offers.");
             return;
         }
         for (Protos.Offer offer : offers) {
-            if (isHostAlreadyRunningTask(offer)) {
-                driver.declineOffer(offer.getId()); // DCOS certification 05
-                LOGGER.info("Declined offer: Host " + offer.getHostname() + " is already running an Elastisearch task");
-            } else if (clusterState.getTaskList().size() == configuration.getElasticsearchNodes()) {
-                driver.declineOffer(offer.getId()); // DCOS certification 05
-                LOGGER.info("Declined offer: Mesos runs already runs " + configuration.getElasticsearchNodes() + " Elasticsearch tasks");
-            } else if (!containsTwoPorts(offer.getResourcesList())) {
-                LOGGER.info("Declined offer: Offer did not contain 2 ports for Elasticsearch client and transport connection");
-                driver.declineOffer(offer.getId());
-            } else if (!isEnoughCPU(offer.getResourcesList())) {
-                LOGGER.info("Declined offer: Not enough CPU resources");
-                driver.declineOffer(offer.getId());
-            } else if (!isEnoughRAM(offer.getResourcesList())) {
-                LOGGER.info("Declined offer: Not enough RAM resources");
-                driver.declineOffer(offer.getId());
-            } else if (!isEnoughDisk(offer.getResourcesList())) {
-                LOGGER.info("Not enough Disk resources");
+            final OfferStrategy.OfferResult result = offerStrategy.evaluate(offer);
+
+            if (!result.acceptable) {
+                LOGGER.debug("Declined offer: " + result.reason.orElse("Unknown"));
                 driver.declineOffer(offer.getId());
             } else {
-                LOGGER.info("Accepted offer: " + offer.getHostname());
                 Protos.TaskInfo taskInfo = taskInfoFactory.createTask(configuration, offer);
                 LOGGER.debug(taskInfo.toString());
                 driver.launchTasks(Collections.singleton(offer.getId()), Collections.singleton(taskInfo));
@@ -126,26 +113,6 @@ public void resourceOffers(SchedulerDriver driver, List<Protos.Offer> offers) {
         }
     }
 
-    private boolean isEnoughDisk(List<Protos.Resource> resourcesList) {
-        ResourceCheck resourceCheck = new ResourceCheck(Resources.RESOURCE_DISK);
-        return resourceCheck.isEnough(resourcesList, configuration.getDisk());
-    }
-
-    private boolean isEnoughCPU(List<Protos.Resource> resourcesList) {
-        ResourceCheck resourceCheck = new ResourceCheck(Resources.RESOURCE_CPUS);
-        return resourceCheck.isEnough(resourcesList, configuration.getCpus());
-    }
-
-    private boolean isEnoughRAM(List<Protos.Resource> resourcesList) {
-        ResourceCheck resourceCheck = new ResourceCheck(Resources.RESOURCE_MEM);
-        return resourceCheck.isEnough(resourcesList, configuration.getMem());
-    }
-
-    private boolean containsTwoPorts(List<Protos.Resource> resources) {
-        int count = Resources.selectTwoPortsFromRange(resources).size();
-        return count == 2;
-    }
-
     @Override
     public void offerRescinded(SchedulerDriver driver, Protos.OfferID offerId) {
         LOGGER.info("Offer " + offerId.getValue() + " rescinded");
diff --git a/scheduler/src/main/java/org/apache/mesos/elasticsearch/scheduler/OfferStrategy.java b/scheduler/src/main/java/org/apache/mesos/elasticsearch/scheduler/OfferStrategy.java
@@ -0,0 +1,111 @@
+package org.apache.mesos.elasticsearch.scheduler;
+
+import org.apache.log4j.Logger;
+import org.apache.mesos.Protos;
+import org.apache.mesos.elasticsearch.scheduler.state.ClusterState;
+
+import java.util.List;
+import java.util.Optional;
+
+import static java.util.Arrays.asList;
+
+/**
+ * Offer strategy
+ */
+public class OfferStrategy {
+    private static final Logger LOGGER = Logger.getLogger(ElasticsearchScheduler.class.toString());
+    private ClusterState clusterState;
+    private Configuration configuration;
+
+    private List<OfferRule> acceptanceRules = asList(
+            new OfferRule("Host already running task", this::isHostAlreadyRunningTask),
+            new OfferRule("Cluster size already fulfilled", offer -> clusterState.getTaskList().size() == configuration.getElasticsearchNodes()),
+            new OfferRule("Offer did not have 2 ports", offer -> !containsTwoPorts(offer.getResourcesList())),
+            new OfferRule("Offer did not have enough CPU resources", offer -> !isEnoughCPU(configuration, offer.getResourcesList())),
+            new OfferRule("Offer did not have enough RAM resources", offer -> !isEnoughRAM(configuration, offer.getResourcesList())),
+            new OfferRule("Offer did not have enough disk resources", offer -> !isEnoughDisk(configuration, offer.getResourcesList()))
+    );
+
+    public OfferStrategy(Configuration configuration, ClusterState clusterState) {
+        this.clusterState = clusterState;
+        this.configuration = configuration;
+    }
+
+    public OfferResult evaluate(Protos.Offer offer) {
+        final Optional<OfferRule> decline = acceptanceRules.stream().filter(offerRule -> offerRule.rule.accepts(offer)).limit(1).findFirst();
+        if (decline.isPresent()) {
+            return OfferResult.decline(decline.get().declineReason);
+        }
+
+        LOGGER.info("Accepted offer: " + offer.getHostname());
+        return OfferResult.accept();
+    }
+
+    /**
+     * Offer result
+     */
+    public static class OfferResult {
+        final boolean acceptable;
+        final Optional<String> reason;
+
+        private OfferResult(boolean acceptable, Optional<String> reason) {
+            this.acceptable = acceptable;
+            this.reason = reason;
+        }
+
+        public static OfferResult accept() {
+            return new OfferResult(true, Optional.<String>empty());
+        }
+
+        public static OfferResult decline(String reason) {
+            return new OfferResult(false, Optional.of(reason));
+        }
+    }
+
+    private boolean isHostAlreadyRunningTask(Protos.Offer offer) {
+        Boolean result = false;
+        List<Protos.TaskInfo> stateList = clusterState.getTaskList();
+        for (Protos.TaskInfo t : stateList) {
+            if (t.getSlaveId().equals(offer.getSlaveId())) {
+                result = true;
+            }
+        }
+        return result;
+    }
+    private boolean isEnoughDisk(Configuration configuration, List<Protos.Resource> resourcesList) {
+        return new ResourceCheck(Resources.RESOURCE_DISK).isEnough(resourcesList, configuration.getDisk());
+    }
+
+    private boolean isEnoughCPU(Configuration configuration, List<Protos.Resource> resourcesList) {
+        return new ResourceCheck(Resources.RESOURCE_CPUS).isEnough(resourcesList, configuration.getCpus());
+    }
+
+    private boolean isEnoughRAM(Configuration configuration, List<Protos.Resource> resourcesList) {
+        return new ResourceCheck(Resources.RESOURCE_MEM).isEnough(resourcesList, configuration.getMem());
+    }
+
+    private boolean containsTwoPorts(List<Protos.Resource> resources) {
+        return Resources.selectTwoPortsFromRange(resources).size() == 2;
+    }
+
+    /**
+     * Rule and reason container object
+     */
+    private static class OfferRule {
+        String declineReason;
+        Rule rule;
+
+        public OfferRule(String declineReason, Rule rule) {
+            this.declineReason = declineReason;
+            this.rule = rule;
+        }
+    }
+
+    /**
+     * Interface for checking offers
+     */
+    @FunctionalInterface
+    private interface Rule {
+        boolean accepts(Protos.Offer offer);
+    }
+}
diff --git a/scheduler/src/test/java/org/apache/mesos/elasticsearch/scheduler/ElasticsearchSchedulerTest.java b/scheduler/src/test/java/org/apache/mesos/elasticsearch/scheduler/ElasticsearchSchedulerTest.java
@@ -12,12 +12,17 @@
 import org.junit.Test;
 import org.mockito.Mockito;
 
+import java.net.InetSocketAddress;
+import java.time.ZonedDateTime;
 import java.util.Date;
 import java.util.UUID;
 
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
 import static org.apache.mesos.elasticsearch.common.Offers.newOfferBuilder;
 import static org.apache.mesos.elasticsearch.scheduler.Resources.*;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 /**
@@ -58,16 +63,9 @@ public class ElasticsearchSchedulerTest {
     private TaskInfoFactory taskInfoFactory;
 
     private org.apache.mesos.elasticsearch.scheduler.Configuration configuration;
-//    private ZonedDateTime now = ZonedDateTime.now();
-//    private InetSocketAddress transportAddress = new InetSocketAddress("localhost", 9300);
-//    private InetSocketAddress clientAddress = new InetSocketAddress("localhost", 9200);
-
 
     @Before
     public void before() {
-        Clock clock = mock(Clock.class);
-        when(clock.now()).thenReturn(TASK1_DATE).thenReturn(TASK2_DATE);
-
         frameworkID = Protos.FrameworkID.newBuilder().setValue(UUID.randomUUID().toString()).build();
         FrameworkState frameworkState = mock(FrameworkState.class);
         when(frameworkState.getFrameworkID()).thenReturn(frameworkID);
@@ -92,6 +90,7 @@ public void before() {
 
         masterInfo = newMasterInfo();
         scheduler.registered(driver, frameworkID, masterInfo);
+        scheduler.offerStrategy = mock(OfferStrategy.class);
     }
 
     @Ignore
@@ -107,7 +106,7 @@ public void shouldCallOberversWhenExecutorLost() {
 
     @Test
     public void testRegistered() {
-        Mockito.verify(driver).requestResources(
+        verify(driver).requestResources(
                 Mockito.argThat(
                         new RequestMatcher(
                                 configuration.getCpus(),
@@ -119,92 +118,33 @@ public void testRegistered() {
         );
     }
 
-    // TODO (pnw): This requires scheduler refactoring to work. Refactor instantiation of objects out of registered method. And use setters.
-//    @Test
-//    public void testResourceOffers_isSlaveAlreadyRunningTask() {
-//        Task task1 = new Task("host1", "1", Protos.TaskState.TASK_RUNNING, now, clientAddress, transportAddress);
-//        Task task2 = new Task("host2", "2", Protos.TaskState.TASK_RUNNING, now, clientAddress, transportAddress);
-//        scheduler.tasks = new HashMap<>();
-//        scheduler.tasks.put("host1", task1);
-//        scheduler.tasks.put("host2", task2);
-//
-//        Protos.Offer.Builder offer = newOffer("host1");
-//
-//        scheduler.resourceOffers(driver, singletonList(offer.build()));
-//
-//        verify(driver).declineOffer(offer.getId());
-//    }
-//
-//    @Test
-//    public void testResourceOffers_enoughNodes() {
-//        Task task1 = new Task("host1", "1", Protos.TaskState.TASK_RUNNING, now, clientAddress, transportAddress);
-//        Task task2 = new Task("host2", "2", Protos.TaskState.TASK_RUNNING, now, clientAddress, transportAddress);
-//        Task task3 = new Task("host3", "3", Protos.TaskState.TASK_RUNNING, now, clientAddress, transportAddress);
-//        scheduler.tasks = new HashMap<>();
-//        scheduler.tasks.put("host1", task1);
-//        scheduler.tasks.put("host2", task2);
-//        scheduler.tasks.put("host3", task3);
-//
-//        Protos.Offer.Builder offer = newOffer("host4");
-//
-//        scheduler.resourceOffers(driver, singletonList(offer.build()));
-//
-//        verify(driver).declineOffer(offer.getId());
-//    }
-//
-//    @Test
-//    public void testResourceOffers_noPorts() {
-//        Task task1 = new Task("host1", "1", Protos.TaskState.TASK_RUNNING, now, clientAddress, transportAddress);
-//        Task task2 = new Task("host2", "2", Protos.TaskState.TASK_RUNNING, now, clientAddress, transportAddress);
-//        scheduler.tasks = new HashMap<>();
-//        scheduler.tasks.put("host1", task1);
-//        scheduler.tasks.put("host2", task2);
-//
-//        Protos.Offer.Builder offer = newOffer("host3");
-//
-//        scheduler.resourceOffers(driver, singletonList(offer.build()));
-//
-//        verify(driver).declineOffer(offer.getId());
-//    }
-//
-//    @SuppressWarnings("unchecked")
-//    @Test
-//    public void testResourceOffers_singlePort() {
-//        Task task = new Task("host1", "task1", Protos.TaskState.TASK_RUNNING, now, clientAddress, transportAddress);
-//        scheduler.tasks = new HashMap<>();
-//        scheduler.tasks.put("host1", task);
-//
-//        Protos.Offer.Builder offerBuilder = newOffer("host3");
-//        offerBuilder.addResources(portRange(9200, 9200));
-//
-//        scheduler.resourceOffers(driver, singletonList(offerBuilder.build()));
-//
-//        Mockito.verify(driver).declineOffer(offerBuilder.build().getId());
-//    }
-//
-//    @Test
-//    public void testResourceOffers_launchTasks() {
-//        scheduler.tasks = new HashMap<>();
-//
-//        Protos.Offer.Builder offerBuilder = newOffer("host3");
-//        offerBuilder.addResources(portRange(9200, 9200));
-//        offerBuilder.addResources(portRange(9300, 9300));
-//
-//        Protos.TaskInfo taskInfo = ProtoTestUtil.getDefaultTaskInfo();
-//
-//        when(taskInfoFactory.createTask(configuration, offerBuilder.build())).thenReturn(taskInfo);
-//
-//        scheduler.resourceOffers(driver, singletonList(offerBuilder.build()));
-//
-//        verify(driver).launchTasks(Collections.singleton(offerBuilder.build().getId()), Collections.singleton(taskInfo));
-//    }
+    @Test
+    public void willDeclineOfferIfStrategyDeclinesOffer() {
+        Protos.Offer offer = newOffer("host1").build();
+
+        when(scheduler.offerStrategy.evaluate(offer)).thenReturn(OfferStrategy.OfferResult.decline("Test"));
+
+        scheduler.resourceOffers(driver, singletonList(offer));
+
+        verify(driver).declineOffer(offer.getId());
+    }
+
+    @Test
+    public void testResourceOffers_launchTasks() {
+        final Protos.Offer offer = newOffer("host3").build();
+        when(scheduler.offerStrategy.evaluate(offer)).thenReturn(OfferStrategy.OfferResult.accept());
+
+        Protos.TaskInfo taskInfo = ProtoTestUtil.getDefaultTaskInfo();
+
+        when(taskInfoFactory.createTask(configuration, offer)).thenReturn(taskInfo);
+
+        scheduler.resourceOffers(driver, singletonList(offer));
+
+        verify(driver).launchTasks(singleton(offer.getId()), singleton(taskInfo));
+    }
 
     private Protos.Offer.Builder newOffer(String hostname) {
-        Protos.Offer.Builder builder = newOfferBuilder(UUID.randomUUID().toString(), hostname, UUID.randomUUID().toString(), frameworkID);
-        builder.addResources(cpus(configuration.getCpus(), configuration.getFrameworkRole()));
-        builder.addResources(mem(configuration.getMem(), configuration.getFrameworkRole()));
-        builder.addResources(disk(configuration.getDisk(), configuration.getFrameworkRole()));
-        return builder;
+        return newOfferBuilder(UUID.randomUUID().toString(), hostname, UUID.randomUUID().toString(), frameworkID);
     }
 
     private Protos.MasterInfo newMasterInfo() {
diff --git a/scheduler/src/test/java/org/apache/mesos/elasticsearch/scheduler/OfferStrategyTest.java b/scheduler/src/test/java/org/apache/mesos/elasticsearch/scheduler/OfferStrategyTest.java
