Skip to content

Commit

Permalink
Merge pull request square#1115 from square/jwilson_1102_mws_reliability
Browse files Browse the repository at this point in the history
Attempt to improve MockWebServer reliability.
  • Loading branch information
swankjesse committed Nov 2, 2014
2 parents ee56def + 12f858c commit 108fd5c
Show file tree
Hide file tree
Showing 4 changed files with 52 additions and 44 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,6 @@
import java.net.Socket;
import java.net.SocketException;
import java.net.URL;
import java.net.UnknownHostException;
import java.security.SecureRandom;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
Expand Down Expand Up @@ -109,25 +108,24 @@ public final class MockWebServer {
private Dispatcher dispatcher = new QueueDispatcher();

private int port = -1;
private InetAddress inetAddress;
private boolean protocolNegotiationEnabled = true;
private List<Protocol> protocols
= Util.immutableList(Protocol.HTTP_2, Protocol.SPDY_3, Protocol.HTTP_1_1);

public int getPort() {
if (port == -1) throw new IllegalStateException("Cannot retrieve port before calling play()");
if (port == -1) throw new IllegalStateException("Call play() before getPort()");
return port;
}

public String getHostName() {
try {
return InetAddress.getLocalHost().getHostName();
} catch (UnknownHostException e) {
throw new AssertionError(e);
}
if (inetAddress == null) throw new IllegalStateException("Call play() before getHostName()");
return inetAddress.getHostName();
}

public Proxy toProxyAddress() {
return new Proxy(Proxy.Type.HTTP, new InetSocketAddress(getHostName(), getPort()));
if (inetAddress == null) throw new IllegalStateException("Call play() before toProxyAddress()");
return new Proxy(Proxy.Type.HTTP, new InetSocketAddress(inetAddress, getPort()));
}

/**
Expand Down Expand Up @@ -284,11 +282,12 @@ public void play() throws IOException {
public void play(int port) throws IOException {
if (executor != null) throw new IllegalStateException("play() already called");
executor = Executors.newCachedThreadPool(Util.threadFactory("MockWebServer", false));
serverSocket = new ServerSocket(port);
inetAddress = InetAddress.getLocalHost();
serverSocket = new ServerSocket(port, 50, inetAddress);
serverSocket.setReuseAddress(true);

this.port = serverSocket.getLocalPort();
executor.execute(new NamedRunnable("MockWebServer %s", port) {
executor.execute(new NamedRunnable("MockWebServer %s", this.port) {
@Override protected void execute() {
try {
acceptConnections();
Expand Down Expand Up @@ -332,8 +331,18 @@ private void acceptConnections() throws Exception {
}

public void shutdown() throws IOException {
if (serverSocket != null) {
serverSocket.close(); // Should cause acceptConnections() to break out.
if (serverSocket == null) return;

// Cause acceptConnections() to break out.
serverSocket.close();

// Await shutdown.
try {
if (!executor.awaitTermination(5, TimeUnit.SECONDS)) {
throw new IOException("Gave up waiting for executor to shut down");
}
} catch (InterruptedException e) {
throw new AssertionError();
}
}

Expand Down Expand Up @@ -440,7 +449,7 @@ private boolean processOneRequest(Socket socket, InputStream in, OutputStream ou
socket.close();
return false;
}
writeResponse(out, response);
writeResponse(socket, out, response);
if (response.getSocketPolicy() == SocketPolicy.DISCONNECT_AT_END) {
in.close();
out.close();
Expand Down Expand Up @@ -525,7 +534,7 @@ private RecordedRequest readRequest(Socket socket, InputStream in, OutputStream
MockResponse throttlePolicy = dispatcher.peek();
if (contentLength != -1) {
hasBody = contentLength > 0;
throttledTransfer(throttlePolicy, in, requestBody, contentLength);
throttledTransfer(throttlePolicy, socket, in, requestBody, contentLength);
} else if (chunked) {
hasBody = true;
while (true) {
Expand All @@ -535,7 +544,7 @@ private RecordedRequest readRequest(Socket socket, InputStream in, OutputStream
break;
}
chunkSizes.add(chunkSize);
throttledTransfer(throttlePolicy, in, requestBody, chunkSize);
throttledTransfer(throttlePolicy, socket, in, requestBody, chunkSize);
readEmptyLine(in);
}
}
Expand All @@ -559,7 +568,8 @@ private RecordedRequest readRequest(Socket socket, InputStream in, OutputStream
requestBody.toByteArray(), sequenceNumber, socket);
}

private void writeResponse(OutputStream out, MockResponse response) throws IOException {
private void writeResponse(Socket socket, OutputStream out, MockResponse response)
throws IOException {
out.write((response.getStatus() + "\r\n").getBytes(Util.US_ASCII));
List<String> headers = response.getHeaders();
for (int i = 0, size = headers.size(); i < size; i++) {
Expand All @@ -571,21 +581,21 @@ private void writeResponse(OutputStream out, MockResponse response) throws IOExc

InputStream in = response.getBodyStream();
if (in == null) return;
throttledTransfer(response, in, out, Long.MAX_VALUE);
throttledTransfer(response, socket, in, out, Long.MAX_VALUE);
}

/**
* Transfer bytes from {@code in} to {@code out} until either {@code length}
* bytes have been transferred or {@code in} is exhausted. The transfer is
* throttled according to {@code throttlePolicy}.
*/
private void throttledTransfer(MockResponse throttlePolicy, InputStream in, OutputStream out,
long limit) throws IOException {
private void throttledTransfer(MockResponse throttlePolicy, Socket socket, InputStream in,
OutputStream out, long limit) throws IOException {
byte[] buffer = new byte[1024];
int bytesPerPeriod = throttlePolicy.getThrottleBytesPerPeriod();
long delayMs = throttlePolicy.getThrottleUnit().toMillis(throttlePolicy.getThrottlePeriod());

while (true) {
while (!socket.isClosed()) {
for (int b = 0; b < bytesPerPeriod; ) {
int toRead = (int) Math.min(Math.min(buffer.length, limit), bytesPerPeriod - b);
int read = in.read(buffer, 0, toRead);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,8 @@ public HttpOverHttp20Draft14Test() {
RecordedRequest request = server.takeRequest();
assertEquals("GET /foo HTTP/1.1", request.getRequestLine());
assertContains(request.getHeaders(), ":scheme: https");
assertContains(request.getHeaders(), hostHeader + ": " + hostName + ":" + server.getPort());
assertContains(request.getHeaders(), hostHeader + ": "
+ server.getHostName() + ":" + server.getPort());

RecordedRequest pushedRequest = server.takeRequest();
assertEquals("GET /foo/bar HTTP/1.1", pushedRequest.getRequestLine());
Expand All @@ -68,7 +69,8 @@ public HttpOverHttp20Draft14Test() {
RecordedRequest request = server.takeRequest();
assertEquals("GET /foo HTTP/1.1", request.getRequestLine());
assertContains(request.getHeaders(), ":scheme: https");
assertContains(request.getHeaders(), hostHeader + ": " + hostName + ":" + server.getPort());
assertContains(request.getHeaders(), hostHeader + ": "
+ server.getHostName() + ":" + server.getPort());

RecordedRequest pushedRequest = server.takeRequest();
assertEquals("HEAD /foo/bar HTTP/1.1", pushedRequest.getRequestLine());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,6 @@ public boolean verify(String hostname, SSLSession session) {

private static final SSLContext sslContext = SslContextBuilder.localhost();
protected final MockWebServer server = new MockWebServer();
protected final String hostName = server.getHostName();
protected final OkUrlFactory client = new OkUrlFactory(new OkHttpClient());
protected HttpURLConnection connection;
protected Cache cache;
Expand Down Expand Up @@ -113,7 +112,8 @@ public boolean verify(String hostname, SSLSession session) {
RecordedRequest request = server.takeRequest();
assertEquals("GET /foo HTTP/1.1", request.getRequestLine());
assertContains(request.getHeaders(), ":scheme: https");
assertContains(request.getHeaders(), hostHeader + ": " + hostName + ":" + server.getPort());
assertContains(request.getHeaders(), hostHeader + ": "
+ server.getHostName() + ":" + server.getPort());
}

@Test public void emptyResponse() throws IOException {
Expand Down Expand Up @@ -406,11 +406,12 @@ public boolean verify(String hostname, SSLSession session) {
@Test public void acceptAndTransmitCookies() throws Exception {
CookieManager cookieManager = new CookieManager();
client.client().setCookieHandler(cookieManager);
server.enqueue(
new MockResponse().addHeader("set-cookie: c=oreo; domain=" + server.getCookieDomain())
.setBody("A"));
server.enqueue(new MockResponse().setBody("B"));
server.play();
server.enqueue(new MockResponse()
.addHeader("set-cookie: c=oreo; domain=" + server.getCookieDomain())
.setBody("A"));
server.enqueue(new MockResponse()
.setBody("B"));

URL url = server.getUrl("/");
assertContent("A", client.open(url), Integer.MAX_VALUE);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -112,10 +112,8 @@ public final class URLConnectionTest {
private final OkUrlFactory client = new OkUrlFactory(new OkHttpClient());
private HttpURLConnection connection;
private Cache cache;
private String hostName;

@Before public void setUp() throws Exception {
hostName = server.getHostName();
server.setProtocolNegotiationEnabled(false);
}

Expand Down Expand Up @@ -1892,8 +1890,8 @@ private void redirectToAnotherOriginServer(boolean https) throws Exception {
assertContent("This is the first server again!", client.open(server.getUrl("/")));
assertContent("This is the 2nd server, again!", client.open(server2.getUrl("/")));

String server1Host = hostName + ":" + server.getPort();
String server2Host = hostName + ":" + server2.getPort();
String server1Host = server.getHostName() + ":" + server.getPort();
String server2Host = server2.getHostName() + ":" + server2.getPort();
assertContains(server.takeRequest().getHeaders(), "Host: " + server1Host);
assertContains(server2.takeRequest().getHeaders(), "Host: " + server2Host);
assertEquals("Expected connection reuse", 1, server.takeRequest().getSequenceNumber());
Expand Down Expand Up @@ -2173,8 +2171,10 @@ private void testRedirect(boolean temporary, String method) throws Exception {
assertContent("DEF", client.open(url));
assertContent("GHI", client.open(url));

assertEquals(Arrays.asList("verify " + hostName), hostnameVerifier.calls);
assertEquals(Arrays.asList("checkServerTrusted [CN=" + hostName + " 1]"), trustManager.calls);
assertEquals(Arrays.asList("verify " + server.getHostName()),
hostnameVerifier.calls);
assertEquals(Arrays.asList("checkServerTrusted [CN=" + server.getHostName() + " 1]"),
trustManager.calls);
}

@Test public void readTimeouts() throws IOException {
Expand Down Expand Up @@ -2203,7 +2203,7 @@ private void testRedirect(boolean temporary, String method) throws Exception {
/** Confirm that an unacknowledged write times out. */
@Test public void writeTimeouts() throws IOException {
server.enqueue(new MockResponse()
.throttleBody(1, 3600, TimeUnit.SECONDS)); // Prevent the server from reading!
.throttleBody(1, 1, TimeUnit.SECONDS)); // Prevent the server from reading!
server.play();

client.client().setWriteTimeout(500, TimeUnit.MILLISECONDS);
Expand Down Expand Up @@ -2365,12 +2365,7 @@ private void testFlushAfterStreamTransmitted(TransferKind transferKind) throws I
}

@Test public void getHeadersThrows() throws IOException {
// Enqueue a response for every IP address held by localhost, because the route selector
// will try each in sequence.
// TODO: use the fake Dns implementation instead of a loop
for (InetAddress inetAddress : InetAddress.getAllByName(server.getHostName())) {
server.enqueue(new MockResponse().setSocketPolicy(DISCONNECT_AT_START));
}
server.enqueue(new MockResponse().setSocketPolicy(DISCONNECT_AT_START));
server.play();

connection = client.open(server.getUrl("/"));
Expand Down Expand Up @@ -3183,7 +3178,7 @@ enum ProxyConfig {
@Override public HttpURLConnection connect(
MockWebServer server, OkUrlFactory streamHandlerFactory, URL url)
throws IOException {
System.setProperty("proxyHost", "localhost");
System.setProperty("proxyHost", server.getHostName());
System.setProperty("proxyPort", Integer.toString(server.getPort()));
return streamHandlerFactory.open(url);
}
Expand All @@ -3193,7 +3188,7 @@ enum ProxyConfig {
@Override public HttpURLConnection connect(
MockWebServer server, OkUrlFactory streamHandlerFactory, URL url)
throws IOException {
System.setProperty("http.proxyHost", "localhost");
System.setProperty("http.proxyHost", server.getHostName());
System.setProperty("http.proxyPort", Integer.toString(server.getPort()));
return streamHandlerFactory.open(url);
}
Expand All @@ -3203,7 +3198,7 @@ enum ProxyConfig {
@Override public HttpURLConnection connect(
MockWebServer server, OkUrlFactory streamHandlerFactory, URL url)
throws IOException {
System.setProperty("https.proxyHost", "localhost");
System.setProperty("https.proxyHost", server.getHostName());
System.setProperty("https.proxyPort", Integer.toString(server.getPort()));
return streamHandlerFactory.open(url);
}
Expand Down

0 comments on commit 108fd5c

Please sign in to comment.