src/examples/java/V1TestFramework.java - mesos - Git at Google

 /**
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */

 import java.io.File;
 import java.io.IOException;

 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Timer;
 import java.util.TimerTask;

 import java.util.concurrent.locks.*;

 import org.apache.mesos.v1.*;

 import org.apache.mesos.v1.Protos.*;

 import org.apache.mesos.v1.scheduler.Mesos;
 import org.apache.mesos.v1.scheduler.Scheduler;
 import org.apache.mesos.v1.scheduler.V0Mesos;
 import org.apache.mesos.v1.scheduler.V1Mesos;

 import org.apache.mesos.v1.scheduler.Protos.Call;
 import org.apache.mesos.v1.scheduler.Protos.Event;

 public class V1TestFramework {
   static class TestScheduler implements Scheduler {

     public TestScheduler(
         String master,
         FrameworkInfo framework,
         ExecutorInfo executor) {
       this(master, framework, executor, 5);
     }

     public TestScheduler(
         String master,
         FrameworkInfo framework,
         ExecutorInfo executor,
         int totalTasks) {
       this.framework = framework;
       this.executor = executor;
       this.totalTasks = totalTasks;
       this.state = State.DISCONNECTED;
     }

     // TODO(anand): Synchronize on `state` instead.
     @Override
     public synchronized void connected(final Mesos mesos) {
       System.out.println("Connected");

       state = State.CONNECTED;

       retryTimer = new Timer();
       retryTimer.schedule(new TimerTask() {
         @Override
         public void run() {
           doReliableRegistration(mesos);
         }
       }, 0, 1000); // Repeat every 1 second
     }

     @Override
     public synchronized void disconnected(Mesos mesos) {
       System.out.println("Disconnected");

       state = state.DISCONNECTED;
       cancelRetryTimer();
     }

     @Override
     public synchronized void received(Mesos mesos, Event event) {
       switch (event.getType()) {
         case SUBSCRIBED: {
           frameworkId = event.getSubscribed().getFrameworkId();
           state = State.SUBSCRIBED;

           System.out.println(
               "Subscribed with ID " + frameworkId + " to master " +
               event.getSubscribed().getMasterInfo().getId());
           break;
         }

         case OFFERS: {
           System.out.println("Received an OFFERS event");

           offers(mesos, event.getOffers().getOffersList());
           break;
         }

         case RESCIND: {
           System.out.println("Received an RESCIND event");
           break;
         }

         case UPDATE: {
           System.out.println("Received an UPDATE event");

           update(mesos, event.getUpdate().getStatus());
           break;
         }

         case MESSAGE: {
           System.out.println("Received a MESSAGE event");
           break;
         }

         case FAILURE: {
           System.out.println("Received a FAILURE event");
           break;
         }

         case ERROR: {
           System.out.println("Received an ERROR event");
           System.exit(1);
         }

         case HEARTBEAT: {
           // TODO(anand): Force reconnection with the master upon lack
           // of heartbeats.
           System.out.println("Received a HEARTBEAT event");
           break;
         }

         case UNKNOWN: {
           System.out.println("Received an UNKNOWN event");
           break;
         }
       }
     }

     public synchronized void doReliableRegistration(Mesos mesos) {
       if (state == State.SUBSCRIBED || state == State.DISCONNECTED) {
         cancelRetryTimer();
         return;
       }

       Call.Builder callBuilder = Call.newBuilder()
           .setType(Call.Type.SUBSCRIBE)
           .setSubscribe(Call.Subscribe.newBuilder()
             .setFrameworkInfo(framework)
             .build());

       mesos.send(callBuilder.build());
     }

     private void cancelRetryTimer() {
       // Cancel previously active timer (if one exists).
       if (retryTimer != null) {
         retryTimer.cancel();
         retryTimer.purge();
       }

       retryTimer = null;
     }

     public void offers(Mesos mesos, List<Offer> offers) {
       double CPUS_PER_TASK = 1;
       double MEM_PER_TASK = 128;

       for (Offer offer : offers) {
         Offer.Operation.Launch.Builder launch =
           Offer.Operation.Launch.newBuilder();

         double offerCpus = 0;
         double offerMem = 0;
         for (Resource resource : offer.getResourcesList()) {
           if (resource.getName().equals("cpus")) {
             offerCpus += resource.getScalar().getValue();
           } else if (resource.getName().equals("mem")) {
             offerMem += resource.getScalar().getValue();
           }
         }

         System.out.println(
             "Received offer " + offer.getId().getValue() + " with cpus: " +
             offerCpus + " and mem: " + offerMem);

         double remainingCpus = offerCpus;
         double remainingMem = offerMem;
         while (launchedTasks < totalTasks &&
                remainingCpus >= CPUS_PER_TASK &&
                remainingMem >= MEM_PER_TASK) {
           TaskID taskId = TaskID.newBuilder()
             .setValue(Integer.toString(launchedTasks++)).build();

           System.out.println("Launching task " + taskId.getValue() +
                              " using offer " + offer.getId().getValue());

           TaskInfo task = TaskInfo.newBuilder()
             .setName("task " + taskId.getValue())
             .setTaskId(taskId)
             .setAgentId(offer.getAgentId())
             .addResources(Resource.newBuilder()
                           .setName("cpus")
                           .setType(Value.Type.SCALAR)
                           .setScalar(Value.Scalar.newBuilder()
                             .setValue(CPUS_PER_TASK)
                             .build()))
             .addResources(Resource.newBuilder()
                           .setName("mem")
                           .setType(Value.Type.SCALAR)
                           .setScalar(Value.Scalar.newBuilder()
                             .setValue(MEM_PER_TASK)
                             .build()))
             .setExecutor(ExecutorInfo.newBuilder(executor))
             .build();

           launch.addTaskInfos(TaskInfo.newBuilder(task));

           remainingCpus -= CPUS_PER_TASK;
           remainingMem -= MEM_PER_TASK;
         }

         mesos.send(Call.newBuilder()
           .setType(Call.Type.ACCEPT)
           .setFrameworkId(frameworkId)
           .setAccept(Call.Accept.newBuilder()
             .addOfferIds(offer.getId())
             .addOperations(Offer.Operation.newBuilder()
               .setType(Offer.Operation.Type.LAUNCH)
               .setLaunch(launch)
               .build())
             .setFilters(Filters.newBuilder()
               .setRefuseSeconds(1)
               .build()))
           .build());
       }
     }

     public void update(Mesos mesos, TaskStatus status) {
       System.out.println(
           "Status update: task " + status.getTaskId().getValue() +
           " is in state " + status.getState().getValueDescriptor().getName());

       if (status.getState() == TaskState.TASK_FINISHED) {
         finishedTasks++;
         System.out.println("Finished tasks: " + finishedTasks);
         if (finishedTasks == totalTasks) {
           lock.lock();
           try {
             finished = true;
             finishedCondtion.signal();
           } finally {
             lock.unlock();
           }
         }
       }

       if (status.getState() == TaskState.TASK_LOST ||
           status.getState() == TaskState.TASK_KILLED ||
           status.getState() == TaskState.TASK_FAILED) {
         System.err.println(
             "Aborting because task " + status.getTaskId().getValue() +
             " is in unexpected state " +
             status.getState().getValueDescriptor().getName() +
             " with reason '" +
             status.getReason().getValueDescriptor().getName() + "'" +
             " from source '" +
             status.getSource().getValueDescriptor().getName() + "'" +
             " with message '" + status.getMessage() + "'");

         System.exit(1);
       }

       mesos.send(Call.newBuilder()
         .setType(Call.Type.ACKNOWLEDGE)
         .setFrameworkId(frameworkId)
         .setAcknowledge(Call.Acknowledge.newBuilder()
           .setAgentId(status.getAgentId())
           .setTaskId(status.getTaskId())
           .setUuid(status.getUuid())
           .build())
         .build());
     }

     private enum State {
       DISCONNECTED,
       CONNECTED,
       SUBSCRIBED
     }

     private FrameworkInfo framework;
     private FrameworkID frameworkId;
     private final ExecutorInfo executor;
     private final int totalTasks;
     private int launchedTasks = 0;
     private int finishedTasks = 0;
     private State state;
     private Timer retryTimer = null;
   }

   private static void usage() {
     String name = V1TestFramework.class.getName();
     System.err.println("Usage: " + name + " master version{0,1}");
   }

   public static void main(String[] args) throws Exception {
     if (args.length < 2 || args.length > 3) {
       usage();
       System.exit(1);
     }

     int version = Integer.parseInt(args[1]);
     if (version != 0 && version != 1) {
       usage();
       System.exit(1);
     }

     String uri = new File("./test-executor").getCanonicalPath();

     ExecutorInfo executor = ExecutorInfo.newBuilder()
       .setExecutorId(ExecutorID.newBuilder().setValue("default"))
       .setCommand(CommandInfo.newBuilder().setValue(uri))
       .setName("Test Executor (Java)")
       .setSource("java_test")
       .build();

     FrameworkInfo.Builder frameworkBuilder = FrameworkInfo.newBuilder()
         .setUser(System.getProperty("user.name", "default-user"))
         .setName("V" + version + " Test Framework (Java)");

     Credential.Builder credentialBuilder = null;

     if (System.getenv("MESOS_EXAMPLE_PRINCIPAL") != null) {
       frameworkBuilder.setPrincipal(System.getenv("MESOS_EXAMPLE_PRINCIPAL"));

       if (System.getenv("MESOS_EXAMPLE_SECRET") != null) {
         credentialBuilder = Credential.newBuilder()
           .setPrincipal(System.getenv("MESOS_EXAMPLE_PRINCIPAL"))
           .setSecret(System.getenv("MESOS_EXAMPLE_SECRET"));
       }
     }

     Scheduler scheduler = args.length == 2
       ? new TestScheduler(args[0], frameworkBuilder.build(), executor)
       : new TestScheduler(args[0],
                           frameworkBuilder.build(),
                           executor,
                           Integer.parseInt(args[2]));

     Mesos mesos;
     if (credentialBuilder != null) {
       mesos = version == 1
                 ? new V1Mesos(scheduler, args[0], credentialBuilder.build())
                 : new V0Mesos(scheduler,
                               frameworkBuilder.build(),
                               args[0],
                               credentialBuilder.build());
     } else {
       mesos = version == 1
                 ? new V1Mesos(scheduler, args[0])
                 : new V0Mesos(scheduler, frameworkBuilder.build(), args[0]);
     }

     lock.lock();
     try {
       while (!finished) {
         finishedCondtion.await();
       }
     } finally {
       lock.unlock();
     }

     // NOTE: Copied from src/examples/java/TestFramework.java:
     // For this test to pass reliably on some platforms, this sleep is
     // required to ensure that the SchedulerDriver teardown is complete
     // before the JVM starts running native object destructors after
     // System.exit() is called. 500ms proved successful in test runs,
     // but on a heavily loaded machine it might not.

     // TODO(greg): Ideally, we would inspect the status of the driver
     // and its associated tasks via the Java API and wait until their
     // teardown is complete to exit.
     Thread.sleep(500);

     System.exit(0);
   }

   static boolean finished = false;
   final static Lock lock = new ReentrantLock();
   final static Condition finishedCondtion = lock.newCondition();
 }
	/**
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/

	import java.io.File;
	import java.io.IOException;

	import java.util.ArrayList;
	import java.util.HashMap;
	import java.util.List;
	import java.util.Map;
	import java.util.Timer;
	import java.util.TimerTask;

	import java.util.concurrent.locks.*;

	import org.apache.mesos.v1.*;

	import org.apache.mesos.v1.Protos.*;

	import org.apache.mesos.v1.scheduler.Mesos;
	import org.apache.mesos.v1.scheduler.Scheduler;
	import org.apache.mesos.v1.scheduler.V0Mesos;
	import org.apache.mesos.v1.scheduler.V1Mesos;

	import org.apache.mesos.v1.scheduler.Protos.Call;
	import org.apache.mesos.v1.scheduler.Protos.Event;

	public class V1TestFramework {
	static class TestScheduler implements Scheduler {

	public TestScheduler(
	String master,
	FrameworkInfo framework,
	ExecutorInfo executor) {
	this(master, framework, executor, 5);
	}

	public TestScheduler(
	String master,
	FrameworkInfo framework,
	ExecutorInfo executor,
	int totalTasks) {
	this.framework = framework;
	this.executor = executor;
	this.totalTasks = totalTasks;
	this.state = State.DISCONNECTED;
	}

	// TODO(anand): Synchronize on `state` instead.
	@Override
	public synchronized void connected(final Mesos mesos) {
	System.out.println("Connected");

	state = State.CONNECTED;

	retryTimer = new Timer();
	retryTimer.schedule(new TimerTask() {
	@Override
	public void run() {
	doReliableRegistration(mesos);
	}
	}, 0, 1000); // Repeat every 1 second
	}

	@Override
	public synchronized void disconnected(Mesos mesos) {
	System.out.println("Disconnected");

	state = state.DISCONNECTED;
	cancelRetryTimer();
	}

	@Override
	public synchronized void received(Mesos mesos, Event event) {
	switch (event.getType()) {
	case SUBSCRIBED: {
	frameworkId = event.getSubscribed().getFrameworkId();
	state = State.SUBSCRIBED;

	System.out.println(
	"Subscribed with ID " + frameworkId + " to master " +
	event.getSubscribed().getMasterInfo().getId());
	break;
	}

	case OFFERS: {
	System.out.println("Received an OFFERS event");

	offers(mesos, event.getOffers().getOffersList());
	break;
	}

	case RESCIND: {
	System.out.println("Received an RESCIND event");
	break;
	}

	case UPDATE: {
	System.out.println("Received an UPDATE event");

	update(mesos, event.getUpdate().getStatus());
	break;
	}

	case MESSAGE: {
	System.out.println("Received a MESSAGE event");
	break;
	}

	case FAILURE: {
	System.out.println("Received a FAILURE event");
	break;
	}

	case ERROR: {
	System.out.println("Received an ERROR event");
	System.exit(1);
	}

	case HEARTBEAT: {
	// TODO(anand): Force reconnection with the master upon lack
	// of heartbeats.
	System.out.println("Received a HEARTBEAT event");
	break;
	}

	case UNKNOWN: {
	System.out.println("Received an UNKNOWN event");
	break;
	}
	}
	}

	public synchronized void doReliableRegistration(Mesos mesos) {
	if (state == State.SUBSCRIBED \|\| state == State.DISCONNECTED) {
	cancelRetryTimer();
	return;
	}

	Call.Builder callBuilder = Call.newBuilder()
	.setType(Call.Type.SUBSCRIBE)
	.setSubscribe(Call.Subscribe.newBuilder()
	.setFrameworkInfo(framework)
	.build());

	mesos.send(callBuilder.build());
	}

	private void cancelRetryTimer() {
	// Cancel previously active timer (if one exists).
	if (retryTimer != null) {
	retryTimer.cancel();
	retryTimer.purge();
	}

	retryTimer = null;
	}

	public void offers(Mesos mesos, List<Offer> offers) {
	double CPUS_PER_TASK = 1;
	double MEM_PER_TASK = 128;

	for (Offer offer : offers) {
	Offer.Operation.Launch.Builder launch =
	Offer.Operation.Launch.newBuilder();

	double offerCpus = 0;
	double offerMem = 0;
	for (Resource resource : offer.getResourcesList()) {
	if (resource.getName().equals("cpus")) {
	offerCpus += resource.getScalar().getValue();
	} else if (resource.getName().equals("mem")) {
	offerMem += resource.getScalar().getValue();
	}
	}

	System.out.println(
	"Received offer " + offer.getId().getValue() + " with cpus: " +
	offerCpus + " and mem: " + offerMem);

	double remainingCpus = offerCpus;
	double remainingMem = offerMem;
	while (launchedTasks < totalTasks &&
	remainingCpus >= CPUS_PER_TASK &&
	remainingMem >= MEM_PER_TASK) {
	TaskID taskId = TaskID.newBuilder()
	.setValue(Integer.toString(launchedTasks++)).build();

	System.out.println("Launching task " + taskId.getValue() +
	" using offer " + offer.getId().getValue());

	TaskInfo task = TaskInfo.newBuilder()
	.setName("task " + taskId.getValue())
	.setTaskId(taskId)
	.setAgentId(offer.getAgentId())
	.addResources(Resource.newBuilder()
	.setName("cpus")
	.setType(Value.Type.SCALAR)
	.setScalar(Value.Scalar.newBuilder()
	.setValue(CPUS_PER_TASK)
	.build()))
	.addResources(Resource.newBuilder()
	.setName("mem")
	.setType(Value.Type.SCALAR)
	.setScalar(Value.Scalar.newBuilder()
	.setValue(MEM_PER_TASK)
	.build()))
	.setExecutor(ExecutorInfo.newBuilder(executor))
	.build();

	launch.addTaskInfos(TaskInfo.newBuilder(task));

	remainingCpus -= CPUS_PER_TASK;
	remainingMem -= MEM_PER_TASK;
	}

	mesos.send(Call.newBuilder()
	.setType(Call.Type.ACCEPT)
	.setFrameworkId(frameworkId)
	.setAccept(Call.Accept.newBuilder()
	.addOfferIds(offer.getId())
	.addOperations(Offer.Operation.newBuilder()
	.setType(Offer.Operation.Type.LAUNCH)
	.setLaunch(launch)
	.build())
	.setFilters(Filters.newBuilder()
	.setRefuseSeconds(1)
	.build()))
	.build());
	}
	}

	public void update(Mesos mesos, TaskStatus status) {
	System.out.println(
	"Status update: task " + status.getTaskId().getValue() +
	" is in state " + status.getState().getValueDescriptor().getName());

	if (status.getState() == TaskState.TASK_FINISHED) {
	finishedTasks++;
	System.out.println("Finished tasks: " + finishedTasks);
	if (finishedTasks == totalTasks) {
	lock.lock();
	try {
	finished = true;
	finishedCondtion.signal();
	} finally {
	lock.unlock();
	}
	}
	}

	if (status.getState() == TaskState.TASK_LOST \|\|
	status.getState() == TaskState.TASK_KILLED \|\|
	status.getState() == TaskState.TASK_FAILED) {
	System.err.println(
	"Aborting because task " + status.getTaskId().getValue() +
	" is in unexpected state " +
	status.getState().getValueDescriptor().getName() +
	" with reason '" +
	status.getReason().getValueDescriptor().getName() + "'" +
	" from source '" +
	status.getSource().getValueDescriptor().getName() + "'" +
	" with message '" + status.getMessage() + "'");

	System.exit(1);
	}

	mesos.send(Call.newBuilder()
	.setType(Call.Type.ACKNOWLEDGE)
	.setFrameworkId(frameworkId)
	.setAcknowledge(Call.Acknowledge.newBuilder()
	.setAgentId(status.getAgentId())
	.setTaskId(status.getTaskId())
	.setUuid(status.getUuid())
	.build())
	.build());
	}

	private enum State {
	DISCONNECTED,
	CONNECTED,
	SUBSCRIBED
	}

	private FrameworkInfo framework;
	private FrameworkID frameworkId;
	private final ExecutorInfo executor;
	private final int totalTasks;
	private int launchedTasks = 0;
	private int finishedTasks = 0;
	private State state;
	private Timer retryTimer = null;
	}

	private static void usage() {
	String name = V1TestFramework.class.getName();
	System.err.println("Usage: " + name + " master version{0,1}");
	}

	public static void main(String[] args) throws Exception {
	if (args.length < 2 \|\| args.length > 3) {
	usage();
	System.exit(1);
	}

	int version = Integer.parseInt(args[1]);
	if (version != 0 && version != 1) {
	usage();
	System.exit(1);
	}

	String uri = new File("./test-executor").getCanonicalPath();

	ExecutorInfo executor = ExecutorInfo.newBuilder()
	.setExecutorId(ExecutorID.newBuilder().setValue("default"))
	.setCommand(CommandInfo.newBuilder().setValue(uri))
	.setName("Test Executor (Java)")
	.setSource("java_test")
	.build();

	FrameworkInfo.Builder frameworkBuilder = FrameworkInfo.newBuilder()
	.setUser(System.getProperty("user.name", "default-user"))
	.setName("V" + version + " Test Framework (Java)");

	Credential.Builder credentialBuilder = null;

	if (System.getenv("MESOS_EXAMPLE_PRINCIPAL") != null) {
	frameworkBuilder.setPrincipal(System.getenv("MESOS_EXAMPLE_PRINCIPAL"));

	if (System.getenv("MESOS_EXAMPLE_SECRET") != null) {
	credentialBuilder = Credential.newBuilder()
	.setPrincipal(System.getenv("MESOS_EXAMPLE_PRINCIPAL"))
	.setSecret(System.getenv("MESOS_EXAMPLE_SECRET"));
	}
	}

	Scheduler scheduler = args.length == 2
	? new TestScheduler(args[0], frameworkBuilder.build(), executor)
	: new TestScheduler(args[0],
	frameworkBuilder.build(),
	executor,
	Integer.parseInt(args[2]));

	Mesos mesos;
	if (credentialBuilder != null) {
	mesos = version == 1
	? new V1Mesos(scheduler, args[0], credentialBuilder.build())
	: new V0Mesos(scheduler,
	frameworkBuilder.build(),
	args[0],
	credentialBuilder.build());
	} else {
	mesos = version == 1
	? new V1Mesos(scheduler, args[0])
	: new V0Mesos(scheduler, frameworkBuilder.build(), args[0]);
	}

	lock.lock();
	try {
	while (!finished) {
	finishedCondtion.await();
	}
	} finally {
	lock.unlock();
	}

	// NOTE: Copied from src/examples/java/TestFramework.java:
	// For this test to pass reliably on some platforms, this sleep is
	// required to ensure that the SchedulerDriver teardown is complete
	// before the JVM starts running native object destructors after
	// System.exit() is called. 500ms proved successful in test runs,
	// but on a heavily loaded machine it might not.

	// TODO(greg): Ideally, we would inspect the status of the driver
	// and its associated tasks via the Java API and wait until their
	// teardown is complete to exit.
	Thread.sleep(500);

	System.exit(0);
	}

	static boolean finished = false;
	final static Lock lock = new ReentrantLock();
	final static Condition finishedCondtion = lock.newCondition();
	}