def background_updating(self):
self.subscription = RedisStreamSubscriber()
for stream in self.streams:
self.stream_writers[stream] = BasicStreamWriter(
os.path.join("/logging/", stream), stream)
self.subscription.subscribe(stream)
self.own_writer.write("Starting to listen")
while True:
with self.lock:
self.log_next_message()
def __init__(self, request, session_key, event, websocket):
self.request = request
self.session_key = session_key
self.event = event
self.websocket = websocket
self.time_connected = time.time()
users_pending_connection.put(self)
self.event.wait()
self.lock = threading.Lock()
# TODO: replace this subscriber with a subscriber pool, to reduce the number of redis connections
self.subscription = RedisStreamSubscriber()
# Keep track of when we last sent a message, to periodically send heartbeats
self.time_last_send = time.time()
self.websocket_file_descriptor = None
self.subscription_file_descriptor = None
def get_next_command(self):
if not self.redis_stream_subscriber:
self.redis_stream_subscriber = RedisStreamSubscriber()
self.redis_stream_subscriber.subscribe(self.redis_stream_name + "-q")
self.redis_stream_publisher = RedisStreamPublisher(self.redis_stream_name + "-a", raw=True)
message, stream_name = self.redis_stream_subscriber.next_message()
if not message:
return message
try:
message = str(message, "utf-8")
except Exception:
self.logger.error("Failed to convert to unicode")
return None
try:
return json.loads(message)
except Exception:
self.logger.error("Failed to parse command " + str(message))
return None
class BaseRedisCommandProcessor(BaseCommandProcessor):
def __init__(self, logger_name, redis_stream_name):
super().__init__(logger_name)
self.redis_stream_name = redis_stream_name
self.redis_stream_subscriber = None
self.redis_stream_publisher = None
def get_next_command(self):
if not self.redis_stream_subscriber:
self.redis_stream_subscriber = RedisStreamSubscriber()
self.redis_stream_subscriber.subscribe(self.redis_stream_name + "-q")
self.redis_stream_publisher = RedisStreamPublisher(self.redis_stream_name + "-a", raw=True)
message, stream_name = self.redis_stream_subscriber.next_message()
if not message:
return message
try:
message = str(message, "utf-8")
except Exception:
self.logger.error("Failed to convert to unicode")
return None
try:
return json.loads(message)
except Exception:
self.logger.error("Failed to parse command " + str(message))
return None
def publish_answer(self, answer, command):
self.redis_stream_publisher.publish_json(answer)
def handle_exception(self, exception):
self.logger.exception("Error processing redis command for " + str(self.__class__.__name__))
self.redis_stream_subscriber = None
self.redis_stream_publisher = None
time.sleep(1.0)
class GreenletRedisStreamListener(GreenletQueueWorker):
def __init__(self,
job_queue,
redis_stream_name,
logger=None,
context=None):
super().__init__(job_queue=job_queue, logger=logger, context=context)
self.redis_stream_name = redis_stream_name
self.redis_stream_subscriber = None
def init(self):
if not self.redis_stream_subscriber:
self.redis_stream_subscriber = RedisStreamSubscriber()
self.redis_stream_subscriber.subscribe(self.redis_stream_name)
def cleanup(self):
self.redis_stream_subscriber = None
def tick(self):
message, stream_name = self.redis_stream_subscriber.next_message()
message = redis_response_to_json(message)
if message:
self.job_queue.put(message)
class StreamWriterManager(object):
def __init__(self, name, streams):
#TODO: this should take in a name and a redis connection
self.name = name
self.own_writer = BasicStreamWriter(os.path.join("/logging/", name),
name)
self.stream_writers = {}
# The redis subscription is created in a separate thread since pubsub is not thread safe (and reads don't need to block)
self.subscription = None
# self.subscription = RedisStreamSubscriber()
self.subscription_queue = queue.Queue()
self.streams = streams
self.lock = threading.Lock()
self.background_updating_thread = ThreadHandler(
"background updating stream " + name, self.background_updating)
def log_next_message(self):
message = None
stream_name = None
try:
message, stream_name = self.subscription.next_message()
if stream_name:
stream_name = stream_name.decode("utf-8")
if stream_name in self.stream_writers and message is not None:
self.stream_writers[stream_name].write(message.decode("utf-8"))
except Exception as e:
self.own_writer.write("Exception in logging:\nStream:" +
str(stream_name) + "\nMessage:" +
str(message) + "\nException: " + str(e) +
"\nMore info: " + str(sys.exc_info()) +
"\n" + traceback.format_exc())
def background_updating(self):
self.subscription = RedisStreamSubscriber()
for stream in self.streams:
self.stream_writers[stream] = BasicStreamWriter(
os.path.join("/logging/", stream), stream)
self.subscription.subscribe(stream)
self.own_writer.write("Starting to listen")
while True:
with self.lock:
self.log_next_message()
def subscribe(self, stream_name, folder_path=""):
if stream_name in self.stream_writers:
print("Stream already added: ", stream_name)
return
if folder_path == "":
folder_path = os.path.join("/logging/", stream_name)
#TODO: ensure folder exists
with self.lock:
self.stream_writers[stream_name] = BasicStreamWriter(
folder_path, stream_name)
self.subscription.subscribe(stream_name)
class UserConnection(object):
"""
Class used to manage a websocket connection to a user
"""
def __init__(self, request, session_key, event, websocket):
self.request = request
self.session_key = session_key
self.event = event
self.websocket = websocket
self.time_connected = time.time()
users_pending_connection.put(self)
self.event.wait()
self.lock = threading.Lock()
# TODO: replace this subscriber with a subscriber pool, to reduce the number of redis connections
self.subscription = RedisStreamSubscriber()
# Keep track of when we last sent a message, to periodically send heartbeats
self.time_last_send = time.time()
self.websocket_file_descriptor = None
self.subscription_file_descriptor = None
def __str__(self):
return "Connection to " + str(self.request.user) + " on IP " + str(
self.user_ip)
@property
def user_ip(self):
# TODO: use the global one
request_meta = self.request.META
if "HTTP_REAL_IP" in request_meta:
return request_meta["HTTP_REAL_IP"]
return request_meta.get("HTTP_REMOTE_ADDR",
request_meta.get("REMOTE_ADDR", "0.0.0.0"))
def to_json(self):
return {
"ip": self.user_ip,
"user": self.user,
"numStreams": self.num_streams,
"timeConnected": self.time_connected
}
def get_user(self):
if self.session_key:
self.request.session = session_engine.SessionStore(
self.session_key)
self.request.user = get_user(self.request)
else:
self.request.user = AnonymousUser()
self.event.set()
@property
def user(self):
return self.request.user
@property
def user_stream_name(self):
if not self.user:
return None
return "user-" + self.user.id
def send(self, message):
self.time_last_send = time.time()
return self.websocket.send(message)
def ensure_heartbeat(self):
"""
This method sees when we last had activity on the stream, and sends a heartbeat just to keep the connection alive
"""
if time.time(
) - self.time_last_send >= settings.WEBSOCKET_HEARTBEAT_INTERVAL:
self.send(settings.WEBSOCKET_HEARTBEAT)
def subscribe(self, stream_name):
can_subscribe, reason = user_can_subscribe_to_stream(
self.user, stream_name)
if can_subscribe:
self.subscription.subscribe(stream_name)
self.send(b"s " + stream_name.encode())
else:
# This is an error we should know about because our clients should know what they are allowed to be subscribed to
logger.error("Registration not allowed for connection " +
str(self) + " to stream " + str(stream_name) +
"\nReason: " + reason)
self.send(b"error invalidSubscription " + stream_name.encode() +
b" " + str(self.user.id).encode() + b" " +
reason.encode())
logger.debug("Subscribed to " + str(self.num_streams) + " streams")
def unsubscribe(self, stream_name):
self.subscription.unsubscribe(stream_name)
@property
def num_streams(self):
return self.subscription.num_streams
def receive_next_message(self):
"""
This should only be called when there's a message to read from a subscribed redis channel
"""
message, stream_name = self.subscription.next_message()
if message:
self.send(b'm ' + stream_name + b' ' + message)
def process_file_descriptors(self, ready_file_descriptors):
for file_descriptor in ready_file_descriptors:
if file_descriptor == self.websocket.get_file_descriptor():
# We've received a message from the user
# TODO: if the message is really large, or we're getting high traffic, ban that asshole
for raw_message in self.websocket.receive_all():
if not raw_message:
continue
user_command = UserCommand(self, raw_message)
logger.debug("Received user websocket message ",
extra={
"user_message": str(raw_message),
"user": str(self.user)
})
# TODO: this should be done on a separate thread, not here!!!!!!
user_command.process()
elif file_descriptor == self.subscription.get_file_descriptor():
self.receive_next_message()
else:
logger.error("Invalid file descriptor: " +
str(file_descriptor))
# At least send a heartbeat at the end
self.ensure_heartbeat()
def close(self):
if self.subscription:
self.subscription.close()
if self.websocket:
self.websocket.close(code=1001, message='Websocket Closed')
self.websocket = None
@property
def connected(self):
"""
:return True if the websocket is connected to a user
"""
return self.websocket and not self.websocket.closed
def init(self):
if not self.redis_stream_subscriber:
self.redis_stream_subscriber = RedisStreamSubscriber()
self.redis_stream_subscriber.subscribe(self.redis_stream_name)