package com.vlkan.log4j2.redis.appender;

import com.google.common.util.concurrent.RateLimiter;

import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

class RedisThrottler implements AutoCloseable {

    private final RedisThrottlerConfig config;

    private final RedisThrottlerReceiver receiver;

    private final boolean ignoreExceptions;

    private final BlockingQueue<byte[]> buffer;

    private final byte[][] batch;

    private final Thread flushTrigger;

    private final RateLimiter rateLimiter;

    private final DebugLogger logger;

    private final AtomicReference<Throwable> lastThrown = new AtomicReference<>(null);

    RedisThrottler(RedisThrottlerConfig config, RedisThrottlerReceiver receiver, boolean ignoreExceptions, boolean debugEnabled) {
        this.config = config;
        this.receiver = receiver;
        this.ignoreExceptions = ignoreExceptions;
        this.buffer = new ArrayBlockingQueue<>(config.getBufferSize());
        this.batch = new byte[config.getBatchSize()][];
        this.flushTrigger = createFlushTrigger();
        this.rateLimiter = config.getMaxEventCountPerSecond() > 0 ? RateLimiter.create(config.getMaxEventCountPerSecond()) : null;
        this.logger = new DebugLogger(RedisThrottler.class, debugEnabled);
    }

    private Thread createFlushTrigger() {
        return new Thread(new Runnable() {
            @Override
            @SuppressWarnings("InfiniteLoopStatement")
            public void run() {
                logger.debug("started");
                do {
                    logger.debug("flushing");
                    try {
                        flush();
                    } catch (InterruptedException ignored) {
                        logger.debug("interrupted");
                        Thread.currentThread().interrupt();
                        break;
                    }
                } while (true);
            }
        });
    }

    private void flush() throws InterruptedException {

        int batchIndex = 0;
        byte[] event;

        // Flush in batches.
        logger.debug("polling");
        while ((event = buffer.poll(config.getFlushPeriodMillis(), TimeUnit.MILLISECONDS)) != null) {
            if (logger.isEnabled()) {
                logger.debug("polled: %s", new String(event));
            }
            batch[batchIndex++] = event;
            if (batchIndex == batch.length) {
                safeConsumeEvents(batch);
                batchIndex = 0;
            }
        }

        // Flush individually.
        if (batchIndex > 0) {
            logger.debug("pushing %d individual events", batchIndex);
            int eventCount = batchIndex;
            batchIndex = 0;
            do {
                event = this.batch[batchIndex];
                safeConsumeEvent(event);
            } while (++batchIndex < eventCount);
        }

    }

    private void safeConsumeEvent(byte[] event) {
        try {
            if (logger.isEnabled()) {
                logger.debug("pushing single event: %s", new String(event));
            }
            receiver.consumeThrottledEvent(event);
        } catch (Throwable thrown) {
            if (logger.isEnabled()) {
                logger.debug("push failure: %s", thrown.getMessage());
                thrown.printStackTrace();
            }
            lastThrown.set(thrown);
        }
    }

    private void safeConsumeEvents(final byte[]... events) {
        try {
            logger.debug("pushing %d events", events.length);
            receiver.consumeThrottledEvents(events);
        } catch (Throwable thrown) {
            if (logger.isEnabled()) {
                logger.debug("push failure: %s", thrown.getMessage());
                thrown.printStackTrace();
            }
            lastThrown.set(thrown);
        }
    }

    public void push(byte[] event) {

        Throwable thrown = lastThrown.getAndSet(null);
        if (thrown != null) {
            tryThrow(thrown.getMessage());
            return;
        }

        if (rateLimiter != null && !rateLimiter.tryAcquire()) {
            tryThrow("failed acquiring rate limiter token");
            return;
        }

        if (!buffer.offer(event)) {
            tryThrow("failed enqueueing");
        }

    }

    private void tryThrow(String error) {
        logger.debug(error);
        if (!ignoreExceptions)
            throw new RuntimeException(error);
    }

    public synchronized void start() {
        logger.debug("starting");
        flushTrigger.start();
    }

    @Override
    public synchronized void close() {
        logger.debug("closing");
        flushTrigger.interrupt();
    }

}

@@ -86,24 +86,49 @@
 teLimiter;%0A%0A
+import java.util.Arrays;%0A
 import java.
@@ -2256,16 +2256,176 @@
     
-while ((
+long waitPeriodMillis = config.getFlushPeriodMillis();%0A        while (waitPeriodMillis %3E 0) %7B%0A            long pollTimeMillis = System.currentTimeMillis();%0A            
 even
@@ -2440,31 +2440,20 @@
 er.poll(
-config.getFlush
+wait
 PeriodMi
@@ -2456,18 +2456,16 @@
 odMillis
-()
 , TimeUn
@@ -2484,20 +2484,79 @@
 NDS)
-) != null) %7B
+;%0A            if (event == null) %7B%0A                break;%0A            %7D
 %0A   
@@ -2854,490 +2854,209 @@
     
-%7D%0A%0A
     
-    // Flush individually.%0A        if (batchIndex %3E 0) %7B%0A            logger.debug(%22pushing %25d individual events%22, batchIndex);%0A            int eventCount = batchIndex;%0A            batchIndex = 0;%0A            do %7B%0A                event = this.batch%5BbatchIndex%5D;%0A                safeConsumeEvent(event);%0A            %7D while (++batchIndex %3C eventCount);%0A        %7D%0A%0A    %7D%0A%0A    private void safeConsumeEvent(byte%5B%5D event) %7B%0A        try %7B%0A            if (logger.isEnabled()) %7B%0A    
+long pollPeriodMillis = System.currentTimeMillis() - pollTimeMillis;%0A            waitPeriodMillis -= pollPeriodMillis;%0A        %7D%0A%0A        // Flush the last remaining.%0A        if (batchIndex %3E 0) %7B%0A
     
@@ -3089,347 +3089,155 @@
 ing 
-single
+remaining %25d
  event
-: %25
 s%22, 
-new String(event));%0A            %7D%0A            receiver.consumeThrottledEvent(event);%0A        %7D catch (Throwable thrown) %7B%0A            if (logger.isEnabled()) %7B%0A                logger.debug(%22push failure: %25s%22, thrown.getMessage());%0A                thrown.printStackTrace();%0A            %7D%0A            lastThrown.set(thrown
+batchIndex);%0A            byte%5B%5D%5B%5D subBatch = Arrays.copyOfRange(batch, 0, batchIndex);%0A            safeConsumeEvents(subBatch
 );%0A 
@@ -3237,32 +3237,33 @@
 tch);%0A        %7D%0A
+%0A
     %7D%0A%0A    priva

Make sure we follow flushPeriodMillis in throttler.