<commit_before>#!/usr/bin/python


import asyncio
import json
import os
import re
import signal
import sys
import time
import functools

import asyncio_redis

from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from redis import StrictRedis
from sqlalchemy import and_, func
from sqlalchemy.orm.exc import NoResultFound
from tornado.gen import coroutine, Task
from tornado.httpserver import HTTPServer
from tornado.ioloop import IOLoop
from tornado.platform.asyncio import AsyncIOMainLoop
from tornado.web import Application, RequestHandler
from tornado.websocket import WebSocketHandler, WebSocketClosedError
from uuid import uuid4

from newparp.helpers.chat import (
    UnauthorizedException,
    BannedException,
    TooManyPeopleException,
    KickedException,
    authorize_joining,
    kick_check,
    join,
    send_userlist,
    get_userlist,
    disconnect,
    send_quit_message,
)
from newparp.model import sm, AnyChat, Ban, ChatUser, Message, User, SearchCharacter
from newparp.model.connections import redis_pool

redis = StrictRedis(connection_pool=redis_pool)
thread_pool = ThreadPoolExecutor()


origin_regex = re.compile("^https?:\/\/%s$" % os.environ["BASE_DOMAIN"].replace(".", "\."))


sockets = set()

DEBUG = "DEBUG" in os.environ or "--debug" in sys.argv

class ChatHandler(WebSocketHandler):
    @property
    def db(self):
        if hasattr(self, "_db") and self._db is not None:
            return self._db
        else:
            self._db = sm()
            return self._db

    @property
    def loop(self):
        return asyncio.get_event_loop()

    def get_chat_user(self):
        return self.db.query(
            ChatUser, User, AnyChat,
        ).join(
            User, ChatUser.user_id == User.id,
        ).join(
            AnyChat, ChatUser.chat_id == AnyChat.id,
        ).filter(and_(
            ChatUser.user_id == self.user_id,
            ChatUser.chat_id == self.chat_id,
        )).one()

    def set_typing(self, is_typing):
        command = redis.sadd if is_typing else redis.srem
        typing_key = "chat:%s:typing" % self.chat_id
        if command(typing_key, self.user_number):
            redis.publish(self.channels["typing"], json.dumps({
                "typing": list(int(_) for _ in redis.smembers(typing_key)),
            }))

    def safe_write(self, message):
        try:
            self.write_message(message)
        except WebSocketClosedError:
            return

    def check_origin(self, origin):
        return origin_regex.match(origin) is not None

    @coroutine
    def prepare(self):
        self.id = str(uuid4())
        self.joined = False
        try:
            self.session_id = self.cookies["newparp"].value
            self.chat_id = int(self.path_args[0])
            self.user_id = int(redis.get("session:%s" % self.session_id))
        except (KeyError, TypeError, ValueError):
            self.send_error(400)
            return
        try:
            self.chat_user, self.user, self.chat = yield thread_pool.submit(self.get_chat_user)
        except NoResultFound:
            self.send_error(404)
            return
        # Remember the user number so typing notifications can refer to it
        # without reopening the database session.
        self.user_number = self.chat_user.number
        self.user.last_online = datetime.now()
        self.user.last_ip = self.request.headers["X-Forwarded-For"]
        if self.user.group == "banned":
            self.send_error(403)
            return
        try:
            yield thread_pool.submit(authorize_joining, redis, self.db, self)
        except (UnauthorizedException, BannedException, TooManyPeopleException):
            self.send_error(403)
            return

    @coroutine
    def open(self, chat_id):

        redis.zadd("sockets_alive", time.time() + 60, "%s/%s/%s" % (self.chat_id, self.session_id, self.id))
        sockets.add(self)
        redis.sadd("chat:%s:sockets:%s" % (self.chat_id, self.session_id), self.id)
        if DEBUG:
            print("socket opened: %s %s %s" % (self.id, self.chat.url, self.user.username))

        try:
            yield thread_pool.submit(kick_check, redis, self)
        except KickedException:
            self.safe_write(json.dumps({"exit": "kick"}))
            self.close()
            return

        # Subscribe
        self.channels = {
            "chat": "channel:%s" % self.chat_id,
            "user": "channel:%s:%s" % (self.chat_id, self.user_id),
            "typing": "channel:%s:typing" % self.chat_id,
        }
        if self.chat.type == "pm":
            self.channels["pm"] = "channel:pm:%s" % self.user_id

        # Send backlog.
        try:
            after = int(self.request.query_arguments["after"][0])
        except (KeyError, IndexError, ValueError):
            after = 0
        messages = redis.zrangebyscore("chat:%s" % self.chat_id, "(%s" % after, "+inf")
        self.safe_write(json.dumps({
            "chat": self.chat.to_dict(),
            "messages": [json.loads(_) for _ in messages],
        }))

        join_message_sent = yield thread_pool.submit(join, redis, self.db, self)
        self.joined = True

        # Send userlist if nothing was sent by join().
        if not join_message_sent:
            userlist = yield thread_pool.submit(get_userlist, self.db, redis, self.chat)
            self.safe_write(json.dumps({"users": userlist}))

        self.db.commit()
        self.db.close()

        yield self.redis_listen()

    def on_message(self, message):
        if redis.zadd("sockets_alive", time.time() + 60, "%s/%s/%s" % (self.chat_id, self.session_id, self.id)):
            # We've been reaped, so disconnect.
            self.close()
            return
        if DEBUG:
            print("message: %s" % message)
        if message in ("typing", "stopped_typing"):
            self.set_typing(message == "typing")

    def on_close(self):
        # Unsubscribe here and let the exit callback handle disconnecting.
        if hasattr(self, "redis_client"):
            self.redis_client.close()

        if hasattr(self, "close_code") and self.close_code in (1000, 1001):
            message_type = "disconnect"
        else:
            message_type = "timeout"
        if self.joined and disconnect(redis, self.chat_id, self.id):
            try:
                send_quit_message(self.db, redis, *self.get_chat_user(), type=message_type)
            except NoResultFound:
                send_userlist(self.db, redis, self.chat)
            self.db.commit()
        # Delete the database connection here and on_finish just to be sure.
        if hasattr(self, "_db"):
            self._db.close()
            del self._db
        self.set_typing(False)
        if DEBUG:
            print("socket closed: %s" % (self.id))
        redis.srem("chat:%s:sockets:%s" % (self.chat_id, self.session_id), self.id)
        redis.zrem("sockets_alive", "%s/%s/%s" % (self.chat_id, self.session_id, self.id))

        try:
            sockets.remove(self)
        except KeyError:
            pass

    def on_finish(self):
        if hasattr(self, "_db"):
            self._db.close()
            del self._db

    async def redis_listen(self):
        self.redis_client = await asyncio_redis.Connection.create(
            host=os.environ["REDIS_HOST"],
            port=int(os.environ["REDIS_PORT"]),
            db=int(os.environ["REDIS_DB"]),
        )
        # Set the connection name, subscribe, and listen.
        await self.redis_client.client_setname("live:%s:%s" % (self.chat_id, self.user_id))

        try:
            subscriber = await self.redis_client.start_subscribe()
            await subscriber.subscribe(list(self.channels.values()))

            while self.ws_connection:
                message = await subscriber.next_published()
                self.on_redis_message(message)
        finally:
            self.redis_client.close()

    def on_redis_message(self, message):
        if DEBUG:
            print("redis message: %s" % str(message))

        self.safe_write(message.value)

        if message.channel == self.channels["user"]:
            data = json.loads(message.value)
            if "exit" in data:
                self.joined = False
                self.close()


