Merge pull request #24 from tkowalcz/21-send-logic

Re #21: Added more elaborate logic to decide when we should send data…

loki-log4j2-appender/src/main/java/pl/tkowalcz/tjahzi/log4j2/LokiAppenderBuilder.java

@@ -16,6 +16,7 @@ import pl.tkowalcz.tjahzi.stats.MutableMonitoringModuleWrapper;
 import pl.tkowalcz.tjahzi.stats.StandardMonitoringModule;
 
 import java.util.HashMap;
+import java.util.concurrent.TimeUnit;
 import java.util.stream.Stream;
 
 import static java.util.Arrays.stream;
@@ -58,6 +59,12 @@ public class LokiAppenderBuilder<B extends LokiAppenderBuilder<B>> extends Abstr
     @PluginBuilderAttribute
     private String logLevelLabel;
 
+    @PluginBuilderAttribute
+    private final long batchSize = 10_2400;
+
+    @PluginBuilderAttribute
+    private final long batchWait = 5;
+
     @PluginBuilderAttribute
     private int maxRequestsInFlight = 100;
 
@@ -109,6 +116,8 @@ public class LokiAppenderBuilder<B extends LokiAppenderBuilder<B>> extends Abstr
                 httpClient,
                 monitoringModuleWrapper,
                 lokiLabels,
+                batchSize,
+                TimeUnit.SECONDS.toMillis(batchWait),
                 bufferSizeBytes,
                 isUseOffHeapBuffer()
         );
@@ -196,6 +205,14 @@ public class LokiAppenderBuilder<B extends LokiAppenderBuilder<B>> extends Abstr
         this.logLevelLabel = logLevelLabel;
     }
 
+    public long getBatchSize() {
+        return batchSize;
+    }
+
+    public long getBatchWait() {
+        return batchWait;
+    }
+
     public void setMaxRequestsInFlight(int maxRequestsInFlight) {
         this.maxRequestsInFlight = maxRequestsInFlight;
     }

tjahzi-core/src/main/java/pl/tkowalcz/tjahzi/LogBufferAgent.java

@@ -11,26 +11,41 @@ import org.agrona.concurrent.ringbuffer.ManyToOneRingBuffer;
 import pl.tkowalcz.tjahzi.http.NettyHttpClient;
 
 import java.io.IOException;
