package com.avast.kafkatests;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Duration;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * Utility functions.
 */
public class Utils {
    private static final Logger LOGGER = LoggerFactory.getLogger(Utils.class);

    public static void logAllUnhandledExceptions() {
        Thread.setDefaultUncaughtExceptionHandler((t, e) -> LOGGER.error("Default uncaught exception handler: {}, {}", t, e, e));
    }

    public static void closeOnShutdown(AutoCloseable closeable) {
        Runtime.getRuntime().addShutdownHook(new Thread() {
            public void run() {
                try {
                    LOGGER.info("====================== Shutdown hook ======================");
                    closeable.close();
                } catch (Exception e) {
                    LOGGER.error("Exception while closing on shutdown: {}", e, e);
                }
            }
        });
    }

    public static void loopWithNoExit() {
        while (true) {
            try {
                new CountDownLatch(1).await();
            } catch (InterruptedException e) {
                LOGGER.error("Interrupted exception");
            }

            LOGGER.error("Never ending loop continues");
        }
    }

    public static void shutdownAndWaitTermination(ExecutorService executor, Duration timeout, String name) {
        try {
            LOGGER.debug("Waiting for termination of executor: {}, max {} ms, {}", name, timeout.toMillis(), executor);
            executor.shutdown();
            executor.awaitTermination(timeout.toMillis(), TimeUnit.MILLISECONDS);
            LOGGER.debug("Waiting for termination of executor finished: {}, max {} ms, {}", name, timeout.toMillis(), executor);
        } catch (InterruptedException e) {
            LOGGER.error("Waiting for termination of executor interrupted: {}, max {} ms, {}, {}", name, timeout.toMillis(), executor, e, e);
        }
    }
}

@@ -1640,24 +1640,25 @@
 shutdown();%0A
+%0A
             
@@ -1649,32 +1649,36 @@
 );%0A%0A            
+if (
 executor.awaitTe
@@ -1729,18 +1729,24 @@
 SECONDS)
-;%0A
+) %7B%0A    
         
@@ -1858,32 +1858,201 @@
 s(), executor);%0A
+            %7D else %7B%0A                LOGGER.error(%22Waiting for termination of executor timed out: %7B%7D, max %7B%7D ms, %7B%7D%22, name, timeout.toMillis(), executor);%0A            %7D%0A
         %7D catch 

Handle return value of executor.awaitTermination() during shutdown.