wyoj-datums/daemon.py

import os
import traceback
import threading
import pymysql
import time
import subprocess

GAP = 2

def execute(s):
    print('Running: ', s)
    try:
        subprocess.run(s, shell=True, check=True)
    except subprocess.CalledProcessError as e:
        raise Exception('Failed to execute the command')

def watch_stream():
    db = pymysql.connect(host='uoj-db', user='root', password='root', database='app_uoj233')
    db.autocommit(True)
    with open('did', 'r') as f:
        did = int(f.read().strip())
    print(f'Watching request id {did}')
    last = ''
    curs = db.cursor()
    curs.execute('insert ignore into datum_streams (id) values (%s)', (did, ))

    while True:
        if not os.path.exists('stream'):
            continue
        with open('stream', 'r') as f:
            d = f.read()
            if last == d:
                if open('flag').read() == 'done':
                    print('Finished!')
                    return
                continue
        curs = db.cursor()
        curs.execute('update datum_streams set stream = %s where id = %s', (d, did))
        last = d
        time.sleep(GAP)

def generate_data_zip(p, stat):
    watcher = threading.Thread(target=watch_stream)
    open('flag', 'w').write('no')
    open('did', 'w').write('')
    open('stream', 'w').write('')

    did, pid, std, timelim, memlim = p
    print(f'Received request for problem {pid}')
    with open('stat.md', 'w') as f:
        f.write(stat)
    with open('std.cpp', 'w') as f:
        f.write(std)
    with open('did', 'w') as f:
        print(did, file=f)
    watcher.start()
    execute(f'sudo -u datum python3 datum.py stat.md --time {timelim} --mem {memlim}')
    open('flag', 'w').write('done')
    watcher.join()
    open('flag', 'w').write('no')
    open('did', 'w').write('')
    open('stream', 'w').write('')

def move_to_upload(pid):
    execute(f'rm -r /var/uoj_data/web/data/upload/{pid} || true')
    execute(f'mkdir /var/uoj_data/web/data/upload/{pid} || true')
    execute(f'mv data.zip /var/uoj_data/web/data/upload/{pid}/')
    execute(f'cd /var/uoj_data/web/data/upload/{pid}/; unzip data.zip; rm data.zip')

if __name__ == '__main__':
    db = pymysql.connect(host='uoj-db', user='root', password='root', database='app_uoj233')
    db.autocommit(True)

    while True:
        curs = db.cursor()
        curs.execute('select id, pid, std, timelimit, memlimit from datum_requests where status = "Waiting" order by id asc limit 1')
        res = curs.fetchall()
        if not res:
            time.sleep(GAP)
            continue

        curs = db.cursor()
        curs.execute(f'update datum_requests set status = "Generating" where id = {res[0][0]}')
        db.commit()

        curs = db.cursor()
        curs.execute(f'select statement_md from problems_contents where id = {res[0][1]}')

        try:
            generate_data_zip(res[0], curs.fetchall()[0][0])
            move_to_upload(res[0][1])
        except Exception as e:
            print(f'Exception {e}')
            traceback.print_exc()
            curs = db.cursor()
            curs.execute(f'update datum_requests set status = "Failed" where id = {res[0][0]}')
            curs.execute(f'update datum_requests set finish_time = now() where id = {res[0][0]}')
            db.commit()
            print(f'Failed {res[0][0]}')
        else:
            curs = db.cursor()
            generator = open('gen.py', 'r').read()
            curs.execute(f'update datum_requests set status = "Done" where id = {res[0][0]}')
            curs.execute('update datum_requests set generator = %s where id = %s', (generator, res[0][0]))
            curs.execute(f'update datum_requests set finish_time = now() where id = {res[0][0]}')
            db.commit()
            print(f'Done {res[0][0]}')