Merge lp:~sidnei/txstatsd/push-producer-ftw into lp:txstatsd

[Sidnei da Silva]

Register the GraphiteProtocol as a producer for the connected transport, so
that if the buffer gets filled we pause message production until the buffer is
flushed.

- Register the GraphiteProtocol as a producer for the connected transport, so
  that if the buffer gets filled we pause message production until the buffer is
  flushed.

[lp:~sidnei/txstatsd/push-producer-ftw]

43. By Sidnei da Silva on 2011-10-07

- Simplify test a bit.

44. By Sidnei da Silva on 2011-10-07

- Turns out numStats was missing a newline and everything else already includes it. So let's get rid of the extra one.

45. By Sidnei da Silva on 2011-10-07

- Fix failures

[Facundo Batista (facundo)]

Looks ok

review: Approve

[Ian Wilkinson (theiw)]

+1

review: Approve

=== modified file 'txstatsd/server/processor.py'
--- txstatsd/server/processor.py	2011-10-01 04:05:34 +0000
+++ txstatsd/server/processor.py	2011-10-07 19:53:28 +0000
@@ -192,7 +192,7 @@
             messages.extend(meter_metrics)
             num_stats += events
 
-        messages.append("statsd.numStats %s %s" % (num_stats, timestamp))
+        messages.append("statsd.numStats %s %s\n" % (num_stats, timestamp))
         return messages
 
     def flush_counter_metrics(self, interval, timestamp):

=== modified file 'txstatsd/server/protocol.py'
--- txstatsd/server/protocol.py	2011-08-31 10:45:16 +0000
+++ txstatsd/server/protocol.py	2011-10-07 19:53:28 +0000
@@ -1,7 +1,8 @@
-from twisted.internet import task, defer
+from zope.interface import implements
+
+from twisted.internet import task, interfaces
 from twisted.internet.protocol import (
-    DatagramProtocol, ReconnectingClientFactory)
-from twisted.protocols.basic import LineOnlyReceiver
+    DatagramProtocol, ReconnectingClientFactory, Protocol)
 
 
 class StatsDServerProtocol(DatagramProtocol):
@@ -27,28 +28,45 @@
             self.processor.process(data)
 
 
-class GraphiteProtocol(LineOnlyReceiver):
+class GraphiteProtocol(Protocol):
     """A client protocol for talking to Graphite.
 
     Messages to Graphite are line-based and C{\n}-separated.
     """
 
-    delimiter = "\n"
+    implements(interfaces.IPushProducer)
 
-    def __init__(self, processor, interval):
+    def __init__(self, processor, interval, clock=None):
+        self.paused = False
         self.processor = processor
         self.interval = interval
         self.flush_task = task.LoopingCall(self.flushProcessor)
+        if clock is not None:
+            self.flush_task.clock = clock
         self.flush_task.start(self.interval / 1000, False)
 
-    @defer.inlineCallbacks
+    def connectionMade(self):
+        """
+        A connection has been made, register ourselves as a producer for the
+        bound transport.
+        """
+        self.transport.registerProducer(self, True)
+
     def flushProcessor(self):
         """Flush messages queued in the processor to Graphite."""
         for message in self.processor.flush(interval=self.interval):
-            for line in message.splitlines():
-                if self.connected:
-                    self.sendLine(line)
-                yield
+            if self.connected and not self.paused:
+                self.transport.write(message)
+
+    def pauseProducing(self):
+        """Pause producing messages, since the buffer is full."""
+        self.paused = True
+
+    stopProducing = pauseProducing
+
+    def resumeProducing(self):
+        """We can write to the transport again. Yay!."""
+        self.paused = False
 
 
 class GraphiteClientFactory(ReconnectingClientFactory):

=== modified file 'txstatsd/tests/test_configurableprocessor.py'
--- txstatsd/tests/test_configurableprocessor.py	2011-09-14 12:01:10 +0000
+++ txstatsd/tests/test_configurableprocessor.py	2011-10-07 19:53:28 +0000
@@ -21,7 +21,7 @@
         self.assertEqual(2, len(messages))
         counters = messages[0].splitlines()
         self.assertEqual("gorets.count 17 42", counters[0])
-        self.assertEqual("statsd.numStats 1 42", messages[1])
+        self.assertEqual("statsd.numStats 1 42", messages[1].splitlines()[0])
 
     def test_flush_counter_with_prefix(self):
         """
@@ -34,7 +34,7 @@
         self.assertEqual(2, len(messages))
         counters = messages[0].splitlines()
         self.assertEqual("test.metric.gorets.count 17 42", counters[0])
-        self.assertEqual("statsd.numStats 1 42", messages[1])
+        self.assertEqual("statsd.numStats 1 42", messages[1].splitlines()[0])
 
     def test_flush_single_timer_single_time(self):
         """
