Merge pull request #243 from jiafu1115/patch-25

support async write point with different rp. Close #162
influxdata · Nov 14, 2016 · 16e79a9 · 16e79a9
2 parents 970aca8 + 3a90c2c
commit 16e79a9
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 9 deletions.
diff --git a/src/main/java/org/influxdb/impl/BatchProcessor.java b/src/main/java/org/influxdb/impl/BatchProcessor.java
@@ -186,8 +186,9 @@ void write() {
       if (this.queue.isEmpty()) {
         return;
       }
-
-      Map<String, BatchPoints> databaseToBatchPoints = Maps.newHashMap();
+      //for batch on HTTP.
+      Map<String, BatchPoints> batchKeyToBatchPoints = Maps.newHashMap();
+      //for batch on UDP.
       Map<Integer, List<String>> udpPortToBatchPoints = Maps.newHashMap();
       List<AbstractBatchEntry> batchEntries = new ArrayList<>(this.queue.size());
       this.queue.drainTo(batchEntries);
@@ -197,12 +198,14 @@ void write() {
         if (batchEntry instanceof HttpBatchEntry) {
             HttpBatchEntry httpBatchEntry = HttpBatchEntry.class.cast(batchEntry);
             String dbName = httpBatchEntry.getDb();
-            if (!databaseToBatchPoints.containsKey(dbName)) {
+            String rp = httpBatchEntry.getRp();
+            String batchKey = dbName + "_" + rp;
+            if (!batchKeyToBatchPoints.containsKey(batchKey)) {
               BatchPoints batchPoints = BatchPoints.database(dbName)
-                                                   .retentionPolicy(httpBatchEntry.getRp()).build();
-              databaseToBatchPoints.put(dbName, batchPoints);
+                                                   .retentionPolicy(rp).build();
+              batchKeyToBatchPoints.put(batchKey, batchPoints);
             }
-            databaseToBatchPoints.get(dbName).point(point);
+            batchKeyToBatchPoints.get(batchKey).point(point);
         } else if (batchEntry instanceof UdpBatchEntry) {
             UdpBatchEntry udpBatchEntry = UdpBatchEntry.class.cast(batchEntry);
             int udpPort = udpBatchEntry.getUdpPort();
@@ -214,8 +217,8 @@ void write() {
         }
       }
 
-      for (BatchPoints batchPoints : databaseToBatchPoints.values()) {
-        BatchProcessor.this.influxDB.write(batchPoints);
+      for (BatchPoints batchPoints : batchKeyToBatchPoints.values()) {
+          BatchProcessor.this.influxDB.write(batchPoints);
       }
       for (Entry<Integer, List<String>> entry : udpPortToBatchPoints.entrySet()) {
           BatchProcessor.this.influxDB.write(entry.getKey(), entry.getValue());

diff --git a/src/test/java/org/influxdb/impl/BatchProcessorTest.java b/src/test/java/org/influxdb/impl/BatchProcessorTest.java
@@ -39,7 +39,25 @@ public void testSchedulerExceptionHandling() throws InterruptedException, IOExce
         // without try catch the 2nd time does not occur
         verify(mockInfluxDB, times(2)).write(any(BatchPoints.class));
     }
-
+
+    @Test
+    public void testBatchWriteWithDifferenctRp() throws InterruptedException, IOException {
+        InfluxDB mockInfluxDB = mock(InfluxDBImpl.class);
+        BatchProcessor batchProcessor = BatchProcessor.builder(mockInfluxDB).actions(Integer.MAX_VALUE)
+            .interval(1, TimeUnit.NANOSECONDS).build();
+
+        Point point = Point.measurement("cpu").field("6", "").build();
+        BatchProcessor.HttpBatchEntry batchEntry1 = new BatchProcessor.HttpBatchEntry(point, "db1", "rp_1");
+        BatchProcessor.HttpBatchEntry batchEntry2 = new BatchProcessor.HttpBatchEntry(point, "db1", "rp_2");
+
+        batchProcessor.put(batchEntry1);
+        batchProcessor.put(batchEntry2);
+
+        Thread.sleep(200); // wait for scheduler
+        // same dbname with different rp should write two batchs instead of only one.
+        verify(mockInfluxDB, times(2)).write(any(BatchPoints.class));
+    }
+
     @Test(expected = IllegalArgumentException.class)
     public void testActionsIsZero() throws InterruptedException, IOException {
         InfluxDB mockInfluxDB = mock(InfluxDBImpl.class);