<commit_before>import uwsgi
import os
import re
import time
import logging
from pathlib import Path

logging.basicConfig(format='%(asctime)s %(levelname)-8s %(message)s', level=logging.INFO)
logger = logging.getLogger('websockets')

project_dir = Path(__file__).resolve().parent.parent.parent.parent


def app(env, start_response):
    uwsgi.websocket_handshake(env['HTTP_SEC_WEBSOCKET_KEY'], env.get('HTTP_ORIGIN', ''))
    logger.info(f"New websocket connection: {env['PATH_INFO']}")
    job_id = env['PATH_INFO'].split("/").pop()

    contest_id, job_id = job_id.rsplit('_', 1)
    contest_id = re.sub('[^0-9a-z_-]', '', contest_id)
    job_id = re.sub('[^0-9a-zA-Z-]', '', job_id)
    log_file = project_dir.joinpath('logs', f'{contest_id}_{job_id}.log')

    if not os.path.isfile(log_file):
        logger.error(f"Log not found: {log_file}")
        uwsgi.websocket_send('The requested log does not exist')
        return

    status_file = project_dir.joinpath('logs', f'{contest_id}.status.json')
    logger.info(f"Opened websocket for {log_file}")

    # Use line-buffering (buffering=1) so that we never send incomplete lines
    close_next_time = False
    with open(log_file, buffering=1, encoding='utf-8') as stream:
        n = 0
        while True:
            try:
                # Handle ping/pong, but don't block
                uwsgi.websocket_recv_nb()

                new_data = stream.read()
                if new_data != '':
                    uwsgi.websocket_send(new_data)
                    if 'Job finished contest' in new_data:
                        logger.info(f"Completed streaming log: {log_file}")
                        return
                
                if close_next_time:
                    logger.info(f"Completed streaming log: {log_file}")
                    return
            except IOError as err:
                logger.info('WebSocket connection closed by client')
                print(err)
                return
            if n % 10 == 0:
                try:
                    with open(status_file) as fp:
                        status = json.load(fp)
                        if int(status['job_id']) == job_id and status['status'] != 'running':
                            close_next_time = True
                except:
                    pass
            n += 1
            time.sleep(1)

<commit_msg>Websocket: Wait one cycle before closing<commit_after>import uwsgi
import os
import re
import time
import logging
from pathlib import Path

logging.basicConfig(format='%(asctime)s %(levelname)-8s %(message)s', level=logging.INFO)
logger = logging.getLogger('websockets')

project_dir = Path(__file__).resolve().parent.parent.parent.parent


def app(env, start_response):
    uwsgi.websocket_handshake(env['HTTP_SEC_WEBSOCKET_KEY'], env.get('HTTP_ORIGIN', ''))
    logger.info(f"New websocket connection: {env['PATH_INFO']}")
    job_id = env['PATH_INFO'].split("/").pop()

    contest_id, job_id = job_id.rsplit('_', 1)
    contest_id = re.sub('[^0-9a-z_-]', '', contest_id)
    job_id = re.sub('[^0-9a-zA-Z-]', '', job_id)
    log_file = project_dir.joinpath('logs', f'{contest_id}_{job_id}.log')

    if not os.path.isfile(log_file):
        logger.error(f"Log not found: {log_file}")
        uwsgi.websocket_send('The requested log does not exist')
        return

    status_file = project_dir.joinpath('logs', f'{contest_id}.status.json')
    logger.info(f"Opened websocket for {log_file}")

    # Use line-buffering (buffering=1) so that we never send incomplete lines
    close_next_time = False
    with open(log_file, buffering=1, encoding='utf-8') as stream:
        n = 0
        while True:
            try:
                # Handle ping/pong, but don't block
                uwsgi.websocket_recv_nb()

                new_data = stream.read()
                if new_data != '':
                    uwsgi.websocket_send(new_data)
                    if 'Job finished contest' in new_data:
                        close_next_time = True
                
                if close_next_time:
                    logger.info(f"Completed streaming log: {log_file}")
                    return
            except IOError as err:
                logger.info('WebSocket connection closed by client')
                print(err)
                return
            if n % 10 == 0:
                try:
                    with open(status_file) as fp:
                        status = json.load(fp)
                        if int(status['job_id']) == job_id and status['status'] != 'running':
                            close_next_time = True
                except:
                    pass
            n += 1
            time.sleep(1)


{'fraction_non_alphanumeric': 0.11712285168682368, 'fraction_numerical': 0.006365372374283896, 'mean_word_length': 1.955485893416928, 'pattern_counts': {'":': 0, '<': 3, '<?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': '1655456', 'n_tokens_mistral': 1470, 'n_tokens_neox': 1395, 'n_words': 310}