TinyApplicationContext.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package net.morimekta.tiny.server;

import com.sun.net.httpserver.HttpExchange;
import com.sun.net.httpserver.HttpHandler;
import com.sun.net.httpserver.HttpServer;
import net.morimekta.tiny.server.http.TinyHealth;
import net.morimekta.tiny.server.http.TinyHealthHttpHandler;
import net.morimekta.tiny.server.http.TinyHttpHandler;
import net.morimekta.tiny.server.http.TinyHttpStatus;
import net.morimekta.tiny.server.http.TinyReadyHttpHandler;
import net.morimekta.tiny.server.utils.TinyThreadFactory;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.util.HashSet;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.Semaphore;
import java.util.concurrent.atomic.AtomicReference;

import static ch.qos.logback.core.util.OptionHelper.isNullOrEmpty;
import static net.morimekta.collect.util.SetOperations.union;
import static net.morimekta.strings.StringUtil.isNotEmpty;

public final class TinyApplicationContext {
    static final String SERVER_READINESS = "tiny-server";

    private final int                                 localPort;
    private final HttpServer                          httpServer;
    private final Map<String, TinyHealth.ReadyCheck>  readyCheckMap;
    private final Map<String, TinyHealth.HealthCheck> healthCheckMap;
    private final AtomicReference<Semaphore> stopping;

    private TinyApplicationContext(
            HttpServer httpServer,
            Map<String, TinyHealth.ReadyCheck> readyCheckMap,
            Map<String, TinyHealth.HealthCheck> healthCheckMap) {
        this.stopping = new AtomicReference<>();
        this.httpServer = httpServer;
        this.localPort = httpServer.getAddress().getPort();
        this.readyCheckMap = readyCheckMap;
        this.healthCheckMap = healthCheckMap;
    }

    /**
     * Get the port the Admin HTTP server is listening to.
     *
     * @return The listening Admin port.
     */
    public int getAdminPort() {
        return localPort;
    }

    /**
     * Add ready check to active checks.
     *
     * @param name       Ready check name. The name must be unique.
     * @param readyCheck The ready check to be done.
     * @return The context.
     */
    public TinyApplicationContext addReadyCheck(String name, TinyHealth.ReadyCheck readyCheck) {
        if (readyCheckMap.containsKey(name)) {
            throw new IllegalArgumentException("Readiness check with name \"" + name + "\" already exists.");
        }
        this.readyCheckMap.put(name, readyCheck);
        return this;
    }

    /**
     * Remove ready check with given name.
     *
     * @param name The ready check name.
     * @return The context.
     */
    public TinyApplicationContext removeReadyCheck(String name) {
        this.readyCheckMap.remove(name);
        return this;
    }

    /**
     * Add health check to active checks.
     *
     * @param name        Ready check name.
     * @param healthCheck The health check to be done.
     * @return The context.
     */
    public TinyApplicationContext addHealthCheck(String name, TinyHealth.HealthCheck healthCheck) {
        if (healthCheckMap.containsKey(name)) {
            throw new IllegalArgumentException("Health check with name \"" + name + "\" already exists.");
        }
        this.healthCheckMap.put(name, healthCheck);
        return this;
    }

    /**
     * Remove health check with given name.
     *
     * @param name The health check name.
     * @return The context.
     */
    public TinyApplicationContext removeHealthCheck(String name) {
        this.healthCheckMap.remove(name);
        return this;
    }

    // -----------------------
    // -----  PROTECTED  -----
    // -----------------------

    void stopServer() {
        this.httpServer.stop(0);
    }

    void setReady() {
        removeReadyCheck(SERVER_READINESS);
    }

    boolean setStopping() {
        if (stopping.getAndUpdate(old -> old != null ? old : new Semaphore(0)) == null) {
            removeReadyCheck(SERVER_READINESS)
                    .addReadyCheck(SERVER_READINESS, () -> TinyHealth.unhealthy("Server is stopping."));
            return true;
        }
        return false;
    }

    Semaphore getStoppingSemaphore() {
        return stopping.get();
    }

