84 lines
3.2 KiB
Python
84 lines
3.2 KiB
Python
import brotli as brotli_
|
|
from flask import request, make_response
|
|
from threading import Thread
|
|
from hashlib import sha256
|
|
import redis
|
|
import os.path
|
|
import mimetypes
|
|
|
|
|
|
class BrotliCache(object):
|
|
def __init__(self, redis_uri='redis://', timeout=0.020, expire=60*60*6):
|
|
self.redis = redis.StrictRedis.from_url(redis_uri)
|
|
self.timeout = timeout
|
|
self.expire = expire
|
|
self.redis.client_setname('brotlicache')
|
|
|
|
def compress(self, cache_key, lock_key, body, mode=brotli_.MODE_GENERIC):
|
|
encbody = brotli_.compress(body, mode=mode)
|
|
self.redis.set(cache_key, encbody, px=int(self.expire*1000))
|
|
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)
|
|
response.headers.set('x-brotli-cache', 'HIT')
|
|
if not encbody:
|
|
response.headers.set('x-brotli-cache', 'MISS')
|
|
lock_key = 'brotlicache:lock:{}'.format(digest)
|
|
if self.redis.set(lock_key, 1, nx=True, ex=10):
|
|
mode = (
|
|
brotli_.MODE_TEXT
|
|
if response.content_type.startswith('text/')
|
|
else brotli_.MODE_GENERIC)
|
|
t = Thread(
|
|
target=self.compress,
|
|
args=(cache_key, lock_key, body, mode))
|
|
t.start()
|
|
if self.timeout > 0:
|
|
t.join(self.timeout)
|
|
encbody = self.redis.get(cache_key)
|
|
if not encbody:
|
|
response.headers.set('x-brotli-cache', 'TIMEOUT')
|
|
else:
|
|
response.headers.set('x-brotli-cache', 'LOCKED')
|
|
if encbody:
|
|
response.headers.set('content-encoding', 'br')
|
|
response.headers.set('vary', 'accept-encoding')
|
|
response.set_data(encbody)
|
|
return response
|
|
|
|
return response
|
|
|
|
|
|
def brotli(app, static=True, dynamic=True, **kwargs):
|
|
original_static = app.view_functions['static']
|
|
|
|
def static_maybe_gzip_brotli(filename=None):
|
|
path = os.path.join(app.static_folder, filename)
|
|
for encoding, extension in (('br', '.br'), ('gzip', '.gz')):
|
|
if encoding not in request.accept_encodings:
|
|
continue
|
|
encpath = path + extension
|
|
if os.path.isfile(encpath):
|
|
resp = make_response(
|
|
original_static(filename=filename + extension))
|
|
resp.headers.set('content-encoding', encoding)
|
|
resp.headers.set('vary', 'accept-encoding')
|
|
mimetype = (mimetypes.guess_type(filename)[0]
|
|
or 'application/octet-stream')
|
|
resp.headers.set('content-type', mimetype)
|
|
return resp
|
|
return original_static(filename=filename)
|
|
if static:
|
|
app.view_functions['static'] = static_maybe_gzip_brotli
|
|
if dynamic:
|
|
cache = BrotliCache(redis_uri=app.config.get('REDIS_URI'), **kwargs)
|
|
app.after_request(cache.wrap_response)
|