267 lines
8.3 KiB
Python
267 lines
8.3 KiB
Python
from collections import defaultdict
|
|
|
|
from channels import Group
|
|
from channels.sessions import session_for_reply_channel
|
|
from django.apps import apps
|
|
from django.core.cache import cache, caches
|
|
|
|
|
|
class BaseWebsocketUserCache:
|
|
"""
|
|
Caches the reply channel names of all open websocket connections. The id of
|
|
the user that that opened the connection is used as reference.
|
|
|
|
This is the Base cache that has to be overriden.
|
|
"""
|
|
cache_key = 'current_websocket_users'
|
|
|
|
def add(self, user_id, channel_name):
|
|
"""
|
|
Adds a channel name to an user id.
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
def remove(self, user_id, channel_name):
|
|
"""
|
|
Removes one channel name from the cache.
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
def get_all(self):
|
|
"""
|
|
Returns all data using a dict where the key is a user id and the value
|
|
is a set of channel_names.
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
def save_data(self, data):
|
|
"""
|
|
Saves the full data set (like created with build_data) to the cache.
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
def build_data(self):
|
|
"""
|
|
Creates all the data, saves it to the cache and returns it.
|
|
"""
|
|
websocket_user_ids = defaultdict(set)
|
|
for channel_name in Group('site').channel_layer.group_channels('site'):
|
|
session = session_for_reply_channel(channel_name)
|
|
user_id = session.get('user_id', None)
|
|
websocket_user_ids[user_id or 0].add(channel_name)
|
|
self.save_data(websocket_user_ids)
|
|
return websocket_user_ids
|
|
|
|
def get_cache_key(self):
|
|
"""
|
|
Returns the cache key.
|
|
"""
|
|
return self.cache_key
|
|
|
|
|
|
class RedisWebsocketUserCache(BaseWebsocketUserCache):
|
|
"""
|
|
Implementation of the WebsocketUserCache that uses redis.
|
|
|
|
This uses one cache key to store all connected user ids in a set and
|
|
for each user another set to save the channel names.
|
|
"""
|
|
|
|
def add(self, user_id, channel_name):
|
|
"""
|
|
Adds a channel name to an user id.
|
|
"""
|
|
redis = get_redis_connection()
|
|
pipe = redis.pipeline()
|
|
pipe.sadd(self.get_cache_key(), user_id)
|
|
pipe.sadd(self.get_user_cache_key(user_id), channel_name)
|
|
pipe.execute()
|
|
|
|
def remove(self, user_id, channel_name):
|
|
"""
|
|
Removes one channel name from the cache.
|
|
"""
|
|
redis = get_redis_connection()
|
|
redis.srem(self.get_user_cache_key(user_id), channel_name)
|
|
|
|
def get_all(self):
|
|
"""
|
|
Returns all data using a dict where the key is a user id and the value
|
|
is a set of channel_names.
|
|
"""
|
|
redis = get_redis_connection()
|
|
user_ids = redis.smembers(self.get_cache_key())
|
|
if user_ids is None:
|
|
websocket_user_ids = self.build_data()
|
|
else:
|
|
websocket_user_ids = dict()
|
|
for user_id in user_ids:
|
|
# Redis returns the id as string. So we have to convert it
|
|
user_id = int(user_id)
|
|
channel_names = redis.smembers(self.get_user_cache_key(user_id))
|
|
if channel_names is not None:
|
|
# If channel name is empty, then we can assume, that the user
|
|
# has no active connection.
|
|
websocket_user_ids[user_id] = set(channel_names)
|
|
return websocket_user_ids
|
|
|
|
def save_data(self, data):
|
|
"""
|
|
Saves the full data set (like created with the method build_data()) to
|
|
the cache.
|
|
"""
|
|
redis = get_redis_connection()
|
|
pipe = redis.pipeline()
|
|
|
|
# Save all user ids
|
|
pipe.delete(self.get_cache_key())
|
|
pipe.sadd(self.get_cache_key(), *data.keys())
|
|
|
|
for user_id, channel_names in data.items():
|
|
pipe.delete(self.get_user_cache_key(user_id))
|
|
pipe.sadd(self.get_user_cache_key(user_id), *channel_names)
|
|
pipe.execute()
|
|
|
|
def get_cache_key(self):
|
|
"""
|
|
Returns the cache key.
|
|
"""
|
|
return cache.make_key(self.cache_key)
|
|
|
|
def get_user_cache_key(self, user_id):
|
|
"""
|
|
Returns a cache key to save the channel names for a specific user.
|
|
"""
|
|
return cache.make_key('{}:{}'.format(self.cache_key, user_id))
|
|
|
|
|
|
class DjangoCacheWebsocketUserCache(BaseWebsocketUserCache):
|
|
"""
|
|
Implementation of the WebsocketUserCache that uses the django cache.
|
|
|
|
If you use this with the inmemory cache, then you should only use one
|
|
worker.
|
|
|
|
This uses only one cache key to save a dict where the key is the user id and
|
|
the value is a set of channel names.
|
|
"""
|
|
|
|
def add(self, user_id, channel_name):
|
|
"""
|
|
Adds a channel name for a user using the django cache.
|
|
"""
|
|
websocket_user_ids = cache.get(self.get_cache_key())
|
|
if websocket_user_ids is None:
|
|
websocket_user_ids = dict()
|
|
|
|
if user_id in websocket_user_ids:
|
|
websocket_user_ids[user_id].add(channel_name)
|
|
else:
|
|
websocket_user_ids[user_id] = set([channel_name])
|
|
cache.set(self.get_cache_key(), websocket_user_ids)
|
|
|
|
def remove(self, user_id, channel_name):
|
|
"""
|
|
Removes one channel name from the django cache.
|
|
"""
|
|
websocket_user_ids = cache.get(self.get_cache_key())
|
|
if websocket_user_ids is not None and user_id in websocket_user_ids:
|
|
websocket_user_ids[user_id].discard(channel_name)
|
|
cache.set(self.get_cache_key(), websocket_user_ids)
|
|
|
|
def get_all(self):
|
|
"""
|
|
Returns the data using the django cache.
|
|
"""
|
|
websocket_user_ids = cache.get(self.get_cache_key())
|
|
if websocket_user_ids is None:
|
|
return self.build_data()
|
|
return websocket_user_ids
|
|
|
|
def save_data(self, data):
|
|
"""
|
|
Saves the data using the django cache.
|
|
"""
|
|
cache.set(self.get_cache_key(), data)
|
|
|
|
|
|
class StartupCache:
|
|
"""
|
|
Cache of all data that are required when a client connects via websocket.
|
|
"""
|
|
cache_key = "full_data_startup_cache"
|
|
|
|
def build(self):
|
|
"""
|
|
Generate the cache by going through all apps. Returns a dict where the
|
|
key is the collection string and the value a list of the full_data from
|
|
the collection elements.
|
|
"""
|
|
cache_data = {}
|
|
for app in apps.get_app_configs():
|
|
try:
|
|
# Get the method get_startup_elements() from an app.
|
|
# This method has to return an iterable of Collection objects.
|
|
get_startup_elements = app.get_startup_elements
|
|
except AttributeError:
|
|
# Skip apps that do not implement get_startup_elements.
|
|
continue
|
|
|
|
for collection in get_startup_elements():
|
|
cache_data[collection.collection_string] = [
|
|
collection_element.get_full_data()
|
|
for collection_element
|
|
in collection.element_generator()]
|
|
|
|
cache.set(self.cache_key, cache_data, 86400)
|
|
return cache_data
|
|
|
|
def clear(self):
|
|
"""
|
|
Clears the cache.
|
|
"""
|
|
cache.delete(self.cache_key)
|
|
|
|
def get_collections(self):
|
|
"""
|
|
Generator that returns all cached Collections.
|
|
|
|
The data is read from the cache if it exists. It builds the cache if it
|
|
does not exists.
|
|
"""
|
|
from .collection import Collection
|
|
data = cache.get(self.cache_key)
|
|
if data is None:
|
|
# The cache does not exist.
|
|
data = self.build()
|
|
for collection_string, full_data in data.items():
|
|
yield Collection(collection_string, full_data)
|
|
|
|
|
|
startup_cache = StartupCache()
|
|
|
|
|
|
def use_redis_cache():
|
|
"""
|
|
Returns True if Redis is used als caching backend.
|
|
"""
|
|
try:
|
|
from django_redis.cache import RedisCache
|
|
except ImportError:
|
|
return False
|
|
return isinstance(caches['default'], RedisCache)
|
|
|
|
|
|
def get_redis_connection():
|
|
"""
|
|
Returns an object that can be used to talk directly to redis.
|
|
"""
|
|
from django_redis import get_redis_connection
|
|
return get_redis_connection("default")
|
|
|
|
|
|
if use_redis_cache():
|
|
websocket_user_cache = RedisWebsocketUserCache()
|
|
else:
|
|
websocket_user_cache = DjangoCacheWebsocketUserCache()
|