<commit_before># -*- coding: utf-8 -*-

import os
import shutil

from monolithe.lib.utils.printer import Printer
from monolithe.lib.utils.constants import Constants
from monolithe.lib.managers import TaskManager
from monolithe.lib.utils.parse import ParsingUtils
from monolithe.generators.lib.writers import TemplateFileWriter

class SDKWriter(object):
    """ Writer of the Python VSD SDK

    """
    IGNORED_FILES = ['__init__.py', 'nuvsdsession.py', 'utils.py', 'nurestuser.py', 'constants.py']
    COMMON_ATTRIBUTES = ['multicast', 'iptype', 'maintenancemode', 'permittedaction', 'connectionstate', 'forwardingclasses']

    def __init__(self, directory):
        """ Initializes a writer to the specific directory

            Args:
                directory: directory where to copy sources

        """
        if not os.path.exists(directory):
            Printer.log("Copying default sources...")
            self.writer_directory = directory
            writer = self.get_writer()
            writer.copy_default_files()

        self.writer_directory = '%s/vsdk' % directory

    def get_writer(self):
        """ Get a writer to write content

        """
        return VSDKFileWriter(directory=self.writer_directory)

    def _prepare_attributes(self, model, constants):
        """ Removes attributes and computes constants

        """
        attributes = list()
        for attribute in model.attributes:

            attributes.append(attribute)

            if attribute.allowed_choices and len(attribute.allowed_choices) > 0:

                if attribute.remote_name.lower() not in SDKWriter.COMMON_ATTRIBUTES:
                    name = '%s%s%s' % (model.name, attribute.remote_name[0].upper(), attribute.remote_name[1:])
                else:
                    name = '%s%s' % (attribute.remote_name[0].upper(), attribute.remote_name[1:])

                constants[name] = self._make_constants_value(attribute.allowed_choices)

        # Another specific case... they really love this!
        constants['ProtocolType'] = Constants.PROTOCOL_TYPES

        model.attributes = attributes

    def _make_constants_value(self, choices):
        """ Create a dictionary for constants

        """
        results = dict()

        for choice in choices:
            results[choice] = choice

        return results

    def write(self, resources, apiversion, revision):
        """ Write all files according to data

            Args:
                resources: A list of all resources to manage
                apiversion: the version of the api
                revision: the revision number of the api

            Returns:
                Writes models and fetchers files

        """

        autogenerate_filenames = dict()
        fetcher_filenames = dict()
        override_filenames = dict()
        constants = dict()

        task_manager = TaskManager()

        for model_name, model in resources.iteritems():
            self._prepare_attributes(model=model, constants=constants)
            task_manager.start_task(method=self._write_autogenerate_file, model=model, filenames=autogenerate_filenames, version=apiversion)
            task_manager.start_task(method=self._write_override_file, model=model, filenames=override_filenames, version=apiversion)
            task_manager.start_task(method=self._write_fetcher_file, model=model, filenames=fetcher_filenames, version=apiversion)

        task_manager.wait_until_exit()
        except_files = autogenerate_filenames.keys() + fetcher_filenames.keys() + override_filenames.keys()
        self._clean_files(except_files=except_files)

        writer = self.get_writer()
        writer.write_setup_file(version=apiversion, revision=revision)

        # Constants
        constants = ParsingUtils.order(constants)
        writer.write_constants_file(constants=constants)

        # VSD Session
        writer.write_vsdsession_file(version=apiversion)

        writer.write_init_autogenerate_files(filenames=autogenerate_filenames)
        writer.write_init_fetcher_files(filenames=fetcher_filenames)
        writer.write_init_override_files(filenames=override_filenames)

    def _write_autogenerate_file(self, model, filenames, version):
        """ Write the autogenerate file for the model

            Args:
                model: the model to write
                filenames: list of generates filenames

        """
        writer = self.get_writer()

        if model.name != Constants.RESTUSER:
            (filename, classname) = writer.write_model(model=model, version=version)
        else:
            (filename, classname) = writer.write_restuser_model(model=model, version=version)

        filenames[filename] = classname

    def _write_override_file(self, model, filenames, version):
        """ Write the override file for the model

            Args:
                model: the model to write

        """
        writer = self.get_writer()
        (filename, classname) = writer.write_model_override(model=model, version=version)

        filenames[filename] = classname

    def _write_fetcher_file(self, model, filenames, version):
        """ Write the fetcher file for the model

            Args:
                model: the model to write
                filenames: list of generates filenames

        """
        writer = self.get_writer()

        if model.name != Constants.RESTUSER:
            (filename, classname) = writer.write_fetcher(model=model, version=version)
            filenames[filename] = classname

    def _clean_files(self, except_files):
        """ Removes not generated files

            Args:
                except_files: dictionary of filenames to avoid removing

        """

        except_files = list(set(except_files + SDKWriter.IGNORED_FILES))

        writer = self.get_writer()
        writer.clean(except_files=except_files)

class VSDKFileWriter(TemplateFileWriter):
    """ Provide usefull method to write Python files.

    """
    def __init__(self, directory):
        """ Initializes a VSDKFileWriter

        """
        super(VSDKFileWriter, self).__init__(directory=directory, package=u'monolithe.generators.vspk')

        self._vanilla_path = '%s/../vanilla/vsdk' % os.path.dirname(os.path.realpath(__file__))
        self._override_path = '%s/overrides' % self._vanilla_path

        self._template_folder = 'templates'
        self._autogenerate_path = '/autogenerates/'
        self._fetchers_path = '/fetchers/'
        self._setup_template = 'setup.py.tpl'
        self._fetcher_template = 'nuobject_fetcher.py.tpl'
        self._model_override_template = 'nuobject_override.py.tpl'
        self._model_template = 'nuobject_autogenerate.py.tpl'
        self._restuser_template = 'nurestuser.py.tpl'
        self._constants_template = 'constants.py.tpl'
        self._vsdsession_template = 'nuvsdsession.py.tpl'
        self._autogenerate_init_template = '__autogenerate_init__.py.tpl'
        self._fetcher_init_template = '__fetcher_init__.py.tpl'
        self._override_init_template = '__override_init__.py.tpl'

    def copy_default_files(self):
        """ Copy default sources to the output directory

        """
        shutil.copytree('%s/base' % self._vanilla_path, self.directory)

    def write_setup_file(self, version, revision):
        """ Write setup.py file

            Will generate a setup.py with version set to version-revision

            Args:
                version: version of the package
                revision: number of the revision

            Example:
                setup.py with version 3.0-1

        """
        destination = '%s/../' % self.directory
        filename = 'setup.py'

        self.write(destination=destination, filename=filename, template_name=self._setup_template,  apiversion=version, revisionnumber=revision)

    def write_constants_file(self, constants):
        """ Write constants file

            Args:
                constants (dict): dict of constants

        """
        destination = self.directory
        filename = 'constants.py'

        self.write(destination=destination, filename=filename, template_name=self._constants_template, constants=constants)

    def write_init_autogenerate_files(self, filenames):
        """ Write constants file

            Args:
                filenames (dict): dict of filename and classes

        """
        destination = '%s%s' % (self.directory, self._autogenerate_path)
        filename = '__init__.py'

        self.write(destination=destination, filename=filename, template_name=self._autogenerate_init_template, filenames=filenames)

    def write_init_fetcher_files(self, filenames):
        """ Write constants file

            Args:
                filenames (dict): dict of filename and classes

        """
        destination = '%s%s' % (self.directory, self._fetchers_path)
        filename = '__init__.py'

        self.write(destination=destination, filename=filename, template_name=self._fetcher_init_template, filenames=filenames)

    def write_init_override_files(self, filenames):
        """ Write constants file

            Args:
                filenames (dict): dict of filename and classes

        """
        destination = self.directory
        filename = '__init__.py'

        self.write(destination=destination, filename=filename, template_name=self._override_init_template, filenames=filenames)

    def write_vsdsession_file(self, version):
        """ Write VSD session file

            Args:
                version (str): the version of the vsd

        """
        destination = self.directory
        filename = 'nuvsdsession.py'

        self.write(destination=destination, filename=filename, template_name=self._vsdsession_template, version=version)

    def write_model(self, model, version):
        """ Write autogenerate model file

        """
        destination = '%s%s' % (self.directory, self._autogenerate_path)
        filename = 'nu%s.py' % model.name.lower()

        self.write(destination=destination, filename=filename, template_name=self._model_template, model=model, version=version)

        return (filename, model.name)

    def write_restuser_model(self, model, version):
        """ Write autogenerate rest user model file

        """
        destination = '%s%s' % (self.directory, self._autogenerate_path)
        filename = 'nu%s.py' % model.name.lower()

        self.write(destination=destination, filename=filename, template_name=self._restuser_template, model=model, version=version)

        return (filename, model.name)

    def write_model_override(self, model, version):
        """ Write model override

        """
        destination = self.directory
        filename = 'nu%s.py' % model.name.lower()

        # Read override from file
        override_path = '%s/%s' % (self._override_path, 'nu%s.override.py' % model.name.lower())
        override_content = None
        if os.path.isfile(override_path):
            override_content = open(override_path).read()

        self.write(destination=destination, filename=filename, template_name=self._model_override_template, model=model, override_content=override_content)
        return (filename, model.name)

    def write_fetcher(self, model, version):
        """ Write fetcher

        """
        destination = '%s%s' % (self.directory, self._fetchers_path)
        filename = 'nu%s_fetcher.py' % model.name.lower()

        self.write(destination=destination, filename=filename, template_name=self._fetcher_template, model=model, version=version)

        return (filename, model.plural_name)

    def clean(self, except_files):
        """ Clean folder to remove all files except those generated

            Args:
                except_files: list of generated files to avoid removing
        """
        self._clean_folder(folder=self._autogenerate_path, except_files=except_files)
        self._clean_folder(folder=self._fetchers_path, except_files=except_files)
        self._clean_folder(folder='', except_files=except_files)

    def _clean_folder(self, folder, except_files):
        """ Removes all files of directory except when file name
            is in except dictionary

            Args:
                folder: the folder
                except_files: dictionary of filenames to avoid removing

        """
        path = self.directory + folder

        for file in os.listdir(path):
            file_path = os.path.join(path, file)

            try:
                if os.path.isfile(file_path) and file not in except_files:
                    Printer.log("Removing file `%s` with path `%s`" % (file, file_path))
                    os.unlink(file_path)
            except:
                Printer.raiseError("An error has been raised on file %s with path %s" % (file, file_path))
