Java Code Examples for org.apache.reef.tang.Injector#bindVolatileInstance()

private void runPhysicalPlan(final TestPlanGenerator.PlanType planType,
                             final Injector injector) throws Exception {
  final MetricMessageHandler metricMessageHandler = mock(MetricMessageHandler.class);
  final PlanRewriter planRewriter = mock(PlanRewriter.class);
  final PhysicalPlan plan = TestPlanGenerator.generatePhysicalPlan(planType, false);

  // Get scheduler
  injector.bindVolatileInstance(MetricMessageHandler.class, metricMessageHandler);
  injector.bindVolatileInstance(PubSubEventHandlerWrapper.class, mock(PubSubEventHandlerWrapper.class));
  injector.bindVolatileInstance(SchedulingConstraintRegistry.class, mock(SchedulingConstraintRegistry.class));
  injector.bindVolatileInstance(PlanRewriter.class, planRewriter);
  planStateManager = injector.getInstance(PlanStateManager.class);
  scheduler = injector.getInstance(BatchScheduler.class);
  blockManagerMaster = injector.getInstance(BlockManagerMaster.class);

  scheduler.schedulePlan(plan, MAX_SCHEDULE_ATTEMPT);
}
 

/**
 * Test {@link LocalFileStore}.
 */
@Test(timeout = 10000)
public void testLocalFileStore() throws Exception {
  FileUtils.deleteDirectory(new File(TMP_FILE_DIRECTORY));
  final Injector injector = Tang.Factory.getTang().newInjector();
  injector.bindVolatileParameter(JobConf.FileDirectory.class, TMP_FILE_DIRECTORY);
  injector.bindVolatileInstance(SerializerManager.class, serializerManager);

  final BlockStore localFileStore = injector.getInstance(LocalFileStore.class);
  shuffle(localFileStore, localFileStore);
  concurrentRead(localFileStore, localFileStore);
  shuffleInHashRange(localFileStore, localFileStore);
  FileUtils.deleteDirectory(new File(TMP_FILE_DIRECTORY));
}
 

@Test(timeout = 10000)
public void testState() throws Exception {
  // Create a simple client endpoint that returns given job state.
  final StateTranslator stateTranslator = mock(StateTranslator.class);
  when(stateTranslator.translateState(any())).then(state -> state.getArgument(0));
  final ClientEndpoint clientEndpoint = new TestClientEndpoint(stateTranslator);
  assertEquals(PlanState.State.READY, clientEndpoint.getPlanState());

  // Wait for connection but not connected.
  assertEquals(PlanState.State.READY, clientEndpoint.waitUntilJobFinish(100, TimeUnit.MILLISECONDS));

  // Create a PlanStateManager of a dag and create a DriverEndpoint with it.
  final PhysicalPlan physicalPlan =
    TestPlanGenerator.generatePhysicalPlan(TestPlanGenerator.PlanType.TwoVerticesJoined, false);
  final Injector injector = Tang.Factory.getTang().newInjector();
  injector.bindVolatileInstance(MetricMessageHandler.class, mock(MetricMessageHandler.class));
  final PlanStateManager planStateManager = injector.getInstance(PlanStateManager.class);
  planStateManager.updatePlan(physicalPlan, MAX_SCHEDULE_ATTEMPT);

  final DriverEndpoint driverEndpoint = new DriverEndpoint(planStateManager, clientEndpoint);

  // Check the current state.
  assertEquals(PlanState.State.EXECUTING, clientEndpoint.getPlanState());

  // Wait for the job to finish but not finished
  assertEquals(PlanState.State.EXECUTING, clientEndpoint.waitUntilJobFinish(100, TimeUnit.MILLISECONDS));

  // Check finish.
  final List<String> tasks = physicalPlan.getStageDAG().getTopologicalSort().stream()
    .flatMap(stage -> planStateManager.getTaskAttemptsToSchedule(stage.getId()).stream())
    .collect(Collectors.toList());
  tasks.forEach(taskId -> planStateManager.onTaskStateChanged(taskId, TaskState.State.EXECUTING));
  tasks.forEach(taskId -> planStateManager.onTaskStateChanged(taskId, TaskState.State.COMPLETE));
  assertEquals(PlanState.State.COMPLETE, clientEndpoint.waitUntilJobFinish());
}
 

/**
 * @param port    a listening port number
 * @param factory an identifier factory
 * @param localAddressProvider a local address provider
 * Constructs a name server
 */
