ControllerTest.java

/*
 * 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.
 */
package org.apache.twill.internal;

import com.google.common.base.Suppliers;
import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.AbstractIdleService;
import com.google.common.util.concurrent.Service;
import com.google.gson.JsonObject;
import org.apache.twill.api.Command;
import org.apache.twill.api.ResourceReport;
import org.apache.twill.api.RunId;
import org.apache.twill.api.ServiceController;
import org.apache.twill.api.TwillController;
import org.apache.twill.api.logging.LogHandler;
import org.apache.twill.common.ServiceListenerAdapter;
import org.apache.twill.common.Threads;
import org.apache.twill.internal.state.StateNode;
import org.apache.twill.internal.zookeeper.InMemoryZKServer;
import org.apache.twill.zookeeper.NodeData;
import org.apache.twill.zookeeper.ZKClient;
import org.apache.twill.zookeeper.ZKClientService;
import org.apache.zookeeper.KeeperException;
import org.junit.Assert;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

/**
 *
 */
public class ControllerTest {

  private static final Logger LOG = LoggerFactory.getLogger(ControllerTest.class);

  @Test
  public void testController() throws ExecutionException, InterruptedException, TimeoutException {
    InMemoryZKServer zkServer = InMemoryZKServer.builder().build();
    zkServer.startAndWait();

    LOG.info("ZKServer: " + zkServer.getConnectionStr());

    try {
      RunId runId = RunIds.generate();
      ZKClientService zkClientService = ZKClientService.Builder.of(zkServer.getConnectionStr()).build();
      zkClientService.startAndWait();

      Service service = createService(zkClientService, runId);
      service.startAndWait();

      TwillController controller = getController(zkClientService, runId);
      controller.sendCommand(Command.Builder.of("test").build()).get(2, TimeUnit.SECONDS);
      controller.stop().get(2, TimeUnit.SECONDS);

      Assert.assertEquals(ServiceController.State.TERMINATED, controller.state());

      final CountDownLatch terminateLatch = new CountDownLatch(1);
      service.addListener(new ServiceListenerAdapter() {
        @Override
        public void terminated(Service.State from) {
          terminateLatch.countDown();
        }
      }, Threads.SAME_THREAD_EXECUTOR);

      Assert.assertTrue(service.state() == Service.State.TERMINATED || terminateLatch.await(2, TimeUnit.SECONDS));

      zkClientService.stopAndWait();

    } finally {
      zkServer.stopAndWait();
    }
  }

  // Test controller created before service starts.
  @Test
  public void testControllerBefore() throws InterruptedException {
    InMemoryZKServer zkServer = InMemoryZKServer.builder().build();
    zkServer.startAndWait();

    LOG.info("ZKServer: " + zkServer.getConnectionStr());
    try {
      RunId runId = RunIds.generate();
      ZKClientService zkClientService = ZKClientService.Builder.of(zkServer.getConnectionStr()).build();
      zkClientService.startAndWait();

      final CountDownLatch runLatch = new CountDownLatch(1);
      final CountDownLatch stopLatch = new CountDownLatch(1);
      TwillController controller = getController(zkClientService, runId);
      controller.addListener(new ServiceListenerAdapter() {
        @Override
        public void running() {
          runLatch.countDown();
        }

        @Override
        public void terminated(Service.State from) {
          stopLatch.countDown();
        }
      }, Threads.SAME_THREAD_EXECUTOR);

      Service service = createService(zkClientService, runId);
      service.start();

      Assert.assertTrue(runLatch.await(2, TimeUnit.SECONDS));
      Assert.assertFalse(stopLatch.await(2, TimeUnit.SECONDS));

      service.stop();

      Assert.assertTrue(stopLatch.await(2, TimeUnit.SECONDS));

    } finally {
      zkServer.stopAndWait();
    }
  }

  // Test controller listener receive first state change without state transition from service
  @Test
  public void testControllerListener() throws InterruptedException {
    InMemoryZKServer zkServer = InMemoryZKServer.builder().build();
    zkServer.startAndWait();

    LOG.info("ZKServer: " + zkServer.getConnectionStr());
    try {
      RunId runId = RunIds.generate();
      ZKClientService zkClientService = ZKClientService.Builder.of(zkServer.getConnectionStr()).build();
      zkClientService.startAndWait();

      Service service = createService(zkClientService, runId);
      service.startAndWait();

      final CountDownLatch runLatch = new CountDownLatch(1);
      TwillController controller = getController(zkClientService, runId);
      controller.addListener(new ServiceListenerAdapter() {
        @Override
        public void running() {
          runLatch.countDown();
        }
      }, Threads.SAME_THREAD_EXECUTOR);

      Assert.assertTrue(runLatch.await(2, TimeUnit.SECONDS));

      service.stopAndWait();

      zkClientService.stopAndWait();
    } finally {
      zkServer.stopAndWait();
    }
  }

  private Service createService(ZKClient zkClient, RunId runId) {
    return new ZKServiceDecorator(
      zkClient, runId, Suppliers.ofInstance(new JsonObject()), new AbstractIdleService() {

      @Override
      protected void startUp() throws Exception {
        LOG.info("Start");
      }

      @Override
      protected void shutDown() throws Exception {
        LOG.info("Stop");
      }
    });
  }

  private TwillController getController(ZKClient zkClient, RunId runId) {
    TwillController controller = new AbstractTwillController(runId, zkClient, ImmutableList.<LogHandler>of()) {

      @Override
      public void kill() {
        // No-op
      }

      @Override
      protected void instanceNodeUpdated(NodeData nodeData) {
        // No-op
      }

      @Override
      protected void instanceNodeFailed(Throwable cause) {
        // Shutdown if the instance node goes away
        if (cause instanceof KeeperException.NoNodeException) {
          forceShutDown();
        }
      }

      @Override
      protected void stateNodeUpdated(StateNode stateNode) {
        // No-op
      }

      @Override
      public ResourceReport getResourceReport() {
        return null;
      }
    };
    controller.startAndWait();
    return controller;
  }
}