    public static final class Builder {
        private final TinyApplication server;

        private int    adminPort   = 0;
        private String adminHost   = "0.0.0.0";
        private int    threads     = 10;
        private String readyPath   = "/ready";
        private String healthyPath = "/healthy";
        private String drainPath   = "/drain";

        private final HttpServer  httpServer;
        private final Set<String> httpContextSet;

        public Builder(TinyApplication server) throws IOException {
            this.server = server;
            this.httpContextSet = new HashSet<>();
            this.httpServer = HttpServer.create();
        }

        public void setReadyPath(String readyPath) {
            Set<String> knownContexts = union(httpContextSet, Set.of(healthyPath, drainPath));
            if (knownContexts.contains(readyPath)) {
                throw new IllegalArgumentException("Context " + readyPath + " already exists.");
            }
            this.readyPath = readyPath;
        }

        public void setHealthyPath(String healthyPath) {
            Set<String> knownContexts = union(httpContextSet, Set.of(readyPath, drainPath));
            if (knownContexts.contains(healthyPath)) {
                throw new IllegalArgumentException("Context " + healthyPath + " already exists.");
            }
            this.healthyPath = healthyPath;
        }

        public void setDrainPath(String drainPath) {
            if (isNullOrEmpty(drainPath)) {
                // In case we want to *not* have the drain servlet.
                this.drainPath = "";
            } else {
                Set<String> knownContexts = union(httpContextSet, Set.of(readyPath, healthyPath));
                if (knownContexts.contains(drainPath)) {
                    throw new IllegalArgumentException("Context " + drainPath + " already exists.");
                }
                this.drainPath = drainPath;
            }
        }

        /**
         * Add a custom HTTP handler to the admin server.
         *
         * @param context The context path for the handler.
         * @param handler The handler.
         */
        public void addHttpHandler(String context, HttpHandler handler) {
            Objects.requireNonNull(context, "context == null");
            Objects.requireNonNull(handler, "handler == null");
            Set<String> knownContexts = union(httpContextSet, Set.of(readyPath, healthyPath, drainPath));
            if (knownContexts.contains(context)) {
                throw new IllegalArgumentException("Context " + context + " already exists.");
            }
            httpServer.createContext(context, handler);
            httpContextSet.add(context);
        }

        void setAdminPort(int localPort) {
            this.adminPort = localPort;
        }

        void setAdminHost(String localHost) {
            this.adminHost = localHost;
        }

        void setAdminServerThreads(int threads) {
            this.threads = threads;
        }

        TinyApplicationContext build() throws IOException {
            var healthChecks = new ConcurrentHashMap<String, TinyHealth.HealthCheck>();
            var readyChecks = new ConcurrentHashMap<String, TinyHealth.ReadyCheck>();
            readyChecks.put(SERVER_READINESS, () -> TinyHealth.unhealthy("Server is starting."));

            var threadFactory = new TinyThreadFactory("TinyHttpServer", true);
            httpServer.setExecutor(threads > 0
                                   ? Executors.newFixedThreadPool(threads, threadFactory)
                                   : Executors.newCachedThreadPool(threadFactory));
            httpServer.createContext(readyPath, new TinyReadyHttpHandler(readyChecks));
            httpServer.createContext(healthyPath, new TinyHealthHttpHandler(healthChecks));
            if (isNotEmpty(drainPath)) {
                httpServer.createContext(drainPath, new TinyHttpHandler() {
                    @Override
                    protected void doGet(HttpExchange exchange) throws IOException {
                        try {
                            server.drain();
                            exchange.sendResponseHeaders(TinyHttpStatus.SC_NO_CONTENT, -1);
                        } catch (Exception e) {
                            exchange.sendResponseHeaders(TinyHttpStatus.SC_INTERNAL, -1);
                        }
                    }
                });
            }
            httpServer.bind(new InetSocketAddress(adminHost, adminPort), 100);
            httpServer.start();
            return new TinyApplicationContext(
                    httpServer,
                    readyChecks,
                    healthChecks);
        }
    }
}