package cucumber.runtime;

import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;

public class Timeout {
    public static <T> T timeout(Callback<T> callback, long timeoutMillis) throws Throwable {
        if (timeoutMillis == 0) {
            return callback.call();
        } else {
            final Thread executionThread = Thread.currentThread();
            final AtomicBoolean done = new AtomicBoolean();
            ScheduledFuture<?> timer = Executors.newSingleThreadScheduledExecutor().schedule(new Runnable() {
                @Override
                public void run() {
                    if (!done.get()) {
                        executionThread.interrupt();
                    }
                }
            }, timeoutMillis, TimeUnit.MILLISECONDS);
            try {
                return callback.call();
            } catch (InterruptedException timeout) {
                throw new TimeoutException("Timed out after " + timeoutMillis + "ms.");
            } finally {
                done.set(true);
                timer.cancel(true);
            }

        }
    }

    public interface Callback<T> {
        T call() throws Throwable;
    }
}

@@ -52,146 +52,9 @@
 ent.
-Executors;%0Aimport java.util.concurrent.ScheduledFuture;%0Aimport java.util.concurrent.TimeUnit;%0Aimport java.util.concurrent.TimeoutException
+*
 ;%0Aim
@@ -428,24 +428,25 @@
 cBoolean();%0A
+%0A
             
@@ -458,23 +458,39 @@
 uled
-Future%3C?%3E timer
+ExecutorService executorService
  = E
@@ -532,16 +532,72 @@
 ecutor()
+;%0A            ScheduledFuture%3C?%3E timer = executorService
 .schedul
@@ -1110,32 +1110,32 @@
 done.set(true);%0A
-
                 
@@ -1146,32 +1146,79 @@
 r.cancel(true);%0A
+                executorService.shutdownNow();%0A
             %7D%0A%0A 

fix the thread leak in timeout class (https://github.com/cucumber/cucumber-jvm/issues/639)