@Inject
private NameServerImpl(
    @Parameter(RemoteConfiguration.HostAddress.class) final String hostAddress,
    @Parameter(NameServerParameters.NameServerPort.class) final int port,
    @Parameter(NameServerParameters.NameServerIdentifierFactory.class) final IdentifierFactory factory,
    final TcpPortProvider portProvider,
    final LocalAddressProvider localAddressProvider) {

  final Injector injector = Tang.Factory.getTang().newInjector();

  this.localAddressProvider = localAddressProvider;
  this.reefEventStateManager = null;
  final Codec<NamingMessage> codec = NamingCodecFactory.createFullCodec(factory);
  final EventHandler<NamingMessage> handler = createEventHandler(codec);

  String host = UNKNOWN_HOST_NAME.equals(hostAddress) ? localAddressProvider.getLocalAddress() : hostAddress;
  injector.bindVolatileParameter(RemoteConfiguration.HostAddress.class, host);
  injector.bindVolatileParameter(RemoteConfiguration.Port.class, port);
  injector.bindVolatileInstance(TcpPortProvider.class, portProvider);
  injector.bindVolatileParameter(RemoteConfiguration.RemoteServerStage.class,
      new SyncStage<>(new NamingServerHandler(handler, codec)));

  try {
    this.transport = injector.getInstance(NettyMessagingTransport.class);
  } catch (final InjectionException e) {
    throw new RuntimeException(e);
  }

  this.port = transport.getListeningPort();
  this.idToAddrMap = Collections.synchronizedMap(new HashMap<Identifier, InetSocketAddress>());

  LOG.log(Level.FINE, "NameServer starting, listening at port {0}", this.port);
}
 

private BlockManagerWorker createWorker(final String executorId, final LocalMessageDispatcher messageDispatcher,
                                        final Injector nameClientInjector) {
  final LocalMessageEnvironment messageEnvironment = new LocalMessageEnvironment(executorId, messageDispatcher);
  final PersistentConnectionToMasterMap conToMaster = new PersistentConnectionToMasterMap(messageEnvironment);
  final Configuration executorConfiguration = TANG.newConfigurationBuilder()
      .bindNamedParameter(JobConf.ExecutorId.class, executorId)
      .bindNamedParameter(MessageParameters.SenderId.class, executorId)
      .build();
  final Injector injector = nameClientInjector.forkInjector(executorConfiguration);
  injector.bindVolatileInstance(MessageEnvironment.class, messageEnvironment);
  injector.bindVolatileInstance(PersistentConnectionToMasterMap.class, conToMaster);
  injector.bindVolatileParameter(JobConf.FileDirectory.class, TMP_LOCAL_FILE_DIRECTORY);
  injector.bindVolatileParameter(JobConf.GlusterVolumeDirectory.class, TMP_REMOTE_FILE_DIRECTORY);
  final BlockManagerWorker blockManagerWorker;
  final MetricManagerWorker metricManagerWorker;
  final SerializerManager serializerManager;
  try {
    blockManagerWorker = injector.getInstance(BlockManagerWorker.class);
    metricManagerWorker =  injector.getInstance(MetricManagerWorker.class);
    serializerManager = injector.getInstance(SerializerManager.class);
    serializerManagers.put(blockManagerWorker, serializerManager);
  } catch (final InjectionException e) {
    throw new RuntimeException(e);
  }

  // Unused, but necessary for wiring up the message environments
  final Executor executor = new Executor(
      executorId,
      EXECUTOR_CAPACITY,
      conToMaster,
      messageEnvironment,
      serializerManager,
      new DataTransferFactory(HASH_RANGE_MULTIPLIER, blockManagerWorker),
      metricManagerWorker);
  injector.bindVolatileInstance(Executor.class, executor);

  return blockManagerWorker;
}
 

@Before
public void setUp() throws Exception {
  final LocalMessageDispatcher messageDispatcher = new LocalMessageDispatcher();
  final LocalMessageEnvironment messageEnvironment =
      new LocalMessageEnvironment(MessageEnvironment.MASTER_COMMUNICATION_ID, messageDispatcher);
  final Injector injector = Tang.Factory.getTang().newInjector();
  injector.bindVolatileInstance(MessageEnvironment.class, messageEnvironment);
  blockManagerMaster = injector.getInstance(BlockManagerMaster.class);
}
 

@Before
public void setUp() throws Exception {
  irDAGBuilder = new DAGBuilder<>();
  final LocalMessageDispatcher messageDispatcher = new LocalMessageDispatcher();
  final LocalMessageEnvironment messageEnvironment =
      new LocalMessageEnvironment(MessageEnvironment.MASTER_COMMUNICATION_ID, messageDispatcher);
  final Injector injector = Tang.Factory.getTang().newInjector();
  injector.bindVolatileInstance(MessageEnvironment.class, messageEnvironment);
  blockManagerMaster = injector.getInstance(BlockManagerMaster.class);
  metricMessageHandler = mock(MetricMessageHandler.class);
  injector.bindVolatileParameter(JobConf.DAGDirectory.class, "");
  physicalPlanGenerator = injector.getInstance(PhysicalPlanGenerator.class);
}
 

