Bug 964370: support memcached in reporter.py
authorDustin J. Mitchell <dustin@mozilla.com>
Tue, 28 Jan 2014 12:02:26 -0500
changeset 236 868f5e3ca772d09095fccd2a804e6714e4f5a071
parent 235 35d7ff385eae63359ae4196e9745239b178d1dc3
child 237 bea9ffb2d17b085c528dca7e4fadd342c01422c3
push id188
push userdmitchell@mozilla.com
push dateTue, 28 Jan 2014 17:02:45 +0000
bugs964370
Bug 964370: support memcached in reporter.py
buildapi/scripts/reporter.py
--- a/buildapi/scripts/reporter.py
+++ b/buildapi/scripts/reporter.py
@@ -1,10 +1,15 @@
 #!/usr/bin/python
-import simplejson as json
+
+import sys, os
+print sys.path
+from optparse import OptionParser
+from ConfigParser import SafeConfigParser
+import json
 
 import sqlalchemy as sa
 from datetime import datetime, timedelta
 
 import calendar
 import time
 
 import buildapi.model.statusdb_orm as model
@@ -77,21 +82,61 @@ def get_request_info(scheduler_db, build
     return (request_ids, first_request_time)
 
 def get_slave_name(cache, session, slave_id):
     slave_key = 'slavenames:%i' % slave_id
     slave_name = cache.get(slave_key)
     if not slave_name:
         s = session.query(model.Slave).get(slave_id)
         slave_name = s.name
-        cache.set(slave_key, slave_name)
-        cache.expire(slave_key, 24*3600)
+        cache.put(slave_key, slave_name, 24*3600)
     return slave_name
 
 
+class Cache(object):
+
+    @staticmethod
+    def new(config):
+        cache_spec = config.get('general', 'cache', None)
+        if not cache_spec:
+            return Cache()
+        if cache_spec.startswith("memcached:"):
+            hosts = cache_spec[10:].split(',')
+            return MemcacheCache(hosts)
+        raise RuntimeError("invalid cache spec %r" % (cache_spec,))
+
+    def get(self, key):
+        return None
+
+    def put(self, key, val, expire=0):
+        pass
+
+
+class MemcacheCache(Cache):
+
+    def __init__(self, hosts):
+        import memcache
+        self.m = memcache.Client(hosts)
+
+    def _utf8(self, s):
+        if isinstance(s, unicode):
+            return s.encode('utf-8')
+        return s
+
+    def get(self, key):
+        return self.m.get(self._utf8(key))
+
+    def put(self, key, val, expire=0):
+        if expire == 0:
+            self.m.set(self._utf8(key), val)
+        else:
+            expire = int(expire - time.time())
+            self.m.set(self._utf8(key), val, expire)
+
+
 def build_report(cache, session, scheduler_db, starttime, endtime, include_steps=False):
     masters = {}
     builders = {}
     builds = []
     slaves = {}
     report = {'masters': masters, 'builders': builders, 'builds': builds, 'slaves': slaves,
             'starttime': starttime, 'endtime': endtime}
 
@@ -140,23 +185,23 @@ def build_report(cache, session, schedul
 
             builder_key = 'builders:%i' % build.builder_id
             builder = cache.get(builder_key)
 
             if builder is not None:
                 try:
                     builder = json.loads(builder.decode("zlib"))
                 except:
-                    cache.set(builder_key, builder.encode("zlib"))
-                    cache.expire(builder_key, 24*3600) # Keep it for a day
+                    cache.put(builder_key, builder.encode("zlib"),
+                              24*3600) # Keep it for a day
                     builder = json.loads(builder)
             else:
                 builder = get_builder()
-                cache.set(builder_key, json.dumps(builder).encode("zlib"))
-                cache.expire(builder_key, 24*3600) # Keep it for a day
+                cache.put(builder_key, json.dumps(builder).encode("zlib"),
+                          24*3600) # Keep it for a day
 
             builders[build.builder_id] = builder
             times['builders'] += time.time() - s0
 
             s0 = time.time()
             for slave_id in builder['slaves']:
                 if slave_id not in slaves:
                     slaves[slave_id] = get_slave_name(cache, session, slave_id)
@@ -208,39 +253,30 @@ def build_report(cache, session, schedul
             return build_dict
 
         build_key = 'builds:%i' % build.id
         build_dict = cache.get(build_key)
         if build_dict is not None:
             try:
                 build_dict = json.loads(build_dict.decode("zlib"))
             except:
-                cache.set(build_key, build_dict.encode("zlib"))
-                cache.expire(build_key, 24*3600)
+                cache.put(build_key, build_dict.encode("zlib"), 24*3600)
                 build_dict = json.loads(build_dict)
         else:
             build_dict = get_build_dict()
-            cache.set(build_key, json.dumps(build_dict).encode("zlib"))
-            cache.expire(build_key, 24*3600)
+            cache.put(build_key, json.dumps(build_dict).encode("zlib"), 24*3600)
         builds.append(build_dict)
 
     e = time.time()
     print "%.2f generate report. %i/%i missed" % (e-s, misses[0], n)
     print times
 
     return report
 
 if __name__ == "__main__":
-    import sys, os
-    try:
-        import simplejson as json
-    except:
-        import json
-    from optparse import OptionParser
-    from ConfigParser import SafeConfigParser
 
     def timedelta_option(option, opt, value, parser, units):
         setattr(parser.values, option.dest, timedelta(seconds=value * units))
 
     def ts_option(option, opt, value, parser):
         setattr(parser.values, option.dest, datetime.utcfromtimestamp(value))
 
     def date_option(option, opt, value, parser):
@@ -308,24 +344,16 @@ if __name__ == "__main__":
 
     if options.compression == "gzip":
         import gzip
         fp = gzip.GzipFile(fileobj=fp)
 
     session_maker = model.connect(options.dburl, pool_recycle=60)
     session = session_maker()
 
-    if config.has_option('general', 'redis'):
-        import redis
-        R = redis.Redis(host=config.get('general', 'redis'))
-    else:
-        import redis
-        R = redis.Redis()
-        # TODO: support memcached
-
     scheduler_db_engine = sa.create_engine(config.get('general', 'scheduler_dburl'), pool_recycle=60)
 
     starttime = options.start_time
     endtime = starttime + options.report_window
 
     start = datetime.now()
     if options.compact:
         encoder = json.JSONEncoder(default=encode_dates, sort_keys=True,
@@ -350,17 +378,18 @@ if __name__ == "__main__":
         if not master:
             parser.error("Couldn't find master %s" % options.list_builders)
 
         for builder in session.query(model.Builder).filter_by(master=master):
             fp.write("%s %s\n" % (builder.id, builder.name))
 
     else:
         print time.asctime()
-        report = build_report(R, session, scheduler_db_engine, starttime, endtime)
+        cache = Cache.new(config)
+        report = build_report(cache, session, scheduler_db_engine, starttime, endtime)
 
         fp.write(encoder.encode(report))
 
         if fp != sys.stdout:
             fp.close()
             os.rename(options.output + ".tmp", options.output)
 
         end = datetime.now()