<commit_msg>Fixed fetcher file name<commit_after># -*- coding: utf-8 -*-

import os
import shutil

from monolithe.lib.utils.printer import Printer
from monolithe.lib.utils.constants import Constants
from monolithe.lib.managers import TaskManager
from monolithe.lib.utils.parse import ParsingUtils
from monolithe.generators.lib.writers import TemplateFileWriter

class SDKWriter(object):
    """ Writer of the Python VSD SDK

    """
    IGNORED_FILES = ['__init__.py', 'nuvsdsession.py', 'utils.py', 'nurestuser.py', 'constants.py']
    COMMON_ATTRIBUTES = ['multicast', 'iptype', 'maintenancemode', 'permittedaction', 'connectionstate', 'forwardingclasses']

    def __init__(self, directory):
        """ Initializes a writer to the specific directory

            Args:
                directory: directory where to copy sources

        """
        if not os.path.exists(directory):
            Printer.log("Copying default sources...")
            self.writer_directory = directory
            writer = self.get_writer()
            writer.copy_default_files()

        self.writer_directory = '%s/vsdk' % directory

    def get_writer(self):
        """ Get a writer to write content

        """
        return VSDKFileWriter(directory=self.writer_directory)

    def _prepare_attributes(self, model, constants):
        """ Removes attributes and computes constants

        """
        attributes = list()
        for attribute in model.attributes:

            attributes.append(attribute)

            if attribute.allowed_choices and len(attribute.allowed_choices) > 0:

                if attribute.remote_name.lower() not in SDKWriter.COMMON_ATTRIBUTES:
                    name = '%s%s%s' % (model.name, attribute.remote_name[0].upper(), attribute.remote_name[1:])
                else:
                    name = '%s%s' % (attribute.remote_name[0].upper(), attribute.remote_name[1:])

                constants[name] = self._make_constants_value(attribute.allowed_choices)

        # Another specific case... they really love this!
        constants['ProtocolType'] = Constants.PROTOCOL_TYPES

        model.attributes = attributes

    def _make_constants_value(self, choices):
        """ Create a dictionary for constants

        """
        results = dict()

        for choice in choices:
            results[choice] = choice

        return results

    def write(self, resources, apiversion, revision):
        """ Write all files according to data

            Args:
                resources: A list of all resources to manage
                apiversion: the version of the api
                revision: the revision number of the api

            Returns:
                Writes models and fetchers files

        """

        autogenerate_filenames = dict()
        fetcher_filenames = dict()
        override_filenames = dict()
        constants = dict()

        task_manager = TaskManager()

        for model_name, model in resources.iteritems():
            self._prepare_attributes(model=model, constants=constants)
            task_manager.start_task(method=self._write_autogenerate_file, model=model, filenames=autogenerate_filenames, version=apiversion)
            task_manager.start_task(method=self._write_override_file, model=model, filenames=override_filenames, version=apiversion)
            task_manager.start_task(method=self._write_fetcher_file, model=model, filenames=fetcher_filenames, version=apiversion)

        task_manager.wait_until_exit()
        except_files = autogenerate_filenames.keys() + fetcher_filenames.keys() + override_filenames.keys()
        self._clean_files(except_files=except_files)

        writer = self.get_writer()
        writer.write_setup_file(version=apiversion, revision=revision)

        # Constants
        constants = ParsingUtils.order(constants)
        writer.write_constants_file(constants=constants)

        # VSD Session
        writer.write_vsdsession_file(version=apiversion)

        writer.write_init_autogenerate_files(filenames=autogenerate_filenames)
        writer.write_init_fetcher_files(filenames=fetcher_filenames)
        writer.write_init_override_files(filenames=override_filenames)

    def _write_autogenerate_file(self, model, filenames, version):
        """ Write the autogenerate file for the model

            Args:
                model: the model to write
                filenames: list of generates filenames

        """
        writer = self.get_writer()

        if model.name != Constants.RESTUSER:
            (filename, classname) = writer.write_model(model=model, version=version)
        else:
            (filename, classname) = writer.write_restuser_model(model=model, version=version)

        filenames[filename] = classname

    def _write_override_file(self, model, filenames, version):
        """ Write the override file for the model

            Args:
                model: the model to write

        """
        writer = self.get_writer()
        (filename, classname) = writer.write_model_override(model=model, version=version)

        filenames[filename] = classname

    def _write_fetcher_file(self, model, filenames, version):
        """ Write the fetcher file for the model

            Args:
                model: the model to write
                filenames: list of generates filenames

        """
        writer = self.get_writer()

        if model.name != Constants.RESTUSER:
            (filename, classname) = writer.write_fetcher(model=model, version=version)
            filenames[filename] = classname

    def _clean_files(self, except_files):
        """ Removes not generated files

            Args:
                except_files: dictionary of filenames to avoid removing

        """

        except_files = list(set(except_files + SDKWriter.IGNORED_FILES))

        writer = self.get_writer()
        writer.clean(except_files=except_files)

