Java Code Examples for org.apache.reef.wake.EventHandler#onNext()

@Override
public void onNext(final TransportEvent transportEvent) {
  final byte[] bytes = transportEvent.getData();
  final ControlMessage.DriverToClientMessage message;
  try {
    message = ControlMessage.DriverToClientMessage.parseFrom(bytes);
  } catch (final InvalidProtocolBufferException e) {
    throw new RuntimeException(e);
  }

  final ControlMessage.DriverToClientMessageType type = message.getType();

  if (type == ControlMessage.DriverToClientMessageType.DriverStarted) {
    link = transportEvent.getLink();
  }

  final EventHandler<ControlMessage.DriverToClientMessage> handler = handlers.get(type);
  if (handler == null) {
    throw new RuntimeException(String.format("Handler for message type %s not registered", type));
  } else {
    handler.onNext(message);
  }
}
 

/**
 * Event handler method for {@link ResourceReleaseEvent}. Sends a TERMINATE command to the appropriate evaluator shim.
 *
 * @param resourceReleaseEvent
 */
public void onResourceReleased(final ResourceReleaseEvent resourceReleaseEvent) {
  String resourceRemoteId = getResourceRemoteId(resourceReleaseEvent.getIdentifier());

  // REEF Common will trigger a ResourceReleaseEvent even if the resource has failed. Since we know that the shim
  // has already failed, we can safely ignore this.
  if (this.failedResources.remove(resourceReleaseEvent.getIdentifier()) != null) {
    LOG.log(Level.INFO, "Received a ResourceReleaseEvent for a failed shim with resourceId = {0}. Ignoring.",
        resourceReleaseEvent.getIdentifier());
  } else if (this.evaluators.get(resourceReleaseEvent.getIdentifier()).isPresent()) {
    EventHandler<EvaluatorShimProtocol.EvaluatorShimControlProto> handler =
        this.remoteManager.getHandler(resourceRemoteId, EvaluatorShimProtocol.EvaluatorShimControlProto.class);

    LOG.log(Level.INFO, "Sending TERMINATE command to the shim with remoteId = {0}.", resourceRemoteId);
    handler.onNext(
        EvaluatorShimProtocol.EvaluatorShimControlProto
            .newBuilder()
            .setCommand(EvaluatorShimProtocol.EvaluatorShimCommand.TERMINATE)
            .build());

    this.updateRuntimeStatus();
  }
}
 

/**
 * Sets the driver restart status to be completed if not yet set and notifies the restart completed event handlers.
 */
private synchronized void onDriverRestartCompleted(final boolean isTimedOut) {
  if (this.state != DriverRestartState.COMPLETED) {
    final Set<String> outstandingEvaluatorIds = getOutstandingEvaluatorsAndMarkExpired();
    driverRuntimeRestartManager.informAboutEvaluatorFailures(outstandingEvaluatorIds);

    this.state = DriverRestartState.COMPLETED;
    final DriverRestartCompleted driverRestartCompleted = new DriverRestartCompletedImpl(
        System.currentTimeMillis(), isTimedOut);

    for (final EventHandler<DriverRestartCompleted> serviceRestartCompletedHandler
        : this.serviceDriverRestartCompletedHandlers) {
      serviceRestartCompletedHandler.onNext(driverRestartCompleted);
    }

    for (final EventHandler<DriverRestartCompleted> restartCompletedHandler : this.driverRestartCompletedHandlers) {
      restartCompletedHandler.onNext(driverRestartCompleted);
    }

    LOG.log(Level.FINE, "Restart completed. Evaluators that have not reported back are: " + outstandingEvaluatorIds);
  }

  restartCompletedTimer.cancel();
}
 

/**
 * Handles message from client.
 *
 * @param message message to process
 */
private void handleMessage(final ControlMessage.ClientToDriverMessage message) {
  final ControlMessage.ClientToDriverMessageType type = message.getType();
  final EventHandler<ControlMessage.ClientToDriverMessage> handler = handlers.get(type);
  if (handler == null) {
    throw new RuntimeException(String.format("Handler for message type %s not registered", type));
  } else {
    handler.onNext(message);
  }
}
 

/**
 * Verify Avro message can be serialized and deserialized
 * between two remote managers.
 */
