<commit_before>import time
from multiprocessing import Process, Value, Lock

from numpy import random
from couchbase.exceptions import ValueFormatError
from logger import logger

from spring.cbgen import CBGen
from spring.docgen import (ExistingKey, KeyForRemoval, SequentialHotKey,
                           NewKey, NewDocument)
from spring.querygen import NewQuery


def with_sleep(method):

    CORRECTION_FACTOR = 0.975  # empiric!

    def wrapper(self, *args, **kwargs):
        if self.target_time is None:
            return method(self, *args, **kwargs)
        else:
            t0 = time.time()
            method(self, *args, **kwargs)
            actual_time = time.time() - t0
            if actual_time < self.target_time:
                time.sleep(CORRECTION_FACTOR * (self.target_time - actual_time))
    return wrapper


class Worker(object):

    BATCH_SIZE = 100

    def __init__(self, workload_settings, target_settings, shutdown_event=None):
        self.ws = workload_settings
        self.ts = target_settings
        self.shutdown_event = shutdown_event

        self.existing_keys = ExistingKey(self.ws.working_set,
                                         self.ws.working_set_access,
                                         self.ts.prefix)
        self.new_keys = NewKey(self.ts.prefix, self.ws.expiration)
        self.keys_for_removal = KeyForRemoval(self.ts.prefix)
        self.docs = NewDocument(self.ws.size)

        self.next_report = 0.05  # report after every 5% of completion

        host, port = self.ts.node.split(':')
        self.init_db({"bucket": self.ts.bucket, "host": host, "port": port,
                      "username": self.ts.bucket, "password": self.ts.password})

    def init_db(self, params):
        try:
            self.cb = CBGen(**params)
        except Exception as e:
            raise SystemExit(e)

    def report_progress(self, curr_ops):
        if self.ws.ops < float('inf') and \
                curr_ops > self.next_report * self.ws.ops:
            progress = 100.0 * curr_ops / self.ws.ops
            self.next_report += 0.05
            logger.info('Current progress: {:.2f} %'.format(progress))

    def time_to_stop(self):
        return (self.shutdown_event is not None and
                self.shutdown_event.is_set())


class KVWorker(Worker):

    def gen_sequence(self):
        ops = \
            ['c'] * self.ws.creates + \
            ['r'] * self.ws.reads + \
            ['u'] * self.ws.updates + \
            ['d'] * self.ws.deletes + \
            ['cas'] * self.ws.cases
        random.shuffle(ops)
        return ops

    @with_sleep
    def do_batch(self):
        curr_items_tmp = self.curr_items.value
        deleted_items_tmp = self.deleted_items.value
        if self.ws.creates:
            with self.lock:
                self.curr_items.value += self.ws.creates
        if self.ws.deletes:
            with self.lock:
                self.deleted_items.value += self.ws.deletes

        for op in self.gen_sequence():
            if op == 'c':
                curr_items_tmp += 1
                key, ttl = self.new_keys.next(curr_items_tmp)
                doc = self.docs.next(key)
                self.cb.create(key, doc, ttl)
            elif op == 'r':
                key = self.existing_keys.next(curr_items_tmp, deleted_items_tmp)
                self.cb.read(key)
            elif op == 'u':
                key = self.existing_keys.next(curr_items_tmp, deleted_items_tmp)
                doc = self.docs.next(key)
                self.cb.update(key, doc)
            elif op == 'd':
                deleted_items_tmp += 1
                key = self.keys_for_removal.next(deleted_items_tmp)
                self.cb.delete(key)
            elif op == 'cas':
                key = self.existing_keys.next(curr_items_tmp, deleted_items_tmp)
                doc = self.docs.next(key)
                self.cb.cas(key, doc)

    def run(self, sid, lock, curr_ops, curr_items, deleted_items):
        if self.ws.throughput < float('inf'):
            self.target_time = float(self.BATCH_SIZE) * self.ws.workers / \
                self.ws.throughput
        else:
            self.target_time = None
        self.lock = lock
        self.curr_items = curr_items
        self.deleted_items = deleted_items

        logger.info('Started: worker-{}'.format(sid))
        try:
            while curr_ops.value < self.ws.ops and not self.time_to_stop():
                with lock:
                    curr_ops.value += self.BATCH_SIZE
                self.do_batch()

                if not sid:  # only first worker
                    self.report_progress(curr_ops.value)
        except (KeyboardInterrupt, ValueFormatError):<commit_msg>generate and run batch separately<commit_after>import time
from multiprocessing import Process, Value, Lock

from numpy import random
from couchbase.exceptions import ValueFormatError
from logger import logger

from spring.cbgen import CBGen
from spring.docgen import (ExistingKey, KeyForRemoval, SequentialHotKey,
                           NewKey, NewDocument)
