exporter and midBuffer items

Author: laullon

proxy/pom.xml

@@ -389,7 +389,7 @@
     <dependency>
       <groupId>org.apache.activemq</groupId>
       <artifactId>artemis-server</artifactId>
-      <version>2.27.0</version>
+      <version>2.27.1</version>
     </dependency>
     <dependency>
       <groupId>com.wavefront</groupId>

proxy/src/main/java/com/wavefront/agent/AbstractAgent.java

@@ -15,6 +15,7 @@
 import com.wavefront.agent.api.APIContainer;
 import com.wavefront.agent.config.LogsIngestionConfig;
 import com.wavefront.agent.core.buffers.BuffersManager;
+import com.wavefront.agent.core.buffers.Exporter;
 import com.wavefront.agent.core.senders.SenderTasksManager;
 import com.wavefront.agent.data.EntityPropertiesFactoryImpl;
 import com.wavefront.agent.logsharvesting.InteractiveLogsTester;
@@ -236,9 +237,17 @@ public void start(String[] args) {
       }
 
       // If we are exporting data from the queue, run export and exit
-      // TODO: queue exporter
-      if (proxyConfig.getExportQueueOutputFile() != null
-          && proxyConfig.getExportQueuePorts() != null) {
+      if (proxyConfig.getExportQueueOutputDir() != null
+          && proxyConfig.getExportQueueAtoms() != null) {
+        try {
+          Exporter.export(
+              proxyConfig.getBufferFile(),
+              proxyConfig.getExportQueueOutputDir(),
+              proxyConfig.getExportQueueAtoms(),
+              proxyConfig.isExportQueueRetainData());
+        } catch (Throwable e) {
+          System.out.println(e.getMessage());
+        }
         System.exit(0);
       }
 

proxy/src/main/java/com/wavefront/agent/ProxyConfig.java

@@ -165,19 +165,19 @@ public class ProxyConfig extends Configuration {
 
   // TODO: review export buffer
   @Parameter(
-      names = {"--exportQueuePorts"},
+      names = {"--exportQueueAtoms"},
       description =
           "Export queued data in plaintext "
-              + "format for specified ports (comma-delimited list) and exit. Set to 'all' to export "
-              + "everything. Default: none")
-  String exportQueuePorts = null;
+              + "format for specified atoms (comma-delimited list) and exit. Set to 'all' to export "
+              + "everything. Default: none, valid values: points, deltaCounters, histograms, sourceTags, spans, spanLogs, events, logs")
+  String exportQueueAtoms = null;
 
   @Parameter(
-      names = {"--exportQueueOutputFile"},
+      names = {"--exportQueueOutputDir"},
       description =
           "Export queued data in plaintext "
               + "format for specified ports (comma-delimited list) and exit. Default: none")
-  String exportQueueOutputFile = null;
+  String exportQueueOutputDir = null;
 
   @Parameter(
       names = {"--exportQueueRetainData"},
@@ -1236,12 +1236,12 @@ public String getSqsQueueIdentifier() {
     return sqsQueueIdentifier;
   }
 
-  public String getExportQueuePorts() {
-    return exportQueuePorts;
+  public String getExportQueueAtoms() {
+    return exportQueueAtoms;
   }
 
-  public String getExportQueueOutputFile() {
-    return exportQueueOutputFile;
+  public String getExportQueueOutputDir() {
+    return exportQueueOutputDir;
   }
 
   public boolean isExportQueueRetainData() {
@@ -2233,8 +2233,8 @@ public void verifyAndInit() {
       customApplicationTags = config.getString("customApplicationTags", customApplicationTags);
       customServiceTags = config.getString("customServiceTags", customServiceTags);
 
-      exportQueuePorts = config.getString("exportQueuePorts", exportQueuePorts);
-      exportQueueOutputFile = config.getString("exportQueueOutputFile", exportQueueOutputFile);
+      exportQueueAtoms = config.getString("exportQueueAtoms", exportQueueAtoms);
+      exportQueueOutputDir = config.getString("exportQueueOutputDir", exportQueueOutputDir);
       exportQueueRetainData = config.getBoolean("exportQueueRetainData", exportQueueRetainData);
       flushThreads = config.getInteger("flushThreads", flushThreads);
       flushThreadsEvents = config.getInteger("flushThreadsEvents", flushThreadsEvents);

proxy/src/main/java/com/wavefront/agent/core/buffers/ActiveMQBuffer.java

@@ -44,7 +44,7 @@ public abstract class ActiveMQBuffer implements Buffer {
   final String name;
   private final Map<String, Session> producers = new ConcurrentHashMap<>();
   private final Map<String, Session> consumers = new ConcurrentHashMap<>();
-  private final Map<String, Gauge<Object>> sizeMetrics = new HashMap<>(); // TODO review
+  private final Map<String, Gauge<Object>> sizeMetrics = new HashMap<>();
   private final Map<String, Histogram> timeMetrics = new HashMap<>();
   private final int serverID;
   protected Buffer nextBuffer;

proxy/src/main/java/com/wavefront/agent/core/buffers/BuffersManager.java

@@ -74,10 +74,7 @@ public static void sendMsg(QueueInfo queue, String strPoint) {
   }
 
   public static void onMsgBatch(
-      QueueInfo handler,
-      int idx,
-      EntityRateLimiter rateLimiter,
-      OnMsgDelegate func) {
+      QueueInfo handler, int idx, EntityRateLimiter rateLimiter, OnMsgDelegate func) {
     memoryBuffer.onMsgBatch(handler, idx, func);
   }
 

proxy/src/main/java/com/wavefront/agent/core/buffers/Exporter.java

@@ -0,0 +1,83 @@
+package com.wavefront.agent.core.buffers;
+
+import com.wavefront.data.ReportableEntityType;
+import java.io.BufferedWriter;
+import java.io.File;
+import java.io.FileWriter;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+import org.apache.activemq.artemis.api.core.SimpleString;
+import org.apache.activemq.artemis.core.message.impl.CoreMessage;
+import org.apache.activemq.artemis.core.server.ActiveMQServer;
+import org.apache.activemq.artemis.core.server.MessageReference;
+import org.apache.activemq.artemis.utils.collections.LinkedListIterator;
+
+public class Exporter {
+  public static void export(String bufferStr, String dirStr, String atomsStr, boolean retainData) {
+    List<String> atomsList = Arrays.asList(atomsStr.split(","));
+    atomsList.replaceAll(String::trim);
+    List<ReportableEntityType> atoms =
+        atomsList.stream()
+            .map(
+                s -> {
+                  ReportableEntityType atom = ReportableEntityType.fromString(s);
+                  if (atom == null) {
+                    throw new IllegalArgumentException("invalid atom '" + s + "'");
+                  }
+                  return atom;
+                })
+            .collect(Collectors.toList());
+    File dir = new File(dirStr);
+
+    DiskBufferConfig config = new DiskBufferConfig();
+    config.buffer = new File(bufferStr);
+    DiskBuffer buffer = new DiskBuffer(1, "disk", config);
+    atoms.forEach(
+        atom -> {
+          ActiveMQServer amq = buffer.activeMQServer;
+          try {
+            File outFile = new File(dir, atom.toString().toLowerCase() + ".txt");
+            System.out.println(
+                "Exporting '" + atom + "' from '" + dirStr + "' to '" + outFile + "'");
+            AtomicInteger c= new AtomicInteger();
+            BufferedWriter out = new BufferedWriter(new FileWriter(outFile));
+            amq.getPostOffice()
+                .listQueuesForAddress(SimpleString.toSimpleString(atom.name()))
+                .forEach(
+                    queue -> {
+                      LinkedListIterator<MessageReference> it = queue.browserIterator();
+                      while (it.hasNext()) {
+                        CoreMessage msg = (CoreMessage) it.next().getMessage();
+                        List<String> points =
+                            Arrays.asList(msg.getReadOnlyBodyBuffer().readString().split("\n"));
+                        try {
+                          out.write(String.join("\n", points));
+                          out.write("\n");
+                        } catch (IOException e) {
+                          throw new RuntimeException("Error writing on the output file.", e);
+                        }
+                        if (!retainData) {
+                          try {
+                            queue.deleteReference(msg.getMessageID());
+                          } catch (Exception e) {
+                            throw new RuntimeException("Error deleting data from the buffer", e);
+                          }
+                        }
+                        if (c.addAndGet(points.size()) %100_000==0){
+                          System.out.print(".");
+                        }
+                      }
+                    });
+            out.flush();
+            out.close();
+            System.out.println(".\nDone, exported "+(retainData?"":"and deleted ")+c+" "+atom.toString()+"\n");
+          } catch (Exception e) {
+            throw new RuntimeException(e);
+          }
+        });
+    buffer.shutdown();
+  }
+}

proxy/src/main/java/com/wavefront/agent/core/buffers/MemoryBuffer.java

@@ -65,7 +65,6 @@ public void shutdown() {
       for (Object obj : queues) {
         QueueControl queue = (QueueControl) obj;
         int c = queue.expireMessages("");
-        System.out.println("-> queue: " + queue.getName() + " - " + c);
         counter += c;
       }
     } catch (Exception e) {
@@ -138,7 +137,7 @@ public void run() {
         boolean done = false;
         while (!done) {
           ArrayList<String> metrics = new ArrayList<>();
-          if (midBuffer.drainTo(metrics, 100) != 0) {
+          if (midBuffer.drainTo(metrics, queue.getMaxItemsPerMessage()) != 0) {
             try {
               sendPoints(queue.getName(), metrics);
             } catch (ActiveMQAddressFullException e) {

proxy/src/main/java/com/wavefront/agent/core/buffers/SQSBuffer.java

@@ -87,7 +87,7 @@ public void onMsgBatch(QueueInfo queue, int idx, OnMsgDelegate func) {
         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);
+        done = !func.checkBatchSize(batch.size(), 0, 0, 0);
       } else {
         done = true;
       }

proxy/src/main/java/com/wavefront/agent/core/queues/Queue.java

@@ -12,12 +12,23 @@ class Queue implements QueueInfo {
   private final String tenant;
   private final int threads;
   private final Map<String, QueueInfo> tenants = new HashMap<>();
+  private final int midBufferItems;
 
   Queue(ReportableEntityType entityType, String tenant, int threads) {
     this.name = entityType + (tenant.equalsIgnoreCase(CENTRAL_TENANT_NAME) ? "" : "." + tenant);
     this.entityType = entityType;
     this.tenant = tenant;
     this.threads = threads;
+    switch (entityType) {
+      case LOGS:
+        midBufferItems = 10;
+        break;
+      case POINT:
+        midBufferItems = 255;
+        break;
+      default:
+        midBufferItems = 100;
+    }
     QueueStats.register(this);
   }
 
@@ -48,6 +59,11 @@ public int getNumberThreads() {
     return threads;
   }
 
+  @Override
+  public int getMaxItemsPerMessage() {
+    return midBufferItems;
+  }
+
   public void addTenant(String tenant, Queue queue) {
     tenants.put(tenant, queue);
   }

proxy/src/main/java/com/wavefront/agent/core/queues/QueueInfo.java

@@ -15,4 +15,6 @@ public interface QueueInfo {
   String getName();
 
   int getNumberThreads();
+
+  int getMaxItemsPerMessage();
 }