Improving remove_unit_duplicate_nevra performance

remove-duplicate-test

#!/usr/bin/env python
"""Benchmark script showing difference between current remove_unit_duplicate_nevra
and implementation with adjusted queries.

Logs for running this over real data:

2019-03-11 03:52:15,745 Testing queries for repo: rhel-7-server-release-e2e-test-1-rpms__x86_64
2019-03-11 03:52:15,752 Start improved for rhel-7-server-release-e2e-test-1-rpms__x86_64 (1584 units)
2019-03-11 03:52:15,755 End improved for rhel-7-server-release-e2e-test-1-rpms__x86_64
2019-03-11 03:52:15,755 Start original for rhel-7-server-release-e2e-test-1-rpms__x86_64 (1584 units)
2019-03-11 03:52:16,278 End original for rhel-7-server-release-e2e-test-1-rpms__x86_64
(...)
2019-03-11 03:52:16,353 Testing queries for repo: rhel-7-server-rpms__7Server__x86_64
2019-03-11 03:52:16,363 Start improved for rhel-7-server-rpms__7Server__x86_64 (23195 units)
2019-03-11 03:52:16,365 End improved for rhel-7-server-rpms__7Server__x86_64
2019-03-11 03:52:16,365 Start original for rhel-7-server-rpms__7Server__x86_64 (23195 units)
2019-03-11 03:52:24,601 End original for rhel-7-server-rpms__7Server__x86_64

=> cuts runtime by ~8 seconds for a realistic large repo
"""

import argparse
import logging
import operator

from mongoengine import Q

from pulp.server.controllers import repository as repo_controller
from pulp.server.db import model, connection
from pulp.plugins.loader import api as plugin_api

from pulp_rpm.plugins.db.models import RPM
from pulp_rpm.plugins.importers.yum import purge


LOG = logging.getLogger()


def remove_unit_duplicate_nevra_improved(unit, repo):
    nevra_filters = unit.unit_key.copy()
    del nevra_filters['checksum']
    del nevra_filters['checksumtype']
    Q_filters = [Q(**{key: value}) for key, value in nevra_filters.iteritems()]
    Q_nevra_filter = reduce(operator.and_, Q_filters)

    _model = plugin_api.get_unit_model_by_id(unit.type_id)
    unit_iterator = _model.objects(q_obj=Q_nevra_filter)

    repo_controller.disassociate_units(repo, unit_iterator)


def pulp_init():
    connection.initialize()
    plugin_api.initialize()


def test_repo(repo_id):
    """Compare performance of remove_unit_duplicate_nevra against adjusted implementation"""

    LOG.info("Testing queries for repo: %s", repo_id)

    repo = model.Repository.objects.get_or_404(repo_id=repo_id)

    # The test unit here is some RPM which doesn't exist; this simulates the typical
    # case where there is no duplicate to be removed.
    unit = RPM(arch='x86_64', epoch='0', name='some-rpm-which-does-not-exist',
               release='1.el7', version='2.0.0')

    units_count = model.RepositoryContentUnit.objects(repo_id=repo_id, unit_type_id='rpm').count()

    LOG.info("Start improved for %s (%d units)", repo_id, units_count)
    remove_unit_duplicate_nevra_improved(unit, repo)
    LOG.info("End improved for %s", repo_id)

    LOG.info("Start original for %s (%d units)", repo_id, units_count)
    purge.remove_unit_duplicate_nevra(unit, repo)
    LOG.info("End original for %s", repo_id)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('repo', nargs='+')

    args = parser.parse_args()

    logging.basicConfig(level=logging.DEBUG, format="%(asctime)s %(message)s")
    pulp_init()

    for repo in args.repo:
        test_repo(repo)


if __name__ == '__main__':
    main()