batch limit control for items and bytes (bytes only on logs for now)

laullon · laullon · commit f3f6f9298d73 · 2022-12-07T18:52:41.000+01:00
diff --git a/proxy/src/main/java/com/wavefront/agent/core/buffers/ActiveMQBuffer.java b/proxy/src/main/java/com/wavefront/agent/core/buffers/ActiveMQBuffer.java
@@ -5,7 +5,6 @@
 
 import com.wavefront.agent.core.queues.QueueInfo;
 import com.wavefront.agent.core.queues.QueueStats;
-import com.wavefront.agent.data.EntityRateLimiter;
 import com.yammer.metrics.Metrics;
 import com.yammer.metrics.core.Gauge;
 import com.yammer.metrics.core.Histogram;
@@ -35,6 +34,8 @@ public abstract class ActiveMQBuffer implements Buffer {
   private static final Logger log =
       LoggerFactory.getLogger(ActiveMQBuffer.class.getCanonicalName());
   private static final Logger slowLog = log;
+  public static final String MSG_ITEMS = "items";
+  public static final String MSG_BYTES = "bytes";
   //      new
   // MessageDedupingLogger(LoggerFactory.getLogger(ActiveMQBuffer.class.getCanonicalName()), 1000,
   // 1);
@@ -221,8 +222,10 @@ public void doSendPoints(String queue, List<String> points) throws ActiveMQAddre
 
     try {
       ClientMessage message = mqCtx.session.createMessage(true);
-      message.writeBodyBufferString(String.join("\n", points));
-      message.putIntProperty("points", points.size());
+      String str = String.join("\n", points);
+      message.writeBodyBufferString(str);
+      message.putIntProperty(MSG_ITEMS, points.size());
+      message.putIntProperty(MSG_BYTES, str.length());
       mqCtx.producer.send(message);
     } catch (ActiveMQAddressFullException e) {
       log.info("queue full: " + e.getMessage());
@@ -253,8 +256,7 @@ private void checkConnection() throws Exception {
   }
 
   @Override
-  public void onMsgBatch(
-      QueueInfo queue, int idx, int batchSize, EntityRateLimiter rateLimiter, OnMsgFunction func) {
+  public void onMsgBatch(QueueInfo queue, int idx, OnMsgDelegate delegate) {
     String sessionKey = "onMsgBatch." + queue.getName() + "." + Thread.currentThread().getName();
     Session mqCtx =
         consumers.computeIfAbsent(
@@ -277,20 +279,29 @@ public void onMsgBatch(
     try {
       long start = System.currentTimeMillis();
       mqCtx.session.start();
-      List<String> batch = new ArrayList<>(batchSize);
+      List<String> batch = new ArrayList<>();
       List<ClientMessage> toACK = new ArrayList<>();
       boolean done = false;
       boolean needRollBack = false;
-      while ((batch.size() < batchSize) && !done && ((System.currentTimeMillis() - start) < 1000)) {
+      int batchBytes = 0;
+      while (!done && ((System.currentTimeMillis() - start) < 1000)) {
         ClientMessage msg = mqCtx.consumer.receive(100);
         if (msg != null) {
           List<String> points = Arrays.asList(msg.getReadOnlyBodyBuffer().readString().split("\n"));
-          boolean ok = rateLimiter.tryAcquire(points.size());
-          if (ok) {
+          boolean ok_size =
+              delegate.checkBatchSize(
+                  batch.size(), batchBytes, points.size(), msg.getIntProperty(MSG_BYTES));
+          boolean ok_rate = delegate.checkRates(points.size(), batchBytes);
+          if (ok_size && ok_rate) {
             toACK.add(msg);
             batch.addAll(points);
+            batchBytes += msg.getIntProperty(MSG_BYTES);
           } else {
-            slowLog.info("rate limit reached on queue '" + queue.getName() + "'");
+            if (!ok_rate) {
+              slowLog.info("rate limit reached on queue '" + queue.getName() + "'");
+            } else {
+              slowLog.info("payload limit reached on queue '" + queue.getName() + "'");
+            }
             done = true;
             needRollBack = true;
           }
@@ -301,7 +312,7 @@ public void onMsgBatch(
 
       try {
         if (batch.size() > 0) {
-          func.run(batch);
+          delegate.processBatch(batch);
         }
         // commit all messages ACKed
         toACK.forEach(
diff --git a/proxy/src/main/java/com/wavefront/agent/core/buffers/Bridge.java b/proxy/src/main/java/com/wavefront/agent/core/buffers/Bridge.java
@@ -1,5 +1,7 @@
 package com.wavefront.agent.core.buffers;
 
+import static com.wavefront.agent.core.buffers.ActiveMQBuffer.MSG_ITEMS;
+
 import com.wavefront.agent.core.queues.QueueStats;
 import java.util.Arrays;
 import java.util.List;
@@ -55,10 +57,10 @@ public void messageAcknowledged(MessageReference ref, AckReason reason, ServerCo
         diskBuffer.sendPoints(queue, points);
         switch (reason) {
           case KILLED:
-            stats.queuedFailed.inc(ref.getMessage().getIntProperty("points"));
+            stats.queuedFailed.inc(ref.getMessage().getIntProperty(MSG_ITEMS));
             break;
           case EXPIRED:
-            stats.queuedExpired.inc(ref.getMessage().getIntProperty("points"));
+            stats.queuedExpired.inc(ref.getMessage().getIntProperty(MSG_ITEMS));
             break;
         }
       } catch (ActiveMQAddressFullException e) {
diff --git a/proxy/src/main/java/com/wavefront/agent/core/buffers/Buffer.java b/proxy/src/main/java/com/wavefront/agent/core/buffers/Buffer.java
@@ -1,15 +1,13 @@
 package com.wavefront.agent.core.buffers;
 
 import com.wavefront.agent.core.queues.QueueInfo;
-import com.wavefront.agent.data.EntityRateLimiter;
 import java.util.List;
 import org.apache.activemq.artemis.api.core.ActiveMQAddressFullException;
 
 public interface Buffer {
   void registerNewQueueInfo(QueueInfo key);
 
-  void onMsgBatch(
-      QueueInfo key, int idx, int batchSize, EntityRateLimiter rateLimiter, OnMsgFunction func);
+  void onMsgBatch(QueueInfo key, int idx, OnMsgDelegate func);
 
   void sendPoints(String queue, List<String> strPoint) throws ActiveMQAddressFullException;
 
diff --git a/proxy/src/main/java/com/wavefront/agent/core/buffers/BuffersManager.java b/proxy/src/main/java/com/wavefront/agent/core/buffers/BuffersManager.java
@@ -76,10 +76,9 @@ public static void sendMsg(QueueInfo queue, String strPoint) {
   public static void onMsgBatch(
       QueueInfo handler,
       int idx,
-      int batchSize,
       EntityRateLimiter rateLimiter,
-      OnMsgFunction func) {
-    memoryBuffer.onMsgBatch(handler, idx, batchSize, rateLimiter, func);
+      OnMsgDelegate func) {
+    memoryBuffer.onMsgBatch(handler, idx, func);
   }
 
   public static void truncateBacklog() {
diff --git a/proxy/src/main/java/com/wavefront/agent/core/buffers/OnMsgDelegate.java b/proxy/src/main/java/com/wavefront/agent/core/buffers/OnMsgDelegate.java
@@ -0,0 +1,11 @@
+package com.wavefront.agent.core.buffers;
+
+import java.util.List;
+
+public interface OnMsgDelegate {
+  void processBatch(List<String> batch) throws Exception;
+
+  boolean checkBatchSize(int items, int bytes, int newItems, int newBytes);
+
+  boolean checkRates(int newItems, int newBytes);
+}
diff --git a/proxy/src/main/java/com/wavefront/agent/core/buffers/OnMsgFunction.java b/proxy/src/main/java/com/wavefront/agent/core/buffers/OnMsgFunction.java
diff --git a/proxy/src/main/java/com/wavefront/agent/core/buffers/PointsGauge.java b/proxy/src/main/java/com/wavefront/agent/core/buffers/PointsGauge.java
@@ -1,5 +1,7 @@
 package com.wavefront.agent.core.buffers;
 
+import static com.wavefront.agent.core.buffers.ActiveMQBuffer.MSG_ITEMS;
+
 import com.wavefront.agent.PushAgent;
 import com.wavefront.agent.core.queues.QueueInfo;
 import com.wavefront.common.NamedThreadFactory;
@@ -47,7 +49,7 @@ long doCount() {
             (QueueControl) amq.getManagementService().getResource(ResourceNames.QUEUE + queueName);
         Map<String, Object>[] messages = queueControl.listMessages("");
         for (Map<String, Object> message : messages) {
-          int p = (int) message.get("points");
+          int p = (int) message.get(MSG_ITEMS);
           count += p;
         }
       }
diff --git a/proxy/src/main/java/com/wavefront/agent/core/buffers/SQSBuffer.java b/proxy/src/main/java/com/wavefront/agent/core/buffers/SQSBuffer.java
@@ -5,7 +5,6 @@
 import com.amazonaws.services.sqs.AmazonSQSClientBuilder;
 import com.amazonaws.services.sqs.model.*;
 import com.wavefront.agent.core.queues.QueueInfo;
-import com.wavefront.agent.data.EntityRateLimiter;
 import java.util.*;
 import org.apache.activemq.artemis.api.core.ActiveMQAddressFullException;
 import org.slf4j.Logger;
@@ -71,15 +70,14 @@ public void registerNewQueueInfo(QueueInfo queue) {
   }
 
   @Override
-  public void onMsgBatch(
-      QueueInfo queue, int idx, int batchSize, EntityRateLimiter rateLimiter, OnMsgFunction func) {
+  public void onMsgBatch(QueueInfo queue, int idx, OnMsgDelegate func) {
 
     String queueUrl = queuesUrls.get(queue.getName());
     long start = System.currentTimeMillis();
-    List<String> batch = new ArrayList<>(batchSize);
+    List<String> batch = new ArrayList<>();
     List<Message> messagesToDelete = new ArrayList<>();
     boolean done = false;
-    while ((batch.size() < batchSize) && !done && ((System.currentTimeMillis() - start) < 1000)) {
+    while (!done && ((System.currentTimeMillis() - start) < 1000)) {
       ReceiveMessageRequest receiveRequest = new ReceiveMessageRequest(queueUrl);
       receiveRequest.setMaxNumberOfMessages(1);
       receiveRequest.setWaitTimeSeconds(1);
@@ -89,14 +87,15 @@ public void onMsgBatch(
         List<String> points = Arrays.asList(messages.get(0).getBody().split("\n"));
         batch.addAll(points);
         messagesToDelete.addAll(messages);
+        done = !func.checkBatchSize(batch.size(), 0,0,0);
       } else {
         done = true;
       }
     }
 
     try {
       if (batch.size() > 0) {
-        func.run(batch);
+        func.processBatch(batch);
       }
       messagesToDelete.forEach(
           message -> {
diff --git a/proxy/src/main/java/com/wavefront/agent/core/senders/LogSenderTask.java b/proxy/src/main/java/com/wavefront/agent/core/senders/LogSenderTask.java
@@ -21,6 +21,7 @@ public class LogSenderTask extends SenderTask {
   private final QueueInfo queue;
   private final LogAPI logAPI;
   private final UUID proxyId;
+  private EntityProperties properties;
 
   /**
    * @param queue handler key, that serves as an identifier of the log pipeline.
@@ -40,6 +41,11 @@ public class LogSenderTask extends SenderTask {
     this.queue = queue;
     this.logAPI = logAPI;
     this.proxyId = proxyId;
+    this.properties = properties;
+  }
+
+  public boolean checkBatchSize(int items, int bytes, int newItems, int newBytes) {
+    return bytes + newBytes <= properties.getDataPerBatch();
   }
 
   protected Response submit(List<String> logs) {
diff --git a/proxy/src/main/java/com/wavefront/agent/core/senders/SenderTask.java b/proxy/src/main/java/com/wavefront/agent/core/senders/SenderTask.java
@@ -3,6 +3,7 @@
 import static com.wavefront.common.Utils.isWavefrontResponse;
 
 import com.wavefront.agent.core.buffers.Buffer;
+import com.wavefront.agent.core.buffers.OnMsgDelegate;
 import com.wavefront.agent.core.queues.QueueInfo;
 import com.wavefront.agent.core.queues.QueueStats;
 import com.wavefront.agent.data.EntityProperties;
@@ -22,7 +23,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-abstract class SenderTask implements Runnable {
+abstract class SenderTask implements Runnable, OnMsgDelegate {
   private static final Logger log = LoggerFactory.getLogger(SenderTask.class.getCanonicalName());
   //      new MessageDedupingLogger(LoggerFactory.getLogger(SenderTask.class.getCanonicalName()),
   // 1000, 1);
@@ -45,18 +46,24 @@ abstract class SenderTask implements Runnable {
   @Override
   public void run() {
     try {
-      buffer.onMsgBatch(
-          queue,
-          idx,
-          properties.getDataPerBatch(),
-          properties.getRateLimiter(),
-          this::processBatch);
+      buffer.onMsgBatch(queue, idx, this);
     } catch (Throwable e) {
       log.error("error sending " + queue.getEntityType().name(), e);
     }
   }
 
-  private void processBatch(List<String> batch) throws SenderTaskException {
+  @Override
+  public boolean checkBatchSize(int items, int bytes, int newItems, int newBytes) {
+    return items + newItems <= properties.getDataPerBatch();
+  }
+
+  @Override
+  public boolean checkRates(int newItems, int newBytes) {
+    return properties.getRateLimiter().tryAcquire(newItems);
+  }
+
+  @Override
+  public void processBatch(List<String> batch) throws SenderTaskException {
     TimerContext timer =
         Metrics.newTimer(
                 new MetricName("push." + queue.getName(), "", "duration"),
diff --git a/proxy/src/test/java/com/wavefront/agent/core/buffers/BufferManagerTest.java b/proxy/src/test/java/com/wavefront/agent/core/buffers/BufferManagerTest.java
