apereo / cas

import com.amazonaws.auth.AWSStaticCredentialsProvider;

import com.amazonaws.auth.BasicAWSCredentials;

import com.amazonaws.services.logs.AWSLogs;

import com.amazonaws.services.logs.AWSLogsClient;

import com.amazonaws.services.logs.AWSLogsClientBuilder;

import com.amazonaws.services.logs.model.CreateLogGroupRequest;

import com.amazonaws.services.logs.model.CreateLogStreamRequest;

import com.amazonaws.services.logs.model.DataAlreadyAcceptedException;

import com.amazonaws.services.logs.model.DescribeLogGroupsRequest;

import com.amazonaws.services.logs.model.DescribeLogGroupsResult;

import com.amazonaws.services.logs.model.DescribeLogStreamsRequest;

import com.amazonaws.services.logs.model.DescribeLogStreamsResult;

import com.amazonaws.services.logs.model.InputLogEvent;

import com.amazonaws.services.logs.model.InvalidSequenceTokenException;

import com.amazonaws.services.logs.model.LogStream;

import com.amazonaws.services.logs.model.PutLogEventsRequest;

import com.amazonaws.services.logs.model.PutLogEventsResult;

import org.apache.commons.lang3.StringUtils;

import org.apache.logging.log4j.core.Layout;

import org.apache.logging.log4j.core.LogEvent;

import org.apache.logging.log4j.core.appender.AbstractAppender;

import org.apache.logging.log4j.core.config.plugins.Plugin;

import org.apache.logging.log4j.core.config.plugins.PluginAttribute;

import org.apache.logging.log4j.core.config.plugins.PluginElement;

import org.apache.logging.log4j.core.config.plugins.PluginFactory;

import org.apache.logging.log4j.core.layout.PatternLayout;

import java.io.Serializable;

import java.nio.charset.StandardCharsets;

import java.util.ArrayList;

import java.util.Comparator;

import java.util.List;

import java.util.concurrent.BlockingQueue;

import java.util.concurrent.LinkedBlockingQueue;

/**

 * This is {@link CloudWatchAppender}.

 *

 * @author Misagh Moayyed

 * @since 5.1.0

 */

@Plugin(name = "CloudWatchAppender", category = "Core", elementType = "appender", printObject = true)

public class CloudWatchAppender extends AbstractAppender {

    private static final long serialVersionUID = 1044758913028847477L;

    private static final int AWS_DRAIN_LIMIT = 256;

    private static final int AWS_LOG_STREAM_MAX_QUEUE_DEPTH = 10000;

    private static final int SHUTDOWN_TIMEOUT_MILLIS = 10000;

    private static final int AWS_LOG_STREAM_FLUSH_PERIOD_IN_SECONDS = 5;

    private final BlockingQueue<InputLogEvent> queue = new LinkedBlockingQueue<>(AWS_LOG_STREAM_MAX_QUEUE_DEPTH);

    private volatile boolean shutdown;

    private int flushPeriodMillis;

    private Thread deliveryThread;

    private final Object monitor = new Object();

    /**

     * Every PutLogEvents request must include the sequenceToken obtained from the response of the previous request.

     */

    private String sequenceTokenCache;

    private long lastReportedTimestamp = -1;

    private String logGroupName;

    private String logStreamName;

    private AWSLogs awsLogsClient;

    private volatile boolean queueFull;

