unrhodecode/rhodecode/tests/load/vcs_performance.py

# Copyright (C) 2016-2024 RhodeCode GmbH
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License, version 3
# (only), as published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# This program is dual-licensed. If you wish to learn more about the
# RhodeCode Enterprise Edition, including its added features, Support services,
# and proprietary license terms, please see https://rhodecode.com/licenses/

"""
VCS Performance measurement tool

Usage:

- Check that required vcs keys can be found in ~/.hgrc and ~/.netrc

- Start a local instance of RhodeCode Enterprise

- Launch the script:

    TMPDIR=/tmp python vcs_performance.py \
        --host=http://vm:5000 \
        --api-key=55c4a33688577da24183dcac5fde4dddfdbf18dc \
        --commits=10 --repositories=100 --log-level=info
"""

import argparse
import functools
import logging
import os
import shutil
import subprocess
import tempfile
import time
from itertools import chain

from api import RCApi, ApiError


log = logging.getLogger(__name__)


def timed(f):
    """Decorator that returns the time it took to execute the function."""

    @functools.wraps(f)
    def wrapped_f(*args, **kwargs):
        start_time = time.time()
        try:
            f(*args, **kwargs)
        finally:
            return time.time() - start_time

    return wrapped_f


def mean(container):
    """Return the mean of the container."""
    if not container:
        return -1.0
    return sum(container) / len(container)


class Config(object):
    args = None

    def __init__(self):
        parser = argparse.ArgumentParser(description="Runs VCS load tests")
        parser.add_argument("--host", dest="host", action="store", required=True, help="RhodeCode Enterprise host")
        parser.add_argument("--api-key", dest="api_key", action="store", required=True, help="API Key")
        parser.add_argument(
            "--file-size", dest="file_size", action="store", required=False, default=1, type=int, help="File size in MB"
        )
        parser.add_argument(
            "--repositories",
            dest="repositories",
            action="store",
            required=False,
            default=1,
            type=int,
            help="Number of repositories",
        )
        parser.add_argument(
            "--commits", dest="commits", action="store", required=False, default=1, type=int, help="Number of commits"
        )
        parser.add_argument(
            "--log-level", dest="log_level", action="store", required=False, default="error", help="Logging level"
        )
        self.args = parser.parse_args()

    def __getattr__(self, attr):
        return getattr(self.args, attr)


class Repository(object):
    FILE_NAME_TEMPLATE = "test_{:09d}.bin"

    def __init__(self, name, base_path, api):
        self.name = name
        self.path = os.path.join(base_path, name)
        self.api = api
        self.url = None

    def create(self):
        self._create_filesystem_repo(self.path)
        try:
            self.url = self.api.create_repo(self.name, self.TYPE, "Performance tests")
        except ApiError as e:
            log.error("api: {}".format(e))

    def delete(self):
        self._delete_filesystem_repo()
        try:
            self.api.delete_repo(self.name)
        except ApiError as e:
            log.error("api: {}".format(e))

    def create_commits(self, number, file_size):
        for i in range(number):
            file_name = self.FILE_NAME_TEMPLATE.format(i)
            log.debug("Create commit[{}] {}".format(self.name, file_name))
            self._create_file(file_name, file_size)
            self._create_commit(file_name)

    @timed
    def push(self):
        raise NotImplementedError()

    @timed
    def clone(self, destination_path):
        raise NotImplementedError()

    @timed
    def pull(self):
        raise NotImplementedError()

    def _run(self, *args):
        command = [self.BASE_COMMAND] + list(args)
        process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        return process.communicate()

    def _create_file(self, name, size):
        file_name = os.path.join(self.path, name)
        with open(file_name, "wb") as f:
            f.write(os.urandom(1024))

    def _delete_filesystem_repo(self):
        shutil.rmtree(self.path)

    def _create_filesystem_repo(self, path):
        raise NotImplementedError()

    def _create_commit(self, file_name):
        raise NotImplementedError()


