From dd7a42f5827a49c59770f5f59cdd0d0bad15cf45 Mon Sep 17 00:00:00 2001
From: Renan DelValle <rdelval1@binghamton.edu>
Date: Tue, 27 Sep 2016 18:12:50 -0400
Subject: [PATCH] Added ability to tie benchmark to single node. Uses offer
 hostname for this, maybe a better soluition would be to look at offer
 attributes. Added shorthand for workloads flag -w

---
 scheduler.go | 23 +++++++++++++++++++----
 task.go      |  1 +
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/scheduler.go b/scheduler.go
index 9735291..6cfdbe5 100644
--- a/scheduler.go
+++ b/scheduler.go
@@ -11,6 +11,7 @@ import (
 	"os"
 	"time"
 	"bitbucket.org/bingcloud/electron/pcp"
+	"strings"
 )
 
 const (
@@ -179,13 +180,21 @@ func (s *electronScheduler) ResourceOffers(driver sched.SchedulerDriver, offers
 		default:
 		}
 
-
 		tasks := []*mesos.TaskInfo{}
 
 		// First fit strategy
 
 		taken := false
 		for i, task := range s.tasks {
+
+			// Check host if it exists
+			if task.Host != "" {
+				// Don't take offer if it doesn't match our task's host requirement
+				if !strings.HasPrefix(*offer.Hostname, task.Host) {
+					continue
+				}
+			}
+
 			// Decision to take the offer or not
 			if TakeOffer(offer, task) {
 
@@ -274,10 +283,16 @@ func (s *electronScheduler) Error(_ sched.SchedulerDriver, err string) {
 	log.Printf("Receiving an error: %s", err)
 }
 
+var master = flag.String("master", "xavier:5050", "Location of leading Mesos master")
+var tasksFile = flag.String("workload", "", "JSON file containing task definitions")
+var ignoreWatts = flag.Bool("ignoreWatts", false, "Don't use watts from offers")
+
+// Short hand args
+func init(){
+	flag.StringVar(tasksFile, "w", "", "JSON file containing task definitions")
+}
+
 func main() {
-	master := flag.String("master", "xavier:5050", "Location of leading Mesos master")
-	tasksFile := flag.String("workload", "", "JSON file containing task definitions")
-	ignoreWatts := flag.Bool("ignoreWatts", false, "Don't use watts from offers")
 	flag.Parse()
 
 	IGNORE_WATTS = *ignoreWatts
diff --git a/task.go b/task.go
index 9d4bd01..cbc0e5a 100644
--- a/task.go
+++ b/task.go
@@ -14,6 +14,7 @@ type Task struct {
 	Image string `json:"image"`
 	CMD string `json:"cmd"`
 	Instances *int `json:"inst"`
+	Host string `json:"host"`
 }
 
 func TasksFromJSON(uri string) ([]Task, error) {