Simple health check server in Java without dependencies

gistfile1.txt

/**
 * @author Shams Imam
 */
public interface HealthCheck {
    interface Result {
        int getStatus();

        String getMessage();
    }

    Result check();

    static Result newResult(int status, String message) {
        return new Result() {
            @Override
            public int getStatus() {
                return status;
            }

            @Override
            public String getMessage() {
                return message;
            }
        };
    }

    static Result healthy(String message) {
        return newResult(200, message);
    }

    static Result healthy() {
        return healthy("Healthy");
    }

    static Result unhealthy(String message) {
        return newResult(500, message);
    }

    static Result unhealthy() {
        return unhealthy("Unhealthy");
    }
}

gistfile2.txt

import java.util.concurrent.CountDownLatch;

/**
 * @author Shams Imam
 */
public class Main {

    public static void main(String[] args) throws Exception {
        int port = args.length > 0 ? Integer.parseInt(args[0]) : 8080;

        HealthCheck alwaysHealthy = new HealthCheck() {
            private final Result healthyResult = HealthCheck.healthy();

            @Override
            public Result check() {
                return healthyResult;
            }

            @Override
            public String toString() {
                return "AlwaysHealthy";
            }
        };

        Server server = Server.startServer(port, alwaysHealthy);

        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
            try {
                server.stop();
            } catch (Throwable th) {
                th.printStackTrace();
            } finally {
                System.out.println("Server shutdown");
            }
        }));

        server.await();
    }
}

gistfile3.txt

package com.twosigma.healthcheck;

import com.sun.net.httpserver.HttpExchange;
import com.sun.net.httpserver.HttpHandler;
import com.sun.net.httpserver.HttpServer;

import java.io.IOException;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;

/**
 * @author Shams Imam
 */
public class Server {

    public static Server startServer(int port, HealthCheck healthCheck) throws IOException {
        return startServer("127.0.0.1", port, healthCheck);
    }

    public static Server startServer(String host, int port, HealthCheck healthCheck) throws IOException {
        return startServer(host, port, healthCheck, "/");
    }

    public static Server startServer(String host, int port, HealthCheck healthCheck, String healthCheckEndpoint) throws IOException {
        System.out.println("Server will listen on " + host + ":" + port);
        HttpServer httpServer = HttpServer.create(new InetSocketAddress(host, port), 0);
        System.out.println("Health check context is " + healthCheckEndpoint);
        httpServer.createContext(healthCheckEndpoint, new HealthCheckHandler(healthCheck));

        Executor threadPoolExecutor = Executors.newFixedThreadPool(2);
        httpServer.setExecutor(threadPoolExecutor);

        Server server = new Server(httpServer);

        server.start();
        System.out.println("Started server");

        return server;
    }

    private static class HealthCheckHandler implements HttpHandler {

        private static final String serverHeader = "Java/" + System.getProperty("java.version");

        private final HealthCheck healthCheck;

        public HealthCheckHandler(HealthCheck healthCheck) {
            if (healthCheck == null) {
                throw new IllegalArgumentException("HealthCheck instance may not be null!");
            }
            this.healthCheck = healthCheck;
        }

        @Override
        public void handle(HttpExchange httpExchange) throws IOException {
            final String requestMethod = httpExchange.getRequestMethod();
            System.out.println("Request: " + requestMethod + " " + httpExchange.getRequestURI() + " " + httpExchange.getRequestHeaders().entrySet());

            if ("GET".equals(requestMethod)) {
                handleHealthCheckResponse(httpExchange);
            } else {
                handleBadRequestMethodResponse(httpExchange);
            }
        }

        private void handleHealthCheckResponse(HttpExchange httpExchange) throws IOException {
            HealthCheck.Result result = healthCheck.check();
            if (result == null) {
                result = HealthCheck.unhealthy("Health check returned no result");
            }
            final int responseStatus = result.getStatus();
            final String textResponse = result.getMessage();
            generateResponse(httpExchange, responseStatus, textResponse);
        }

        private void handleBadRequestMethodResponse(HttpExchange httpExchange) throws IOException {
            generateResponse(httpExchange, 405, "Method Not Allowed");
        }

        private void generateResponse(HttpExchange httpExchange, int responseStatus, String textResponse) throws IOException {
            OutputStream outputStream = httpExchange.getResponseBody();
            attachCorrelationIdResponseHeader(httpExchange);
            httpExchange.getResponseHeaders().add("content-type", "text/plain");
            httpExchange.getResponseHeaders().add("server", serverHeader);

            System.out.println("Response: " + responseStatus + " " + httpExchange.getResponseHeaders().entrySet());
            httpExchange.sendResponseHeaders(responseStatus, textResponse.length());
            outputStream.write(textResponse.getBytes());
            outputStream.flush();
            outputStream.close();
        }

        private void attachCorrelationIdResponseHeader(HttpExchange httpExchange) {
            final String cidHeaderName = "x-cid";
            String cid = httpExchange.getRequestHeaders().getFirst(cidHeaderName);
            if (cid == null) {
                cid = "r" + System.nanoTime() + "." + Thread.currentThread().getId();
            }
            httpExchange.getResponseHeaders().add(cidHeaderName, cid);
        }
    }

    private final CountDownLatch latch;
    private final HttpServer server;

    public Server(final HttpServer server) {
        this.latch = new CountDownLatch(1);
        this.server = server;
    }

    public void await() throws InterruptedException {
        latch.await();
    }

    protected void start() {
        System.out.println("Starting server...");
        server.start();
    }

    public void stop() {
        System.out.println("Shutting down server...");
        server.stop(0);
        latch.countDown();
    }
}