class VSDKFileWriter(TemplateFileWriter):
    """ Provide usefull method to write Python files.

    """
    def __init__(self, directory):
        """ Initializes a VSDKFileWriter

        """
        super(VSDKFileWriter, self).__init__(directory=directory, package=u'monolithe.generators.vspk')

        self._vanilla_path = '%s/../vanilla/vsdk' % os.path.dirname(os.path.realpath(__file__))
        self._override_path = '%s/overrides' % self._vanilla_path

        self._template_folder = 'templates'
        self._autogenerate_path = '/autogenerates/'
        self._fetchers_path = '/fetchers/'
        self._setup_template = 'setup.py.tpl'
        self._fetcher_template = 'nuobject_fetcher.py.tpl'
        self._model_override_template = 'nuobject_override.py.tpl'
        self._model_template = 'nuobject_autogenerate.py.tpl'
        self._restuser_template = 'nurestuser.py.tpl'
        self._constants_template = 'constants.py.tpl'
        self._vsdsession_template = 'nuvsdsession.py.tpl'
        self._autogenerate_init_template = '__autogenerate_init__.py.tpl'
        self._fetcher_init_template = '__fetcher_init__.py.tpl'
        self._override_init_template = '__override_init__.py.tpl'

    def copy_default_files(self):
        """ Copy default sources to the output directory

        """
        shutil.copytree('%s/base' % self._vanilla_path, self.directory)

    def write_setup_file(self, version, revision):
        """ Write setup.py file

            Will generate a setup.py with version set to version-revision

            Args:
                version: version of the package
                revision: number of the revision

            Example:
                setup.py with version 3.0-1

        """
        destination = '%s/../' % self.directory
        filename = 'setup.py'

        self.write(destination=destination, filename=filename, template_name=self._setup_template,  apiversion=version, revisionnumber=revision)

    def write_constants_file(self, constants):
        """ Write constants file

            Args:
                constants (dict): dict of constants

        """
        destination = self.directory
        filename = 'constants.py'

        self.write(destination=destination, filename=filename, template_name=self._constants_template, constants=constants)

    def write_init_autogenerate_files(self, filenames):
        """ Write constants file

            Args:
                filenames (dict): dict of filename and classes

        """
        destination = '%s%s' % (self.directory, self._autogenerate_path)
        filename = '__init__.py'

        self.write(destination=destination, filename=filename, template_name=self._autogenerate_init_template, filenames=filenames)

    def write_init_fetcher_files(self, filenames):
        """ Write constants file

            Args:
                filenames (dict): dict of filename and classes

        """
        destination = '%s%s' % (self.directory, self._fetchers_path)
        filename = '__init__.py'

        self.write(destination=destination, filename=filename, template_name=self._fetcher_init_template, filenames=filenames)

    def write_init_override_files(self, filenames):
        """ Write constants file

            Args:
                filenames (dict): dict of filename and classes

        """
        destination = self.directory
        filename = '__init__.py'

        self.write(destination=destination, filename=filename, template_name=self._override_init_template, filenames=filenames)

    def write_vsdsession_file(self, version):
        """ Write VSD session file

            Args:
                version (str): the version of the vsd

        """
        destination = self.directory
        filename = 'nuvsdsession.py'

        self.write(destination=destination, filename=filename, template_name=self._vsdsession_template, version=version)

    def write_model(self, model, version):
        """ Write autogenerate model file

        """
        destination = '%s%s' % (self.directory, self._autogenerate_path)
        filename = 'nu%s.py' % model.name.lower()

        self.write(destination=destination, filename=filename, template_name=self._model_template, model=model, version=version)

        return (filename, model.name)

    def write_restuser_model(self, model, version):
        """ Write autogenerate rest user model file

        """
        destination = '%s%s' % (self.directory, self._autogenerate_path)
        filename = 'nu%s.py' % model.name.lower()

        self.write(destination=destination, filename=filename, template_name=self._restuser_template, model=model, version=version)

        return (filename, model.name)

    def write_model_override(self, model, version):
        """ Write model override

        """
        destination = self.directory
        filename = 'nu%s.py' % model.name.lower()

        # Read override from file
        override_path = '%s/%s' % (self._override_path, 'nu%s.override.py' % model.name.lower())
        override_content = None
        if os.path.isfile(override_path):
            override_content = open(override_path).read()

        self.write(destination=destination, filename=filename, template_name=self._model_override_template, model=model, override_content=override_content)
        return (filename, model.name)

    def write_fetcher(self, model, version):
        """ Write fetcher

        """
        destination = '%s%s' % (self.directory, self._fetchers_path)
        filename = 'nu%s_fetcher.py' % model.plural_name.lower()

        self.write(destination=destination, filename=filename, template_name=self._fetcher_template, model=model, version=version)

        return (filename, model.plural_name)

    def clean(self, except_files):
        """ Clean folder to remove all files except those generated

            Args:
                except_files: list of generated files to avoid removing
        """
        self._clean_folder(folder=self._autogenerate_path, except_files=except_files)
        self._clean_folder(folder=self._fetchers_path, except_files=except_files)
        self._clean_folder(folder='', except_files=except_files)

    def _clean_folder(self, folder, except_files):
        """ Removes all files of directory except when file name
            is in except dictionary

            Args:
                folder: the folder
                except_files: dictionary of filenames to avoid removing

        """
        path = self.directory + folder

        for file in os.listdir(path):
            file_path = os.path.join(path, file)

            try:
                if os.path.isfile(file_path) and file not in except_files:
                    Printer.log("Removing file `%s` with path `%s`" % (file, file_path))
                    os.unlink(file_path)
            except:
                Printer.raiseError("An error has been raised on file %s with path %s" % (file, file_path))
<|endoftext|>


<commit_before>from django.utils import timezone
from django.db import transaction
from .models import Contest, ContestParticipant
from account.models import User
from submission.models import SubmissionStatus
from functools import cmp_to_key

