TinyApplication.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 ch.qos.logback.classic.LoggerContext;
import ch.qos.logback.core.ConsoleAppender;
import net.morimekta.collect.UnmodifiableList;
import net.morimekta.terminal.args.ArgException;
import net.morimekta.terminal.args.ArgHelp;
import net.morimekta.terminal.args.ArgNameFormat;
import net.morimekta.terminal.args.ArgParser;
import net.morimekta.terminal.args.annotations.ArgNaming;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.bridge.SLF4JBridgeHandler;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.PrintStream;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.logging.Level;

import static java.nio.charset.StandardCharsets.UTF_8;
import static java.util.logging.Logger.getLogger;
import static net.morimekta.terminal.args.Flag.flag;
import static net.morimekta.terminal.args.Option.optionLong;
import static net.morimekta.terminal.args.ValueParser.i32;

/**
 * Tiny microservice application base class. Extend this class to set up the
 * server itself, and use the static {@link #start(TinyApplication, String...)}
 * method to actually start it.
 */
public abstract class TinyApplication {
    private static final Logger LOGGER = LoggerFactory.getLogger(TinyApplication.class);

    private final String                                  applicationName;
    private final AtomicReference<TinyApplicationContext> context;

    protected TinyApplication(String applicationName) {
        setUpUncaughtExceptionHandler();
        this.applicationName = Objects.requireNonNull(applicationName, "applicationName == null");
        this.context = new AtomicReference<>();
    }

    /**
     * Override this method if you want to have a special uncaught exception handler,
     * or if you need to keep some other default uncaught exception handler.
     */
    protected void setUpUncaughtExceptionHandler() {
        // Ensure there is an uncaught exception handler. But do not replace any
        // the user may have already set.
        Thread.setDefaultUncaughtExceptionHandler((thread, throwable) -> LOGGER.error(
                "Uncaught exception in thread {}: {}",
                thread.getName(),
                throwable.getMessage(),
                throwable));
    }

    /**
     * This method is called during the initialization phase of setting up the
     * tiny server. This is done before the HTTP server is started.l
     *
     * @param argBuilder     The argument parser builder to parse command line arguments.
     * @param contextBuilder The application context builder.
     */
    protected abstract void initialize(ArgParser.Builder argBuilder, TinyApplicationContext.Builder contextBuilder);

    /**
     * This method is called after the HTTP server is started, but before the
     * service is considered "ready".
     *
     * @param context The application context.
     * @throws Exception If starting the app failed.
     */
    protected abstract void onStart(TinyApplicationContext context) throws Exception;

    /**
     * This method is called immediately when the service should start shutting
     * down. The service is already considered "not ready", but the HTTP server
     * will stay alive as long as this method does not return.
     *
     * @param context The application context of the service.
     * @throws Exception If stopping the service failed.
     */
    protected void beforeStop(TinyApplicationContext context) throws Exception {
        // not implemented.
    }

    /**
     * This method is called after the HTTP service has been stopped, but before
     * the application exits.
     *
     * @param context The application context of the service.
     * @throws Exception If stopping the service failed.
     */
    protected void afterStop(TinyApplicationContext context) throws Exception {
        // not implemented.
    }

    /**
     * @return The application name.
     */
    public final String getApplicationName() {
        return applicationName;
    }

    /**
     * @return The application version.
     */
    public String getApplicationVersion() {
        return "latest";
    }

    /**
     * @return The application description.
     */
    public String getApplicationDescription() {
        return "A Tiny-Server Application";
    }

