How to give deadline to okio via okHttp

Question

From looking at okHttp source code, when call.execute() is called the body being transferred from server to the client. It doesn't make sense because it makes impossible to set deadline to okio which means i cannot give timeout to the whole request but only readTimeout and connectTimeout which have effect only until the first byte is ready for read.

Am i missing something here?

Answer 1

There's no way to give a deadline to the entire request. You should open a feature request on this! OkHttp's use of Okio is one of it's differentiating features, and exposing more Okio functionality through OkHttp's API is a great way to put more power in OkHttp's users.

Answer 2

This is on the schedule for the next version of okhttp ( https://github.com/square/okhttp/issues/2840 ), but for now we successfully implemented a deadline for both the request and response body reading by subclassing Call in our application in production:

package com.pushd.util;

import android.support.annotation.NonNull;
import android.support.annotation.Nullable;

import java.io.IOException;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.logging.Logger;

import okhttp3.Call;
import okhttp3.Callback;
import okhttp3.MediaType;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.Response;
import okhttp3.ResponseBody;
import okhttp3.internal.http2.StreamResetException;
import okio.Buffer;
import okio.BufferedSource;
import okio.ForwardingSource;
import okio.Okio;

/**
 * An okhttp3.Call with a deadline timeout from the start of isExecuted until ResponseBody.source() is closed or unused.
 */
public class DeadlineCall implements Call {
    private final static Logger LOGGER = Logger.getLogger(DeadlineCall.class.getName());

    private static AtomicInteger sFutures = new AtomicInteger();
    private static final ScheduledExecutorService sHTTPCancelExecutorService = Executors.newSingleThreadScheduledExecutor(new ThreadFactory() {
        @Override
        public Thread newThread(Runnable r) {
            Thread t = new Thread(r, "DeadlineCallCancel");
            t.setDaemon(true);
            t.setPriority(Thread.NORM_PRIORITY);
            return t;
        }
    });

    private final Call mUnderlying;
    private final int mDeadlineTimeout;
    private volatile ScheduledFuture mDeadline;
    private volatile boolean mDeadlineHit;
    private volatile boolean mCancelled;
    private volatile BufferedSource mBodySource;

    DeadlineCall(Call underlying, int deadlineTimeout) {
        mUnderlying = underlying;
        mDeadlineTimeout = deadlineTimeout;
    }

    /**
     * Factory wrapper for OkHttpClient.newCall(request) to create a new DeadlineCall scheduled to cancel its underlying Call after the deadline. 
     * @param client
     * @param request
     * @param deadlineTimeout in ms
     * @return Call
     */
    public static DeadlineCall newDeadlineCall(@NonNull OkHttpClient client, @NonNull Request request, int deadlineTimeout) {
        final Call underlying = client.newCall(request);
        return new DeadlineCall(underlying, deadlineTimeout);
    }

    /**
     * Shuts down thread that cancels calls when their deadline is hit.
     */
    public static void shutdownNow() {
        sHTTPCancelExecutorService.shutdownNow();
    }

    @Override
    public Request request() {
        return mUnderlying.request();
    }

    /**
     * Response MUST be closed to clean up deadline even if body is not read, e.g. on !isSuccessful
     * @return
     * @throws IOException
     */
    @Override
    public Response execute() throws IOException {
        startDeadline();

        try {
            return wrapResponse(mUnderlying.execute());
        } catch (IOException e) {
            cancelDeadline();
            throw wrapIfDeadline(e);
        }
    }

    /**
     * Deadline is removed when onResponse returns unless response.body().source() or a method using
     * it is called synchronously from onResponse to indicate caller's committment to close it themselves.
     * This includes peekBody so prefer DeadlineResponseBody.peek unless you explicitly close after peekBody.
     * @param responseCallback
     */
    @Override
    public void enqueue(final Callback responseCallback) {
        startDeadline();

        mUnderlying.enqueue(new Callback() {
            @Override
            public void onFailure(Call underlying, IOException e) {
                cancelDeadline(); // there is no body to read so no need for deadline anymore
                responseCallback.onFailure(DeadlineCall.this, wrapIfDeadline(e));
            }

            @Override
            public void onResponse(Call underlying, Response response) throws IOException {
                try {
                    responseCallback.onResponse(DeadlineCall.this, wrapResponse(response));
                    if (mBodySource == null) {
                        cancelDeadline(); // remove deadline if body was never opened
                    }
                } catch (IOException e) {
                    cancelDeadline();
                    throw wrapIfDeadline(e);
                }
            }
        });
    }

    private IOException wrapIfDeadline(IOException e) {
        if (mDeadlineHit && isCancellationException(e)) {
            return new DeadlineException(e);
        }

        return e;
    }

    public class DeadlineException extends IOException {
        public DeadlineException(Throwable cause) {
            super(cause);
        }
    }

    /**
     * Wraps response to cancelDeadline when response closed and throw correct DeadlineException when deadline happens during response reading.
     * @param response
     * @return
     */
    private Response wrapResponse(final Response response) {
        return response.newBuilder().body(new DeadlineResponseBody(response)).build();
    }

    public class DeadlineResponseBody extends ResponseBody {
        private final Response mResponse;

        DeadlineResponseBody(final Response response) {
            mResponse = response;
        }

        @Override
        public MediaType contentType() {
            return mResponse.body().contentType();
        }

        @Override
        public long contentLength() {
            return mResponse.body().contentLength();
        }

        /**
         * @return the body source indicating it will be closed later by the caller to cancel the deadline
         */
        @Override
        public BufferedSource source() {
            if (mBodySource == null) {
                mBodySource = Okio.buffer(new ForwardingSource(mResponse.body().source()) {
                    @Override
                    public long read(Buffer sink, long byteCount) throws IOException {
                        try {
                            return super.read(sink, byteCount);
                        } catch (IOException e) {
                            throw wrapIfDeadline(e);
                        }
                    }

                    @Override
                    public void close() throws IOException {
                        cancelDeadline();
                        super.close();
                    }
                });
            }

            return mBodySource;
        }

        /**
         * @return the body source without indicating it will be closed later by caller, e.g. to peekBody on unsucessful requests
         */
        public BufferedSource peekSource() {
            return mResponse.body().source();
        }

        /**
         * Copy of https://square.github.io/okhttp/3.x/okhttp/okhttp3/Response.html#peekBody-long- that uses peekSource() since Response class is final
         * @param byteCount
         * @return
         * @throws IOException
         */
        public ResponseBody peek(long byteCount) throws IOException {
            BufferedSource source = peekSource();
            source.request(byteCount);
            Buffer copy = source.buffer().clone();

            // There may be more than byteCount bytes in source.buffer(). If there is, return a prefix.
            Buffer result;
            if (copy.size() > byteCount) {
                result = new Buffer();
                result.write(copy, byteCount);
                copy.clear();
            } else {
                result = copy;
            }

            return ResponseBody.create(mResponse.body().contentType(), result.size(), result);
        }
    }

    private void startDeadline() {
        mDeadline = sHTTPCancelExecutorService.schedule(new Runnable() {
            @Override
            public void run() {
                mDeadlineHit = true;
                mUnderlying.cancel(); // calls onFailure or causes body read to throw
                LOGGER.fine("Deadline hit for " + request()); // should trigger a subsequent wrapIfDeadline but if we see this log line without that it means     the caller orphaned us without closing
            }
        }, mDeadlineTimeout, TimeUnit.MILLISECONDS);

        LOGGER.fine("started deadline for " + request());

        if (sFutures.incrementAndGet() == 1000) {
            LOGGER.warning("1000 pending DeadlineCalls, may be leaking due to not calling close()");
        }
    }

    private void cancelDeadline() {
        if (mDeadline != null) {
            mDeadline.cancel(false);
            mDeadline = null;
            sFutures.decrementAndGet();
            LOGGER.fine("canceled deadline for " + request());
        } else {
            LOGGER.info("deadline already canceled for " + request());
        }
    }

    @Override
    public void cancel() {
        mCancelled = true;

        // should trigger onFailure or raise from execute or responseCallback.onResponse which will cancelDeadline
        mUnderlying.cancel();
    }

    @Override
    public boolean isExecuted() {
        return mUnderlying.isExecuted();
    }

    @Override
    public boolean isCanceled() {
        return mCancelled;
    }

    @Override
    public Call clone() {
        return new DeadlineCall(mUnderlying.clone(), mDeadlineTimeout);
    }

    private static boolean isCancellationException(IOException e) {
        // okhttp cancel from HTTP/2 calls
        if (e instanceof StreamResetException) {
            switch (((StreamResetException) e).errorCode) {
                case CANCEL:
                    return true;
            }
        }

        // https://android.googlesource.com/platform/external/okhttp/+/master/okhttp/src/main/java/com/squareup/okhttp/Call.java#281
        if (e instanceof IOException &&
                e.getMessage() != null && e.getMessage().equals("Canceled")) {
            return true;
        }

        return false;
    }
}

Note that we also have a separate interceptor to timeout DNS since even our deadline doesn't cover that:

/**
 * Based on http://stackoverflow.com/questions/693997/how-to-set-httpresponse-timeout-for-android-in-java/31643186#31643186
 * as per https://github.com/square/okhttp/issues/95
 */
private static class DNSTimeoutInterceptor implements Interceptor {
    long mTimeoutMillis;

    public DNSTimeoutInterceptor(long timeoutMillis) {
        mTimeoutMillis = timeoutMillis;
    }

    @Override
    public Response intercept(final Chain chain) throws IOException {
        Request request = chain.request();
        Log.SplitTimer timer = (request.tag() instanceof RequestTag ? ((RequestTag) request.tag()).getTimer() : null);

        // underlying call should timeout after 2 tries of 5s:  https://android.googlesource.com/platform/bionic/+/android-5.1.1_r38/libc/dns/include/resolv_private.h#137
        // could use our own Dns implementation that falls back to public DNS servers:  https://garage.easytaxi.com/tag/dns-android-okhttp/
        if (!DNSResolver.isDNSReachable(request.url().host(), mTimeoutMillis)) {
            throw new UnknownHostException("DNS timeout");
        }
        return chain.proceed(request);
    }

    private static class DNSResolver implements Runnable {
        private String mDomain;
        private InetAddress mAddress;

        public static boolean isDNSReachable(String domain, long timeoutMillis) {
            try {
                DNSResolver dnsRes = new DNSResolver(domain);

                Thread t = new Thread(dnsRes, "DNSResolver");
                t.start();
                t.join(timeoutMillis);
                return dnsRes.get() != null;
            }  catch(Exception e)  {
                return false;
            }
        }

        public DNSResolver(String domain) {
            this.mDomain = domain;
        }

        public void run() {
            try {
                InetAddress addr = InetAddress.getByName(mDomain);
                set(addr);
            } catch (UnknownHostException e) {
            }
        }

        public synchronized void set(InetAddress inetAddr) {
            this.mAddress = inetAddr;
        }
        public synchronized InetAddress get() {
            return mAddress;
        }
    }
}