@Test(timeout = 3000)
public void testState() throws Exception {
  // Create a simple client endpoint that returns given job state.
  final StateTranslator stateTranslator = mock(StateTranslator.class);
  when(stateTranslator.translateState(any())).then(state -> state.getArgument(0));
  final ClientEndpoint clientEndpoint = new TestClientEndpoint(stateTranslator);
  assertEquals(clientEndpoint.getJobState(), JobState.State.READY);

  // Wait for connection but not connected.
  assertEquals(clientEndpoint.waitUntilJobFinish(100, TimeUnit.MILLISECONDS), JobState.State.READY);

  // Create a JobStateManager of an empty dag and create a DriverEndpoint with it.
  final DAGBuilder<IRVertex, IREdge> irDagBuilder = new DAGBuilder<>();
  final DAG<IRVertex, IREdge> irDAG = irDagBuilder.build();
  final Injector injector = Tang.Factory.getTang().newInjector();
  injector.bindVolatileParameter(JobConf.DAGDirectory.class, "");
  final PhysicalPlanGenerator physicalPlanGenerator = injector.getInstance(PhysicalPlanGenerator.class);
  final DAG<PhysicalStage, PhysicalStageEdge> physicalDAG = irDAG.convert(physicalPlanGenerator);

  final LocalMessageDispatcher messageDispatcher = new LocalMessageDispatcher();
  final LocalMessageEnvironment messageEnvironment =
      new LocalMessageEnvironment(MessageEnvironment.MASTER_COMMUNICATION_ID, messageDispatcher);
  injector.bindVolatileInstance(MessageEnvironment.class, messageEnvironment);
  final BlockManagerMaster pmm = injector.getInstance(BlockManagerMaster.class);
  final JobStateManager jobStateManager = new JobStateManager(
      new PhysicalPlan("TestPlan", physicalDAG, physicalPlanGenerator.getTaskIRVertexMap()),
      pmm, metricMessageHandler, MAX_SCHEDULE_ATTEMPT);

  final DriverEndpoint driverEndpoint = new DriverEndpoint(jobStateManager, clientEndpoint);

  // Check the current state.
  assertEquals(clientEndpoint.getJobState(), JobState.State.EXECUTING);

  // Wait for the job to finish but not finished
  assertEquals(clientEndpoint.waitUntilJobFinish(100, TimeUnit.MILLISECONDS), JobState.State.EXECUTING);

  // Check finish.
  jobStateManager.onJobStateChanged(JobState.State.COMPLETE);
  assertEquals(clientEndpoint.waitUntilJobFinish(), JobState.State.COMPLETE);
}
 

private GlusterFileStore createGlusterFileStore(final String executorId)
  throws InjectionException {
  final Injector injector = LocalMessageEnvironment.forkInjector(baseInjector, executorId);
  injector.bindVolatileParameter(JobConf.GlusterVolumeDirectory.class, TMP_FILE_DIRECTORY);
  injector.bindVolatileParameter(JobConf.JobId.class, "GFS test");
  injector.bindVolatileParameter(JobConf.ExecutorId.class, executorId);
  injector.bindVolatileInstance(SerializerManager.class, serializerManager);
  injector.bindVolatileParameter(JobConf.ExecutorMemoryMb.class, 640);
  injector.bindVolatileParameter(JobConf.MaxOffheapRatio.class, 0.2);
  return injector.getInstance(GlusterFileStore.class);
}
 

/**
 * Test {@link LocalFileStore}.
 *
 * @throws Exception exception on the way.
 */
@Test(timeout = 10000)
public void testLocalFileStore() throws Exception {
  FileUtils.deleteDirectory(new File(TMP_FILE_DIRECTORY));
  final Injector injector = Tang.Factory.getTang().newInjector();
  injector.bindVolatileParameter(JobConf.FileDirectory.class, TMP_FILE_DIRECTORY);
  injector.bindVolatileInstance(SerializerManager.class, serializerManager);
  injector.bindVolatileParameter(JobConf.ExecutorMemoryMb.class, 640);
  injector.bindVolatileParameter(JobConf.MaxOffheapRatio.class, 0.2);
  final BlockStore localFileStore = injector.getInstance(LocalFileStore.class);
  shuffle(localFileStore, localFileStore);
  concurrentRead(localFileStore, localFileStore);
  shuffleInHashRange(localFileStore, localFileStore);
  FileUtils.deleteDirectory(new File(TMP_FILE_DIRECTORY));
}
 