class GitRepository(Repository):
    TYPE = "git"
    BASE_COMMAND = "git"

    @timed
    def push(self):
        os.chdir(self.path)
        self._run("push", "--set-upstream", self.url, "master")

    @timed
    def clone(self, destination_path):
        self._run("clone", self.url, os.path.join(destination_path, self.name))

    @timed
    def pull(self, destination_path):
        path = os.path.join(destination_path, self.name)
        self._create_filesystem_repo(path)
        os.chdir(path)
        self._run("remote", "add", "origin", self.url)
        self._run("pull", "origin", "master")

    def _create_filesystem_repo(self, path):
        self._run("init", path)

    def _create_commit(self, file_name):
        os.chdir(self.path)
        self._run("add", file_name)
        self._run("commit", file_name, "-m", '"Add {}"'.format(file_name))


class HgRepository(Repository):
    TYPE = "hg"
    BASE_COMMAND = "hg"

    @timed
    def push(self):
        os.chdir(self.path)
        self._run("push", self.url)

    @timed
    def clone(self, destination_path):
        self._run("clone", self.url, os.path.join(destination_path, self.name))

    @timed
    def pull(self, destination_path):
        path = os.path.join(destination_path, self.name)
        self._create_filesystem_repo(path)
        os.chdir(path)
        self._run("pull", "-r", "tip", self.url)

    def _create_filesystem_repo(self, path):
        self._run("init", path)

    def _create_commit(self, file_name):
        os.chdir(self.path)
        self._run("add", file_name)
        self._run("commit", file_name, "-m", '"Add {}"'.format(file_name))


class Benchmark(object):
    REPO_CLASSES = {"git": GitRepository, "hg": HgRepository}
    REPO_NAME = "{}_performance_{:03d}"

    def __init__(self, config):
        self.api = RCApi(api_key=config.api_key, rc_endpoint=config.host)
        self.source_path = tempfile.mkdtemp(suffix="vcsperformance")

        self.config = config
        self.git_repos = []
        self.hg_repos = []

        self._set_log_level()

    def start(self):
        self._create_repos()
        repos = {"git": self.git_repos, "hg": self.hg_repos}

        clone_destination_path = tempfile.mkdtemp(suffix="clone")
        pull_destination_path = tempfile.mkdtemp(suffix="pull")
        operations = [("push",), ("clone", clone_destination_path), ("pull", pull_destination_path)]

        for operation in operations:
            for type_ in repos:
                times = self._measure(repos[type_], *operation)
                print("Mean[of {}] {:5s} {:5s} time: {:.3f} sec.".format(len(times), type_, operation[0], mean(times)))

    def cleanup(self):
        log.info("Cleaning up...")
        for repo in chain(self.git_repos, self.hg_repos):
            repo.delete()

    def _measure(self, repos, operation, *args):
        times = []
        for repo in repos:
            method = getattr(repo, operation)
            times.append(method(*args))
        return times

    def _create_repos(self):
        log.info("Creating repositories...")
        for i in range(self.config.repositories):
            self.git_repos.append(self._create_repo("git", i))
            self.hg_repos.append(self._create_repo("hg", i))

    def _create_repo(self, type_, id_):
        RepoClass = self.REPO_CLASSES[type_]
        repo = RepoClass(self.REPO_NAME.format(type_, id_), self.source_path, self.api)
        repo.create()
        repo.create_commits(self.config.commits, self.config.file_size)
        return repo

    def _set_log_level(self):
        try:
            log_level = getattr(logging, config.log_level.upper())
        except:
            log_level = logging.ERROR
        handler = logging.StreamHandler()
        log.addHandler(handler)
        log.setLevel(log_level)


if __name__ == "__main__":
    config = Config()
    benchmark = Benchmark(config)
    try:
        benchmark.start()
    finally:
        benchmark.cleanup()