from spring.querygen import NewQuery


def with_sleep(method):

    CORRECTION_FACTOR = 0.975  # empiric!

    def wrapper(self, *args, **kwargs):
        if self.target_time is None:
            return method(self, *args, **kwargs)
        else:
            t0 = time.time()
            method(self, *args, **kwargs)
            actual_time = time.time() - t0
            if actual_time < self.target_time:
                time.sleep(CORRECTION_FACTOR * (self.target_time - actual_time))
    return wrapper


class Worker(object):

    BATCH_SIZE = 100

    def __init__(self, workload_settings, target_settings, shutdown_event=None):
        self.ws = workload_settings
        self.ts = target_settings
        self.shutdown_event = shutdown_event

        self.existing_keys = ExistingKey(self.ws.working_set,
                                         self.ws.working_set_access,
                                         self.ts.prefix)
        self.new_keys = NewKey(self.ts.prefix, self.ws.expiration)
        self.keys_for_removal = KeyForRemoval(self.ts.prefix)
        self.docs = NewDocument(self.ws.size)

        self.next_report = 0.05  # report after every 5% of completion

        host, port = self.ts.node.split(':')
        self.init_db({"bucket": self.ts.bucket, "host": host, "port": port,
                      "username": self.ts.bucket, "password": self.ts.password})

    def init_db(self, params):
        try:
            self.cb = CBGen(**params)
        except Exception as e:
            raise SystemExit(e)

    def report_progress(self, curr_ops):
        if self.ws.ops < float('inf') and \
                curr_ops > self.next_report * self.ws.ops:
            progress = 100.0 * curr_ops / self.ws.ops
            self.next_report += 0.05
            logger.info('Current progress: {:.2f} %'.format(progress))

    def time_to_stop(self):
        return (self.shutdown_event is not None and
                self.shutdown_event.is_set())


class KVWorker(Worker):

    def gen_sequence(self):
        ops = \
            ['c'] * self.ws.creates + \
            ['r'] * self.ws.reads + \
            ['u'] * self.ws.updates + \
            ['d'] * self.ws.deletes + \
            ['cas'] * self.ws.cases
        random.shuffle(ops)
        return ops

    @with_sleep
    def do_batch(self):
        curr_items_tmp = self.curr_items.value
        deleted_items_tmp = self.deleted_items.value
        if self.ws.creates:
            with self.lock:
                self.curr_items.value += self.ws.creates
        if self.ws.deletes:
            with self.lock:
                self.deleted_items.value += self.ws.deletes

        cmds = []
        for op in self.gen_sequence():
            if op == 'c':
                curr_items_tmp += 1
                key, ttl = self.new_keys.next(curr_items_tmp)
                doc = self.docs.next(key)
                cmds.append((self.cb.create, (key, doc, ttl)))
            elif op == 'r':
                key = self.existing_keys.next(curr_items_tmp, deleted_items_tmp)
                cmds.append((self.cb.read, (key, )))
            elif op == 'u':
                key = self.existing_keys.next(curr_items_tmp, deleted_items_tmp)
                doc = self.docs.next(key)
                cmds.append((self.cb.update, (key, doc)))
            elif op == 'd':
                deleted_items_tmp += 1
                key = self.keys_for_removal.next(deleted_items_tmp)
                cmds.append((self.cb.delete, (key, )))
            elif op == 'cas':
                key = self.existing_keys.next(curr_items_tmp, deleted_items_tmp)
                doc = self.docs.next(key)
                cmds.append((self.cb.cas, (key, doc)))
        for cmd, args in cmds:
            cmd(*args)

    def run(self, sid, lock, curr_ops, curr_items, deleted_items):
        if self.ws.throughput < float('inf'):
            self.target_time = float(self.BATCH_SIZE) * self.ws.workers / \
                self.ws.throughput
        else:
            self.target_time = None
        self.lock = lock
        self.curr_items = curr_items
        self.deleted_items = deleted_items

        logger.info('Started: worker-{}'.format(sid))
        try:
            while curr_ops.value < self.ws.ops and not self.time_to_stop():
                with lock:
                    curr_ops.value += self.BATCH_SIZE
                self.do_batch()

                if not sid:  # only first worker
                    self.report_progress(curr_ops.value)
        except (KeyboardInterrupt, ValueFormatError):
{'fraction_non_alphanumeric': 0.10576425540722342, 'fraction_numerical': 0.004760426368622581, 'mean_word_length': 1.789838337182448, 'pattern_counts': {'":': 10, '<': 11, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5747342', 'n_tokens_mistral': 2953, 'n_tokens_neox': 2826, 'n_words': 603}