Merge pull request influxdata#123 from mnuessler/timeout

Allow configuration of connection timeout parameters

Author: majst01

pom.xml

@@ -89,6 +89,18 @@
 			<version>1.10.19</version>
 			<scope>test</scope>
 		</dependency>
+		<dependency>
+			<groupId>com.github.tomakehurst</groupId>
+			<artifactId>wiremock</artifactId>
+			<version>1.58</version>
+			<scope>test</scope>
+		</dependency>
+		<dependency>
+			<groupId>org.slf4j</groupId>
+			<artifactId>slf4j-simple</artifactId>
+			<version>1.7.5</version>
+			<scope>test</scope>
+		</dependency>
 		<dependency>
 			<groupId>com.google.guava</groupId>
 			<artifactId>guava</artifactId>

src/main/java/org/influxdb/InfluxDB.java

@@ -1,5 +1,6 @@
 package org.influxdb;
 
+import java.net.URLConnection;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 
@@ -22,6 +23,9 @@
  * 
  */
 public interface InfluxDB {
+	int CONNECT_TIMEOUT_SECONDS_DEFAULT = 15;
+	int READ_TIMEOUT_SECONDS_DEFAULT = 20;
+	int WRITE_TIMEOUT_SECONDS_DEFAULT = 20;
 
 	/** Controls the level of logging of the REST layer. */
 	public enum LogLevel {
@@ -171,4 +175,41 @@ public String value() {
 	 */
 	public List<String> describeDatabases();
 
-}
+	/**
+	 * Sets the default connect timeout for new connections. A value of 0 means no timeout. If not
+	 * set explicitly, a default timeout of CONNECT_TIMEOUT_SECONDS_DEFAULT seconds will be used.
+	 *
+	 * @param connectTimeout
+	 *            the connect timeout
+	 * @param timeUnit
+	 *            the time unit for the connect timeout
+	 *
+	 * @see java.net.URLConnection#setConnectTimeout(int)
+	 */
+	void setConnectTimeout(long connectTimeout, TimeUnit timeUnit);
+
+	/**
+	 * Sets the default read timeout for new connections. A value of 0 means no timeout. If not
+	 * set explicitly, a default timeout of READ_TIMEOUT_SECONDS_DEFAULT seconds will be used.
+	 *
+	 * @param readTimeout
+	 *            the read timeout
+	 * @param timeUnit
+	 *            the time unit for the read timeout
+	 *
+	 * @see java.net.URLConnection#setReadTimeout(int)
+	 */
+	void setReadTimeout(long readTimeout, TimeUnit timeUnit);
+
+	/**
+	 * Sets the default write timeout for new connections. A value of 0 means no timeout. If not
+	 * set explicitly, a default timeout of WRITE_TIMEOUT_SECONDS_DEFAULT seconds will be used.
+	 *
+	 * @param writeTimeout
+	 *            the write timeout
+	 * @param timeUnit
+	 *            the time unit for the write timeout
+	 */
+	void setWriteTimeout(long writeTimeout, TimeUnit timeUnit);
+
+}

src/main/java/org/influxdb/impl/InfluxDBImpl.java

@@ -42,6 +42,7 @@ public class InfluxDBImpl implements InfluxDB {
 	private final AtomicLong unBatchedCount = new AtomicLong();
 	private final AtomicLong batchedCount = new AtomicLong();
 	private LogLevel logLevel = LogLevel.NONE;
+	private OkHttpClient okHttpClient;
 
 	/**
 	 * Constructor which should only be used from the InfluxDBFactory.
@@ -57,7 +58,11 @@ public InfluxDBImpl(final String url, final String username, final String passwo
 		super();
 		this.username = username;
 		this.password = password;
-		Client client = new OkClient(new OkHttpClient());
+		okHttpClient = new OkHttpClient();
+		okHttpClient.setConnectTimeout(CONNECT_TIMEOUT_SECONDS_DEFAULT, TimeUnit.SECONDS);
+		okHttpClient.setReadTimeout(READ_TIMEOUT_SECONDS_DEFAULT, TimeUnit.SECONDS);
+		okHttpClient.setWriteTimeout(WRITE_TIMEOUT_SECONDS_DEFAULT, TimeUnit.SECONDS);
+		Client client = new OkClient(okHttpClient);
 		this.restAdapter = new RestAdapter.Builder()
 				.setEndpoint(url)
 				.setErrorHandler(new InfluxDBErrorHandler())
@@ -219,4 +224,28 @@ public List<String> describeDatabases() {
 		return databases;
 	}
 
+	/**
+	 * {@inheritDoc}
+	 */
+	@Override
+	public void setConnectTimeout(long connectTimeout, TimeUnit timeUnit) {
+		okHttpClient.setConnectTimeout(connectTimeout, timeUnit);
+	}
+
+	/**
+	 * {@inheritDoc}
+	 */
+	@Override
+	public void setReadTimeout(long readTimeout, TimeUnit timeUnit) {
+		okHttpClient.setReadTimeout(readTimeout, timeUnit);
+	}
+
+	/**
+	 * {@inheritDoc}
+	 */
+	@Override
+	public void setWriteTimeout(long writeTimeout, TimeUnit timeUnit) {
+		okHttpClient.setWriteTimeout(writeTimeout, timeUnit);
+	}
+
 }

src/test/java/org/influxdb/InfluxDBTimeoutTest.java

@@ -0,0 +1,67 @@
+package org.influxdb;
+
+import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
+import static com.github.tomakehurst.wiremock.client.WireMock.get;
+import static com.github.tomakehurst.wiremock.client.WireMock.stubFor;
+import static com.github.tomakehurst.wiremock.client.WireMock.urlPathEqualTo;
+import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig;
+
+import java.util.concurrent.TimeUnit;
+
+import org.influxdb.dto.Query;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import retrofit.RetrofitError;
+
+import com.github.tomakehurst.wiremock.WireMockServer;
+import com.github.tomakehurst.wiremock.client.WireMock;
+
+/**
+ * Test the timeout settings.
+ */
+public class InfluxDBTimeoutTest {
+	private static final int WIREMOCK_PORT = 8080;
+
+	private WireMockServer wireMockServer;
+
+	@BeforeClass
+	public void setUp() {
+		wireMockServer = new WireMockServer(wireMockConfig().port(WIREMOCK_PORT));
+		wireMockServer.start();
+		WireMock.addRequestProcessingDelay(2 * 1000);
+		WireMock.configureFor(WIREMOCK_PORT);
+
+		stubFor(get(urlPathEqualTo("/query")).willReturn(
+				aResponse()
+						.withStatus(200)
+						.withHeader("Content-Type", "application/json")
+						.withBody("{}")));
+	}
+
+	@AfterClass
+	public void tearDown() {
+		if (wireMockServer != null) {
+			wireMockServer.stop();
+		}
+	}
+
+	@Test(timeOut = 3 * 1000, expectedExceptions = RetrofitError.class, expectedExceptionsMessageRegExp = "connect timed out")
+	public void testConnectTimeout() {
+		// connect to any non-routable IP address will result in a connect timeout
+		InfluxDB influxDB = InfluxDBFactory.connect("http://10.0.0.0:" + WIREMOCK_PORT, "user", "password");
+		influxDB.setConnectTimeout(1, TimeUnit.SECONDS);
+
+		influxDB.query(new Query("SELECT value FROM cpu", "test_db"));
+	}
+
+	@Test(timeOut = 3 * 1000, expectedExceptions = RetrofitError.class, expectedExceptionsMessageRegExp = "timeout")
+	public void testReadTimeout() {
+		InfluxDB influxDB = InfluxDBFactory.connect("http://localhost:" + WIREMOCK_PORT, "user", "password");
+		influxDB.setReadTimeout(1, TimeUnit.SECONDS);
+
+		influxDB.query(new Query("SELECT value FROM cpu", "test_db"));
+	}
+
+}