/**
 * Test {@link SerializedMemoryStore}.
 *
 * @throws Exception exception on the way.
 */
@Test(timeout = 10000)
public void testSerMemoryStore() throws Exception {
  final Injector injector = Tang.Factory.getTang().newInjector();
  injector.bindVolatileInstance(SerializerManager.class, serializerManager);
  injector.bindVolatileParameter(JobConf.ExecutorMemoryMb.class, 640);
  injector.bindVolatileParameter(JobConf.MaxOffheapRatio.class, 0.2);
  final BlockStore serMemoryStore = injector.getInstance(SerializedMemoryStore.class);
  shuffle(serMemoryStore, serMemoryStore);
  concurrentRead(serMemoryStore, serMemoryStore);
  shuffleInHashRange(serMemoryStore, serMemoryStore);
}
 

/**
 * Test {@link MemoryStore}.
 *
 * @throws Exception exception on the way.
 */
@Test(timeout = 10000)
public void testMemoryStore() throws Exception {
  final Injector injector = Tang.Factory.getTang().newInjector();
  injector.bindVolatileInstance(SerializerManager.class, serializerManager);
  injector.bindVolatileParameter(JobConf.ExecutorMemoryMb.class, 640);
  injector.bindVolatileParameter(JobConf.MaxOffheapRatio.class, 0.2);
  final BlockStore memoryStore = injector.getInstance(MemoryStore.class);
  shuffle(memoryStore, memoryStore);
  concurrentRead(memoryStore, memoryStore);
  shuffleInHashRange(memoryStore, memoryStore);
}
 

private Pair<BlockManagerWorker, IntermediateDataIOFactory> createWorker(
  final String executorId,
  final Injector dispatcherInjector,
  final Injector nameClientInjector) throws InjectionException {
  final Injector messageEnvironmentInjector = LocalMessageEnvironment.forkInjector(dispatcherInjector, executorId);
  final MessageEnvironment messageEnvironment = messageEnvironmentInjector.getInstance(MessageEnvironment.class);
  final PersistentConnectionToMasterMap conToMaster = messageEnvironmentInjector
    .getInstance(PersistentConnectionToMasterMap.class);
  final Configuration executorConfiguration = TANG.newConfigurationBuilder()
    .bindNamedParameter(JobConf.ExecutorId.class, executorId)
    .bindNamedParameter(MessageParameters.SenderId.class, executorId)
    .bindNamedParameter(JobConf.ExecutorMemoryMb.class, "640")
    .bindNamedParameter(JobConf.MaxOffheapRatio.class, "0.2")
    .build();
  final Injector injector = nameClientInjector.forkInjector(executorConfiguration);
  injector.bindVolatileInstance(MessageEnvironment.class, messageEnvironment);
  injector.bindVolatileInstance(PersistentConnectionToMasterMap.class, conToMaster);
  injector.bindVolatileParameter(JobConf.FileDirectory.class, TMP_LOCAL_FILE_DIRECTORY);
  injector.bindVolatileParameter(JobConf.GlusterVolumeDirectory.class, TMP_REMOTE_FILE_DIRECTORY);
  final BlockManagerWorker blockManagerWorker;
  final SerializerManager serializerManager;
  final IntermediateDataIOFactory intermediateDataIOFactory;
  try {
    blockManagerWorker = injector.getInstance(BlockManagerWorker.class);
    serializerManager = injector.getInstance(SerializerManager.class);
    serializerManagers.put(blockManagerWorker, serializerManager);
    intermediateDataIOFactory = injector.getInstance(IntermediateDataIOFactory.class);
  } catch (final InjectionException e) {
    throw new RuntimeException(e);
  }

  // Unused, but necessary for wiring up the message environments
  injector.getInstance(Executor.class);

  return Pair.of(blockManagerWorker, intermediateDataIOFactory);
}
 

@Before
public void setUp() throws InjectionException {

  final MessageEnvironment mockMsgEnv = mock(MessageEnvironment.class);
  when(mockMsgEnv.asyncConnect(anyString(), anyString())).thenReturn(mock(Future.class));
  final Configuration configuration = Tang.Factory.getTang().newConfigurationBuilder()
    .bindNamedParameter(JobConf.ScheduleSerThread.class, "1")
    .build();
  final Injector injector = Tang.Factory.getTang().newInjector(configuration);
  injector.bindVolatileInstance(EvaluatorRequestor.class, mock(EvaluatorRequestor.class));
  injector.bindVolatileInstance(MessageEnvironment.class, mockMsgEnv);
  containerManager = injector.getInstance(ContainerManager.class);
}
 