def recalculate_for_participant(contest, submissions, problems):
    """
    :param contest: contest
    :param submissions: submissions
    :param problems: contest problems
    :return: (calculated score, penalty, cache)

    Penalty is the same for all rules: Every failed submission till the accepted one
    will add 1200 to it. And the accepted one will add the time (seconds) to it

    Score is different for rules:
    1. ACM Rule: solved problems
    2. OI Rule: The sum of passed percentage for every problem

    Cache is different for rules:
    1. ACM Rule: solved, penalty,
        +
        or
        -1

i:s(if accepted)
        Note that CE is ignored on the board.
    2. OI Rule: score, penalty,
        100
        or
        30
        or
        0

i:s(if accepted)
    """

    def get_penalty(start_time, submit_time):
        """
        :param start_time: time in DateTimeField
        :param submit_time: time in DateTimeField
        :return: penalty in seconds
        """
        return int((submit_time - start_time).total_seconds())

    def get_time_display(time):
        """
        :param time: time in seconds
        :return: mm:ss
        """
        time_in_minutes = time // 60
        return "%.2d:%.2d" % (time_in_minutes // 60, time_in_minutes % 60)

    # Get the cache: problem => contest_problem identifier
    from collections import Counter, defaultdict
    wrong = Counter()
    subcnt = Counter()
    accept = set()
    waiting = set()
    max_score = defaultdict(int)
    latest_score = defaultdict(int)
    accept_time = dict()
    first_blood = set()
    penalty = 0

    identify_problem = dict()
    for problem in problems:
        identify_problem[problem.problem.pk] = problem.identifier
    subs = [(identify_problem[submission.problem.pk], submission.status, submission.status_percent,
             submission.create_time, submission.problem.pk, submission.pk)
            for submission in submissions if identify_problem.get(submission.problem.pk)]

    # pre-processing: find all waiting problems
    for sub in reversed(subs):
        problem, status, _, create_time, _, _ = sub
        # After freeze time, everything becomes waiting
        if contest.get_status() == 'running' and contest.freeze and create_time >= contest.freeze_time:
            status = SubmissionStatus.WAITING
        if contest.get_frozen() == 'frozen2':
            status = SubmissionStatus.WAITING
        if status == SubmissionStatus.WAITING or status == SubmissionStatus:
            waiting.add(problem)

    # From beginning to the end
    for sub in reversed(subs):
        problem, status, score, create_time, original_problem, submission_id = sub
        subcnt[problem] += 1

        # If problem is waiting, then nothing can be done to this problem
        # It is waiting forever......
        if problem not in waiting:
            max_score[problem] = max(max_score[problem], score)
            latest_score[problem] = score

            if status == SubmissionStatus.WAITING or status == SubmissionStatus.JUDGING:
                waiting.add(problem)
                if problem in accept:
                    accept.remove(problem)
            if problem not in accept:
                if status == SubmissionStatus.ACCEPTED:
                    accept.add(problem)
                    accept_time[problem] = get_penalty(contest.start_time, create_time)
                    if contest.rule != 'oi2':
                        penalty += accept_time[problem] + wrong[problem] * 1200
                    # check if it is first blood
                    if contest.submission_set.filter(problem_id=original_problem,
                                                     status=SubmissionStatus.ACCEPTED).last().pk == submission_id:
                        first_blood.add(problem)
                elif SubmissionStatus.is_penalty(status):
                    wrong[problem] += 1

    score = 0
    cache = ''

    # HTML Caching and global score
    html_danger = '{text}'
    html_success = '{text}'
    html_first_blood = '{text}'
    html_warning = '{text}'
    html_info = '{text}'
    html_small = '{text}'
    html_column = '{column}	'

    # ACM Rule
    if contest.rule == 'acm':
        score = len(accept)
        for problem in sorted(identify_problem.values()):
            if problem in waiting:
                sub_cache = html_info.format(text='?')
            elif problem in accept:
                success_cache = '+' + (str(wrong[problem]) if wrong[problem] > 0 else '')
                if problem in first_blood:
                    sub_cache = html_first_blood.format(text=success_cache)
                else:
                    sub_cache = html_success.format(text=success_cache)
                sub_cache += '
' + html_small.format(text=get_time_display(accept_time[problem]))
            elif wrong[problem] > 0:
                danger_cache = '-' + (str(wrong[problem]))
                sub_cache = html_danger.format(text=danger_cache)
            else:
                sub_cache = ''
            cache += html_column.format(column=sub_cache)
    elif contest.rule == 'oi':
        score = sum(max_score.values())
        for problem in sorted(identify_problem.values()):
            local_score = max_score[problem]
            if problem in waiting:
                sub_cache = html_info.format(text='?')
            elif local_score == 100:
                if problem in first_blood:
                    sub_cache = html_first_blood.format(text=local_score)
                else:
                    sub_cache = html_success.format(text=local_score)
                sub_cache += '
' + html_small.format(text=get_time_display(accept_time[problem]))
            elif local_score > 0:
                sub_cache = html_warning.format(text=local_score)
            elif wrong[problem] > 0 and local_score == 0:
                sub_cache = html_danger.format(text=local_score)
            else:
                sub_cache = ''
            cache += html_column.format(column=sub_cache)
    elif contest.rule == 'oi2':
        score = sum(latest_score.values())
        for problem in sorted(identify_problem.values()):
            local_score = latest_score[problem]
            if problem in waiting:
                sub_cache = html_info.format(text='?')
            elif local_score == 100:
                sub_cache = html_success.format(text=local_score)
                sub_cache += '
' + html_small.format(text=get_time_display(accept_time[problem]))
            elif local_score > 0:
                sub_cache = html_warning.format(text=local_score)
            elif subcnt[problem] > 0 and local_score == 0:
                sub_cache = html_danger.format(text=local_score)
            else:
                sub_cache = ''
            cache += html_column.format(column=sub_cache)
    if contest.rule == 'oi2':
        cache = html_column.format(column=score) + cache
    else:
        cache = html_column.format(column=score) + html_column.format(column=int(penalty // 60)) + cache

    # print(score, penalty, cache)
    return score, penalty, cache

def update_problem_and_participant(contest_id, problem_id, user_id, accept_increment=0):
    """
    This function is used when judge result changed

    The three things we are going to do:
    1. update AC / Submit for ContestProblem
    2. update User (Participant) score
    3. update standings
    :param contest_id: the contest that is going to be affected
    :param problem_id: the problem that is going to be affected
    :param user_id: the user that is going to be affected
    :param accept_increment: does the submission increase AC or decrease?
    """
    contest = Contest.objects.get(pk=contest_id)
    with transaction.atomic():
        if accept_increment != 0:
            problem = contest.contestproblem_set.select_for_update().get(problem_id=problem_id)
            problem.add_accept(accept_increment)
            problem.save()

        participant, _ = contest.contestparticipant_set.select_for_update().\
            get_or_create(user__pk=user_id, defaults={'user': User.objects.get(pk=user_id), 'contest': contest})
        _update_participant(contest, participant)

    _recalculate_rank(contest)
    contest.standings_update_time = timezone.now()
    contest.save(update_fields=["standings_update_time"])

def _update_participant(contest, participant):
    submissions = contest.submission_set.filter(author=participant.user).all()
    problems = contest.contestproblem_set.all()
    participant.score, participant.penalty, participant.html_cache = \
        recalculate_for_participant(contest, submissions, problems)
    participant.save(update_fields=["score", "penalty", "html_cache"])

def _update_header(contest):

    def _get_standings_header(rule, contest_problems):
        template = '{info}	'

        if rule == 'oi' or rule == 'oi2':
            res = template.format(width=3, info='=')
        else:
            res = template.format(width=2, info='=')
        if rule != 'oi2':
            res += template.format(width=5, info='Penalty')
        for contest_problem in contest_problems:
            res += template.format(width=4, info=contest_problem.identifier)
        return res

    contest.contest_header = _get_standings_header(contest.rule, contest.contestproblem_set.all())
    contest.save()

def update_contest(contest):
    with transaction.atomic():
        _update_header(contest)
        participants = contest.contestparticipant_set.select_for_update().all()
        for participant in participants:
            _update_participant(contest, participant)
    _recalculate_rank(contest)

    contest.standings_update_time = timezone.now()
    contest.save(update_fields=["standings_update_time"])

def _recalculate_rank(contest):

    def compare(a, b):
        if a.score == b.score:
            if a.penalty == b.penalty:
                return 0
            return -1 if a.penalty  b.score else 1

    with transaction.atomic():
        participants = contest.contestparticipant_set.select_for_update().all()
        for ind, par in enumerate(sorted(participants, key=cmp_to_key(compare)), start=1):
            new_rank = ind
            if ind > 1 and compare(participants[ind - 2], par) == 0:
                new_rank = participants[ind - 1].rank
            if new_rank != par.rank:
                par.rank = new_rank
                par.save(update_fields=["rank"])

def add_participant_with_invitation(contest_pk, invitation_pk, user):
    with transaction.atomic():
        contest = Contest.objects.get(pk=contest_pk)
        invitation = contest.contestinvitation_set.get(pk=invitation_pk)
        ContestParticipant.objects.create(user=user, comment=invitation.comment, contest=contest)
        invitation.delete()
        update_contest(contest)
<commit_msg>fix: rank issue<commit_after>from django.utils import timezone
from django.db import transaction
from .models import Contest, ContestParticipant
from account.models import User
from submission.models import SubmissionStatus
from functools import cmp_to_key

def recalculate_for_participant(contest, submissions, problems):
    """
    :param contest: contest
    :param submissions: submissions
    :param problems: contest problems
    :return: (calculated score, penalty, cache)

    Penalty is the same for all rules: Every failed submission till the accepted one
    will add 1200 to it. And the accepted one will add the time (seconds) to it

    Score is different for rules:
    1. ACM Rule: solved problems
    2. OI Rule: The sum of passed percentage for every problem

    Cache is different for rules:
    1. ACM Rule: solved, penalty,
        +
        or
        -1

i:s(if accepted)
        Note that CE is ignored on the board.
    2. OI Rule: score, penalty,
        100
        or
        30
        or
        0

i:s(if accepted)
    """

    def get_penalty(start_time, submit_time):
        """
        :param start_time: time in DateTimeField
        :param submit_time: time in DateTimeField
        :return: penalty in seconds
        """
        return int((submit_time - start_time).total_seconds())

    def get_time_display(time):
        """
        :param time: time in seconds
        :return: mm:ss
        """
        time_in_minutes = time // 60
        return "%.2d:%.2d" % (time_in_minutes // 60, time_in_minutes % 60)

    # Get the cache: problem => contest_problem identifier
    from collections import Counter, defaultdict
    wrong = Counter()
    subcnt = Counter()
    accept = set()
    waiting = set()
    max_score = defaultdict(int)
    latest_score = defaultdict(int)
    accept_time = dict()
    first_blood = set()
    penalty = 0

    identify_problem = dict()
    for problem in problems:
        identify_problem[problem.problem.pk] = problem.identifier
    subs = [(identify_problem[submission.problem.pk], submission.status, submission.status_percent,
             submission.create_time, submission.problem.pk, submission.pk)
            for submission in submissions if identify_problem.get(submission.problem.pk)]

    # pre-processing: find all waiting problems
    for sub in reversed(subs):
        problem, status, _, create_time, _, _ = sub
        # After freeze time, everything becomes waiting
        if contest.get_status() == 'running' and contest.freeze and create_time >= contest.freeze_time:
            status = SubmissionStatus.WAITING
        if contest.get_frozen() == 'frozen2':
            status = SubmissionStatus.WAITING
        if status == SubmissionStatus.WAITING or status == SubmissionStatus:
            waiting.add(problem)

    # From beginning to the end
    for sub in reversed(subs):
        problem, status, score, create_time, original_problem, submission_id = sub
        subcnt[problem] += 1

        # If problem is waiting, then nothing can be done to this problem
        # It is waiting forever......
        if problem not in waiting:
            max_score[problem] = max(max_score[problem], score)
            latest_score[problem] = score

            if status == SubmissionStatus.WAITING or status == SubmissionStatus.JUDGING:
                waiting.add(problem)
                if problem in accept:
                    accept.remove(problem)
            if problem not in accept:
                if status == SubmissionStatus.ACCEPTED:
                    accept.add(problem)
                    accept_time[problem] = get_penalty(contest.start_time, create_time)
                    if contest.rule != 'oi2':
                        penalty += accept_time[problem] + wrong[problem] * 1200
                    # check if it is first blood
                    if contest.submission_set.filter(problem_id=original_problem,
                                                     status=SubmissionStatus.ACCEPTED).last().pk == submission_id:
                        first_blood.add(problem)
                elif SubmissionStatus.is_penalty(status):
                    wrong[problem] += 1

    score = 0
    cache = ''

    # HTML Caching and global score
    html_danger = '{text}'
    html_success = '{text}'
    html_first_blood = '{text}'
    html_warning = '{text}'
    html_info = '{text}'
    html_small = '{text}'
    html_column = '{column}	'

    # ACM Rule
    if contest.rule == 'acm':
        score = len(accept)
        for problem in sorted(identify_problem.values()):
            if problem in waiting:
                sub_cache = html_info.format(text='?')
            elif problem in accept:
                success_cache = '+' + (str(wrong[problem]) if wrong[problem] > 0 else '')
                if problem in first_blood:
                    sub_cache = html_first_blood.format(text=success_cache)
                else:
                    sub_cache = html_success.format(text=success_cache)
                sub_cache += '
' + html_small.format(text=get_time_display(accept_time[problem]))
            elif wrong[problem] > 0:
                danger_cache = '-' + (str(wrong[problem]))
                sub_cache = html_danger.format(text=danger_cache)
            else:
                sub_cache = ''
            cache += html_column.format(column=sub_cache)
    elif contest.rule == 'oi':
        score = sum(max_score.values())
        for problem in sorted(identify_problem.values()):
            local_score = max_score[problem]
            if problem in waiting:
                sub_cache = html_info.format(text='?')
            elif local_score == 100:
                if problem in first_blood:
                    sub_cache = html_first_blood.format(text=local_score)
                else:
                    sub_cache = html_success.format(text=local_score)
                sub_cache += '
' + html_small.format(text=get_time_display(accept_time[problem]))
            elif local_score > 0:
                sub_cache = html_warning.format(text=local_score)
            elif wrong[problem] > 0 and local_score == 0:
                sub_cache = html_danger.format(text=local_score)
            else:
                sub_cache = ''
            cache += html_column.format(column=sub_cache)
    elif contest.rule == 'oi2':
        score = sum(latest_score.values())
        for problem in sorted(identify_problem.values()):
            local_score = latest_score[problem]
            if problem in waiting:
                sub_cache = html_info.format(text='?')
            elif local_score == 100:
                sub_cache = html_success.format(text=local_score)
                sub_cache += '
' + html_small.format(text=get_time_display(accept_time[problem]))
            elif local_score > 0:
                sub_cache = html_warning.format(text=local_score)
            elif subcnt[problem] > 0 and local_score == 0:
                sub_cache = html_danger.format(text=local_score)
            else:
                sub_cache = ''
            cache += html_column.format(column=sub_cache)
    if contest.rule == 'oi2':
        cache = html_column.format(column=score) + cache
    else:
        cache = html_column.format(column=score) + html_column.format(column=int(penalty // 60)) + cache

    # print(score, penalty, cache)
    return score, penalty, cache

def update_problem_and_participant(contest_id, problem_id, user_id, accept_increment=0):
    """
    This function is used when judge result changed

    The three things we are going to do:
    1. update AC / Submit for ContestProblem
    2. update User (Participant) score
    3. update standings
    :param contest_id: the contest that is going to be affected
    :param problem_id: the problem that is going to be affected
    :param user_id: the user that is going to be affected
    :param accept_increment: does the submission increase AC or decrease?
    """
    contest = Contest.objects.get(pk=contest_id)
    with transaction.atomic():
        if accept_increment != 0:
            problem = contest.contestproblem_set.select_for_update().get(problem_id=problem_id)
            problem.add_accept(accept_increment)
            problem.save()

        participant, _ = contest.contestparticipant_set.select_for_update().\
            get_or_create(user__pk=user_id, defaults={'user': User.objects.get(pk=user_id), 'contest': contest})
        _update_participant(contest, participant)

    _recalculate_rank(contest)
    contest.standings_update_time = timezone.now()
    contest.save(update_fields=["standings_update_time"])

def _update_participant(contest, participant):
    submissions = contest.submission_set.filter(author=participant.user).all()
    problems = contest.contestproblem_set.all()
    participant.score, participant.penalty, participant.html_cache = \
        recalculate_for_participant(contest, submissions, problems)
    participant.save(update_fields=["score", "penalty", "html_cache"])

def _update_header(contest):

    def _get_standings_header(rule, contest_problems):
        template = '{info}	'

        if rule == 'oi' or rule == 'oi2':
            res = template.format(width=3, info='=')
        else:
            res = template.format(width=2, info='=')
        if rule != 'oi2':
            res += template.format(width=5, info='Penalty')
        for contest_problem in contest_problems:
            res += template.format(width=4, info=contest_problem.identifier)
        return res

    contest.contest_header = _get_standings_header(contest.rule, contest.contestproblem_set.all())
    contest.save()

def update_contest(contest):
    with transaction.atomic():
        _update_header(contest)
        participants = contest.contestparticipant_set.select_for_update().all()
        for participant in participants:
            _update_participant(contest, participant)
    _recalculate_rank(contest)

    contest.standings_update_time = timezone.now()
    contest.save(update_fields=["standings_update_time"])

def _recalculate_rank(contest):

    def compare(a, b):
        if a.score == b.score:
            if a.penalty == b.penalty:
                return 0
            return -1 if a.penalty  b.score else 1

    with transaction.atomic():
        participants = contest.contestparticipant_set.select_for_update().all()
        for ind, par in enumerate(sorted(participants, key=cmp_to_key(compare)), start=1):
            new_rank = ind
            if ind > 1 and compare(participants[ind - 2], par) == 0:
                new_rank = participants[ind - 2].rank
            if new_rank != par.rank:
                par.rank = new_rank
                par.save(update_fields=["rank"])

def add_participant_with_invitation(contest_pk, invitation_pk, user):
    with transaction.atomic():
        contest = Contest.objects.get(pk=contest_pk)
        invitation = contest.contestinvitation_set.get(pk=invitation_pk)
        ContestParticipant.objects.create(user=user, comment=invitation.comment, contest=contest)
        invitation.delete()
        update_contest(contest)
<|endoftext|>


<commit_before>import random
from queue import Queue
import json
from threading import Thread, Lock, Semaphore
import time
import uuid
import datetime
import logging
import boto3
from motorway.messages import Message
from motorway.ramp import Ramp
from boto3.dynamodb.conditions import Attr
shard_election_logger = logging.getLogger("motorway.contrib.amazon_kinesis.shard_election")

logger = logging.getLogger(__name__)

class NoItemsReturned(Exception):
    pass

class KinesisRamp(Ramp):
    stream_name = None
    heartbeat_timeout = 30  # Wait 10 seconds for a heartbeat update, or kill it
    MAX_UNCOMPLETED_ITEMS = 3000
    GET_RECORDS_LIMIT = 1000

    def __init__(self, shard_threads_enabled=True, **kwargs):
        super(KinesisRamp, self).__init__(**kwargs)
        self.conn = boto3.client(**self.connection_parameters('kinesis'))
        assert self.stream_name, "Please define attribute stream_name on your KinesisRamp"

        control_table_name = self.get_control_table_name()

        self.worker_id = str(uuid.uuid4())
        self.semaphore = Semaphore()
        self.uncompleted_ids = {}
        self.dynamodb_client = boto3.client(**self.connection_parameters('dynamodb'))

        if shard_threads_enabled:
            self.dynamodb = boto3.resource(**self.connection_parameters('dynamodb'))

            try:
                self.dynamodb_client.describe_table(TableName=control_table_name)
            except self.dynamodb_client.exceptions.ResourceNotFoundException:
                self.dynamodb_client.create_table(
                    TableName=control_table_name,
                    KeySchema=[
                        {
                            'AttributeName': 'shard_id',
                            'KeyType': 'HASH'
                        },
                    ],
                    ProvisionedThroughput={
                        'ReadCapacityUnits': 10,
                        'WriteCapacityUnits': 10
                    },
                    AttributeDefinitions=[
                        {
                            'AttributeName': 'shard_id',
                            'AttributeType': 'S'
                        },
                    ],
                )
                self.dynamodb_client.get_waiter('table_exists').wait(TableName=control_table_name)

            self.control_table = self.dynamodb.Table(control_table_name)
            shards = self.conn.describe_stream(StreamName=self.stream_name)['StreamDescription']['Shards']
            random.shuffle(shards)  # Start the threads in random order, in case of bulk restart
            threads = []
            self.insertion_queue = Queue()
            for i, shard in enumerate(shards):
                self.uncompleted_ids[shard['ShardId']] = set()
                t = Thread(target=self.process_shard, name="%s-%s" % (self.__class__.__name__, i), args=(shard['ShardId'], ))
                threads.append(t)
                t.start()

    def get_control_table_name(self):
        return 'pipeline-control-%s' % self.stream_name

    def claim_shard(self, shard_id):
        """
        Atomically update the shard in DynamoDB

        :param shard_id:
        :return: bool
        """
        shard_election_logger.info("Claiming shard %s" % shard_id)
        try:
            control_record = self.control_table.get_item(Key={'shard_id': shard_id})['Item']
        except KeyError:
            raise NoItemsReturned()
        control_record['worker_id'] = self.worker_id
        control_record['heartbeat'] = 0
        try:
            self.control_table.put_item(Item=control_record,
                                        ConditionExpression=Attr('shard_id').eq(shard_id) & Attr('checkpoint').eq(control_record['checkpoint'])
                                        # ensure that the record was not changed between the get and put.
                                        )
        except self.dynamodb_client.exceptions.ConditionalCheckFailedException:  # Someone else edited the record
            shard_election_logger.debug("Failed to claim %s to %s" % (shard_id, self.worker_id))
            return False
        return True

    def can_claim_shard(self, shard_id):
        """
        Determine whether or not a given shard can be claimed because of

        1) It's currently not being processed by another process

        2) It's unevenly balanced between the consuming nodes/workers

        :param shard_id:
        :return: bool
        """

        # =====================================
        # | C L A I M  S T A L E  S H A R D S |
        # =====================================

        worker_heartbeats = {}  # Store all heartbeats so we can compare them easily to track changes

        control_record = None
        shards = self.control_table.scan()['Items']
        for shard in shards:
            if shard['shard_id'] == shard_id:
                control_record = dict(shard)
            worker_heartbeats[shard['worker_id']] = shard['heartbeat']

        if control_record is None:
            raise NoItemsReturned()

        worker_heartbeats[self.worker_id] = 0
        time.sleep(self.heartbeat_timeout)
        updated_control_record = self.control_table.get_item(Key={'shard_id': shard_id})['Item']

        # Continue sleeping if heartbeat or worker id has changed
        if control_record['heartbeat'] == updated_control_record['heartbeat'] and control_record['worker_id'] == updated_control_record['worker_id']:
            # if both the heartbeat and the worker_id is the same
            shard_election_logger.debug("Shard %s - heartbeat and worker id remained unchanged for defined time, taking over" % shard_id)
            return True
        elif updated_control_record['worker_id'] != control_record['worker_id']:
            shard_election_logger.debug("Shard %s - Worker id changed to %s, continue sleeping" % (shard_id, updated_control_record['worker_id']))
        else:
            shard_election_logger.debug("Shard %s - Heartbeat changed, continue sleeping" % shard_id)

        # =====================
        # | B A L A N C I N G |
        # =====================

        # Balance, if possible
        active_workers = {
            self.worker_id: True
        }

        # re-fetch the shards and compare the heartbeat
        shards = self.control_table.scan()['Items']
        for shard in shards:  # Update active worker cache
            if shard['worker_id'] in worker_heartbeats and worker_heartbeats[shard['worker_id']] == shard['heartbeat']:
                active_workers[shard['worker_id']] = False
            else:
                active_workers[shard['worker_id']] = True

        number_of_active_workers = sum([1 for is_active in active_workers.values() if is_active])
        number_of_shards = len(shards)
        optimal_number_of_shards_per_worker = number_of_shards // number_of_active_workers
        workers = set([shard['worker_id'] for shard in shards])
        shards_per_worker = {worker: sum([1 for shard in shards if shard['worker_id'] == worker]) for worker in workers}
        for shard in shards:
            if shard['shard_id'] == shard_id:
                if (
                    # Check if the shards current worker has too many, or if the worker has no workers, then take
                    # the shard if the current worker has more than one shard!
                        shards_per_worker.get(shard['worker_id'], 0) > optimal_number_of_shards_per_worker or (
                                    shards_per_worker.get(self.worker_id, 0) == 0 and
                                    shards_per_worker.get(shard['worker_id'], 0) > 1
                                )
                ) and (
                    # Only get shards for balancing purposes, if we have too little
                        shards_per_worker.get(self.worker_id, 0) < optimal_number_of_shards_per_worker
                ):
                        shard_election_logger.debug("Taking over %s from %s" % (shard_id, shard['worker_id']))
                        return True
        return False

    def process_shard(self, shard_id):
        """
        Every shard (at startup) has an active thread that runs this function to either consume or wait to be
        ready to consume data from a shard

        :param shard_id:
        :return:
        """
        while True:
            try:
                # try to claim the shard
                try:
                    # Continuously try to claim until broken out of
                    while True:
                        # First check if we can claim outside the semaphore in parallel on all shards
                        if self.can_claim_shard(shard_id):
                            # If we can claim it, try again with a semaphore ensuring we fully check the entire
                            # table of workers/shards before we take the final decision
                            with self.semaphore:
                                if self.can_claim_shard(shard_id):
                                    if self.claim_shard(shard_id):
                                        break
                        # Wait a bit until we check if it's available to claim again
                        time.sleep(random.randrange(2, 15))
                except NoItemsReturned:
                    # no record for this shard found, nobody ever claimed the shard yet, so claim it
                    self.control_table.put_item(Item={
                        'shard_id': shard_id,
                        'checkpoint': 0,
                        'worker_id': self.worker_id,
                        'heartbeat': 0,
                    })

                # get initial iterator
                control_record = self.control_table.get_item(Key={'shard_id': shard_id})['Item']
                if control_record['checkpoint'] > 0:
                    # if we have a checkpoint, start from the checkpoint
                    iterator = self.conn.get_shard_iterator(
                        StreamName=self.stream_name,
                        ShardId=shard_id,
                        ShardIteratorType="AT_SEQUENCE_NUMBER",
                        StartingSequenceNumber=str(control_record['checkpoint'])
                    )['ShardIterator']
                else:
                    # we have no checkpoint stored, start from the latest item in Kinesis
                    iterator = self.conn.get_shard_iterator(
                        StreamName=self.stream_name,
                        ShardId=shard_id,
                        ShardIteratorType="LATEST",
                    )['ShardIterator']

                cloudwatch = boto3.client(**self.connection_parameters('cloudwatch'))
                current_minute = lambda: datetime.datetime.now().minute
                minute = None
                latest_item = None
                while True:
                    control_record = self.control_table.get_item(Key={'shard_id': shard_id})['Item']  # always retrieve this at the top of the loop
                    current_checkpoint = control_record['checkpoint']
                    current_heartbeat = control_record['heartbeat']
                    # if the shard was claimed by another worker, break out of the loop
                    if not control_record['worker_id'] == self.worker_id:
                        shard_election_logger.info("Lost shard %s, going back to standby" % shard_id)
                        break

                    # update the heartbeat and the checkpoint
                    control_record['heartbeat'] += 1
                    if len(self.uncompleted_ids[shard_id]):
                        # Get the "youngest" uncompleted sequence number
                        control_record['checkpoint'] = min(self.uncompleted_ids[shard_id])
                    elif latest_item:
                        # or the latest item we yielded
                        control_record['checkpoint'] = latest_item

                    self.control_table.put_item(Item=control_record,
                                                ConditionExpression=Attr('shard_id').eq(shard_id) & Attr('checkpoint').eq(current_checkpoint) & Attr('worker_id').eq(self.worker_id) & Attr('heartbeat').eq(current_heartbeat)
                                                # Will fail if someone else modified it - ConditionalCheckFailedException
                                                )

                    if len(self.uncompleted_ids[shard_id]) < self.MAX_UNCOMPLETED_ITEMS:
                        # get records from Kinesis, using the previously created iterator
                        result = self.conn.get_records(ShardIterator=iterator, Limit=self.GET_RECORDS_LIMIT)

                        # insert the records into the queue, and use the provided iterator for the next loop
                        for record in result['Records']:
                            self.uncompleted_ids[shard_id].add(record['SequenceNumber'])
                            latest_item = record['SequenceNumber']
                            self.insertion_queue.put(record)
                        iterator = result['NextShardIterator']
                    else:
                        logger.debug("Pausing, too many uncompleted items (%s/%s)" % (len(self.uncompleted_ids[shard_id]), self.MAX_UNCOMPLETED_ITEMS))
                        # we have too many uncompleted items, so back off for a while
                        # however, the iterator needs to be updated, because it expires after a while
                        # use the latest record we added to the queue as the starting point
                        next_iterator_number = latest_item if latest_item else str(control_record['checkpoint'])
                        iterator = self.conn.get_shard_iterator(
                            StreamName=self.stream_name,
                            ShardId=shard_id,
                            ShardIteratorType="AT_SEQUENCE_NUMBER",
                            StartingSequenceNumber=next_iterator_number
                        )['ShardIterator']

                        # get just one item to update the MillisBehindLatest below
                        result = self.conn.get_records(ShardIterator=iterator, Limit=1)

                    # Push metrics to CloudWatch
                    delay = result['MillisBehindLatest']
                    if minute != current_minute():  # push once per minute to CloudWatch.
                        minute = current_minute()
                        cloudwatch.put_metric_data(Namespace='Motorway/Kinesis',
                                                   MetricData=[{'MetricName': 'MillisecondsBehind',
                                                                'Dimensions': [{
                                                                    'Name': 'Stream',
                                                                    'Value': self.stream_name
                                                                }, {
                                                                    'Name': 'Shard',
                                                                    'Value': shard_id
                                                                }],
                                                                'Value': delay,
                                                                'Unit': 'Milliseconds'
                                                                }])

                    # recommended pause between fetches from AWS
                    time.sleep(1)
            except self.dynamodb_client.exceptions.ConditionalCheckFailedException as e:
                logger.warning(e)
                pass  # we're no longer worker for this shard
            except (self.dynamodb_client.exceptions.ProvisionedThroughputExceededException, self.dynamodb_client.exceptions.ProvisionedThroughputExceededException,
                    self.conn.exceptions.LimitExceededException, self.conn.exceptions.ProvisionedThroughputExceededException) as e:
                logger.warning(e)
                time.sleep(random.randrange(5, self.heartbeat_timeout//2))  # back off for a while

    def connection_parameters(self, service_name):
        return {
            'region_name': 'eu-west-1',
            'service_name': service_name,
            # Add this or use ENV VARS
            # 'aws_access_key_id': '',
            # 'aws_secret_access_key': ''
        }

    def next(self):
        msg = self.insertion_queue.get()
        try:
            yield Message(msg['SequenceNumber'], json.loads(msg['Data']), grouping_value=msg['PartitionKey'])
        except ValueError as e:
            logger.exception(e)

    def success(self, _id):
        for uncompleted_ids in self.uncompleted_ids.values():
            if _id in uncompleted_ids:
                uncompleted_ids.remove(_id)
<commit_msg>fixed if str to int comparison<commit_after>import random
from queue import Queue
import json
from threading import Thread, Lock, Semaphore
import time
import uuid
import datetime
import logging
import boto3
from motorway.messages import Message
from motorway.ramp import Ramp
from boto3.dynamodb.conditions import Attr
shard_election_logger = logging.getLogger("motorway.contrib.amazon_kinesis.shard_election")

logger = logging.getLogger(__name__)

class NoItemsReturned(Exception):
    pass

class KinesisRamp(Ramp):
    stream_name = None
    heartbeat_timeout = 30  # Wait 10 seconds for a heartbeat update, or kill it
    MAX_UNCOMPLETED_ITEMS = 3000
    GET_RECORDS_LIMIT = 1000

    def __init__(self, shard_threads_enabled=True, **kwargs):
        super(KinesisRamp, self).__init__(**kwargs)
        self.conn = boto3.client(**self.connection_parameters('kinesis'))
        assert self.stream_name, "Please define attribute stream_name on your KinesisRamp"

        control_table_name = self.get_control_table_name()

        self.worker_id = str(uuid.uuid4())
        self.semaphore = Semaphore()
        self.uncompleted_ids = {}
        self.dynamodb_client = boto3.client(**self.connection_parameters('dynamodb'))

        if shard_threads_enabled:
            self.dynamodb = boto3.resource(**self.connection_parameters('dynamodb'))

            try:
                self.dynamodb_client.describe_table(TableName=control_table_name)
            except self.dynamodb_client.exceptions.ResourceNotFoundException:
                self.dynamodb_client.create_table(
                    TableName=control_table_name,
                    KeySchema=[
                        {
                            'AttributeName': 'shard_id',
                            'KeyType': 'HASH'
                        },
                    ],
                    ProvisionedThroughput={
                        'ReadCapacityUnits': 10,
                        'WriteCapacityUnits': 10
                    },
                    AttributeDefinitions=[
                        {
                            'AttributeName': 'shard_id',
                            'AttributeType': 'S'
                        },
                    ],
                )
                self.dynamodb_client.get_waiter('table_exists').wait(TableName=control_table_name)

            self.control_table = self.dynamodb.Table(control_table_name)
            shards = self.conn.describe_stream(StreamName=self.stream_name)['StreamDescription']['Shards']
            random.shuffle(shards)  # Start the threads in random order, in case of bulk restart
            threads = []
            self.insertion_queue = Queue()
            for i, shard in enumerate(shards):
                self.uncompleted_ids[shard['ShardId']] = set()
                t = Thread(target=self.process_shard, name="%s-%s" % (self.__class__.__name__, i), args=(shard['ShardId'], ))
                threads.append(t)
                t.start()

    def get_control_table_name(self):
        return 'pipeline-control-%s' % self.stream_name

    def claim_shard(self, shard_id):
        """
        Atomically update the shard in DynamoDB

        :param shard_id:
        :return: bool
        """
        shard_election_logger.info("Claiming shard %s" % shard_id)
        try:
            control_record = self.control_table.get_item(Key={'shard_id': shard_id})['Item']
        except KeyError:
            raise NoItemsReturned()
        control_record['worker_id'] = self.worker_id
        control_record['heartbeat'] = 0
        try:
            self.control_table.put_item(Item=control_record,
                                        ConditionExpression=Attr('shard_id').eq(shard_id) & Attr('checkpoint').eq(control_record['checkpoint'])
                                        # ensure that the record was not changed between the get and put.
                                        )
        except self.dynamodb_client.exceptions.ConditionalCheckFailedException:  # Someone else edited the record
            shard_election_logger.debug("Failed to claim %s to %s" % (shard_id, self.worker_id))
            return False
        return True

    def can_claim_shard(self, shard_id):
        """
        Determine whether or not a given shard can be claimed because of

        1) It's currently not being processed by another process

        2) It's unevenly balanced between the consuming nodes/workers

        :param shard_id:
        :return: bool
        """

        # =====================================
        # | C L A I M  S T A L E  S H A R D S |
        # =====================================

        worker_heartbeats = {}  # Store all heartbeats so we can compare them easily to track changes

        control_record = None
        shards = self.control_table.scan()['Items']
        for shard in shards:
            if shard['shard_id'] == shard_id:
                control_record = dict(shard)
            worker_heartbeats[shard['worker_id']] = shard['heartbeat']

        if control_record is None:
            raise NoItemsReturned()

        worker_heartbeats[self.worker_id] = 0
        time.sleep(self.heartbeat_timeout)
        updated_control_record = self.control_table.get_item(Key={'shard_id': shard_id})['Item']

        # Continue sleeping if heartbeat or worker id has changed
        if control_record['heartbeat'] == updated_control_record['heartbeat'] and control_record['worker_id'] == updated_control_record['worker_id']:
            # if both the heartbeat and the worker_id is the same
            shard_election_logger.debug("Shard %s - heartbeat and worker id remained unchanged for defined time, taking over" % shard_id)
            return True
        elif updated_control_record['worker_id'] != control_record['worker_id']:
            shard_election_logger.debug("Shard %s - Worker id changed to %s, continue sleeping" % (shard_id, updated_control_record['worker_id']))
        else:
            shard_election_logger.debug("Shard %s - Heartbeat changed, continue sleeping" % shard_id)

        # =====================
        # | B A L A N C I N G |
        # =====================

        # Balance, if possible
        active_workers = {
            self.worker_id: True
        }

        # re-fetch the shards and compare the heartbeat
        shards = self.control_table.scan()['Items']
        for shard in shards:  # Update active worker cache
            if shard['worker_id'] in worker_heartbeats and worker_heartbeats[shard['worker_id']] == shard['heartbeat']:
                active_workers[shard['worker_id']] = False
            else:
                active_workers[shard['worker_id']] = True

        number_of_active_workers = sum([1 for is_active in active_workers.values() if is_active])
        number_of_shards = len(shards)
        optimal_number_of_shards_per_worker = number_of_shards // number_of_active_workers
        workers = set([shard['worker_id'] for shard in shards])
        shards_per_worker = {worker: sum([1 for shard in shards if shard['worker_id'] == worker]) for worker in workers}
        for shard in shards:
            if shard['shard_id'] == shard_id:
                if (
                    # Check if the shards current worker has too many, or if the worker has no workers, then take
                    # the shard if the current worker has more than one shard!
                        shards_per_worker.get(shard['worker_id'], 0) > optimal_number_of_shards_per_worker or (
                                    shards_per_worker.get(self.worker_id, 0) == 0 and
                                    shards_per_worker.get(shard['worker_id'], 0) > 1
                                )
                ) and (
                    # Only get shards for balancing purposes, if we have too little
                        shards_per_worker.get(self.worker_id, 0) < optimal_number_of_shards_per_worker
                ):
                        shard_election_logger.debug("Taking over %s from %s" % (shard_id, shard['worker_id']))
                        return True
        return False

    def process_shard(self, shard_id):
        """
        Every shard (at startup) has an active thread that runs this function to either consume or wait to be
        ready to consume data from a shard

        :param shard_id:
        :return:
        """
        while True:
            try:
                # try to claim the shard
                try:
                    # Continuously try to claim until broken out of
                    while True:
                        # First check if we can claim outside the semaphore in parallel on all shards
                        if self.can_claim_shard(shard_id):
                            # If we can claim it, try again with a semaphore ensuring we fully check the entire
                            # table of workers/shards before we take the final decision
                            with self.semaphore:
                                if self.can_claim_shard(shard_id):
                                    if self.claim_shard(shard_id):
                                        break
                        # Wait a bit until we check if it's available to claim again
                        time.sleep(random.randrange(2, 15))
                except NoItemsReturned:
                    # no record for this shard found, nobody ever claimed the shard yet, so claim it
                    self.control_table.put_item(Item={
                        'shard_id': shard_id,
                        'checkpoint': 0,
                        'worker_id': self.worker_id,
                        'heartbeat': 0,
                    })

                # get initial iterator
                control_record = self.control_table.get_item(Key={'shard_id': shard_id})['Item']
                if control_record['checkpoint']:
                    # if we have a checkpoint, start from the checkpoint
                    iterator = self.conn.get_shard_iterator(
                        StreamName=self.stream_name,
                        ShardId=shard_id,
                        ShardIteratorType="AT_SEQUENCE_NUMBER",
                        StartingSequenceNumber=str(control_record['checkpoint'])
                    )['ShardIterator']
                else:
                    # we have no checkpoint stored, start from the latest item in Kinesis
                    iterator = self.conn.get_shard_iterator(
                        StreamName=self.stream_name,
                        ShardId=shard_id,
                        ShardIteratorType="LATEST",
                    )['ShardIterator']

                cloudwatch = boto3.client(**self.connection_parameters('cloudwatch'))
                current_minute = lambda: datetime.datetime.now().minute
                minute = None
                latest_item = None
                while True:
                    control_record = self.control_table.get_item(Key={'shard_id': shard_id})['Item']  # always retrieve this at the top of the loop
                    current_checkpoint = control_record['checkpoint']
                    current_heartbeat = control_record['heartbeat']
                    # if the shard was claimed by another worker, break out of the loop
                    if not control_record['worker_id'] == self.worker_id:
                        shard_election_logger.info("Lost shard %s, going back to standby" % shard_id)
                        break

                    # update the heartbeat and the checkpoint
                    control_record['heartbeat'] += 1
                    if len(self.uncompleted_ids[shard_id]):
                        # Get the "youngest" uncompleted sequence number
                        control_record['checkpoint'] = min(self.uncompleted_ids[shard_id])
                    elif latest_item:
                        # or the latest item we yielded
                        control_record['checkpoint'] = latest_item

                    self.control_table.put_item(Item=control_record,
                                                ConditionExpression=Attr('shard_id').eq(shard_id) & Attr('checkpoint').eq(current_checkpoint) & Attr('worker_id').eq(self.worker_id) & Attr('heartbeat').eq(current_heartbeat)
                                                # Will fail if someone else modified it - ConditionalCheckFailedException
                                                )

                    if len(self.uncompleted_ids[shard_id]) < self.MAX_UNCOMPLETED_ITEMS:
                        # get records from Kinesis, using the previously created iterator
                        result = self.conn.get_records(ShardIterator=iterator, Limit=self.GET_RECORDS_LIMIT)

                        # insert the records into the queue, and use the provided iterator for the next loop
                        for record in result['Records']:
                            self.uncompleted_ids[shard_id].add(record['SequenceNumber'])
                            latest_item = record['SequenceNumber']
                            self.insertion_queue.put(record)
                        iterator = result['NextShardIterator']
                    else:
                        logger.debug("Pausing, too many uncompleted items (%s/%s)" % (len(self.uncompleted_ids[shard_id]), self.MAX_UNCOMPLETED_ITEMS))
                        # we have too many uncompleted items, so back off for a while
                        # however, the iterator needs to be updated, because it expires after a while
                        # use the latest record we added to the queue as the starting point
                        next_iterator_number = latest_item if latest_item else str(control_record['checkpoint'])
                        iterator = self.conn.get_shard_iterator(
                            StreamName=self.stream_name,
                            ShardId=shard_id,
                            ShardIteratorType="AT_SEQUENCE_NUMBER",
                            StartingSequenceNumber=next_iterator_number
                        )['ShardIterator']

                        # get just one item to update the MillisBehindLatest below
                        result = self.conn.get_records(ShardIterator=iterator, Limit=1)

                    # Push metrics to CloudWatch
                    delay = result['MillisBehindLatest']
                    if minute != current_minute():  # push once per minute to CloudWatch.
                        minute = current_minute()
                        cloudwatch.put_metric_data(Namespace='Motorway/Kinesis',
                                                   MetricData=[{'MetricName': 'MillisecondsBehind',
                                                                'Dimensions': [{
                                                                    'Name': 'Stream',
                                                                    'Value': self.stream_name
                                                                }, {
                                                                    'Name': 'Shard',
                                                                    'Value': shard_id
                                                                }],
                                                                'Value': delay,
                                                                'Unit': 'Milliseconds'
                                                                }])

                    # recommended pause between fetches from AWS
                    time.sleep(1)
            except self.dynamodb_client.exceptions.ConditionalCheckFailedException as e:
                logger.warning(e)
                pass  # we're no longer worker for this shard
            except (self.dynamodb_client.exceptions.ProvisionedThroughputExceededException, self.dynamodb_client.exceptions.ProvisionedThroughputExceededException,
                    self.conn.exceptions.LimitExceededException, self.conn.exceptions.ProvisionedThroughputExceededException) as e:
                logger.warning(e)
                time.sleep(random.randrange(5, self.heartbeat_timeout//2))  # back off for a while

    def connection_parameters(self, service_name):
        return {
            'region_name': 'eu-west-1',
            'service_name': service_name,
            # Add this or use ENV VARS
            # 'aws_access_key_id': '',
            # 'aws_secret_access_key': ''
        }

    def next(self):
        msg = self.insertion_queue.get()
        try:
            yield Message(msg['SequenceNumber'], json.loads(msg['Data']), grouping_value=msg['PartitionKey'])
        except ValueError as e:
            logger.exception(e)

    def success(self, _id):
        for uncompleted_ids in self.uncompleted_ids.values():
            if _id in uncompleted_ids:
                uncompleted_ids.remove(_id)
<|endoftext|>


<commit_before>from core import app
app()<commit_msg>main.py: set window configuration<commit_after>from core.app import PocketLordApp
from kivy.core.window import Window
from kivy.config import Config

Config.set("graphics", "width", 1000)
Config.set("graphics", "height", 800)
PocketLordApp().run()<|endoftext|>