SocketTimeoutTest.java revision e78f117bcbd6b57d783737107f445ef75ecb474a 
1/*
2 * Copyright (C) 2014 Square, Inc.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *      http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16package okio;
17
18import java.io.EOFException;
19import java.io.IOException;
20import java.io.InputStream;
21import java.io.InterruptedIOException;
22import java.io.OutputStream;
23import java.net.ServerSocket;
24import java.net.Socket;
25import java.util.concurrent.TimeUnit;
26import org.junit.Test;
27
28import static org.junit.Assert.assertTrue;
29import static org.junit.Assert.fail;
30
31public class SocketTimeoutTest {
32
33  // The size of the socket buffers to use. Less than half the data transferred during tests to
34  // ensure send and receive buffers are flooded and any necessary blocking behavior takes place.
35  private static final int SOCKET_BUFFER_SIZE = 256 * 1024;
36  private static final int ONE_MB = 1024 * 1024;
37
38  @Test public void readWithoutTimeout() throws Exception {
39    Socket socket = socket(ONE_MB, 0);
40    BufferedSource source = Okio.buffer(Okio.source(socket));
41    source.timeout().timeout(5000, TimeUnit.MILLISECONDS);
42    source.require(ONE_MB);
43    socket.close();
44  }
45
46  @Test public void readWithTimeout() throws Exception {
47    Socket socket = socket(0, 0);
48    BufferedSource source = Okio.buffer(Okio.source(socket));
49    source.timeout().timeout(250, TimeUnit.MILLISECONDS);
50    try {
51      source.require(ONE_MB);
52      fail();
53    } catch (InterruptedIOException expected) {
54    }
55    socket.close();
56  }
57
58  @Test public void writeWithoutTimeout() throws Exception {
59    Socket socket = socket(0, ONE_MB);
60    Sink sink = Okio.buffer(Okio.sink(socket));
61    sink.timeout().timeout(500, TimeUnit.MILLISECONDS);
62    byte[] data = new byte[ONE_MB];
63    sink.write(new Buffer().write(data), data.length);
64    sink.flush();
65    socket.close();
66  }
67
68  @Test public void writeWithTimeout() throws Exception {
69    Socket socket = socket(0, 0);
70    Sink sink = Okio.sink(socket);
71    sink.timeout().timeout(500, TimeUnit.MILLISECONDS);
72    byte[] data = new byte[ONE_MB];
73    long start = System.nanoTime();
74    try {
75      sink.write(new Buffer().write(data), data.length);
76      sink.flush();
77      fail();
78    } catch (InterruptedIOException expected) {
79    }
80    long elapsed = System.nanoTime() - start;
81    socket.close();
82
83    assertTrue("elapsed: " + elapsed, TimeUnit.NANOSECONDS.toMillis(elapsed) >= 500);
84    assertTrue("elapsed: " + elapsed, TimeUnit.NANOSECONDS.toMillis(elapsed) <= 750);
85  }
86
87  /**
88   * Returns a socket that can read {@code readableByteCount} incoming bytes and
89   * will accept {@code writableByteCount} written bytes. The socket will idle
90   * for 5 seconds when the required data has been read and written.
91   */
92  static Socket socket(final int readableByteCount, final int writableByteCount) throws IOException {
93    final ServerSocket serverSocket = new ServerSocket(0);
94    serverSocket.setReuseAddress(true);
95    serverSocket.setReceiveBufferSize(SOCKET_BUFFER_SIZE);
96
97    Thread peer = new Thread("peer") {
98      @Override public void run() {
99        Socket socket = null;
100        try {
101          socket = serverSocket.accept();
102          socket.setSendBufferSize(SOCKET_BUFFER_SIZE);
103          writeFully(socket.getOutputStream(), readableByteCount);
104          readFully(socket.getInputStream(), writableByteCount);
105          Thread.sleep(5000); // Sleep 5 seconds so the peer can close the connection.
106        } catch (Exception ignored) {
107        } finally {
108          try {
109            if (socket != null) socket.close();
110          } catch (IOException ignored) {
111          }
112        }
113      }
114    };
115    peer.start();
116
117    Socket socket = new Socket(serverSocket.getInetAddress(), serverSocket.getLocalPort());
118    socket.setReceiveBufferSize(SOCKET_BUFFER_SIZE);
119    socket.setSendBufferSize(SOCKET_BUFFER_SIZE);
120    return socket;
121  }
122
123  private static void writeFully(OutputStream out, int byteCount) throws IOException {
124    out.write(new byte[byteCount]);
125    out.flush();
126  }
127
128  private static byte[] readFully(InputStream in, int byteCount) throws IOException {
129    int count = 0;
130    byte[] result = new byte[byteCount];
131    while (count < byteCount) {
132      int read = in.read(result, count, result.length - count);
133      if (read == -1) throw new EOFException();
134      count += read;
135    }
136    return result;
137  }
138}
139