@Test
public void testProtocolSerializerTest() throws InterruptedException {

  final int[] numbers = {12, 25};
  final String[] strings = {"The first string", "The second string"};

  // Queues for storing messages byte messages.
  final BlockingQueue<byte[]> queue1 = new LinkedBlockingQueue<>();
  final BlockingQueue<byte[]> queue2 = new LinkedBlockingQueue<>();

  // Remote managers for sending and receiving byte messages.
  final RemoteManager remoteManager1 = remoteManagerFactory.getInstance("RemoteManagerOne");
  final RemoteManager remoteManager2 = remoteManagerFactory.getInstance("RemoteManagerTwo");

  // Register message handlers for byte level messages.
  remoteManager1.registerHandler(byte[].class, new ByteMessageObserver(queue1));
  remoteManager2.registerHandler(byte[].class, new ByteMessageObserver(queue2));

  final EventHandler<byte[]> sender1 = remoteManager1.getHandler(remoteManager2.getMyIdentifier(), byte[].class);
  final EventHandler<byte[]> sender2 = remoteManager2.getHandler(remoteManager1.getMyIdentifier(), byte[].class);

  sender1.onNext(serializer.write(new AvroTestMessage(numbers[0], strings[0]), 1));
  sender2.onNext(serializer.write(new AvroTestMessage(numbers[1], strings[1]), 2));

  final AvroMessageObserver avroObserver1 = new AvroMessageObserver();
  final AvroMessageObserver avroObserver2 = new AvroMessageObserver();

  serializer.read(queue1.take(), avroObserver1);
  serializer.read(queue2.take(), avroObserver2);

  assertEquals(numbers[0], avroObserver2.getNumber());
  assertEquals(strings[0], avroObserver2.getDataString());

  assertEquals(numbers[1], avroObserver1.getNumber());
  assertEquals(strings[1], avroObserver1.getDataString());
}
 

/**
 * Invokes a specific handler for the event class type if it exists.
 *
 * @param event an event
 * @throws WakeRuntimeException
 */
@Override
public void onNext(final T event) {
  final EventHandler<T> handler = (EventHandler<T>) map.get(event.getClass());
  if (handler == null) {
    throw new WakeRuntimeException("No event " + event.getClass() + " handler");
  }
  handler.onNext(event);
}
 

/**
 * We must implement this synchronously in order to catch exceptions and
 * forward them back via the bridge before the server shuts down, after
 * this method returns.
 * @param stopTime stop time
 */
@SuppressWarnings("checkstyle:illegalCatch")
public Throwable dispatch(final StopTime stopTime) {
  try {
    for (final EventHandler<StopTime> handler : stopHandlers) {
      handler.onNext(stopTime);
    }
    return null;
  } catch (Throwable t) {
    return t;
  }
}
 

@Test
public void testRemoteManagerPBufTest() throws Exception {
  System.out.println(LOG_PREFIX + name.getMethodName());
  LoggingUtils.setLoggingLevel(Level.INFO);

  final Monitor monitor = new Monitor();
  final TimerStage timer = new TimerStage(new TimeoutHandler(monitor), 2000, 2000);

  final Map<Class<?>, Codec<?>> clazzToCodecMap = new HashMap<>();
  clazzToCodecMap.put(TestEvent.class, new TestEventCodec());
  final Codec<?> codec = new MultiCodec<Object>(clazzToCodecMap);

  final String hostAddress = localAddressProvider.getLocalAddress();

  final RemoteManager rm = this.remoteManagerFactory.getInstance(
      "name", hostAddress, 0, codec, new LoggingEventHandler<Throwable>(), false, 3, 10000,
      localAddressProvider, Tang.Factory.getTang().newInjector().getInstance(TcpPortProvider.class));

  final RemoteIdentifier remoteId = rm.getMyIdentifier();

  final EventHandler<TestEvent> proxyHandler = rm.getHandler(remoteId, TestEvent.class);

  final AtomicInteger counter = new AtomicInteger(0);
  final int finalSize = 0;
  rm.registerHandler(TestEvent.class, new MessageTypeEventHandler<TestEvent>(rm, monitor, counter, finalSize));

  proxyHandler.onNext(new TestEvent("hello", 0.0));

  monitor.mwait();

  Assert.assertEquals(finalSize, counter.get());

  rm.close();
  timer.close();
}
 

/**
 * Sends the TaskStop event to the handlers for it.
 */
@SuppressWarnings("checkstyle:illegalcatch")
public void afterTaskExit() throws TaskStopHandlerFailure {
  LOG.log(Level.FINEST, "Sending TaskStop event to the registered event handlers.");
  for (final EventHandler<TaskStop> stopHandler : this.taskStopHandlers) {
    try {
      stopHandler.onNext(this.taskStop);
    } catch (final Throwable throwable) {
      throw new TaskStopHandlerFailure(stopHandler, throwable);
    }
  }
  LOG.log(Level.FINEST, "Done sending TaskStop event to the registered event handlers.");
}
 

/**
 * Sends the TaskStart event to the handlers for it.
 */