/**
 * Initializes injector by copying needed handlers.
 * @param runtimeInjector The injector to initialize
 * @throws InjectionException on configuration error.
 */
private void initializeInjector(final Injector runtimeInjector) throws InjectionException {

  copyEventHandler(runtimeInjector, RuntimeParameters.ResourceStatusHandler.class);
  copyEventHandler(runtimeInjector, RuntimeParameters.NodeDescriptorHandler.class);
  copyEventHandler(runtimeInjector, RuntimeParameters.ResourceAllocationHandler.class);
  copyEventHandler(runtimeInjector, RuntimeParameters.RuntimeStatusHandler.class);

  try {
    runtimeInjector.bindVolatileInstance(HttpServer.class, this.originalInjector.getInstance(HttpServer.class));
    LOG.log(Level.INFO, "Binding http server for the runtime implementation");
  } catch (final InjectionException e) {
    LOG.log(Level.INFO, "Http Server is not configured for the runtime", e);
  }
}
 

/**
 * NetworkService messaging test.
 */
@Test
public void testMessagingNetworkService() throws Exception {
  LOG.log(Level.FINEST, name.getMethodName());

  final IdentifierFactory factory = new StringIdentifierFactory();

  final Injector injector = Tang.Factory.getTang().newInjector();
  injector.bindVolatileParameter(NameServerParameters.NameServerIdentifierFactory.class, factory);
  injector.bindVolatileInstance(LocalAddressProvider.class, this.localAddressProvider);

  try (NameServer server = injector.getInstance(NameServer.class)) {
    final int nameServerPort = server.getPort();

    final int numMessages = 10;
    final Monitor monitor = new Monitor();

    // network service
    final String name2 = "task2";
    final String name1 = "task1";
    final Configuration nameResolverConf =
        Tang.Factory.getTang().newConfigurationBuilder(NameResolverConfiguration.CONF
        .set(NameResolverConfiguration.NAME_SERVER_HOSTNAME, this.localAddress)
        .set(NameResolverConfiguration.NAME_SERVICE_PORT, nameServerPort)
        .build())
        .build();

    final Injector injector2 = Tang.Factory.getTang().newInjector(nameResolverConf);

    LOG.log(Level.FINEST, "=== Test network service receiver start");
    LOG.log(Level.FINEST, "=== Test network service sender start");
    try (NameResolver nameResolver = injector2.getInstance(NameResolver.class)) {
      injector2.bindVolatileParameter(NetworkServiceParameters.NetworkServiceIdentifierFactory.class, factory);
      injector2.bindVolatileInstance(NameResolver.class, nameResolver);
      injector2.bindVolatileParameter(NetworkServiceParameters.NetworkServiceCodec.class, new StringCodec());
      injector2.bindVolatileParameter(NetworkServiceParameters.NetworkServiceTransportFactory.class,
          injector.getInstance(MessagingTransportFactory.class));
      injector2.bindVolatileParameter(NetworkServiceParameters.NetworkServiceExceptionHandler.class,
          new ExceptionHandler());

      final Injector injectorNs2 = injector2.forkInjector();
      injectorNs2.bindVolatileParameter(NetworkServiceParameters.NetworkServiceHandler.class,
          new MessageHandler<String>(name2, monitor, numMessages));
      final NetworkService<String> ns2 = injectorNs2.getInstance(NetworkService.class);

      final Injector injectorNs1 = injector2.forkInjector();
      injectorNs1.bindVolatileParameter(NetworkServiceParameters.NetworkServiceHandler.class,
          new MessageHandler<String>(name1, null, 0));
      final NetworkService<String> ns1 = injectorNs1.getInstance(NetworkService.class);

      ns2.registerId(factory.getNewInstance(name2));
      final int port2 = ns2.getTransport().getListeningPort();
      server.register(factory.getNewInstance("task2"), new InetSocketAddress(this.localAddress, port2));

      ns1.registerId(factory.getNewInstance(name1));
      final int port1 = ns1.getTransport().getListeningPort();
      server.register(factory.getNewInstance("task1"), new InetSocketAddress(this.localAddress, port1));

      final Identifier destId = factory.getNewInstance(name2);

      try (Connection<String> conn = ns1.newConnection(destId)) {
        conn.open();
        for (int count = 0; count < numMessages; ++count) {
          conn.write("hello! " + count);
        }
        monitor.mwait();

      } catch (final NetworkException e) {
        e.printStackTrace();
        throw new RuntimeException(e);
      }
    }
  }
}
 

