Refactor the YumFinder (#38)

* Refactor the YumFinder

Implement our own XML fetch/parse routines for consuming repomd
data, and create a simple lookup dict from the results rather than
trying to shoehorn in use of the `repomd.findall` library call.

This yields a 30%-50% speedup in the functional tests.

* Drop yum.lookup_in_repomd function

Stepping through the refactored workflow showed that the YumFinder was
deep-copying a dict into a subprocess, looking up a value, and
returning the result.  Surely, there must be a better way.

Co-Authored-By: Nicolas Simonds <nisimond@cisco.com>

Author: 0xDEC0DE

soufi/finders/yum.py

@@ -2,15 +2,18 @@
 # All rights reserved.
 
 import abc
-import lzma
+import gzip
+import pathlib
 import pickle  # nosec
 import sys
 import textwrap
 import warnings
 from multiprocessing import Process, Queue
 from types import SimpleNamespace
 
+import defusedxml.lxml
 import repomd
+import requests
 from dogpile.cache.backends.null import NullBackend
 
 from soufi import exceptions, finder
@@ -116,12 +119,12 @@ def _walk_source_repos(self, name, version=None):
             version = self.version
         locations = set()
         for repo_url in self.generate_source_repos():
-            baseurl, repo_xml = self._cache.get_or_create(
+            baseurl, repo = self._cache.get_or_create(
                 f"repo-{repo_url}",
                 do_task,
                 creator_args=([get_repomd, repo_url], {}),
             )
-            for package in do_task(lookup_in_repomd, baseurl, repo_xml, name):
+            for package in repo.get(name, []):
                 # If the package version in the repomd is our version,
                 # it's easy.  Note that we want to match epoch-full and
                 # epoch-less version formats.
@@ -144,12 +147,12 @@ def _walk_source_repos(self, name, version=None):
     def _walk_binary_repos(self, name):
         packages = set()
         for repo_url in self.generate_binary_repos():
-            baseurl, repo_xml = self._cache.get_or_create(
+            _, repo = self._cache.get_or_create(
                 f"repo-{repo_url}",
                 do_task,
                 creator_args=([get_repomd, repo_url], {}),
             )
-            for package in do_task(lookup_in_repomd, baseurl, repo_xml, name):
+            for package in repo.get(name, []):
                 # If we have a binary package matching our version, but
                 # with a different name than the corresponding source
                 # package, return the NVR fields
@@ -263,18 +266,50 @@ def do_task(target, *args):
     return response
 
 
-# NOTE(nic): To allow for serializing repomd objects, We will instead
-#  re-serialize the object into its original XML, then re-create and use a
-#  new Repo object from the cached XML on every cache hit.  This makes cache
-#  hits relatively expensive, but they're still orders of magnitude faster
-#  than the alternative.  The XML is LZMA-compressed here to keep cache
-#  storage utilization low, and in the case of Redis caching, to reduce the
-#  amount of cache traffic sent over the wire.
+# NOTE(nic): stolen almost verbatim from repomd.load, except this one:
+#  - has timeouts
+#  - uses requests instead of urllib to do the heavy lifting
+#  - uses `lxml.parse` and file objects instead of `lxml.fromstring`
+#  - returns a plain dict instead of a Repo object
+def load_repomd(url):
+    timeout = YumFinder.timeout
+    baseurl = requests.utils.parse_url(url)
+    path = pathlib.PurePosixPath(baseurl.path)
+
+    # first we must get the repomd.xml file
+    repomd_path = path / 'repodata' / 'repomd.xml'
+    repomd_url = baseurl._replace(path=str(repomd_path))
+
+    with requests.get(repomd_url, stream=True, timeout=timeout) as r:
+        r.raw.decode_content = True
+        repomd_xml = defusedxml.lxml.parse(r.raw)
+
+    # determine the location of *primary.xml.gz from the repomd.xml
+    primary_element = repomd_xml.find(
+        'repo:data[@type="primary"]/repo:location', namespaces=repomd._ns
+    )
+    primary_path = path / primary_element.get('href')
+    primary_url = baseurl._replace(path=str(primary_path))
+
+    # download and consume *-primary.xml into a dict object for fast lookups.
+    # We will use repomd.Package objects rather than reimplement them,
+    # but only as an intermediate step (see `serialize_package`, below)
+    repo = {}
+    with requests.get(primary_url, stream=True, timeout=timeout) as r:
+        r.raw.decode_content = True
+        with gzip.GzipFile(fileobj=r.raw) as uncompressed:
+            for element in defusedxml.lxml.parse(uncompressed).getroot():
+                package = repomd.Package(element)
+                repo.setdefault(package.name, [])
+                repo[package.name].append(serialize_package(package))
+    return repo
+
+
 def get_repomd(queue, url):
     if not url.endswith('/'):
         url += '/'
     try:
-        repo = repomd.load(url)
+        repo = load_repomd(url)
     except Exception as e:
         # Try and send exceptions back to the caller, just like anything
         # else.  It's up to the receiver to inspect and re-raise.  If the
@@ -289,17 +324,18 @@ def get_repomd(queue, url):
 
         queue.put((e,), timeout=YumFinder.timeout)
         return
-    payload = repomd.defusedxml.lxml.tostring(repo._metadata)
-    queue.put((str(repo.baseurl), lzma.compress(payload)))
+    queue.put((url, repo))
 
 
 # NOTE(nic): repomd.Package object properties do XPath lookups into the
 #  ElementTree and other similar tricks on the backend, which makes them
 #  unsuitable for the IPC-based workflow we're trying to use, so we'll
-#  convert them into simple objects with identical names so that they can be
-#  easily pickled and passed around.  The upside is that we only need to
-#  carry over what we need to do package lookups.  The downside is that we only
-#  get what we've carried over.
+#  convert them into simple objects with identical field names so that they
+#  can be easily pickled and passed around.  The upside is that we only need
+#  to carry over what we need to do package lookups, reducing the space
+#  requirements by at least two orders of magnitude.  The downside is that we
+#  only get what we've carried over, but in practice that's the opposite of
+#  a problem.
 def serialize_package(package):
     return SimpleNamespace(
         name=str(package.name),
@@ -309,12 +345,3 @@ def serialize_package(package):
         location=str(package.location),
         sourcerpm=str(package.sourcerpm),
     )
-
-
-def lookup_in_repomd(queue, baseurl, repomd_xml, name):
-    if None in (baseurl, repomd_xml):
-        queue.put([], timeout=YumFinder.timeout)
-        return
-    payload = lzma.decompress(repomd_xml)
-    repo = repomd.Repo(baseurl, repomd.defusedxml.lxml.fromstring(payload))
-    queue.put([serialize_package(p) for p in repo.findall(name)])

soufi/tests/finders/test_yum_finder.py

@@ -1,9 +1,7 @@
 # Copyright (c) 2021 Cisco Systems, Inc. and its affiliates
 # All rights reserved.
 import io
-import lzma
 import string
-import urllib
 from itertools import repeat
 from unittest import mock
 
@@ -204,7 +202,7 @@ def test__walk_source_repos(self):
         finder = self.make_finder(source_repos=src, binary_repos=bin)
         package = self.FakePackage(vr=finder.version)
         do_task = self.patch(yum, 'do_task')
-        do_task.side_effect = ([baseurl, None], [package])
+        do_task.return_value = (baseurl, {finder.name: [package]})
         self.patch(finder, 'test_url').return_value = False
         url = finder._walk_source_repos(finder.name)
         self.assertEqual(baseurl + package.location, url)
@@ -216,7 +214,7 @@ def test__walk_source_repos_different_version_hit(self):
         finder = self.make_finder(source_repos=src, binary_repos=bin)
         package = self.FakePackage()
         do_task = self.patch(yum, 'do_task')
-        do_task.side_effect = ([baseurl, None], [package])
+        do_task.return_value = (baseurl, {finder.name: [package]})
         self.patch(finder, 'test_url').return_value = True
         url = finder._walk_source_repos(finder.name)
         self.assertEqual(baseurl + package.location, url)
@@ -228,7 +226,7 @@ def test__walk_source_repos_different_version_miss(self):
         finder = self.make_finder(source_repos=src, binary_repos=bin)
         package = self.FakePackage()
         do_task = self.patch(yum, 'do_task')
-        do_task.side_effect = ([baseurl, None], [package])
+        do_task.return_value = (baseurl, {finder.name: [package]})
         self.patch(finder, 'test_url').return_value = False
         url = finder._walk_source_repos(finder.name)
         self.assertIsNone(url)
@@ -243,7 +241,7 @@ def test__walk_binary_repos(self):
         srcrpm = self.make_package(n=n, v=v, r=r)
         package = self.FakePackage(vr=finder.version, sourcerpm=srcrpm)
         do_task = self.patch(yum, 'do_task')
-        do_task.side_effect = ([None, None], [package])
+        do_task.return_value = (None, {finder.name: [package]})
         name, version = finder._walk_binary_repos(finder.name)
         self.expectThat(name, Equals(n))
         self.expectThat(version, Equals(f"{v}-{r}"))
@@ -254,7 +252,7 @@ def test__walk_binary_repos_no_sourcerpm(self):
         finder = self.make_finder(source_repos=src, binary_repos=bin)
         package = self.FakePackage(vr=finder.version, sourcerpm='')
         do_task = self.patch(yum, 'do_task')
-        do_task.side_effect = ([None, None], [package])
+        do_task.return_value = (None, {finder.name: [package]})
         name, version = finder._walk_binary_repos(finder.name)
         self.assertIsNone(name)
         self.assertIsNone(version)
@@ -266,7 +264,7 @@ def test__walk_binary_repos_different_name_multiple_versions(self):
         package1 = self.FakePackage()
         package2 = self.FakePackage()
         do_task = self.patch(yum, 'do_task')
-        do_task.side_effect = ([None, None], [package1, package2])
+        do_task.return_value = (None, {finder.name: [package1, package2]})
         name, version = finder._walk_binary_repos(finder.name)
         self.assertIsNone(name)
         self.assertIsNone(version)
@@ -281,7 +279,7 @@ def test__walk_binary_repos_different_name_different_version(self):
         srcrpm = self.make_package(n=n, v=v, r=r)
         package = self.FakePackage(sourcerpm=srcrpm)
         do_task = self.patch(yum, 'do_task')
-        do_task.side_effect = ([None, None], [package])
+        do_task.return_value = (None, {finder.name: [package]})
         name, version = finder._walk_binary_repos(finder.name)
         self.expectThat(name, Equals(n))
         self.expectThat(version, Equals(f"{v}-{r}"))
@@ -355,97 +353,48 @@ def setUp(self):
     def test_get_repomd(self):
         # Mock up a successful repomd fetch
         url = self.factory.make_url()
-        repo = self.patch(repomd, 'load')
-        lxml = self.patch(repomd.defusedxml.lxml, 'tostring')
-        lxml.return_value = b"<xml>Test</xml>"
-        compress = self.patch(lzma, 'compress')
+        self.patch(requests, 'get')
+        lxml = self.patch(repomd.defusedxml.lxml, 'parse')
+        package = mock.MagicMock()
+        lxml.return_value.getroot.return_value = [package]
 
-        # Ensure that get_repomd re-serializes the object XML and returns a
-        # compressed payload
         yum.get_repomd(self.queue, url)
-        compress.assert_called_once_with(lxml.return_value)
+        # Mocking out all the various and sundry Package properties would be
+        # tedious, and we don't really care about anything past the package
+        # name anyhow
         self.queue.put.assert_called_once_with(
-            (str(repo.return_value.baseurl), compress.return_value)
+            (url + "/", {package.findtext.return_value: [mock.ANY]})
         )
 
     def test_get_repomd_http_error(self):
         # Mock up a failure to fetch the repomd
         url = self.factory.make_url()
-        load = self.patch(repomd, 'load')
-        load.side_effect = urllib.error.HTTPError(None, None, None, None, None)
-        lxml = self.patch(repomd.defusedxml.lxml, 'tostring')
-        compress = self.patch(lzma, 'compress')
+        load = self.patch(requests, 'get')
+        load.side_effect = requests.exceptions.HTTPError()
+        lxml = self.patch(repomd.defusedxml.lxml, 'parse')
 
         # Ensure that get_repomd won't fill the cache with garbage
         yum.get_repomd(self.queue, url)
         lxml.assert_not_called()
-        compress.assert_not_called()
         self.queue.put.assert_called_once_with((load.side_effect,), timeout=30)
 
     def test_get_repomd_unserializable_http_error(self):
-        # Ibid, but initializing the HTTPError with a live file pointer will
+        # Ibid, but initializing the exception with a live file pointer will
         # make it refuse to serialize
         url = self.factory.make_url()
         fp = io.BufferedReader(io.StringIO())
-        load = self.patch(repomd, 'load')
-        load.side_effect = urllib.error.HTTPError(None, None, None, None, fp)
-        lxml = self.patch(repomd.defusedxml.lxml, 'tostring')
-        compress = self.patch(lzma, 'compress')
+        load = self.patch(requests, 'get')
+        load.side_effect = requests.exceptions.RequestException(fp)
+        lxml = self.patch(repomd.defusedxml.lxml, 'parse')
 
         # Ensure that we get a re-raised plain Exception
         yum.get_repomd(self.queue, url)
         lxml.assert_not_called()
-        compress.assert_not_called()
         self.queue.put.assert_called_once_with((mock.ANY,), timeout=30)
         self.assertIn(
             're-raising as plain Exception', str(self.queue.put.call_args)
         )
 
-    def test_lookup_in_repomd(self):
-        # Mock up a successful package lookup
-        baseurl = self.factory.make_url()
-        repomd_xml = self.factory.make_string()
-        name = self.factory.make_string()
-        decompress = self.patch(lzma, 'decompress')
-        lxml = self.patch(repomd.defusedxml.lxml, 'fromstring')
-        repo = self.patch(repomd, 'Repo')
-        package = mock.MagicMock()
-        repo.return_value.findall.return_value = [package]
-
-        # lookup_in_repomd should decompress/rehydrate a repomd object
-        yum.lookup_in_repomd(self.queue, baseurl, repomd_xml, name)
-        decompress.assert_called_once_with(repomd_xml)
-        lxml.assert_called_once_with(decompress.return_value)
-        repo.assert_called_once_with(baseurl, lxml.return_value)
-        # The response should be converted into a "simplified package"
-        self.queue.put.assert_called_once_with(
-            [yum.serialize_package(package)]
-        )
-
-    def test_lookup_in_repomd_no_baseurl(self):
-        # Ensure that calling lookup_in_repomd without valid args does no work
-        repomd_xml = self.factory.make_string()
-        name = self.factory.make_string()
-        decompress = self.patch(lzma, 'decompress')
-        repo = self.patch(repomd, 'Repo')
-
-        yum.lookup_in_repomd(self.queue, None, repomd_xml, name)
-        decompress.assert_not_called()
-        repo.assert_not_called()
-        self.queue.put.assert_called_once_with([], timeout=30)
-
-    def test_lookup_in_repomd_no_repoxml(self):
-        # Ibid.
-        baseurl = self.factory.make_url()
-        name = self.factory.make_string()
-        decompress = self.patch(lzma, 'decompress')
-        repo = self.patch(repomd, 'Repo')
-
-        yum.lookup_in_repomd(self.queue, baseurl, None, name)
-        decompress.assert_not_called()
-        repo.assert_not_called()
-        self.queue.put.assert_called_once_with([], timeout=30)
-
     def test_do_task(self):
         # Mock up a process that does not exit upon return
         data = self.factory.make_string('response')