    public CloudWatchAppender(final String name,

                              final String awsLogGroupName,

                              final String awsLogStreamName,

                              final String awsLogStreamFlushPeriodInSeconds,

                              final String credentialAccessKey,

                              final String credentialSecretKey,

                              final String awsLogRegionName,

                              final Layout<Serializable> layout) {

        super(name, null, layout == null ? PatternLayout.createDefaultLayout() : layout, false);

        try {

            int flushPeriod = AWS_LOG_STREAM_FLUSH_PERIOD_IN_SECONDS;

            if (awsLogStreamFlushPeriodInSeconds != null) {

                flushPeriod = Integer.parseInt(awsLogStreamFlushPeriodInSeconds);

            }

            flushPeriodMillis = flushPeriod * 1_000;

            LOGGER.debug("Connecting to AWS CloudWatch...");

            final AWSLogsClientBuilder builder = AWSLogsClient.builder();

            final BasicAWSCredentials credentials = new BasicAWSCredentials(credentialAccessKey, credentialSecretKey);

            builder.setCredentials(new AWSStaticCredentialsProvider(credentials));

            builder.setRegion(awsLogRegionName);

            this.awsLogsClient = builder.build();

            this.logGroupName = awsLogGroupName;

            this.logStreamName = awsLogStreamName;

            this.sequenceTokenCache = createLogGroupAndLogStreamIfNeeded();

        } catch (final Exception e) {

            LOGGER.error(e.getMessage(), e);

        }

    }

    private void flush() {

        int drained;

        final List<InputLogEvent> logEvents = new ArrayList<>(AWS_DRAIN_LIMIT);

        do {

            drained = queue.drainTo(logEvents, AWS_DRAIN_LIMIT);

            if (logEvents.isEmpty()) {

                break;

            }

            logEvents.sort(Comparator.comparing(InputLogEvent::getTimestamp));

            if (lastReportedTimestamp > 0) {

                for (final InputLogEvent event : logEvents) {

                    if (event.getTimestamp() < lastReportedTimestamp) {

                        event.setTimestamp(lastReportedTimestamp);

                    }

                }

            }

            lastReportedTimestamp = logEvents.get(logEvents.size() - 1).getTimestamp();

            final PutLogEventsRequest putLogEventsRequest = new PutLogEventsRequest(logGroupName, logStreamName, logEvents);

            putLogEventsRequest.setSequenceToken(sequenceTokenCache);

            try {

                final PutLogEventsResult putLogEventsResult = awsLogsClient.putLogEvents(putLogEventsRequest);

                sequenceTokenCache = putLogEventsResult.getNextSequenceToken();

            } catch (final DataAlreadyAcceptedException daae) {

                sequenceTokenCache = daae.getExpectedSequenceToken();

            } catch (final InvalidSequenceTokenException iste) {

                sequenceTokenCache = iste.getExpectedSequenceToken();

            } catch (final Exception e) {

                LOGGER.error(e.getMessage(), e);

            }

            logEvents.clear();

        } while (drained >= AWS_DRAIN_LIMIT);

    }

    @Override

    public void append(final LogEvent logEvent) {

        final LogEvent event = LoggingUtils.prepareLogEvent(logEvent);

        final InputLogEvent awsLogEvent = new InputLogEvent();

        final long timestamp = event.getTimeMillis();

        final String message = new String(getLayout().toByteArray(event), StandardCharsets.UTF_8);

        awsLogEvent.setTimestamp(timestamp);

        awsLogEvent.setMessage(message);

        if (!queue.offer(awsLogEvent) && !queueFull) {

            queueFull = true;

        } else if (queueFull) {

            queueFull = false;

        }

    }

