From dec00fe2f4e987573323126b49f26bf83a4a46c8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lubom=C3=ADr=20Sedl=C3=A1=C5=99?= <lsedlar@redhat.com>
Date: Tue, 3 Oct 2017 15:09:53 +0200
Subject: [PATCH] createrepo: Allow customizing number of threads
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The default is now to use one thread per CPU.

Signed-off-by: Lubomír Sedlář <lsedlar@redhat.com>
---
 doc/configuration.rst         |  4 ++++
 pungi/checks.py               | 12 ++++++++++++
 pungi/phases/createrepo.py    |  2 +-
 tests/test_createrepophase.py | 12 ++++++++----
 4 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/doc/configuration.rst b/doc/configuration.rst
index 55dccc40..c3a7987a 100644
--- a/doc/configuration.rst
+++ b/doc/configuration.rst
@@ -399,6 +399,10 @@ Options
     (*bool*) -- whether to pass ``--xz`` to the createrepo command. This will
     cause the SQLite databases to be compressed with xz.
 
+**createrepo_num_threads**
+    (*int*) -- how many concurrent ``createrepo`` process to run. The default
+    is to use one thread per CPU available on the machine.
+
 **product_id** = None
     (:ref:`scm_dict <scm_support>`) -- If specified, it should point to a
     directory with certificates ``<variant_uid>-<arch>-*.pem``. This
diff --git a/pungi/checks.py b/pungi/checks.py
index 3796c01c..8a805ef8 100644
--- a/pungi/checks.py
+++ b/pungi/checks.py
@@ -45,6 +45,7 @@ import jsonschema
 import six
 from kobo.shortcuts import force_list
 from productmd.composeinfo import COMPOSE_TYPES
+import multiprocessing
 
 from . import util
 
@@ -605,6 +606,10 @@ def make_schema():
                 "type": "boolean",
                 "default": False,
             },
+            "createrepo_num_threads": {
+                "type": "number",
+                "default": get_num_cpus(),
+            },
             "repoclosure_strictness": _variant_arch_mapping({
                 "type": "string",
                 "default": "lenient",
@@ -1117,6 +1122,13 @@ def _one_or_list(value):
     }
 
 
+def get_num_cpus():
+    try:
+        return multiprocessing.cpu_count()
+    except NotImplementedError:
+        return 3
+
+
 # This is a mapping of configuration option dependencies and conflicts.
 #
 # The key in this mapping is the trigger for the check. When the option is
diff --git a/pungi/phases/createrepo.py b/pungi/phases/createrepo.py
index 1324f9ca..10a12539 100644
--- a/pungi/phases/createrepo.py
+++ b/pungi/phases/createrepo.py
@@ -63,7 +63,7 @@ class CreaterepoPhase(PhaseBase):
 
     def run(self):
         get_productids_from_scm(self.compose)
-        for i in range(3):
+        for i in range(self.compose.conf['createrepo_num_threads']):
             self.pool.add(CreaterepoThread(self.pool))
 
         for variant in self.compose.get_variants():
diff --git a/tests/test_createrepophase.py b/tests/test_createrepophase.py
index 0eba0e4f..0c47ba02 100644
--- a/tests/test_createrepophase.py
+++ b/tests/test_createrepophase.py
@@ -59,8 +59,10 @@ class TestCreaterepoPhase(PungiTestCase):
 
         self.assertIn('deltas', str(ctx.exception))
 
+    @mock.patch('pungi.checks.get_num_cpus')
     @mock.patch('pungi.phases.createrepo.ThreadPool')
-    def test_starts_jobs(self, ThreadPoolCls):
+    def test_starts_jobs(self, ThreadPoolCls, get_num_cpus):
+        get_num_cpus.return_value = 5
         compose = DummyCompose(self.topdir, {})
 
         pool = ThreadPoolCls.return_value
@@ -69,7 +71,7 @@ class TestCreaterepoPhase(PungiTestCase):
         phase.run()
         self.maxDiff = None
 
-        self.assertEqual(len(pool.add.mock_calls), 3)
+        self.assertEqual(len(pool.add.mock_calls), 5)
         self.assertItemsEqual(
             pool.queue_put.mock_calls,
             [mock.call((compose, 'x86_64', compose.variants['Server'], 'rpm')),
@@ -86,8 +88,10 @@ class TestCreaterepoPhase(PungiTestCase):
              mock.call((compose, 'amd64', compose.variants['Client'], 'debuginfo')),
              mock.call((compose, None, compose.variants['Client'], 'srpm'))])
 
+    @mock.patch('pungi.checks.get_num_cpus')
     @mock.patch('pungi.phases.createrepo.ThreadPool')
-    def test_skips_empty_variants(self, ThreadPoolCls):
+    def test_skips_empty_variants(self, ThreadPoolCls, get_num_cpus):
+        get_num_cpus.return_value = 5
         compose = DummyCompose(self.topdir, {})
         compose.variants['Client'].is_empty = True
 
@@ -97,7 +101,7 @@ class TestCreaterepoPhase(PungiTestCase):
         phase.run()
         self.maxDiff = None
 
-        self.assertEqual(len(pool.add.mock_calls), 3)
+        self.assertEqual(len(pool.add.mock_calls), 5)
         self.assertItemsEqual(
             pool.queue_put.mock_calls,
             [mock.call((compose, 'x86_64', compose.variants['Server'], 'rpm')),