class HealthHandler(RequestHandler):
    @property
    def loop(self):
        return asyncio.get_event_loop()

    def test_sql(self):
        db = sm()

        try:
            db.query(SearchCharacter).first()
        finally:
            db.close()
            del db

    def test_redis(self):
        redis.set("health", 1)

    async def get(self):
        try:
            await self.loop.run_in_executor(thread_pool, self.test_sql)
            await self.loop.run_in_executor(thread_pool, self.test_redis)
        except:
            self.send_error(500)

        self.write("ok")


def sig_handler(sig, frame):
    print("Caught signal %s." % sig)
    ioloop.add_callback_from_signal(shutdown)


def shutdown():
    print("Shutting down.")

    for socket in sockets:
        ioloop.add_callback(socket.close)

    for task in asyncio.Task.all_tasks():
        task.cancel()

    deadline = time.time() + 5

    def stop_loop():
        now = time.time()
        if now < deadline and len(sockets) != 0:
            ioloop.add_timeout(now + 0.1, stop_loop)
        else:
            ioloop.stop()

    stop_loop()

if __name__ == "__main__":

    AsyncIOMainLoop().install()
    ioloop = IOLoop.instance()

    application = Application([
        (r"/(\d+)", ChatHandler),
        (r"/health", HealthHandler)
    ])

    http_server = HTTPServer(application)
    http_server.listen(int(os.environ.get("LISTEN_PORT", 5000)))

    signal.signal(signal.SIGTERM, sig_handler)
    signal.signal(signal.SIGINT, sig_handler)

    ioloop.start()
<commit_msg>Ensure future for the pubsub and on_redis_message functions.<commit_after>#!/usr/bin/python


import asyncio
import json
import os
import re
import signal
import sys
import time
import functools

import asyncio_redis

