Add g.stats.count_string and start counting user agents.

r2/r2/controllers/reddit_base.py

@@ -707,6 +707,8 @@ class MinimalController(BaseController):
         # True/False forces. None updates for most non-POST requests
         c.update_last_visit = None
 
+        g.stats.count_string('user_agents', request.user_agent)
+
     def try_pagecache(self):
         #check content cache
         if request.method.upper() == 'GET' and not c.user_is_loggedin:

r2/r2/lib/stats.py

@@ -21,6 +21,7 @@
 ###############################################################################
 
 import collections
+import functools
 import os
 import random
 import socket
@@ -99,6 +100,34 @@ class CountingStatBuffer:
             yield k, str(v) + '|c'
 
 
+class StringCountBuffer:
+    """Dictionary of keys to counts of various values."""
+
+    def __init__(self):
+        self.data = collections.defaultdict(
+            functools.partial(collections.defaultdict, int))
+
+    @staticmethod
+    def _encode_string(string):
+        # escape \ -> \\, | -> \&, : -> \;, and newline -> \n
+        return (
+            string.replace('\\', '\\\\')
+                .replace('\n', '\\n')
+                .replace('|', '\\&')
+                .replace(':', '\\;'))
+
+    def record(self, key, value, count=1):
+        self.data[key][value] += count
+
+    def flush(self):
+        new_data = collections.defaultdict(
+            functools.partial(collections.defaultdict, int))
+        data, self.data = self.data, new_data
+        for k, counts in data.iteritems():
+            for v, count in counts.iteritems():
+                yield k, str(count) + '|s|' + self._encode_string(v)
+
+
 class StatsdConnection:
     def __init__(self, addr, compress=True):
         if addr:
@@ -150,6 +179,7 @@ class StatsdClient:
         self.sample_rate = sample_rate
         self.timing_stats = TimingStatBuffer()
         self.counting_stats = CountingStatBuffer()
+        self.string_counts = StringCountBuffer()
         self.connect(addr)
 
     def connect(self, addr):
@@ -161,6 +191,7 @@ class StatsdClient:
     def flush(self):
         data = list(self.timing_stats.flush())
         data.extend(self.counting_stats.flush())
+        data.extend(self.string_counts.flush())
         self.conn.send(self._data_iterator(data))
 
 
@@ -365,6 +396,9 @@ class Stats:
             return
         key = '.'.join(['pg', db_server.replace('.', '-'), db_name])
         self.client.timing_stats.record(key, start, end)
+
+    def count_string(self, key, value, count=1):
+        self.client.string_counts.record(key, value, count=count)
    
 
 class CacheStats:

r2/r2/tests/unit/lib/stats_test.py

@@ -56,6 +56,29 @@ class CountingStatBufferTest(unittest.TestCase):
                  ('3', '6|c')]),
             set(csb.flush()))
 
+class StringCountBufferTest(unittest.TestCase):
+    def test_encode_string(self):
+        enc = stats.StringCountBuffer._encode_string
+        self.assertEquals('test', enc('test'))
+        self.assertEquals('\\n\\&\\\\&', enc('\n|\\&'))
+
+    def test_scb(self):
+        scb = stats.StringCountBuffer()
+        self.assertEquals([], list(scb.flush()))
+
+        for i in xrange(1, 4):
+            for j in xrange(i):
+                for k in xrange(j + 1):
+                    scb.record(str(i), str(j))
+        self.assertEquals(
+            set([('1', '1|s|0'),
+                 ('2', '1|s|0'),
+                 ('2', '2|s|1'),
+                 ('3', '1|s|0'),
+                 ('3', '2|s|1'),
+                 ('3', '3|s|2')]),
+            set(scb.flush()))
+
 class FakeUdpSocket:
     def __init__(self, *ignored_args):
         self.host = None