Re #21: Added more elaborate logic to decide when we should send data to Loki. Using terminology from promtail introduced 'batchSize' and 'batchWait' parameters. Fixed ByteBuf and HttpMessage leak due to missing release.

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
         );