from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from redis import StrictRedis
from sqlalchemy import and_, func
from sqlalchemy.orm.exc import NoResultFound
from tornado.gen import coroutine, Task
from tornado.httpserver import HTTPServer
from tornado.ioloop import IOLoop
from tornado.platform.asyncio import AsyncIOMainLoop
from tornado.web import Application, RequestHandler
from tornado.websocket import WebSocketHandler, WebSocketClosedError
from uuid import uuid4

from newparp.helpers.chat import (
    UnauthorizedException,
    BannedException,
    TooManyPeopleException,
    KickedException,
    authorize_joining,
    kick_check,
    join,
    send_userlist,
    get_userlist,
    disconnect,
    send_quit_message,
)
from newparp.model import sm, AnyChat, Ban, ChatUser, Message, User, SearchCharacter
from newparp.model.connections import redis_pool

redis = StrictRedis(connection_pool=redis_pool)
thread_pool = ThreadPoolExecutor()


origin_regex = re.compile("^https?:\/\/%s$" % os.environ["BASE_DOMAIN"].replace(".", "\."))


sockets = set()

DEBUG = "DEBUG" in os.environ or "--debug" in sys.argv

class ChatHandler(WebSocketHandler):
    @property
    def db(self):
        if hasattr(self, "_db") and self._db is not None:
            return self._db
        else:
            self._db = sm()
            return self._db

    @property
    def loop(self):
        return asyncio.get_event_loop()

    def get_chat_user(self):
        return self.db.query(
            ChatUser, User, AnyChat,
        ).join(
            User, ChatUser.user_id == User.id,
        ).join(
            AnyChat, ChatUser.chat_id == AnyChat.id,
        ).filter(and_(
            ChatUser.user_id == self.user_id,
            ChatUser.chat_id == self.chat_id,
        )).one()

    def set_typing(self, is_typing):
        command = redis.sadd if is_typing else redis.srem
        typing_key = "chat:%s:typing" % self.chat_id
        if command(typing_key, self.user_number):
            redis.publish(self.channels["typing"], json.dumps({
                "typing": list(int(_) for _ in redis.smembers(typing_key)),
            }))

    def safe_write(self, message):
        try:
            self.write_message(message)
        except WebSocketClosedError:
            return

    def check_origin(self, origin):
        return origin_regex.match(origin) is not None

    @coroutine
    def prepare(self):
        self.id = str(uuid4())
        self.joined = False
        try:
            self.session_id = self.cookies["newparp"].value
            self.chat_id = int(self.path_args[0])
            self.user_id = int(redis.get("session:%s" % self.session_id))
        except (KeyError, TypeError, ValueError):
            self.send_error(400)
            return
        try:
            self.chat_user, self.user, self.chat = yield thread_pool.submit(self.get_chat_user)
        except NoResultFound:
            self.send_error(404)
            return
        # Remember the user number so typing notifications can refer to it
        # without reopening the database session.
        self.user_number = self.chat_user.number
        self.user.last_online = datetime.now()
        self.user.last_ip = self.request.headers["X-Forwarded-For"]
        if self.user.group == "banned":
            self.send_error(403)
            return
        try:
            yield thread_pool.submit(authorize_joining, redis, self.db, self)
        except (UnauthorizedException, BannedException, TooManyPeopleException):
            self.send_error(403)
            return

    @coroutine
    def open(self, chat_id):

        redis.zadd("sockets_alive", time.time() + 60, "%s/%s/%s" % (self.chat_id, self.session_id, self.id))
        sockets.add(self)
        redis.sadd("chat:%s:sockets:%s" % (self.chat_id, self.session_id), self.id)
        if DEBUG:
            print("socket opened: %s %s %s" % (self.id, self.chat.url, self.user.username))

        try:
            yield thread_pool.submit(kick_check, redis, self)
        except KickedException:
            self.safe_write(json.dumps({"exit": "kick"}))
            self.close()
            return

        # Subscribe
        self.channels = {
            "chat": "channel:%s" % self.chat_id,
            "user": "channel:%s:%s" % (self.chat_id, self.user_id),
            "typing": "channel:%s:typing" % self.chat_id,
        }

        if self.chat.type == "pm":
            self.channels["pm"] = "channel:pm:%s" % self.user_id

        asyncio.ensure_future(self.redis_listen())

        # Send backlog.
        try:
            after = int(self.request.query_arguments["after"][0])
        except (KeyError, IndexError, ValueError):
            after = 0
        messages = redis.zrangebyscore("chat:%s" % self.chat_id, "(%s" % after, "+inf")
        self.safe_write(json.dumps({
            "chat": self.chat.to_dict(),
            "messages": [json.loads(_) for _ in messages],
        }))

        join_message_sent = yield thread_pool.submit(join, redis, self.db, self)
        self.joined = True

        # Send userlist if nothing was sent by join().
        if not join_message_sent:
            userlist = yield thread_pool.submit(get_userlist, self.db, redis, self.chat)
            self.safe_write(json.dumps({"users": userlist}))

        self.db.commit()
        self.db.close()

    def on_message(self, message):
        if redis.zadd("sockets_alive", time.time() + 60, "%s/%s/%s" % (self.chat_id, self.session_id, self.id)):
            # We've been reaped, so disconnect.
            self.close()
            return
        if DEBUG:
            print("message: %s" % message)
        if message in ("typing", "stopped_typing"):
            self.set_typing(message == "typing")

    def on_close(self):
        # Unsubscribe here and let the exit callback handle disconnecting.
        if hasattr(self, "redis_client"):
            self.redis_client.close()

        if hasattr(self, "close_code") and self.close_code in (1000, 1001):
            message_type = "disconnect"
        else:
            message_type = "timeout"
        if self.joined and disconnect(redis, self.chat_id, self.id):
            try:
                send_quit_message(self.db, redis, *self.get_chat_user(), type=message_type)
            except NoResultFound:
                send_userlist(self.db, redis, self.chat)
            self.db.commit()
        # Delete the database connection here and on_finish just to be sure.
        if hasattr(self, "_db"):
            self._db.close()
            del self._db
        self.set_typing(False)
        if DEBUG:
            print("socket closed: %s" % (self.id))
        redis.srem("chat:%s:sockets:%s" % (self.chat_id, self.session_id), self.id)
        redis.zrem("sockets_alive", "%s/%s/%s" % (self.chat_id, self.session_id, self.id))

        try:
            sockets.remove(self)
        except KeyError:
            pass

    def on_finish(self):
        if hasattr(self, "_db"):
            self._db.close()
            del self._db

    async def redis_listen(self):
        self.redis_client = await asyncio_redis.Connection.create(
            host=os.environ["REDIS_HOST"],
            port=int(os.environ["REDIS_PORT"]),
            db=int(os.environ["REDIS_DB"]),
        )
        # Set the connection name, subscribe, and listen.
        await self.redis_client.client_setname("live:%s:%s" % (self.chat_id, self.user_id))

        try:
            subscriber = await self.redis_client.start_subscribe()
            await subscriber.subscribe(list(self.channels.values()))

            while self.ws_connection:
                message = await subscriber.next_published()
                asyncio.ensure_future(self.on_redis_message(message))
        finally:
            self.redis_client.close()

    async def on_redis_message(self, message):
        if DEBUG:
            print("redis message: %s" % str(message))

        self.safe_write(message.value)

        if message.channel == self.channels["user"]:
            data = json.loads(message.value)
            if "exit" in data:
                self.joined = False
                self.close()