    /**
     * Stop the server and trigger the internal stop mechanisms.
     */
    public final void stop() {
        LOGGER.info("Stopping server.");

        var context = this.context.getAndSet(null);
        if (context == null) {
            return;
        }

        // Triggers readiness to start failing.
        if (context.setStopping()) {
            try {
                beforeStop(context);
            } catch (Exception e) {
                LOGGER.warn("Exception in beforeStop(): {}", e.getMessage(), e);
            } finally {
                context.getStoppingSemaphore().release(100);
            }
        } else {
            try {
                context.getStoppingSemaphore().acquire();
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }

        context.stopServer();

        try {
            afterStop(context);
        } catch (Exception e) {
            LOGGER.warn("Exception in afterStop(): {}", e.getMessage(), e);
        }

        flushConsoleLogging();
    }

    public final void drain() {
        LOGGER.info("Draining server.");

        var context = this.context.get();
        if (context == null) {
            return;
        }

        // Triggers readiness to start failing.
        if (context.setStopping()) {
            try {
                beforeStop(context);
            } catch (Exception e) {
                LOGGER.warn("Exception in beforeStop(): {}", e.getMessage(), e);
            } finally {
                context.getStoppingSemaphore().release(100);
            }
        }
    }

    /**
     * Start the server. If the server start failed for any reason, will forcefully exit
     * the program.
     *
     * <pre>{@code
     * public class MyServer extends TinyServer {
     *     // implement...
     *
     *     public static void main(String[] args) {
     *         TinyApplication.start(new MyServer(), args);
     *     }
     * }
     * }</pre>
     *
     * @param app  The tiny server application to start.
     * @param args Arguments form command line.
     */
    public static void start(TinyApplication app, String... args) {
        try {
            startUnsafe(app, args);
        } catch (Exception e) {
            // Untestable, as it will abort the test ...
            System.exit(1);
        }
    }

    /**
     * Same as the {@link #start(TinyApplication, String...)} method, but will throw
     * the exception. Visible for testing.
     *
     * @param app  The tiny server application to start.
     * @param args Arguments form command line.
     * @throws Exception If startup failed for any reason.
     */
    public static void startUnsafe(TinyApplication app, String... args) throws Exception {
        // Send JUL & java System logger to SLF4J.
        SLF4JBridgeHandler.removeHandlersForRootLogger();
        SLF4JBridgeHandler.install();
        getLogger("").setLevel(Level.FINEST);

        ArgParser parser = null;
        ArgNaming argNaming = app.getClass().getAnnotation(ArgNaming.class);
        ArgNameFormat defaultFormat = argNaming != null ? argNaming.value() : ArgNameFormat.SNAKE;

        try {
            var argParserBuilder = ArgParser.argParser(
                    app.getApplicationName(),
                    app.getApplicationVersion(),
                    app.getApplicationDescription());
            argParserBuilder.generateArgsNameFormat(defaultFormat);

            var help = new AtomicBoolean();
            argParserBuilder.add(flag("--help", "?h", "Show help", help::set));

            var contextBuilder = new TinyApplicationContext.Builder(app);
            argParserBuilder.add(optionLong("--" + defaultFormat.format("admin_host"),
                                            "Set IP address the admin server should listen to.",
                                            contextBuilder::setAdminHost)
                                         .defaultValue("0.0.0.0").metaVar("IP"));
            argParserBuilder.add(optionLong("--" + defaultFormat.format("admin_port"),
                                            "Set HTTP port the admin server should listen to.",
                                            i32(contextBuilder::setAdminPort))
                                         .defaultValue("0").metaVar("PORT"));
            argParserBuilder.add(optionLong("--" + defaultFormat.format("admin_threads"),
                                            "Number of threads to use for admin server.",
                                            i32(contextBuilder::setAdminServerThreads))
                                         .defaultValue(10).metaVar("THR"));

            app.initialize(argParserBuilder, contextBuilder);

            parser = argParserBuilder.parse(args);
            if (help.get()) {
                ArgHelp.argHelp(parser).printHelp(System.out);
                return;
            }
            parser.validate();
            app.startInternal(contextBuilder.build());
        } catch (ArgException e) {
            if (e.getParser() != null) {
                LOGGER.error("{}\n{}", e.getMessage(), getHelp(e.getParser()), e);
            } else if (parser != null) {
                LOGGER.error("{}\n{}", e.getMessage(), getHelp(parser), e);
            }
            flushConsoleLogging();
            throw e;
        } catch (RuntimeException e) {
            if (parser != null) {
                // Exception validating arguments.
                LOGGER.error("{}\n{}", e.getMessage(), getHelp(parser), e);
            } else {
                LOGGER.error("{}", e.getMessage(), e);
            }
            flushConsoleLogging();
            throw e;
        } catch (IOException e) {
            LOGGER.error("IO Exception: {}", e.getMessage(), e);
            flushConsoleLogging();
            throw e;
        } catch (Exception e) {
            LOGGER.error("Exception: {}", e.getMessage(), e);
            flushConsoleLogging();
            throw e;
        }
    }

    // ---- PRIVATE ----

    private void startInternal(TinyApplicationContext context) throws Exception {
        try {
            onStart(context);
        } catch (Exception e) {
            LOGGER.error("Exception in onStarted(): {}", e.getMessage(), e);
            context.stopServer();
            throw e;
        }

        this.context.set(context);

        Thread stopServerThread = new Thread(this::stop);
        stopServerThread.setDaemon(false);
        stopServerThread.setName("TinyServerShutdownHook");
        Runtime.getRuntime().addShutdownHook(stopServerThread);

        context.setReady();
    }

    private static void flushConsoleLogging() {
        // Force console loggers to flush output. Otherwise, some logging may
        // be missing just before application exits.
        ((LoggerContext) LoggerFactory.getILoggerFactory())
                .getLoggerList()
                .stream()
                .map(ch.qos.logback.classic.Logger::iteratorForAppenders)
                .flatMap(it -> UnmodifiableList.asList(it).stream())
                .filter(app -> app instanceof ConsoleAppender)
                .map(app -> (ConsoleAppender<?>) app)
                .forEach(app -> {
                    try {
                        app.getOutputStream().flush();
                    } catch (IOException ignore) {
                    }
                });
    }

    private static String getHelp(ArgParser parser) {
        try (var out = new ByteArrayOutputStream();
             var writer = new PrintStream(out, true, UTF_8)) {
            ArgHelp.argHelp(parser)
                   .usageWidth(80)
                   .printPreamble(writer);
            return out.toString(UTF_8);
        } catch (IOException e) {
            return "";  // should be impossible.
        }
    }
}