From 10abb9cd2f34d878656de041a3b3af070cf10ce4 Mon Sep 17 00:00:00 2001
From: codl <codl@codl.fr>
Date: Fri, 11 Aug 2017 19:13:37 +0200
Subject: [PATCH] working brotli cache

---
 app.py            |  2 +-
 config.example.py | 24 +++++++++++++++++-------
 lib/brotli.py     | 43 ++++++++++++++++++++++++++++++++++---------
 routes.py         |  4 +++-
 4 files changed, 55 insertions(+), 18 deletions(-)

diff --git a/app.py b/app.py
index 20f3064..e99c375 100644
--- a/app.py
+++ b/app.py
@@ -14,7 +14,7 @@ default_config = {
         "SQLALCHEMY_TRACK_MODIFICATIONS": False,
         "SQLALCHEMY_DATABASE_URI": "postgresql+psycopg2:///forget",
         "SECRET_KEY": "hunter2",
-        "CELERY_BROKER": "amqp://",
+        "CELERY_BROKER": "redis://",
         "HTTPS": True,
         "SENTRY_CONFIG": {},
         "RATELIMIT_STORAGE_URL": "redis://",
diff --git a/config.example.py b/config.example.py
index be7e9a9..0f472b1 100644
--- a/config.example.py
+++ b/config.example.py
@@ -11,7 +11,7 @@ determines where to connect to the database
 see http://docs.sqlalchemy.org/en/latest/core/engines.html#database-urls for syntax
 only postgresql with psycopg2 driver is officially supported
 """
-SQLALCHEMY_DATABASE_URI='postgresql+psycopg2:///forget'
+# SQLALCHEMY_DATABASE_URI='postgresql+psycopg2:///forget'
 
 """
 TWITTER CREDENTIALS
@@ -19,17 +19,17 @@ TWITTER CREDENTIALS
 get these at apps.twitter.com
 blah
 """
-TWITTER_CONSUMER_KEY='vdsvdsvds'
-TWITTER_CONSUMER_SECRET='hjklhjklhjkl'
+# TWITTER_CONSUMER_KEY='vdsvdsvds'
+# TWITTER_CONSUMER_SECRET='hjklhjklhjkl'
 
 """
 this will be necessary so we can tell twitter where to redirect
 """
-SERVER_NAME="localhost:5000"
+# SERVER_NAME="localhost:5000"
 
-CELERY_BROKER='amqp://'
+# CELERY_BROKER='redis://'
 
-HTTPS=True
+# HTTPS=True
 
 # SENTRY_DSN='https://foo:bar@sentry.io/69420'
 
@@ -41,7 +41,17 @@ requests and uploading hundreds of bogus tweet archives
 
 docs here <https://flask-limiter.readthedocs.io/en/stable/#configuration>
 '''
-RATELIMIT_STORAGE_URL='redis://'
+# RATELIMIT_STORAGE_URL='redis://'
+
+# REDIS=dict(
+#        db=0
+#
+#        host='localhost'
+#        port=6379
+#        # or...
+#        unix_socket_path='/var/run/redis/redis.sock'
+#        # see `pydoc redis.StrictRedis.__init__` for full list of arguments
+#     )
 
 """
 you can also use any config variable that flask expects here, such as
diff --git a/lib/brotli.py b/lib/brotli.py
index 9976b1a..b41ac2a 100644
--- a/lib/brotli.py
+++ b/lib/brotli.py
@@ -1,12 +1,37 @@
 import brotli
+from flask import request
+from functools import wraps
+from threading import Thread
+from hashlib import sha256
+import redis
+
+class BrotliCache(object):
+    def __init__(self, redis_kwargs={}):
+        self.redis = redis.StrictRedis(**redis_kwargs)
+
+    def compress(self, cache_key, lock_key, body):
+        encbody = brotli.compress(body)
+        self.redis.set(cache_key, encbody, ex=3600)
+        self.redis.delete(lock_key)
+
+    def wrap_response(self, response):
+        if 'br' not in request.accept_encodings or response.is_streamed:
+            return response
+
+        body = response.get_data()
+        digest = sha256(body).hexdigest()
+        cache_key = 'brotlicache:{}'.format(digest)
+
+        encbody = self.redis.get(cache_key)
+        if encbody:
+            response.headers.set('content-encoding', 'br')
+            response.headers.set('vary', 'content-encoding')
+            response.set_data(encbody)
+            return response
+        else:
+            lock_key = 'brotlicache:lock:{}'.format(digest)
+            if self.redis.set(lock_key, 1, nx=True, ex=60):
+                t = Thread(target=self.compress, args=(cache_key, lock_key, body))
+                t.run()
 
-def compress_response(response):
-    if response.is_streamed:
         return response
-    mode = brotli.MODE_GENERIC
-    if response.headers.get('content-type', '').startswith('text/'):
-        mode = brotli.MODE_TEXT
-    response.set_data(brotli.compress(response.get_data(), mode=mode))
-    response.headers.set('content-encoding', 'br')
-    response.headers.set('vary', 'content-encoding')
-    return response
diff --git a/routes.py b/routes.py
index e3ed564..34a1226 100644
--- a/routes.py
+++ b/routes.py
@@ -36,7 +36,9 @@ def touch_viewer(resp):
         db.session.commit()
     return resp
 
-app.after_request(lib.brotli.compress_response)
+brcache = lib.brotli.BrotliCache(app.config.get('REDIS', {}))
+
+app.after_request(brcache.wrap_response)
 
 @app.route('/')
 def index():