+import java.time.Clock;
 import java.util.Map;
 
 public class LogBufferAgent implements Agent, MessageHandler {
 
-    public static final int MESSAGE_COUNT_LIMIT = 100;
+    public static final int MAX_MESSAGES_TO_RETRIEVE = 100;
+
+    private final Clock clock;
 
     private final ManyToOneRingBuffer logBuffer;
     private final NettyHttpClient httpClient;
+    private final long batchSize;
+    private final long batchWaitMillis;
 
     private final Map<String, String> staticLabels;
     private final LogBufferDeserializer logBufferDeserializer = new LogBufferDeserializer();
 
     private Logproto.PushRequest.Builder request = Logproto.PushRequest.newBuilder();
 
+    private int estimatedBytesPending;
+    private long timeoutDeadline;
+
     public LogBufferAgent(
+            Clock clock,
             ManyToOneRingBuffer logBuffer,
             NettyHttpClient httpClient,
+            long batchSizeBytes,
+            long batchWaitMillis,
             Map<String, String> staticLabels
     ) {
+        this.clock = clock;
+
         this.logBuffer = logBuffer;
+        this.batchSize = batchSizeBytes;
+        this.batchWaitMillis = batchWaitMillis;
         this.httpClient = httpClient;
 
         this.staticLabels = staticLabels;
@@ -38,9 +53,10 @@ public class LogBufferAgent implements Agent, MessageHandler {
 
     @Override
     public int doWork() throws IOException {
-        int workDone = logBuffer.read(this, MESSAGE_COUNT_LIMIT);
+        int workDone = logBuffer.read(this, MAX_MESSAGES_TO_RETRIEVE);
 
-        if (workDone > 0) {
+        long currentTimeMillis = clock.millis();
+        if (exceededBatchSizeThreshold() || exceededWaitTimeThreshold(currentTimeMillis)) {
             ByteBuf buffer = PooledByteBufAllocator.DEFAULT.buffer();
             request.build().writeTo(
                     new ByteBufOutputStream(buffer)
@@ -48,6 +64,9 @@ public class LogBufferAgent implements Agent, MessageHandler {
 
             httpClient.log(buffer);
             request = Logproto.PushRequest.newBuilder();
+
+            estimatedBytesPending = 0;
+            timeoutDeadline = currentTimeMillis + batchWaitMillis;
         }
 
         return workDone;
@@ -61,14 +80,28 @@ public class LogBufferAgent implements Agent, MessageHandler {
             int length
     ) {
         if (msgTypeId == TjahziLogger.LOG_MESSAGE_TYPE_ID) {
+            estimatedBytesPending += length;
             processMessage(buffer, index);
         } else {
             // Ignore
         }
     }
 
+    @Override
+    public String roleName() {
+        return "ReadingLogBufferAndSendingHttp";
+    }
+
+    private boolean exceededWaitTimeThreshold(long currentTimeMillis) {
+        return currentTimeMillis > timeoutDeadline & estimatedBytesPending > 0;
+    }
+
+    private boolean exceededBatchSizeThreshold() {
+        return estimatedBytesPending > batchSize;
+    }
+
     private void processMessage(MutableDirectBuffer buffer, int index) {
-        Logproto.StreamAdapter stream = logBufferDeserializer.deserialize(
+        Logproto.StreamAdapter stream = logBufferDeserializer.deserializeIntoProtobuf(
                 buffer,
                 index,
                 staticLabels
@@ -76,9 +109,4 @@ public class LogBufferAgent implements Agent, MessageHandler {
 
         request.addStreams(stream);
     }
-
-    @Override
-    public String roleName() {
-        return "ReadingLogBufferAndSendingHttp";
-    }
 }

tjahzi-core/src/main/java/pl/tkowalcz/tjahzi/LogBufferDeserializer.java

@@ -11,7 +11,7 @@ import java.util.Map;
 
 public class LogBufferDeserializer {
 
-    public Logproto.StreamAdapter deserialize(
+    public Logproto.StreamAdapter deserializeIntoProtobuf(
             DirectBuffer buffer,
             int index,
             Map<String, String> staticLabels

tjahzi-core/src/main/java/pl/tkowalcz/tjahzi/TjahziInitializer.java

@@ -19,6 +19,8 @@ public class TjahziInitializer {
             NettyHttpClient httpClient,
             MonitoringModule monitoringModule,
             Map<String, String> staticLabels,
+            long batchSizeBytes,
+            long batchWaitMillis,
             int bufferSizeBytes,
             boolean offHeap) {
         bufferSizeBytes = findNearestPowerOfTwo(bufferSizeBytes);
@@ -29,8 +31,11 @@ public class TjahziInitializer {
         );
 
         LogBufferAgent agent = new LogBufferAgent(
+                monitoringModule.getClock(),
                 logBuffer,
                 httpClient,
+                batchSizeBytes,
+                batchWaitMillis,
                 staticLabels
         );
 

tjahzi-core/src/main/java/pl/tkowalcz/tjahzi/http/NettyHttpClient.java

@@ -34,6 +34,7 @@ public class NettyHttpClient implements Closeable {
     public void log(ByteBuf dataBuffer) {
         ByteBuf output = PooledByteBufAllocator.DEFAULT.buffer();
         snappy.encode(dataBuffer, output, dataBuffer.readableBytes());
+        dataBuffer.release();
 
         FullHttpRequest request = new DefaultFullHttpRequest(
                 HttpVersion.HTTP_1_1,

tjahzi-core/src/main/java/pl/tkowalcz/tjahzi/http/RequestAndResponseHandler.java

@@ -8,6 +8,7 @@ import io.netty.handler.codec.http.FullHttpRequest;
 import io.netty.handler.codec.http.FullHttpResponse;
 import io.netty.handler.codec.http.HttpStatusClass;
 import io.netty.handler.codec.http.HttpUtil;
+import io.netty.util.ReferenceCountUtil;
 import pl.tkowalcz.tjahzi.stats.MonitoringModule;
 
 @ChannelHandler.Sharable
@@ -41,6 +42,8 @@ class RequestAndResponseHandler extends ChannelDuplexHandler {
         if (!HttpUtil.isKeepAlive(msg)) {
             ctx.close();
         }
+
+        ReferenceCountUtil.release(msg);
     }
 
     @Override

tjahzi-core/src/test/java/pl/tkowalcz/tjahzi/HeadersTest.java

@@ -17,7 +17,6 @@ import java.util.concurrent.TimeUnit;
 import static com.github.tomakehurst.wiremock.client.WireMock.*;
 import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig;
 import static org.awaitility.Awaitility.await;
-import static org.mockito.Mockito.mock;
 
 class HeadersTest {
 
@@ -38,7 +37,7 @@ class HeadersTest {
         wireMockServer.start();
 
         initializer = new TjahziInitializer();
-        monitoringModule = mock(StandardMonitoringModule.class);
+        monitoringModule = new StandardMonitoringModule();
     }
 
     @AfterEach
@@ -83,6 +82,8 @@ class HeadersTest {
                 httpClient,
                 monitoringModule,
                 Map.of(),
+                0,
+                0,
                 1024 * 1024,
                 false
         );
@@ -134,6 +135,8 @@ class HeadersTest {
                 httpClient,
                 monitoringModule,
                 Map.of(),
+                0,
+                0,
                 1024 * 1024,
                 false
         );
@@ -187,6 +190,8 @@ class HeadersTest {
                 httpClient,
                 monitoringModule,
                 Map.of(),
+                0,
+                0,
                 1024 * 1024,
                 false
         );

tjahzi-core/src/test/java/pl/tkowalcz/tjahzi/LogBufferAgentTest.java

@@ -0,0 +1,184 @@
+package pl.tkowalcz.tjahzi;
+
+import org.agrona.concurrent.UnsafeBuffer;
+import org.agrona.concurrent.ringbuffer.ManyToOneRingBuffer;
+import org.agrona.concurrent.ringbuffer.RingBufferDescriptor;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import pl.tkowalcz.tjahzi.http.NettyHttpClient;
+import pl.tkowalcz.tjahzi.stats.SettableClock;
+import pl.tkowalcz.tjahzi.stats.StandardMonitoringModule;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Map;
+
+import static org.mockito.Matchers.any;
+import static org.mockito.Mockito.*;
+
+class LogBufferAgentTest {
+
+    private ManyToOneRingBuffer logBuffer;
+    private TjahziLogger logger;
+    private SettableClock clock;
+
+    @BeforeEach
+    private void setUp() {
+        logBuffer = new ManyToOneRingBuffer(
+                new UnsafeBuffer(new byte[1024 * 1024 + RingBufferDescriptor.TRAILER_LENGTH])
+        );
+
+        logger = new TjahziLogger(logBuffer, new StandardMonitoringModule());
+        clock = new SettableClock();
+    }
+
+    @Test
+    void shouldSendDataIfOverSizeLimit() throws IOException {
+        // Given
+        long waitForever = Long.MAX_VALUE;
+        int sentBatchAfter5kb = 5 * 1024;
+
+        NettyHttpClient httpClient = mock(NettyHttpClient.class);
+        LogBufferAgent agent = new LogBufferAgent(
+                clock,
+                logBuffer,
+                httpClient,
+                sentBatchAfter5kb,
+                waitForever,
+                Map.of()
+        );
+
+        for (int i = 0; i < 100; i++) {
+            logger.log(
+                    42L,
+                    Map.of(),
+                    null,
+                    null,
+                    ByteBuffer.wrap((
+                            "Cupcake ipsum dolor sit amet cake wafer. " +
+                                    "Soufflé jelly beans biscuit topping. " +
+                                    "Danish bonbon gummies powder caramels. " +
+                                    "Danish jelly beans sweet roll topping jelly beans oat cake toffee. " +
+                                    "Chocolate cake sesame snaps brownie biscuit cheesecake. " +
+                                    "Ice cream dessert sweet donut marshmallow. " +
+                                    "Muffin bear claw cookie jelly-o sugar plum jelly beans apple pie fruitcake cookie. " +
+                                    "Tootsie roll carrot cake pastry jujubes jelly beans chupa chups. " +
+                                    "Soufflé cake muffin liquorice tart soufflé pie sesame snaps."
+                    ).getBytes())
+            );
+        }
+
+        // When
+        agent.doWork();
+
+        // Then
+        verify(httpClient).log(any());
+    }
+
+    @Test
+    void shouldNotSendDataBelowSizeLimit() throws IOException {
+        // Given
+        int wait5s = 5000;
+        int sentBatchAfter5kb = 5 * 1024;
+
+        NettyHttpClient httpClient = mock(NettyHttpClient.class);
+        LogBufferAgent agent = new LogBufferAgent(
+                clock,
+                logBuffer,
+                httpClient,
+                sentBatchAfter5kb,
+                wait5s,
+                Map.of()
+        );
+
+        logger.log(
+                42L,
+                Map.of(),
+                null,
+                null,
+                ByteBuffer.wrap("Test".getBytes())
+        );
+
+        // When
+        agent.doWork();
+
+        // Then
+        verifyZeroInteractions(httpClient);
+    }
+
+    @Test
+    void shouldNotAttemptSendingDataIfThereIsNothingToSend() throws IOException {
+        // Given
+        int wait5s = 5000;
+        int sentBatchAfter5kb = 5 * 1024;
+
+        NettyHttpClient httpClient = mock(NettyHttpClient.class);
+        LogBufferAgent agent = new LogBufferAgent(
+                clock,
+                logBuffer,
+                httpClient,
+                sentBatchAfter5kb,
+                wait5s,
+                Map.of()
+        );
+
+        // When
+        agent.doWork();
+        clock.tick(wait5s + 1);
+        agent.doWork();
+        clock.tick(wait5s + 1);
+        agent.doWork();
+
+        // Then
+        verifyZeroInteractions(httpClient);
+    }
+
+    @Test
+    void shouldSendDataBelowSizeLimitIfTimeoutExpires() throws IOException {
+        // Given
+        int wait5s = 5000;
+        int sentBatchAfter5kb = 5 * 1024;
+
+        ManyToOneRingBuffer logBuffer = new ManyToOneRingBuffer(
+                new UnsafeBuffer(new byte[1024 + RingBufferDescriptor.TRAILER_LENGTH])
+        );
+
+        TjahziLogger logger = new TjahziLogger(logBuffer, new StandardMonitoringModule());
+
+        NettyHttpClient httpClient = mock(NettyHttpClient.class);
+        SettableClock clock = new SettableClock();
+
+        LogBufferAgent agent = new LogBufferAgent(
+                clock,
+                logBuffer,
+                httpClient,
+                sentBatchAfter5kb,
+                wait5s,
+                Map.of()
+        );
+
+        // When
+        logger.log(
+                42L,
+                Map.of(),
+                null,
+                null,
+                ByteBuffer.wrap("Test".getBytes())
+        );
+        clock.tick(wait5s + 1);
+        agent.doWork();
+
+        logger.log(
+                42L,
+                Map.of(),
+                null,
+                null,
+                ByteBuffer.wrap("Test".getBytes())
+        );
+        clock.tick(wait5s + 1);
+        agent.doWork();
+
+        // Then
+        verify(httpClient, times(2)).log(any());
+    }
+}

tjahzi-core/src/test/java/pl/tkowalcz/tjahzi/LogBufferDeserializerTest.java

@@ -46,7 +46,7 @@ class LogBufferDeserializerTest {
 
         // When
         LogBufferDeserializer deserializer = new LogBufferDeserializer();
-        Logproto.StreamAdapter stream = deserializer.deserialize(
+        Logproto.StreamAdapter stream = deserializer.deserializeIntoProtobuf(
                 buffer,
                 0,
                 staticLabels
@@ -101,7 +101,7 @@ class LogBufferDeserializerTest {
 
         // When
         LogBufferDeserializer deserializer = new LogBufferDeserializer();
-        Logproto.StreamAdapter stream = deserializer.deserialize(
+        Logproto.StreamAdapter stream = deserializer.deserializeIntoProtobuf(
                 buffer,
                 0,
                 staticLabels

tjahzi-core/src/test/java/pl/tkowalcz/tjahzi/LogBufferSerializerDeserializerTest.java

@@ -90,7 +90,7 @@ class LogBufferSerializerDeserializerTest {
 
         // Then
         LogBufferDeserializer deserializer = new LogBufferDeserializer();
-        Logproto.StreamAdapter stream = deserializer.deserialize(buffer, 0, Map.of());
+        Logproto.StreamAdapter stream = deserializer.deserializeIntoProtobuf(buffer, 0, Map.of());
 
         assertThat(stream.getEntriesList()).hasSize(1);
         assertThat(stream.getEntriesList().get(0).getTimestamp()).isEqualTo(

tjahzi-core/src/test/java/pl/tkowalcz/tjahzi/LoggingSystemSanityCheckTest.java

@@ -25,7 +25,6 @@ import java.util.concurrent.TimeUnit;
 import static io.restassured.RestAssured.given;
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.CoreMatchers.hasItems;
-import static org.mockito.Mockito.mock;
 
 @Testcontainers
 class LoggingSystemSanityCheckTest {
@@ -43,7 +42,7 @@ class LoggingSystemSanityCheckTest {
             .withExposedPorts(3100);
 
     private LoggingSystem loggingSystem;
-    private final StandardMonitoringModule monitoringModule = mock(StandardMonitoringModule.class);
+    private final StandardMonitoringModule monitoringModule = new StandardMonitoringModule();
 
     @BeforeEach
     void setUp() {
@@ -65,6 +64,8 @@ class LoggingSystemSanityCheckTest {
                 httpClient,
                 monitoringModule,
                 Map.of("version", "0.43", "server", "127.0.0.1"),
+                0,
+                0,
                 1024 * 1024,
                 false
         );

tjahzi-core/src/test/java/pl/tkowalcz/tjahzi/LoggingSystemTest.java

@@ -28,7 +28,6 @@ import java.util.concurrent.TimeUnit;
 import static io.restassured.RestAssured.given;
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.CoreMatchers.hasItems;
-import static org.mockito.Mockito.mock;
 
 @Testcontainers
 class LoggingSystemTest {
@@ -45,7 +44,7 @@ class LoggingSystemTest {
             )
             .withExposedPorts(3100);
 
-    private final StandardMonitoringModule monitoringModule = mock(StandardMonitoringModule.class);
+    private final StandardMonitoringModule monitoringModule = new StandardMonitoringModule();
     private LoggingSystem loggingSystem;
 
     @BeforeEach
@@ -69,6 +68,8 @@ class LoggingSystemTest {
                 monitoringModule,
                 Map.of("version", "0.43", "server", "127.0.0.1"),
                 1024 * 1024,
+                0,
+                0,
                 false
         );
 

tjahzi-core/src/test/java/pl/tkowalcz/tjahzi/ReconnectTest.java

@@ -17,7 +17,6 @@ import java.util.concurrent.TimeUnit;
 import static com.github.tomakehurst.wiremock.client.WireMock.*;
 import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig;
 import static org.awaitility.Awaitility.await;
-import static org.mockito.Mockito.mock;
 
 public class ReconnectTest {
 
@@ -41,7 +40,7 @@ public class ReconnectTest {
                         ));
 
         initializer = new TjahziInitializer();
-        monitoringModule = mock(StandardMonitoringModule.class);
+        monitoringModule = new StandardMonitoringModule();
     }
 
     @AfterEach
@@ -78,6 +77,8 @@ public class ReconnectTest {
                 httpClient,
                 monitoringModule,
                 Map.of(),
+                0,
+                0,
                 1024 * 1024,
                 false
         );
@@ -127,6 +128,8 @@ public class ReconnectTest {
                 httpClient,
                 monitoringModule,
                 Map.of(),
+                0,
+                0,
                 1024 * 1024,
                 false
         );

tjahzi-core/src/test/java/pl/tkowalcz/tjahzi/ResourcesCleanupOnCloseTest.java

@@ -68,6 +68,8 @@ class ResourcesCleanupOnCloseTest {
                 httpClient,
                 monitoringModule,
                 Map.of(),
+                0,
+                0,
                 1024 * 1024,
                 false
         );