exec/java-exec/src/test/java/org/apache/drill/test/TestGracefulShutdown.java - drill - 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.
  */
 package org.apache.drill.test;

 import org.apache.commons.lang3.reflect.FieldUtils;
 import org.apache.drill.categories.SlowTest;
 import org.apache.drill.common.exceptions.UserException;
 import org.apache.drill.exec.ExecConstants;
 import org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint;
 import org.apache.drill.exec.server.Drillbit;
 import org.apache.drill.exec.server.rest.WebServer;
 import org.junit.Assert;
 import org.junit.BeforeClass;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 import org.junit.rules.TestRule;

 import java.io.BufferedReader;
 import java.io.BufferedWriter;
 import java.io.File;
 import java.io.FileWriter;
 import java.io.InputStreamReader;
 import java.io.PrintWriter;
 import java.lang.reflect.Field;
 import java.net.HttpURLConnection;
 import java.net.URL;
 import java.nio.file.Path;
 import java.util.Collection;

 import static org.hamcrest.CoreMatchers.containsString;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;

 @Category({SlowTest.class})
 public class TestGracefulShutdown extends ClusterTest {

   @Rule
   public final TestRule TIMEOUT = TestTools.getTimeoutRule(120_000);

   @BeforeClass
   public static void setUpTestData() throws Exception {
     for (int i = 0; i < 300; i++) {
       setupFile(i);
     }
   }

   private static ClusterFixtureBuilder builderWithEnabledWebServer() {
     return builderWithEnabledPortHunting()
         .configProperty(ExecConstants.HTTP_ENABLE, true)
         .configProperty(ExecConstants.HTTP_PORT_HUNT, true)
         .configProperty(ExecConstants.SLICE_TARGET, 10);
   }

   private static ClusterFixtureBuilder builderWithEnabledPortHunting() {
     return ClusterFixture.builder(dirTestWatcher)
         .configProperty(ExecConstants.DRILL_PORT_HUNT, true)
         .configProperty(ExecConstants.GRACE_PERIOD, 500)
         .configProperty(ExecConstants.ALLOW_LOOPBACK_ADDRESS_BINDING, true);
   }

   /*
   Start multiple drillbits and then shutdown a drillbit. Query the online
   endpoints and check if the drillbit still exists.
    */
   @Test
   public void testOnlineEndPoints() throws Exception {

     String[] drillbits = {"db1", "db2", "db3"};
     ClusterFixtureBuilder builder = builderWithEnabledPortHunting()
         .withLocalZk()
         .withBits(drillbits);

     try (ClusterFixture cluster = builder.build()) {

       Drillbit drillbit = cluster.drillbit("db2");
       int zkRefresh = drillbit.getContext().getConfig().getInt(ExecConstants.ZK_REFRESH);
       DrillbitEndpoint drillbitEndpoint = drillbit.getRegistrationHandle().getEndPoint();
       cluster.closeDrillbit("db2");

       while (true) {
         Collection<DrillbitEndpoint> drillbitEndpoints = cluster.drillbit()
             .getContext()
             .getClusterCoordinator()
             .getOnlineEndPoints();

         if (!drillbitEndpoints.contains(drillbitEndpoint)) {
           // Success
           return;
         }

         Thread.sleep(zkRefresh);
       }
     }
   }

   /*
    Test shutdown through RestApi
    */
   @Test
   public void testRestApi() throws Exception {

     String[] drillbits = {"db1", "db2", "db3"};
     ClusterFixtureBuilder builder = builderWithEnabledWebServer()
         .withLocalZk()
         .withBits(drillbits);

     QueryBuilder.QuerySummaryFuture listener;
     final String sql = "select * from dfs.root.`.`";
     try (ClusterFixture cluster = builder.build();
          ClientFixture client = cluster.clientFixture()) {
       Drillbit drillbit = cluster.drillbit("db1");
       int port = drillbit.getWebServerPort();
       int zkRefresh = drillbit.getContext().getConfig().getInt(ExecConstants.ZK_REFRESH);
       listener = client.queryBuilder().sql(sql).futureSummary();
       URL url = new URL("http://localhost:" + port + "/gracefulShutdown");
       HttpURLConnection conn = (HttpURLConnection) url.openConnection();
       conn.setRequestMethod("POST");
       if (conn.getResponseCode() != 200) {
         throw new RuntimeException("Failed : HTTP error code : "
                 + conn.getResponseCode());
       }
       while (!listener.isDone()) {
         Thread.sleep(100L);
       }

       if (waitAndAssertDrillbitCount(cluster, zkRefresh, drillbits.length)) {
         return;
       }
       Assert.fail("Timed out");
     }
   }

   /*
    Test default shutdown through RestApi
    */
   @Test
   public void testRestApiShutdown() throws Exception {

     String[] drillbits = {"db1", "db2", "db3"};
     ClusterFixtureBuilder builder = builderWithEnabledWebServer().withLocalZk().withBits(drillbits);

     QueryBuilder.QuerySummaryFuture listener;
     final String sql = "select * from dfs.root.`.`";
     try (ClusterFixture cluster = builder.build();
          final ClientFixture client = cluster.clientFixture()) {
       Drillbit drillbit = cluster.drillbit("db1");
       int port = drillbit.getWebServerPort();
       int zkRefresh = drillbit.getContext().getConfig().getInt(ExecConstants.ZK_REFRESH);
       listener = client.queryBuilder().sql(sql).futureSummary();
       while (!listener.isDone()) {
         Thread.sleep(100L);
       }
       URL url = new URL("http://localhost:" + port + "/shutdown");
       HttpURLConnection conn = (HttpURLConnection) url.openConnection();
       conn.setRequestMethod("POST");
       if (conn.getResponseCode() != 200) {
         throw new RuntimeException("Failed : HTTP error code : "
                 + conn.getResponseCode());
       }
       if (waitAndAssertDrillbitCount(cluster, zkRefresh, drillbits.length)) {
         return;
       }
       Assert.fail("Timed out");
     }
   }

   @Test // DRILL-6912
   public void testDrillbitWithSamePortContainsShutdownThread() throws Exception {
     ClusterFixtureBuilder fixtureBuilder = ClusterFixture.builder(dirTestWatcher)
         .withLocalZk()
         .configProperty(ExecConstants.ALLOW_LOOPBACK_ADDRESS_BINDING, true)
         .configProperty(ExecConstants.INITIAL_USER_PORT, QueryTestUtil.getFreePortNumber(31170, 300))
         .configProperty(ExecConstants.INITIAL_BIT_PORT, QueryTestUtil.getFreePortNumber(31180, 300));

     try (ClusterFixture fixture = fixtureBuilder.build();
          Drillbit drillbitWithSamePort = new Drillbit(fixture.config(),
              fixtureBuilder.configBuilder().getDefinitions(), fixture.serviceSet())) {
       // Assert preconditions :
       //      1. First drillbit instance should be started normally
       //      2. Second instance startup should fail, because ports are occupied by the first one
       assertNotNull("First drillbit instance should be initialized", fixture.drillbit());
       try {
         drillbitWithSamePort.run();
         fail("Invocation of 'drillbitWithSamePort.run()' should throw UserException");
       } catch (UserException e) {
         assertThat(e.getMessage(), containsString("RESOURCE ERROR: Drillbit could not bind to port"));
         // Ensure that drillbit with failed startup may be safely closed
         assertNotNull("Drillbit.gracefulShutdownThread shouldn't be null, otherwise close() may throw NPE (if so, check suppressed exception).",
             drillbitWithSamePort.getGracefulShutdownThread());
       }
     }
   }

   @Test // DRILL-7056
   public void testDrillbitTempDir() throws Exception {
     File originalDrillbitTempDir;
     ClusterFixtureBuilder fixtureBuilder = ClusterFixture.builder(dirTestWatcher).withLocalZk()
         .configProperty(ExecConstants.ALLOW_LOOPBACK_ADDRESS_BINDING, true)
         .configProperty(ExecConstants.INITIAL_USER_PORT, QueryTestUtil.getFreePortNumber(31170, 300))
         .configProperty(ExecConstants.INITIAL_BIT_PORT, QueryTestUtil.getFreePortNumber(31180, 300));

     try (ClusterFixture fixture = fixtureBuilder.build();
          Drillbit twinDrillbitOnSamePort = new Drillbit(fixture.config(),
             fixtureBuilder.configBuilder().getDefinitions(), fixture.serviceSet())) {
       // Assert preconditions :
       //      1. First drillbit instance should be started normally
       //      2. Second instance startup should fail, because ports are occupied by the first one
       Drillbit originalDrillbit = fixture.drillbit();
       assertNotNull("First drillbit instance should be initialized", originalDrillbit);
       originalDrillbitTempDir = getWebServerTempDirPath(originalDrillbit);
       assertTrue("First drillbit instance should have a temporary Javascript dir initialized", originalDrillbitTempDir.exists());
       try {
         twinDrillbitOnSamePort.run();
         fail("Invocation of 'twinDrillbitOnSamePort.run()' should throw UserException");
       } catch (UserException userEx) {
         assertThat(userEx.getMessage(), containsString("RESOURCE ERROR: Drillbit could not bind to port"));
       }
     }
     // Verify deletion
     assertFalse("First drillbit instance should have a temporary Javascript dir deleted", originalDrillbitTempDir.exists());
   }

   private File getWebServerTempDirPath(Drillbit drillbit) throws IllegalAccessException {
     Field webServerField = FieldUtils.getField(drillbit.getClass(), "webServer", true);
     WebServer webServerHandle = (WebServer) FieldUtils.readField(webServerField, drillbit, true);
     return webServerHandle.getOrCreateTmpJavaScriptDir();
   }

   private boolean waitAndAssertDrillbitCount(ClusterFixture cluster, int zkRefresh, int bitsNum)
       throws InterruptedException {

     while (true) {
       Collection<DrillbitEndpoint> drillbitEndpoints = cluster.drillbit()
               .getContext()
               .getClusterCoordinator()
               .getAvailableEndpoints();
       if (drillbitEndpoints.size() == bitsNum - 1) {
         return true;
       }
       Thread.sleep(zkRefresh);
     }
   }

   private static void setupFile(int file_num) throws Exception {
     String file = "employee" + file_num + ".json";
     Path path = dirTestWatcher.getRootDir().toPath().resolve(file);
     StringBuilder stringBuilder = new StringBuilder();
     int rowsCount = 7;
     try (BufferedReader reader = new BufferedReader(new InputStreamReader(TestGracefulShutdown.class.getResourceAsStream("/employee.json")))) {
       for (int i = 0; i < rowsCount; i++) {
         stringBuilder.append(reader.readLine());
       }
     }
     String content = stringBuilder.toString();
     try (PrintWriter out = new PrintWriter(new BufferedWriter(new FileWriter(path.toFile(), true)))) {
       out.println(content);
     }
   }

 }
	/*
	* 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.drill.test;

	import org.apache.commons.lang3.reflect.FieldUtils;
	import org.apache.drill.categories.SlowTest;
	import org.apache.drill.common.exceptions.UserException;
	import org.apache.drill.exec.ExecConstants;
	import org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint;
	import org.apache.drill.exec.server.Drillbit;
	import org.apache.drill.exec.server.rest.WebServer;
	import org.junit.Assert;
	import org.junit.BeforeClass;
	import org.junit.Rule;
	import org.junit.Test;
	import org.junit.experimental.categories.Category;
	import org.junit.rules.TestRule;

	import java.io.BufferedReader;
	import java.io.BufferedWriter;
	import java.io.File;
	import java.io.FileWriter;
	import java.io.InputStreamReader;
	import java.io.PrintWriter;
	import java.lang.reflect.Field;
	import java.net.HttpURLConnection;
	import java.net.URL;
	import java.nio.file.Path;
	import java.util.Collection;

	import static org.hamcrest.CoreMatchers.containsString;
	import static org.junit.Assert.assertFalse;
	import static org.junit.Assert.assertNotNull;
	import static org.junit.Assert.assertThat;
	import static org.junit.Assert.assertTrue;
	import static org.junit.Assert.fail;

	@Category({SlowTest.class})
	public class TestGracefulShutdown extends ClusterTest {

	@Rule
	public final TestRule TIMEOUT = TestTools.getTimeoutRule(120_000);

	@BeforeClass
	public static void setUpTestData() throws Exception {
	for (int i = 0; i < 300; i++) {
	setupFile(i);
	}
	}

	private static ClusterFixtureBuilder builderWithEnabledWebServer() {
	return builderWithEnabledPortHunting()
	.configProperty(ExecConstants.HTTP_ENABLE, true)
	.configProperty(ExecConstants.HTTP_PORT_HUNT, true)
	.configProperty(ExecConstants.SLICE_TARGET, 10);
	}

	private static ClusterFixtureBuilder builderWithEnabledPortHunting() {
	return ClusterFixture.builder(dirTestWatcher)
	.configProperty(ExecConstants.DRILL_PORT_HUNT, true)
	.configProperty(ExecConstants.GRACE_PERIOD, 500)
	.configProperty(ExecConstants.ALLOW_LOOPBACK_ADDRESS_BINDING, true);
	}

	/*
	Start multiple drillbits and then shutdown a drillbit. Query the online
	endpoints and check if the drillbit still exists.
	*/
	@Test
	public void testOnlineEndPoints() throws Exception {

	String[] drillbits = {"db1", "db2", "db3"};
	ClusterFixtureBuilder builder = builderWithEnabledPortHunting()
	.withLocalZk()
	.withBits(drillbits);

	try (ClusterFixture cluster = builder.build()) {

	Drillbit drillbit = cluster.drillbit("db2");
	int zkRefresh = drillbit.getContext().getConfig().getInt(ExecConstants.ZK_REFRESH);
	DrillbitEndpoint drillbitEndpoint = drillbit.getRegistrationHandle().getEndPoint();
	cluster.closeDrillbit("db2");

	while (true) {
	Collection<DrillbitEndpoint> drillbitEndpoints = cluster.drillbit()
	.getContext()
	.getClusterCoordinator()
	.getOnlineEndPoints();

	if (!drillbitEndpoints.contains(drillbitEndpoint)) {
	// Success
	return;
	}

	Thread.sleep(zkRefresh);
	}
	}
	}

	/*
	Test shutdown through RestApi
	*/
	@Test
	public void testRestApi() throws Exception {

	String[] drillbits = {"db1", "db2", "db3"};
	ClusterFixtureBuilder builder = builderWithEnabledWebServer()
	.withLocalZk()
	.withBits(drillbits);

	QueryBuilder.QuerySummaryFuture listener;
	final String sql = "select * from dfs.root.`.`";
	try (ClusterFixture cluster = builder.build();
	ClientFixture client = cluster.clientFixture()) {
	Drillbit drillbit = cluster.drillbit("db1");
	int port = drillbit.getWebServerPort();
	int zkRefresh = drillbit.getContext().getConfig().getInt(ExecConstants.ZK_REFRESH);
	listener = client.queryBuilder().sql(sql).futureSummary();
	URL url = new URL("http://localhost:" + port + "/gracefulShutdown");
	HttpURLConnection conn = (HttpURLConnection) url.openConnection();
	conn.setRequestMethod("POST");
	if (conn.getResponseCode() != 200) {
	throw new RuntimeException("Failed : HTTP error code : "
	+ conn.getResponseCode());
	}
	while (!listener.isDone()) {
	Thread.sleep(100L);
	}

	if (waitAndAssertDrillbitCount(cluster, zkRefresh, drillbits.length)) {
	return;
	}
	Assert.fail("Timed out");
	}
	}

	/*
	Test default shutdown through RestApi
	*/
	@Test
	public void testRestApiShutdown() throws Exception {

	String[] drillbits = {"db1", "db2", "db3"};
	ClusterFixtureBuilder builder = builderWithEnabledWebServer().withLocalZk().withBits(drillbits);

	QueryBuilder.QuerySummaryFuture listener;
	final String sql = "select * from dfs.root.`.`";
	try (ClusterFixture cluster = builder.build();
	final ClientFixture client = cluster.clientFixture()) {
	Drillbit drillbit = cluster.drillbit("db1");
	int port = drillbit.getWebServerPort();
	int zkRefresh = drillbit.getContext().getConfig().getInt(ExecConstants.ZK_REFRESH);
	listener = client.queryBuilder().sql(sql).futureSummary();
	while (!listener.isDone()) {
	Thread.sleep(100L);
	}
	URL url = new URL("http://localhost:" + port + "/shutdown");
	HttpURLConnection conn = (HttpURLConnection) url.openConnection();
	conn.setRequestMethod("POST");
	if (conn.getResponseCode() != 200) {
	throw new RuntimeException("Failed : HTTP error code : "
	+ conn.getResponseCode());
	}
	if (waitAndAssertDrillbitCount(cluster, zkRefresh, drillbits.length)) {
	return;
	}
	Assert.fail("Timed out");
	}
	}

	@Test // DRILL-6912
	public void testDrillbitWithSamePortContainsShutdownThread() throws Exception {
	ClusterFixtureBuilder fixtureBuilder = ClusterFixture.builder(dirTestWatcher)
	.withLocalZk()
	.configProperty(ExecConstants.ALLOW_LOOPBACK_ADDRESS_BINDING, true)
	.configProperty(ExecConstants.INITIAL_USER_PORT, QueryTestUtil.getFreePortNumber(31170, 300))
	.configProperty(ExecConstants.INITIAL_BIT_PORT, QueryTestUtil.getFreePortNumber(31180, 300));

	try (ClusterFixture fixture = fixtureBuilder.build();
	Drillbit drillbitWithSamePort = new Drillbit(fixture.config(),
	fixtureBuilder.configBuilder().getDefinitions(), fixture.serviceSet())) {
	// Assert preconditions :
	// 1. First drillbit instance should be started normally
	// 2. Second instance startup should fail, because ports are occupied by the first one
	assertNotNull("First drillbit instance should be initialized", fixture.drillbit());
	try {
	drillbitWithSamePort.run();
	fail("Invocation of 'drillbitWithSamePort.run()' should throw UserException");
	} catch (UserException e) {
	assertThat(e.getMessage(), containsString("RESOURCE ERROR: Drillbit could not bind to port"));
	// Ensure that drillbit with failed startup may be safely closed
	assertNotNull("Drillbit.gracefulShutdownThread shouldn't be null, otherwise close() may throw NPE (if so, check suppressed exception).",
	drillbitWithSamePort.getGracefulShutdownThread());
	}
	}
	}

	@Test // DRILL-7056
	public void testDrillbitTempDir() throws Exception {
	File originalDrillbitTempDir;
	ClusterFixtureBuilder fixtureBuilder = ClusterFixture.builder(dirTestWatcher).withLocalZk()
	.configProperty(ExecConstants.ALLOW_LOOPBACK_ADDRESS_BINDING, true)
	.configProperty(ExecConstants.INITIAL_USER_PORT, QueryTestUtil.getFreePortNumber(31170, 300))
	.configProperty(ExecConstants.INITIAL_BIT_PORT, QueryTestUtil.getFreePortNumber(31180, 300));

	try (ClusterFixture fixture = fixtureBuilder.build();
	Drillbit twinDrillbitOnSamePort = new Drillbit(fixture.config(),
	fixtureBuilder.configBuilder().getDefinitions(), fixture.serviceSet())) {
	// Assert preconditions :
	// 1. First drillbit instance should be started normally
	// 2. Second instance startup should fail, because ports are occupied by the first one
	Drillbit originalDrillbit = fixture.drillbit();
	assertNotNull("First drillbit instance should be initialized", originalDrillbit);
	originalDrillbitTempDir = getWebServerTempDirPath(originalDrillbit);
	assertTrue("First drillbit instance should have a temporary Javascript dir initialized", originalDrillbitTempDir.exists());
	try {
	twinDrillbitOnSamePort.run();
	fail("Invocation of 'twinDrillbitOnSamePort.run()' should throw UserException");
	} catch (UserException userEx) {
	assertThat(userEx.getMessage(), containsString("RESOURCE ERROR: Drillbit could not bind to port"));
	}
	}
	// Verify deletion
	assertFalse("First drillbit instance should have a temporary Javascript dir deleted", originalDrillbitTempDir.exists());
	}

	private File getWebServerTempDirPath(Drillbit drillbit) throws IllegalAccessException {
	Field webServerField = FieldUtils.getField(drillbit.getClass(), "webServer", true);
	WebServer webServerHandle = (WebServer) FieldUtils.readField(webServerField, drillbit, true);
	return webServerHandle.getOrCreateTmpJavaScriptDir();
	}

	private boolean waitAndAssertDrillbitCount(ClusterFixture cluster, int zkRefresh, int bitsNum)
	throws InterruptedException {

	while (true) {
	Collection<DrillbitEndpoint> drillbitEndpoints = cluster.drillbit()
	.getContext()
	.getClusterCoordinator()
	.getAvailableEndpoints();
	if (drillbitEndpoints.size() == bitsNum - 1) {
	return true;
	}
	Thread.sleep(zkRefresh);
	}
	}

	private static void setupFile(int file_num) throws Exception {
	String file = "employee" + file_num + ".json";
	Path path = dirTestWatcher.getRootDir().toPath().resolve(file);
	StringBuilder stringBuilder = new StringBuilder();
	int rowsCount = 7;
	try (BufferedReader reader = new BufferedReader(new InputStreamReader(TestGracefulShutdown.class.getResourceAsStream("/employee.json")))) {
	for (int i = 0; i < rowsCount; i++) {
	stringBuilder.append(reader.readLine());
	}
	}
	String content = stringBuilder.toString();
	try (PrintWriter out = new PrintWriter(new BufferedWriter(new FileWriter(path.toFile(), true)))) {
	out.println(content);
	}
	}

	}