@SuppressWarnings("checkstyle:illegalcatch")
public void beforeTaskStart() throws TaskStartHandlerFailure {
  LOG.log(Level.FINEST, "Sending TaskStart event to the registered event handlers.");
  for (final EventHandler<TaskStart> startHandler : this.taskStartHandlers) {
    try {
      startHandler.onNext(this.taskStart);
    } catch (final Throwable throwable) {
      throw new TaskStartHandlerFailure(startHandler, throwable);
    }
  }
  LOG.log(Level.FINEST, "Done sending TaskStart event to the registered event handlers.");
}
 

/**
 * Fires ContextStop to all registered event handlers.
 */
void close() {
  final ContextStop contextStop = new ContextStopImpl(this.identifier);
  for (final EventHandler<ContextStop> stopHandler : this.contextStopHandlers) {
    stopHandler.onNext(contextStop);
  }
}
 

@Override
public Integer call() throws Exception {

  final Monitor monitor = new Monitor();
  final TimerStage timer = new TimerStage(new TimeoutHandler(monitor), timeout, timeout);

  final String hostAddress = localAddressProvider.getLocalAddress();
  final RemoteIdentifierFactory factory = new DefaultRemoteIdentifierFactoryImplementation();
  final RemoteIdentifier remoteId = factory.getNewInstance("socket://" + hostAddress + ":" + remotePort);

  final EventHandler<StartEvent> proxyConnection = rm.getHandler(remoteId, StartEvent.class);
  final EventHandler<TestEvent1> proxyHandler1 = rm.getHandler(remoteId, TestEvent1.class);
  final EventHandler<TestEvent2> proxyHandler2 = rm.getHandler(remoteId, TestEvent2.class);

  final AtomicInteger counter = new AtomicInteger(0);
  final int finalSize = 0;
  rm.registerHandler(StartEvent.class, new MessageTypeEventHandler<StartEvent>(rm, monitor, counter, finalSize));

  proxyConnection.onNext(new StartEvent());

  monitor.mwait();

  proxyHandler1.onNext(new TestEvent1("hello1", 0.0)); // registration after send expected to fail
  proxyHandler2.onNext(new TestEvent2("hello2", 0.0)); // registration after send expected to fail
  timer.close();

  return counter.get();
}
 

/**
 * Recovers the list of alive and failed evaluators and inform the driver restart handlers and inform the
 * evaluator failure handlers based on the specific runtime. Also sets the expected amount of evaluators to report
 * back as alive to the job driver.
 */
public synchronized void onRestart(final StartTime startTime,
                                   final List<EventHandler<DriverRestarted>> orderedHandlers) {
  if (this.state == DriverRestartState.BEGAN) {
    restartEvaluators = driverRuntimeRestartManager.getPreviousEvaluators();
    final DriverRestarted restartedInfo = new DriverRestartedImpl(resubmissionAttempts, startTime, restartEvaluators);

    for (final EventHandler<DriverRestarted> handler : orderedHandlers) {
      handler.onNext(restartedInfo);
    }

    this.state = DriverRestartState.IN_PROGRESS;
  } else {
    final String errMsg = "Should not be setting the set of expected alive evaluators more than once.";
    LOG.log(Level.SEVERE, errMsg);
    throw new DriverFatalRuntimeException(errMsg);
  }

  driverRuntimeRestartManager.informAboutEvaluatorFailures(getFailedEvaluators());

  if (driverRestartEvaluatorRecoverySeconds != Integer.MAX_VALUE) {
    // Don't use Clock here because if there is an event scheduled, the driver will not be idle, even if
    // driver restart has already completed, and we cannot cancel the event.
    restartCompletedTimer.schedule(new TimerTask() {
      @Override
      public void run() {
        onDriverRestartCompleted(true);
      }
    }, driverRestartEvaluatorRecoverySeconds * 1000L);
  }
}
 

@Override
public void onNext(final TransportEvent transportEvent) {
  final NetworkConnectionServiceMessage nsMessage = codec.decode(transportEvent.getData());
  nsMessage.setRemoteAddress(transportEvent.getRemoteAddress());
  final NetworkConnectionFactory connFactory = connFactoryMap.get(nsMessage.getConnectionFactoryId());
  final EventHandler eventHandler = connFactory.getEventHandler();
  eventHandler.onNext(nsMessage);
}
 

@Override
public void onNext(final E event) {
  for (final EventHandler<E> handler : handlers) {
    handler.onNext(event);
  }
}
 

private void onStart(final StartTime startTime) {
  for (final EventHandler<StartTime> startHandler : this.startHandlers) {
    startHandler.onNext(startTime);
  }
}
 