@Test
public void testMultithreadedSharedConnMessagingNetworkServiceRate() throws Exception {

  Assume.assumeFalse("Use log level INFO to run benchmarking", LOG.isLoggable(Level.FINEST));

  LOG.log(Level.FINEST, name.getMethodName());

  final IdentifierFactory factory = new StringIdentifierFactory();

  final Injector injector = Tang.Factory.getTang().newInjector();
  injector.bindVolatileParameter(NameServerParameters.NameServerIdentifierFactory.class, factory);
  injector.bindVolatileInstance(LocalAddressProvider.class, this.localAddressProvider);
  try (NameServer server = injector.getInstance(NameServer.class)) {
    final int nameServerPort = server.getPort();

    final int[] messageSizes = {2000}; // {1,16,32,64,512,64*1024,1024*1024};

    for (final int size : messageSizes) {
      final int numMessages = 300000 / (Math.max(1, size / 512));
      final int numThreads = 2;
      final int totalNumMessages = numMessages * numThreads;
      final Monitor monitor = new Monitor();


      // network service
      final String name2 = "task2";
      final String name1 = "task1";
      final Configuration nameResolverConf =
          Tang.Factory.getTang().newConfigurationBuilder(NameResolverConfiguration.CONF
          .set(NameResolverConfiguration.NAME_SERVER_HOSTNAME, this.localAddress)
          .set(NameResolverConfiguration.NAME_SERVICE_PORT, nameServerPort)
          .build())
          .build();

      final Injector injector2 = Tang.Factory.getTang().newInjector(nameResolverConf);

      LOG.log(Level.FINEST, "=== Test network service receiver start");
      LOG.log(Level.FINEST, "=== Test network service sender start");
      try (NameResolver nameResolver = injector2.getInstance(NameResolver.class)) {
        injector2.bindVolatileParameter(NetworkServiceParameters.NetworkServiceIdentifierFactory.class, factory);
        injector2.bindVolatileInstance(NameResolver.class, nameResolver);
        injector2.bindVolatileParameter(NetworkServiceParameters.NetworkServiceCodec.class, new StringCodec());
        injector2.bindVolatileParameter(NetworkServiceParameters.NetworkServiceTransportFactory.class,
            injector.getInstance(MessagingTransportFactory.class));
        injector2.bindVolatileParameter(NetworkServiceParameters.NetworkServiceExceptionHandler.class,
            new ExceptionHandler());

        final Injector injectorNs2 = injector2.forkInjector();
        injectorNs2.bindVolatileParameter(NetworkServiceParameters.NetworkServiceHandler.class,
            new MessageHandler<String>(name2, monitor, totalNumMessages));
        final NetworkService<String> ns2 = injectorNs2.getInstance(NetworkService.class);

        final Injector injectorNs1 = injector2.forkInjector();
        injectorNs1.bindVolatileParameter(NetworkServiceParameters.NetworkServiceHandler.class,
            new MessageHandler<String>(name1, null, 0));
        final NetworkService<String> ns1 = injectorNs1.getInstance(NetworkService.class);

        ns2.registerId(factory.getNewInstance(name2));
        final int port2 = ns2.getTransport().getListeningPort();
        server.register(factory.getNewInstance("task2"), new InetSocketAddress(this.localAddress, port2));

        ns1.registerId(factory.getNewInstance(name1));
        final int port1 = ns1.getTransport().getListeningPort();
        server.register(factory.getNewInstance("task1"), new InetSocketAddress(this.localAddress, port1));

        final Identifier destId = factory.getNewInstance(name2);

        try (Connection<String> conn = ns1.newConnection(destId)) {
          conn.open();

          final String message = StringUtils.repeat('1', size);
          final ExecutorService e = Executors.newCachedThreadPool();

          final long start = System.currentTimeMillis();
          for (int i = 0; i < numThreads; i++) {
            e.submit(new Runnable() {

              @Override
              public void run() {
                for (int i = 0; i < numMessages; i++) {
                  conn.write(message);
                }
              }
            });
          }


          e.shutdown();
          e.awaitTermination(30, TimeUnit.SECONDS);
          monitor.mwait();

          final long end = System.currentTimeMillis();
          final double runtime = ((double) end - start) / 1000;

          LOG.log(Level.FINEST, "size: " + size + "; messages/s: " + totalNumMessages / runtime + 
              " bandwidth(bytes/s): " + ((double) totalNumMessages * 2 * size) / runtime); // x2 for unicode chars
        }
      }
    }
  }
}
 

/**
 * Test method for {@link org.apache.reef.io.network.naming.NameClient#lookup()}.
 * To check caching behavior with expireAfterAccess & expireAfterWrite
 * Changing NameCache's pattern to expireAfterAccess causes this test to fail
 *
 * @throws Exception
 */
