Merge lp:~sidnei/txstatsd/dup-flag into lp:txstatsd

[Sidnei da Silva]

Add dup flag to rewrite/set_metric_type targets, so it yields the original message in addition to the rewritten one.

Add dup flag to rewrite/set_metric_type targets, so it yields the original
message in addition to the rewritten one.

[Martin Albisetti (beuno)]

review: Approve

=== modified file 'txstatsd/server/router.py'
--- txstatsd/server/router.py	2012-02-07 12:28:35 +0000
+++ txstatsd/server/router.py	2012-03-14 20:15:23 +0000
@@ -230,21 +230,25 @@
     def build_target_drop(self):
         """Returns a target that stops the processing of a message."""
         def drop(*args):
-            raise StopProcessingException()
+            return
         return drop
 
-    def build_target_rewrite(self, pattern, repl):
+    def build_target_rewrite(self, pattern, repl, dup="no-dup"):
         rexp = re.compile(pattern)
 
         def rewrite_target(metric_type, key, fields):
+            if dup == "dup" and rexp.match(key) is not None:
+                yield metric_type, key, fields
             key = rexp.sub(repl, key)
-            return (metric_type, key, fields)
+            yield metric_type, key, fields
 
         return rewrite_target
 
-    def build_target_set_metric_type(self, metric_type):
+    def build_target_set_metric_type(self, metric_type, dup="no-dup"):
         def set_metric_type(_, key, fields):
-            return metric_type, key, fields
+            if dup == "dup":
+                yield _, key, fields
+            yield metric_type, key, fields
         return set_metric_type
 
     def build_target_redirect_udp(self, host, port):
@@ -262,7 +266,7 @@
         def redirect_udp_target(metric_type, key, fields):
             message = self.rebuild_message(metric_type, key, fields)
             protocol.write(message)
-            return metric_type, key, fields
+            yield metric_type, key, fields
         return redirect_udp_target
 
     def build_target_redirect_tcp(self, host, port):
@@ -280,17 +284,26 @@
         def redirect_tcp_target(metric_type, key, fields):
             message = self.rebuild_message(metric_type, key, fields)
             factory.write(message)
-            return metric_type, key, fields
+            yield metric_type, key, fields
         return redirect_tcp_target
 
     def process_message(self, message, metric_type, key, fields):
-        try:
+        metrics = [(metric_type, key, fields)]
+        if self.rules:
+            pending, metrics = metrics, []
             for condition, target in self.rules:
-                if condition(metric_type, key, fields):
-                    metric_type, key, fields = target(metric_type, key, fields)
+                if not pending:
+                    return
+                for metric_type, key, fields in pending:
+                    if not condition(metric_type, key, fields):
+                        metrics.append((metric_type, key, fields))
+                        continue
+                    result = target(metric_type, key, fields)
+                    if result is not None:
+                        metrics.extend(result)
+                pending = metrics
 
+        for (metric_type, key, fields) in metrics:
             message = self.rebuild_message(metric_type, key, fields)
             self.message_processor.process_message(message, metric_type,
                                                    key, fields)
-        except StopProcessingException:
-            pass

=== modified file 'txstatsd/tests/test_router.py'
--- txstatsd/tests/test_router.py	2012-02-07 12:28:35 +0000
+++ txstatsd/tests/test_router.py	2012-03-14 20:15:23 +0000
@@ -85,6 +85,31 @@
         self.assertEqual(self.processor.messages[0][2], "glork.gorets")
         self.assertEqual(self.processor.messages[1][2], "nomatch")
 
+    def test_rewrite_and_dup(self):
+        """
+        Process all messages but only rewrite matching ones. If dup flag is set
+        then duplicate original message without rewriting it.
+        """
+        self.update_rules(r"any => rewrite (gorets) glork.\1 dup")
+        self.router.process("gorets:1|c")
+        self.router.process("nomatch:1|d")
+        self.assertEqual(len(self.processor.messages), 3)
+        self.assertEqual(self.processor.messages[0][2], "gorets")
+        self.assertEqual(self.processor.messages[1][2], "glork.gorets")
+        self.assertEqual(self.processor.messages[2][2], "nomatch")
+
+    def test_rewrite_and_no_dup(self):
+        """
+        Process all messages but only rewrite matching ones. If dup flag is set
+        to no-dup, then the original message is not duplicated.
+        """
+        self.update_rules(r"any => rewrite (gorets) glork.\1 no-dup")
+        self.router.process("gorets:1|c")
+        self.router.process("nomatch:1|d")
+        self.assertEqual(len(self.processor.messages), 2)
+        self.assertEqual(self.processor.messages[0][2], "glork.gorets")
+        self.assertEqual(self.processor.messages[1][2], "nomatch")
+
     def test_set_metric_type(self):
         """
         Set metric type to something else.
@@ -94,6 +119,28 @@
         self.assertEqual(self.processor.messages[0][1], "d")
         self.assertEqual(self.processor.messages[0][2], "gorets")
 
+    def test_set_metric_type_dup(self):
+        """
+        Set metric type to something else. If the dup flag is set, duplicate
+        the original message.
+        """
+        self.update_rules(r"any => set_metric_type d dup")
+        self.router.process("gorets:1|c")
+        self.assertEqual(self.processor.messages[0][1], "c")
+        self.assertEqual(self.processor.messages[0][2], "gorets")
+        self.assertEqual(self.processor.messages[1][1], "d")
+        self.assertEqual(self.processor.messages[1][2], "gorets")
+
+    def test_set_metric_type_no_dup(self):
+        """
+        Set metric type to something else. If the dup flag is set to no-dup
+        then do not duplicate the original message.
+        """
+        self.update_rules(r"any => set_metric_type d no-dup")
+        self.router.process("gorets:1|c")
+        self.assertEqual(self.processor.messages[0][1], "d")
+        self.assertEqual(self.processor.messages[0][2], "gorets")
+
 
 class TestUDPRedirect(TxTestCase):