    private String createLogGroupAndLogStreamIfNeeded() {

        LOGGER.debug("Attempting to locate the log group [{}]", logGroupName);

        final DescribeLogGroupsResult describeLogGroupsResult =

                awsLogsClient.describeLogGroups(new DescribeLogGroupsRequest().withLogGroupNamePrefix(logGroupName));

        boolean createLogGroup = true;

        if (describeLogGroupsResult != null && describeLogGroupsResult.getLogGroups() != null && !describeLogGroupsResult.getLogGroups().isEmpty()) {

            createLogGroup = !describeLogGroupsResult.getLogGroups().stream().anyMatch(g -> g.getLogGroupName().equals(logGroupName));

        }

        if (createLogGroup) {

            LOGGER.debug("Creating log group [{}]", logGroupName);

            final CreateLogGroupRequest createLogGroupRequest = new CreateLogGroupRequest(logGroupName);

            awsLogsClient.createLogGroup(createLogGroupRequest);

        }

        String logSequenceToken = null;

        boolean createLogStream = true;

        LOGGER.debug("Attempting to locate the log stream [{}] for group [{}]", logStreamName, logGroupName);

        final DescribeLogStreamsRequest describeLogStreamsRequest = new DescribeLogStreamsRequest(logGroupName).withLogStreamNamePrefix(logStreamName);

        final DescribeLogStreamsResult describeLogStreamsResult = awsLogsClient.describeLogStreams(describeLogStreamsRequest);

        if (describeLogStreamsResult != null && describeLogStreamsResult.getLogStreams() != null && !describeLogStreamsResult.getLogStreams().isEmpty()) {

            for (final LogStream ls : describeLogStreamsResult.getLogStreams()) {

                if (logStreamName.equals(ls.getLogStreamName())) {

                    createLogStream = false;

                    logSequenceToken = ls.getUploadSequenceToken();

                    LOGGER.debug("Found log stream [{}] with sequence token [{}]", logStreamName, logSequenceToken);

                    break;

                }

            }

        }

        if (createLogStream) {

            LOGGER.debug("Creating log stream [{}] for group [{}]", logStreamName, logGroupName);

            final CreateLogStreamRequest createLogStreamRequest = new CreateLogStreamRequest(logGroupName, logStreamName);

            awsLogsClient.createLogStream(createLogStreamRequest);

        }

        return logSequenceToken;

    }

    @Override

    public void start() {

        super.start();

        this.deliveryThread = new Thread(() -> {

            while (!shutdown) {

                try {

                    flush();

                } catch (final Exception e) {

                    LOGGER.error(e.getMessage(), e);

                }

                if (!shutdown && queue.size() < AWS_DRAIN_LIMIT) {

                    try {

                        synchronized (monitor) {

                            monitor.wait(flushPeriodMillis);

                        }

                    } catch (final InterruptedException e) {

                        LOGGER.error(e.getMessage(), e);

                    }

                }

            }

            while (!queue.isEmpty()) {

                flush();

            }

        }, "CloudWatchAppenderDeliveryThread");

        deliveryThread.start();

    }

    }

    /**

     * Create appender cloud watch appender.

     *

     * @param name                             the name

     * @param awsLogStreamName                 the aws log stream name

     * @param awsLogGroupName                  the aws log group name

     * @param awsLogStreamFlushPeriodInSeconds the aws log stream flush period in seconds

     * @param credentialAccessKey              the credential access key

     * @param credentialSecretKey              the credential secret key

     * @param awsLogRegionName                 the aws log region name

     * @param layout                           the layout

     * @return the cloud watch appender

     */

    @PluginFactory

    public static CloudWatchAppender createAppender(@PluginAttribute("name") final String name,

                                                    @PluginAttribute("awsLogStreamName") final String awsLogStreamName,

                                                    @PluginAttribute("awsLogGroupName") final String awsLogGroupName,

                                                    @PluginAttribute("awsLogStreamFlushPeriodInSeconds") final String awsLogStreamFlushPeriodInSeconds,

                                                    @PluginAttribute("credentialAccessKey") final String credentialAccessKey,

                                                    @PluginAttribute("credentialSecretKey") final String credentialSecretKey,

                                                    @PluginAttribute("awsLogRegionName") final String awsLogRegionName,

                                                    @PluginElement("Layout") final Layout<Serializable> layout) {

        return new CloudWatchAppender(

                name,

                awsLogGroupName,

                awsLogStreamName,

                awsLogStreamFlushPeriodInSeconds,

                StringUtils.defaultIfBlank(credentialAccessKey, System.getProperty("AWS_ACCESS_KEY")),

                StringUtils.defaultIfBlank(credentialSecretKey, System.getProperty("AWS_SECRET_KEY")),

                StringUtils.defaultIfBlank(awsLogRegionName, System.getProperty("AWS_REGION_NAME")),

                layout);

    }

}