StreamingLogsController.java

/*
 * Copyright 2020 Global Crop Diversity Trust
 *
 * Licensed 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 org.gringlobal.mvc.admin;

import java.io.BufferedWriter;
import java.io.OutputStreamWriter;
import java.io.Serializable;
import java.io.Writer;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.atomic.AtomicBoolean;

import javax.servlet.http.HttpServletResponse;

import lombok.extern.slf4j.Slf4j;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.core.Layout;
import org.apache.logging.log4j.core.LogEvent;
import org.apache.logging.log4j.core.Logger;
import org.apache.logging.log4j.core.LoggerContext;
import org.apache.logging.log4j.core.appender.AbstractAppender;
import org.apache.logging.log4j.core.config.Property;
import org.apache.logging.log4j.core.config.plugins.Plugin;
import org.apache.logging.log4j.core.layout.PatternLayout;
import org.apache.logging.log4j.core.util.Throwables;
import org.apache.logging.log4j.util.Strings;
import org.bouncycastle.util.Arrays;
import org.springframework.http.MediaType;
import org.springframework.security.access.prepost.PreAuthorize;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;

/**
 * The Class StreamingLogsController.
 *
 * @author Maxym Borodenko
 * @author Matija Obreza
 */
@Controller
@RequestMapping("/admin")
@PreAuthorize("hasAuthority('GROUP_ADMINS')")
@Slf4j
public class StreamingLogsController {

	private static final String DEFAULT_LAYOUT_PATTERN = "%d{yyyy-MM-dd HH:mm:ss} %t %-5p %c{1}:%L - %m%n";

	@PostMapping(value = "/action", params = "action=streamlogs", produces = { MediaType.TEXT_PLAIN_VALUE })
	public void streamLogs(HttpServletResponse response) {
		response.setContentType(MediaType.TEXT_PLAIN_VALUE);

		LoggerContext loggerContext = (LoggerContext) LogManager.getContext(false);
		Logger logger = loggerContext.getRootLogger();

		LogAppenderImpl logAppender = null;
		try {
			BufferedWriter writer = new BufferedWriter(new OutputStreamWriter(response.getOutputStream(), StandardCharsets.UTF_8));
			writer.append("Streaming log messages...").append("\n");
			writer.flush();

			PatternLayout patternLayout = PatternLayout.newBuilder().withPattern(DEFAULT_LAYOUT_PATTERN).build();
			var monitor = new AtomicBoolean(true);
			logAppender = new LogAppenderImpl("streamLogger" + System.currentTimeMillis(), writer, patternLayout, monitor);

			log.warn("Registering streaming log appender {}", logAppender.getName());
			logAppender.start();
			logger.addAppender(logAppender);
			loggerContext.updateLoggers();

			while (monitor.get()) {
				Thread.sleep(1000);
				writer.flush();
			}

		} catch (Throwable e) {
			log.warn(e.getMessage(), e);
		} finally {
			if (logAppender != null) {
				var la = logAppender;
				loggerContext.getLoggers().forEach(l -> {
					log.info("Removing appender from {}", l.getName());
					l.removeAppender(la);
				});
				loggerContext.getConfiguration().getAppenders().remove(logAppender.getName());
				loggerContext.updateLoggers();
				logAppender.stop();
				log.warn("Removed streaming logger {}", logAppender.getName());
			}
		}
	}

	@Plugin(name = "LogAppender", category = "Core", elementType = "appender", printObject = true)
	private static final class LogAppenderImpl extends AbstractAppender {

		private Writer out;
		private AtomicBoolean running;

		LogAppenderImpl(String name, Writer out, Layout<? extends Serializable> layout, AtomicBoolean runningMonitor) {
			super(name, null, layout, false, Property.EMPTY_ARRAY);
			this.out = out;
			this.running = runningMonitor;
		}

		@Override
		public void append(LogEvent event) {
			if (! running.get()) {
				return;
			}

			String[] throwableStr = Throwables.toStringList(event.getThrown()).toArray(Strings.EMPTY_ARRAY);
			StringBuilder builder = new StringBuilder();
			if (!Arrays.isNullOrEmpty(throwableStr)) {
				for (String str : throwableStr) {
					builder.append(str).append('\n');
				}
			}
			final String mes = new String(this.getLayout().toByteArray(event), StandardCharsets.UTF_8) + (Arrays.isNullOrEmpty(throwableStr) ? "" : builder.toString());

			try {
				// 1 thread at a time
				synchronized (this) {
					// this is blocking, what happens if there's really slow connection
					out.append(mes);
					out.flush();
				}
			} catch (Throwable e) {
				running.set(false);
			}
		}
	}
}