@Test
public final void testLookup() throws Exception {
  final String localAddress = localAddressProvider.getLocalAddress();
  final IdentifierFactory factory = new StringIdentifierFactory();
  final Injector injector = Tang.Factory.getTang().newInjector();
  injector.bindVolatileParameter(NameServerParameters.NameServerIdentifierFactory.class, factory);
  injector.bindVolatileInstance(LocalAddressProvider.class, this.localAddressProvider);

  try (NameServer server = injector.getInstance(NameServer.class)) {
    final int serverPort = server.getPort();
    final Configuration nameResolverConf = NameResolverConfiguration.CONF
        .set(NameResolverConfiguration.NAME_SERVER_HOSTNAME, localAddress)
        .set(NameResolverConfiguration.NAME_SERVICE_PORT, serverPort)
        .set(NameResolverConfiguration.CACHE_TIMEOUT, 150)
        .set(NameResolverConfiguration.RETRY_TIMEOUT, RETRY_TIMEOUT)
        .set(NameResolverConfiguration.RETRY_COUNT, RETRY_COUNT)
        .build();

    try (NameResolver client =
             Tang.Factory.getTang().newInjector(nameResolverConf).getInstance(NameClient.class)) {
      final Identifier id = factory.getNewInstance("Task1");
      client.register(id, new InetSocketAddress(localAddress, 7001));
      client.lookup(id); // caches the entry
      client.unregister(id);
      Thread.sleep(100);
      try {
        InetSocketAddress addr = client.lookup(id);
        Thread.sleep(100);
        //With expireAfterAccess, the previous lookup would reset expiry to 150ms
        //more and 100ms wait will not expire the item and will return the cached value
        //With expireAfterWrite, the extra wait of 100 ms will expire the item
        //resulting in NamingException and the test passes
        addr = client.lookup(id);
        Assert.assertNull("client.lookup(id)", addr);
      } catch (final Exception e) {
        if (e instanceof ExecutionException) {
          Assert.assertTrue("Execution Exception cause is instanceof NamingException",
              e.getCause() instanceof NamingException);
        } else {
          throw e;
        }
      }
    }
  }
}
 

/**
 * NetworkService messaging rate benchmark.
 */
@Test
public void testMessagingNetworkServiceRate() throws Exception {

  Assume.assumeFalse("Use log level INFO to run benchmarking", LOG.isLoggable(Level.FINEST));

  LOG.log(Level.FINEST, name.getMethodName());

  final IdentifierFactory factory = new StringIdentifierFactory();

  final Injector injector = Tang.Factory.getTang().newInjector();
  injector.bindVolatileParameter(NameServerParameters.NameServerIdentifierFactory.class, factory);
  injector.bindVolatileInstance(LocalAddressProvider.class, this.localAddressProvider);

  try (NameServer server = injector.getInstance(NameServer.class)) {
    final int nameServerPort = server.getPort();

    final int[] messageSizes = {1, 16, 32, 64, 512, 64 * 1024, 1024 * 1024};

    for (final int size : messageSizes) {
      final int numMessages = 300000 / (Math.max(1, size / 512));
      final Monitor monitor = new Monitor();

      // network service
      final String name2 = "task2";
      final String name1 = "task1";
      final Configuration nameResolverConf =
          Tang.Factory.getTang().newConfigurationBuilder(NameResolverConfiguration.CONF
          .set(NameResolverConfiguration.NAME_SERVER_HOSTNAME, this.localAddress)
          .set(NameResolverConfiguration.NAME_SERVICE_PORT, nameServerPort)
          .build())
          .build();

      final Injector injector2 = Tang.Factory.getTang().newInjector(nameResolverConf);

      LOG.log(Level.FINEST, "=== Test network service receiver start");
      LOG.log(Level.FINEST, "=== Test network service sender start");
      try (NameResolver nameResolver = injector2.getInstance(NameResolver.class)) {
        injector2.bindVolatileParameter(NetworkServiceParameters.NetworkServiceIdentifierFactory.class, factory);
        injector2.bindVolatileInstance(NameResolver.class, nameResolver);
        injector2.bindVolatileParameter(NetworkServiceParameters.NetworkServiceCodec.class, new StringCodec());
        injector2.bindVolatileParameter(NetworkServiceParameters.NetworkServiceTransportFactory.class,
            injector.getInstance(MessagingTransportFactory.class));
        injector2.bindVolatileParameter(NetworkServiceParameters.NetworkServiceExceptionHandler.class,
            new ExceptionHandler());

        final Injector injectorNs2 = injector2.forkInjector();
        injectorNs2.bindVolatileParameter(NetworkServiceParameters.NetworkServiceHandler.class,
            new MessageHandler<String>(name2, monitor, numMessages));
        final NetworkService<String> ns2 = injectorNs2.getInstance(NetworkService.class);

        final Injector injectorNs1 = injector2.forkInjector();
        injectorNs1.bindVolatileParameter(NetworkServiceParameters.NetworkServiceHandler.class,
            new MessageHandler<String>(name1, null, 0));
        final NetworkService<String> ns1 = injectorNs1.getInstance(NetworkService.class);

        ns2.registerId(factory.getNewInstance(name2));
        final int port2 = ns2.getTransport().getListeningPort();
        server.register(factory.getNewInstance("task2"), new InetSocketAddress(this.localAddress, port2));

        ns1.registerId(factory.getNewInstance(name1));
        final int port1 = ns1.getTransport().getListeningPort();
        server.register(factory.getNewInstance("task1"), new InetSocketAddress(this.localAddress, port1));

        final Identifier destId = factory.getNewInstance(name2);
        final String message = StringUtils.repeat('1', size);

        final long start = System.currentTimeMillis();
        try (Connection<String> conn = ns1.newConnection(destId)) {
          conn.open();
          for (int i = 0; i < numMessages; i++) {
            conn.write(message);
          }
          monitor.mwait();
        } catch (final NetworkException e) {
          e.printStackTrace();
          throw new RuntimeException(e);
        }
        final long end = System.currentTimeMillis();
        final double runtime = ((double) end - start) / 1000;
        LOG.log(Level.FINEST, "size: " + size + "; messages/s: " + numMessages / runtime +
            " bandwidth(bytes/s): " + ((double) numMessages * 2 * size) / runtime); // x2 for unicode chars
      }
    }
  }
}
 