@Test
public void testRemoteManagerOrderingGuaranteeTest() throws Exception {
  System.out.println(LOG_PREFIX + name.getMethodName());
  LoggingUtils.setLoggingLevel(Level.INFO);

  final Monitor monitor = new Monitor();
  final TimerStage timer = new TimerStage(new TimeoutHandler(monitor), 2000, 2000);

  final Map<Class<?>, Codec<?>> clazzToCodecMap = new HashMap<>();
  clazzToCodecMap.put(StartEvent.class, new ObjectSerializableCodec<StartEvent>());
  clazzToCodecMap.put(TestEvent.class, new ObjectSerializableCodec<TestEvent>());
  clazzToCodecMap.put(TestEvent1.class, new ObjectSerializableCodec<TestEvent1>());
  clazzToCodecMap.put(TestEvent2.class, new ObjectSerializableCodec<TestEvent2>());
  final Codec<?> codec = new MultiCodec<Object>(clazzToCodecMap);

  final String hostAddress = localAddressProvider.getLocalAddress();

  final RemoteManager rm = this.remoteManagerFactory.getInstance(
      "name", hostAddress, 0, codec, new LoggingEventHandler<Throwable>(), true, 3, 10000,
      localAddressProvider, Tang.Factory.getTang().newInjector().getInstance(TcpPortProvider.class));

  final RemoteIdentifier remoteId = rm.getMyIdentifier();

  final EventHandler<StartEvent> proxyConnection = rm.getHandler(remoteId, StartEvent.class);
  final EventHandler<TestEvent1> proxyHandler1 = rm.getHandler(remoteId, TestEvent1.class);
  final EventHandler<TestEvent2> proxyHandler2 = rm.getHandler(remoteId, TestEvent2.class);

  final AtomicInteger counter = new AtomicInteger(0);
  final int finalSize = 2;
  rm.registerHandler(StartEvent.class, new MessageTypeEventHandler<StartEvent>(rm, monitor, counter, finalSize));

  proxyConnection.onNext(new StartEvent());

  monitor.mwait();

  proxyHandler1.onNext(new TestEvent1("hello1", 0.0));
  proxyHandler2.onNext(new TestEvent2("hello2", 1.0));

  monitor.mwait();

  Assert.assertEquals(finalSize, counter.get());

  rm.close();
  timer.close();
}
 

@Override
public void publish(final ContextMessage contextMessage) {
  for (final EventHandler<ContextMessage> handler : this.contextMessageHandlers) {
    handler.onNext(contextMessage);
  }
}
 

@Override
public void onNext(final T msg) {
  for (final EventHandler<T> handler : handlers) {
    handler.onNext(msg);
  }
}
 

@Test
public void testRemoteManagerTest() throws Exception {
  System.out.println(LOG_PREFIX + name.getMethodName());
  LoggingUtils.setLoggingLevel(Level.INFO);

  final Monitor monitor = new Monitor();
  final TimerStage timer = new TimerStage(new TimeoutHandler(monitor), 2000, 2000);

  final Map<Class<?>, Codec<?>> clazzToCodecMap = new HashMap<>();
  clazzToCodecMap.put(StartEvent.class, new ObjectSerializableCodec<StartEvent>());
  clazzToCodecMap.put(TestEvent.class, new ObjectSerializableCodec<TestEvent>());
  clazzToCodecMap.put(TestEvent1.class, new ObjectSerializableCodec<TestEvent1>());
  clazzToCodecMap.put(TestEvent2.class, new ObjectSerializableCodec<TestEvent2>());
  final Codec<?> codec = new MultiCodec<Object>(clazzToCodecMap);

  final String hostAddress = localAddressProvider.getLocalAddress();

  final RemoteManager rm = this.remoteManagerFactory.getInstance(
      "name", hostAddress, 0, codec, new LoggingEventHandler<Throwable>(), false, 3, 10000,
      localAddressProvider, Tang.Factory.getTang().newInjector().getInstance(TcpPortProvider.class));

  final RemoteIdentifier remoteId = rm.getMyIdentifier();
  Assert.assertTrue(rm.getMyIdentifier().equals(remoteId));

  final EventHandler<StartEvent> proxyConnection = rm.getHandler(remoteId, StartEvent.class);
  final EventHandler<TestEvent1> proxyHandler1 = rm.getHandler(remoteId, TestEvent1.class);
  final EventHandler<TestEvent2> proxyHandler2 = rm.getHandler(remoteId, TestEvent2.class);

  final AtomicInteger counter = new AtomicInteger(0);
  final int finalSize = 2;
  rm.registerHandler(StartEvent.class, new MessageTypeEventHandler<StartEvent>(rm, monitor, counter, finalSize));

  proxyConnection.onNext(new StartEvent());

  monitor.mwait();
  proxyHandler1.onNext(new TestEvent1("hello1", 0.0)); // registration after send expected to fail
  proxyHandler2.onNext(new TestEvent2("hello2", 1.0));

  monitor.mwait();

  Assert.assertEquals(finalSize, counter.get());

  rm.close();
  timer.close();
}