Search code examples
javamultithreadingjava-11

Java: How to timeout the current thread


I have a multithreaded Java application that uses ThreadLocal fields to keep the threads isolated from each other. As part of this application I also have a requirement to implement timeouts on certain functions to prevent DOS attacks.

I'm looking for a way to time out a Java function that is running in the current thread

I've seen plenty of solutions such as How to timeout a thread which will create a Future to execute some code, launch it in a new thread and and wait for it to complete. I want to make it work the other way round.

Consider the following code, which will be run in a multi-threaded environment:

class MyClass {
    // ThreadLocal is not private so callback can access it
    ThreadLocal<AtomicInteger> counter = ThreadLocal.withInitial(AtomicInteger::new);

    public void entry(Function<?, ?> callback) {
        counter.get().set(10);                      // Calling thread performs set up
        I_need_a_timeout(callback, 110);            // Call a function which might take a long time
        int result = counter.get().get();           // If there is no time out this will be 110
    }

    private void I_need_a_timeout(Function<?, ?> callback, int loop) {
        while (loop-- >= 0) {
            counter.get().incrementAndGet();
            callback.apply(null);                   // This may take some time
        }
    }
}

I need to be able to terminate I_need_a_timeout if it runs for too long, but if I were to execute it in a future then it would have it's own thread and therefore it's own instance of AtomicInteger so the value read by the calling code would always be the value I initialise it to (in this case 10)

Update: I've updated the sample code to be closer to my real application. The client passes a function to I_need_a_timeout that could take any amount of time to return (or potentially may never return), so polling solutions won't work


Solution

  • You can straight-forwardly create a future whose operation is executed in the current thread and cancel it from another thread. E.g.

    public class CancelCurrentThread {
        public static void main(String[] args) {
            try(var canceler = Executors.newScheduledThreadPool(1)) {
                for(long timeout: List.of(500, 1500)) {
                    FutureTask<String> task
                        = new FutureTask<>(CancelCurrentThread::operation);
                    System.out.println("Timeout: " + timeout + "ms");
                    canceler.schedule(
                        () -> task.cancel(true), timeout, TimeUnit.MILLISECONDS);
                    task.run();
                    try {
                        System.out.println("Result: " + task.get());
                    }
                    catch(CancellationException ex) {
                        System.out.println("Canceled");
                    }
                    catch(InterruptedException|ExecutionException e) {
                        System.out.println("Failed: " + e);
                    }
                }
            }
        }
      
        static String operation() throws InterruptedException {
            System.out.println("performing operation in " + Thread.currentThread());
            Thread.sleep(1000);
            return "result";
        }
    }
    
    Timeout: 500ms
    performing operation in Thread[#1,main,5,main]
    Canceled
    Timeout: 1500ms
    performing operation in Thread[#1,main,5,main]
    Result: result
    

    It’s important to keep in mind that the same restrictions as with futures from an ExecutorService apply. To return quickly upon cancelation, the operation must respond to interruption, either by using operations which abort on interruption (like the sleep in the example) or by polling Thread.interrupted() in reasonable intervals.