@Before
public void setUp() throws Exception {
  final Injector injector = Tang.Factory.getTang().newInjector();
  final PlanRewriter planRewriter = mock(PlanRewriter.class);
  injector.bindVolatileInstance(PlanRewriter.class, planRewriter);
  injector.bindVolatileParameter(JobConf.DAGDirectory.class, "");

  executorRegistry = injector.getInstance(ExecutorRegistry.class);
  injector.bindVolatileInstance(BlockManagerMaster.class, mock(BlockManagerMaster.class));
  injector.bindVolatileInstance(PubSubEventHandlerWrapper.class, mock(PubSubEventHandlerWrapper.class));
  injector.bindVolatileInstance(MetricMessageHandler.class, mock(MetricMessageHandler.class));
  planStateManager = injector.getInstance(PlanStateManager.class);
  scheduler = injector.getInstance(BatchScheduler.class);

  final ActiveContext activeContext = mock(ActiveContext.class);
  Mockito.doThrow(new RuntimeException()).when(activeContext).close();

  final ExecutorService serializationExecutorService = Executors.newSingleThreadExecutor();
  final ResourceSpecification computeSpec =
    new ResourceSpecification(ResourcePriorityProperty.COMPUTE, EXECUTOR_CAPACITY, 0);
  final Function<String, ExecutorRepresenter> computeSpecExecutorRepresenterGenerator = executorId ->
    new DefaultExecutorRepresenter(executorId, computeSpec, mockMsgSender, activeContext, serializationExecutorService,
      executorId);
  final ExecutorRepresenter a3 = computeSpecExecutorRepresenterGenerator.apply("a3");
  final ExecutorRepresenter a2 = computeSpecExecutorRepresenterGenerator.apply("a2");
  final ExecutorRepresenter a1 = computeSpecExecutorRepresenterGenerator.apply("a1");

  final ResourceSpecification storageSpec =
    new ResourceSpecification(ResourcePriorityProperty.TRANSIENT, EXECUTOR_CAPACITY, 0);
  final Function<String, ExecutorRepresenter> storageSpecExecutorRepresenterGenerator = executorId ->
    new DefaultExecutorRepresenter(executorId, storageSpec, mockMsgSender, activeContext, serializationExecutorService,
      executorId);
  final ExecutorRepresenter b2 = storageSpecExecutorRepresenterGenerator.apply("b2");
  final ExecutorRepresenter b1 = storageSpecExecutorRepresenterGenerator.apply("b1");

  // Add compute nodes
  scheduler.onExecutorAdded(a1);
  scheduler.onExecutorAdded(a2);
  scheduler.onExecutorAdded(a3);

  // Add storage nodes
  scheduler.onExecutorAdded(b1);
  scheduler.onExecutorAdded(b2);
}