@@ -59,7 +59,7 @@
         self.assertEqual("glork.98percentile 24.0 42", timers[7])
         self.assertEqual("glork.99percentile 24.0 42", timers[8])
         self.assertEqual("glork.999percentile 24.0 42", timers[9])
-        self.assertEqual("statsd.numStats 1 42", messages[1])
+        self.assertEqual("statsd.numStats 1 42", messages[1].splitlines()[0])
 
     def test_flush_single_timer_multiple_times(self):
         """
@@ -94,7 +94,7 @@
         self.assertEqual("glork.98percentile 42.0 42", timers[7])
         self.assertEqual("glork.99percentile 42.0 42", timers[8])
         self.assertEqual("glork.999percentile 42.0 42", timers[9])
-        self.assertEqual("statsd.numStats 1 42", messages[1])
+        self.assertEqual("statsd.numStats 1 42", messages[1].splitlines()[0])
 
 
 class FlushMeterMetricMessagesTest(TestCase):
@@ -138,4 +138,5 @@
             "test.metric.gorets.15min_rate 0.0 %s" % self.time_now,
             meter_metric[4])
         self.assertEqual(
-            "statsd.numStats 1 %s" % self.time_now, messages[1])
+            "statsd.numStats 1 %s" % self.time_now,
+            messages[1].splitlines()[0])

=== modified file 'txstatsd/tests/test_processor.py'
--- txstatsd/tests/test_processor.py	2011-09-14 12:01:10 +0000
+++ txstatsd/tests/test_processor.py	2011-10-07 19:53:28 +0000
@@ -117,7 +117,7 @@
         Flushing the message processor when there are no stats available should
         still produce one message where C{statsd.numStats} is set to zero.
         """