class HealthHandler(RequestHandler):
    @property
    def loop(self):
        return asyncio.get_event_loop()

    def test_sql(self):
        db = sm()

        try:
            db.query(SearchCharacter).first()
        finally:
            db.close()
            del db

    def test_redis(self):
        redis.set("health", 1)

    async def get(self):
        try:
            await self.loop.run_in_executor(thread_pool, self.test_sql)
            await self.loop.run_in_executor(thread_pool, self.test_redis)
        except:
            self.send_error(500)

        self.write("ok")


def sig_handler(sig, frame):
    print("Caught signal %s." % sig)
    ioloop.add_callback_from_signal(shutdown)


def shutdown():
    print("Shutting down.")

    for socket in sockets:
        ioloop.add_callback(socket.close)

    for task in asyncio.Task.all_tasks():
        task.cancel()

    deadline = time.time() + 5

    def stop_loop():
        now = time.time()
        if now < deadline and len(sockets) != 0:
            ioloop.add_timeout(now + 0.1, stop_loop)
        else:
            ioloop.stop()

    stop_loop()

if __name__ == "__main__":

    AsyncIOMainLoop().install()
    ioloop = IOLoop.instance()

    application = Application([
        (r"/(\d+)", ChatHandler),
        (r"/health", HealthHandler)
    ])

    http_server = HTTPServer(application)
    http_server.listen(int(os.environ.get("LISTEN_PORT", 5000)))

    signal.signal(signal.SIGTERM, sig_handler)
    signal.signal(signal.SIGINT, sig_handler)

    ioloop.start()

{'fraction_non_alphanumeric': 0.11482602118003025, 'fraction_numerical': 0.004135148764498235, 'mean_word_length': 2.2670510708401976, 'pattern_counts': {'":': 22, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2451305', 'n_tokens_mistral': 6244, 'n_tokens_neox': 6055, 'n_words': 1315}