-        self.assertEqual(["statsd.numStats 0 42"], self.processor.flush())
+        self.assertEqual(["statsd.numStats 0 42\n"], self.processor.flush())
 
     def test_flush_counter(self):
         """
@@ -130,7 +130,7 @@
         counters = messages[0].splitlines()
         self.assertEqual("stats.gorets 4 42", counters[0])
         self.assertEqual("stats_counts.gorets 42 42", counters[1])
-        self.assertEqual("statsd.numStats 1 42", messages[1])
+        self.assertEqual("statsd.numStats 1 42", messages[1].splitlines()[0])
         self.assertEqual(0, self.processor.counter_metrics["gorets"])
 
     def test_flush_counter_one_second_interval(self):
@@ -144,7 +144,7 @@
         counters = messages[0].splitlines()
         self.assertEqual("stats.gorets 42 42", counters[0])
         self.assertEqual("stats_counts.gorets 42 42", counters[1])
-        self.assertEqual("statsd.numStats 1 42", messages[1])
+        self.assertEqual("statsd.numStats 1 42", messages[1].splitlines()[0])
         self.assertEqual(0, self.processor.counter_metrics["gorets"])
 
     def test_flush_single_timer_single_time(self):
@@ -162,7 +162,7 @@
         self.assertEqual("stats.timers.glork.upper_90 24 42", timers[2])
         self.assertEqual("stats.timers.glork.lower 24 42", timers[3])
         self.assertEqual("stats.timers.glork.count 1 42", timers[4])
-        self.assertEqual("statsd.numStats 1 42", messages[1])
+        self.assertEqual("statsd.numStats 1 42", messages[1].splitlines()[0])
         self.assertEqual([], self.processor.timer_metrics["glork"])
 
     def test_flush_single_timer_multiple_times(self):
@@ -183,7 +183,7 @@
         self.assertEqual("stats.timers.glork.upper_90 23 42", timers[2])
         self.assertEqual("stats.timers.glork.lower 4 42", timers[3])
         self.assertEqual("stats.timers.glork.count 6 42", timers[4])
-        self.assertEqual("statsd.numStats 1 42", messages[1])
+        self.assertEqual("statsd.numStats 1 42", messages[1].splitlines()[0])
         self.assertEqual([], self.processor.timer_metrics["glork"])
 
     def test_flush_single_timer_50th_percentile(self):
@@ -207,7 +207,7 @@
         self.assertEqual("stats.timers.glork.upper_50 15 42", timers[2])
         self.assertEqual("stats.timers.glork.lower 4 42", timers[3])
         self.assertEqual("stats.timers.glork.count 6 42", timers[4])
-        self.assertEqual("statsd.numStats 1 42", messages[1])
+        self.assertEqual("statsd.numStats 1 42", messages[1].splitlines()[0])
         self.assertEqual([], self.processor.timer_metrics["glork"])
 
     def test_flush_gauge_metric(self):
@@ -224,7 +224,7 @@
         self.assertEqual(
             "stats.gauge.gorets.value 9.6 42", gauge_metric[0])
         self.assertEqual(
-            "statsd.numStats 1 42", messages[1])
+            "statsd.numStats 1 42", messages[1].splitlines()[0])
         self.assertEqual(0, len(self.processor.gauge_metrics))
 
 
@@ -268,7 +268,8 @@
             "stats.meter.gorets.15min_rate 0.0 %s" % self.time_now,
             meter_metric[4])
         self.assertEqual(
-            "statsd.numStats 1 %s" % self.time_now, messages[1])
+            "statsd.numStats 1 %s" % self.time_now,
+            messages[1].splitlines()[0])
 
         # As we are employing the expected results from test_ewma.py
         # we perform the initial tick(), before advancing the clock 60sec.
@@ -295,4 +296,4 @@
             meter_metric[4].startswith(
                 "stats.meter.gorets.15min_rate 0.5613041"))
         self.assertEqual(
-            "statsd.numStats 1 %s" % self.time_now, messages[1])
+            "statsd.numStats 1 %s" % self.time_now, messages[1].splitlines()[0])

=== added file 'txstatsd/tests/test_protocol.py'
--- txstatsd/tests/test_protocol.py	1970-01-01 00:00:00 +0000
+++ txstatsd/tests/test_protocol.py	2011-10-07 19:53:28 +0000
@@ -0,0 +1,104 @@
+"""Tests for the Graphite Protocol classes."""
+
+from twisted.trial.unittest import TestCase
+
+from txstatsd.server.protocol import GraphiteProtocol, GraphiteClientFactory
+from twisted.internet import task
+from twisted.test import proto_helpers
+
+
+class FakeProcessor(object):
+
+    def __init__(self):
+        self.sequence = 0
+
+    def flush(self, interval):
+        """Always produce a sequence number followed by 9 lines of output"""
+        self.sequence += 1
+        return [str(self.sequence)]
+
+
+class FakeTransport(object):
+
+    def __init__(self):
+        self.messages = []
+
+    def write(self, data):
+        self.messages.append(data)
+
+
+class TestGraphiteProtocol(TestCase):
+
+    def setUp(self):
+        super(TestGraphiteProtocol, self).setUp()
+        self.processor = FakeProcessor()
+        self.transport = FakeTransport()
+        self.clock = task.Clock()
+        self.protocol = GraphiteProtocol(self.processor, 1000, clock=self.clock)
+        self.protocol.transport = self.transport
+        self.protocol.connected = True
+
+    def test_write_unless_paused(self):
+        """
+        If the producer isn't paused, then write to the transport. Once the
+        producer is paused, nothing is written to the transport anymore.
+        """
+        self.assertEqual(0, len(self.transport.messages))
+        self.clock.advance(1)
+        self.assertEqual(1, len(self.transport.messages))
+        self.clock.advance(1)
+        self.assertEqual(2, len(self.transport.messages))
+        self.protocol.pauseProducing()
+        self.clock.advance(1)
+        self.assertEqual(2, len(self.transport.messages))
+        
+    def test_paused_producer_discards_everything_until_resumed(self):
+        """
+        If the producer is paused, everything is discarded until the producer
+        is resumed.
+        """
+        self.protocol.pauseProducing()
+        self.assertEqual(0, len(self.transport.messages))
+        self.clock.advance(1)
+        self.assertEqual(0, len(self.transport.messages))
+        self.clock.advance(1)
+        self.assertEqual(0, len(self.transport.messages))
+        self.protocol.resumeProducing()
+        self.clock.advance(1)
+        self.assertEqual(1, len(self.transport.messages))
+        self.assertEqual("3", self.transport.messages[-1])
+
+    def test_stopped_producer_discards_everything(self):
+        """
+        If the producer is stopped, everything is discarded.
+        """
+        self.assertEqual(0, len(self.transport.messages))
+        self.clock.advance(1)
+        self.assertEqual(1, len(self.transport.messages))
+        self.protocol.stopProducing()
+        self.clock.advance(1)
+        self.assertEqual(1, len(self.transport.messages))
+        self.clock.advance(1)
+        self.assertEqual(1, len(self.transport.messages))
+
+
+class TestProducerRegistration(TestCase):
+
+    def test_register_producer(self):
+        """
+        The Graphite protocol client registers itself as a producer of the
+        transport it's connected to.
+        """
+        processor = FakeProcessor()
+
+        # We don't care about the address argument here.
+        client = GraphiteClientFactory(processor, 1).buildProtocol(None)
+        clientTransport = proto_helpers.StringTransport()
+        client.makeConnection(clientTransport)
+        
+        # check that the producer is registered
+        self.assertTrue(clientTransport.producer is client)
+
+        # check the streaming attribute
+        self.assertTrue(clientTransport.streaming)
+        client.flush_task.stop()