[Plugin] Allow for manual collections during collect phase

Historically, sos has limited collections to existing files and command
output. While there have been many reasons for this, there have also
been several exceptions made to current plugins as well as ongoing
requests for data that is not currently provided via existing commands.
As sos evolves, it should in turn be more capable to provide diagnostic
data beyond what is strictly available via command outputs.

As such, add a new step to the collection phase that allows plugins to
perform these manual data collections. Plugins may now define their own
`collect()` method to do so, thus moving the existing exceptions out of
`setup()` phase execution.

To aide in writing these collections to the plugin directory, a new
`collection_file()` generator has been added which will handle creating,
managing, and closing the new file so that plugins only need to be
concerned with the content being written to such files.

Plugin contributors should note that these manual collections are
executed at the end of the collection phase - meaning they are more
likely to be skipped or interrupted due to plugin timeouts.

Closes: #2992

Signed-off-by: Jake Hunsaker <jhunsake@redhat.com>

4 files changed, 83 insertions, 17 deletions

diff --git a/sos/report/__init__.py b/sos/report/__init__.py
index 4770ec12..bbf73d21 100644
--- a/sos/report/__init__.py
+++ b/sos/report/__init__.py
@@ -1328,7 +1328,7 @@ class SoSReport(SoSComponent):
         )
         self.ui_progress(status_line)
         try:
-            plug.collect()
+            plug.collect_plugin()
             # certain exceptions can cause either of these lists to no
             # longer contain the plugin, which will result in sos hanging
             # so we can't blindly call remove() on these two.
diff --git a/sos/report/plugins/__init__.py b/sos/report/plugins/__init__.py
index c77627ab..42032d50 100644
--- a/sos/report/plugins/__init__.py
+++ b/sos/report/plugins/__init__.py
@@ -17,6 +17,7 @@ from sos.utilities import (sos_get_command_output, import_module, grep,
                            recursive_dict_values_by_key)
 
 from sos.archive import P_FILE
+import contextlib
 import os
 import glob
 import re
@@ -615,6 +616,7 @@ class Plugin():
         self.manifest.add_list('files', [])
         self.manifest.add_field('strings', {})
         self.manifest.add_field('containers', {})
+        self.manifest.add_field('collections', {})
 
     def set_default_cmd_environment(self, env_vars):
         """
@@ -3076,7 +3078,7 @@ class Plugin():
 
     def _collect_tailed_files(self):
         for _file, _size in self._tail_files_list:
-            self._log_info("collecting tail of '%s' due to size limit" % _file)
+            self._log_info(f"collecting tail of '{_file}' due to size limit")
             file_name = _file
             if file_name[0] == os.sep:
                 file_name = file_name.lstrip(os.sep)
@@ -3105,7 +3107,71 @@ class Plugin():
                 self._log_debug("could not add string '%s': %s"
                                 % (file_name, e))
 
+    def _collect_manual(self):
+        """Kick off manual collections performed by the plugin. These manual
+        collections are anything the plugin collects outside of existing
+        files and/or command output. Anything the plugin manually compiles or
+        constructs for data that is included in the final archive.
+
+        Plugins will need to define these collections by overriding the
+        ``collect()`` method, similar to how plugins define their own
+        ``setup()`` methods.
+        """
+        try:
+            self.collect()
+        except Exception as err:
+            self._log_error(f"Error during plugin collections: {err}")
+
     def collect(self):
+        """If a plugin needs to manually compile some data for a collection,
+        that should be specified here by overriding this method.
+
+        These collections are run last during a plugin's execution, and as such
+        are more likely to be interrupted by timeouts than file or command
+        output collections.
+        """
+        pass
+
+    @contextlib.contextmanager
+    def collection_file(self, fname, subdir=None, tags=[]):
+        """Handles creating and managing files within a plugin's subdirectory
+        within the archive, and is intended to be used to save manually
+        compiled data generated during a plugin's ``_collect_manual()`` step
+        of the collection phase.
+
+        Plugins should call this method using a ``with`` context manager.
+
+        :param fname:       The name of the file within the plugin directory
+        :type fname:        ``str``
+
+        :param subdir:      If needed, specify a subdir to write the file to
+        :type subdir:       ``str``
+
+        :param tags:        Tags to be added to this file in the manifest
+        :type tags:         ``str`` or ``list`` of ``str``s
+        """
+        try:
+            start = time()
+            _pfname = self._make_command_filename(fname, subdir=subdir)
+            self.archive.check_path(_pfname, P_FILE)
+            _name = self.archive.dest_path(_pfname)
+            _file = open(_name, 'w')
+            self._log_debug(f"manual collection file opened: {_name}")
+            yield _file
+            _file.close()
+            end = time()
+            run = end - start
+            self._log_info(f"manual collection '{fname}' finished in {run}")
+            if isinstance(tags, str):
+                tags = [tags]
+            self.manifest.collections[fname.split('.')[0]] = {
+                'filepath': _pfname,
+                'tags': tags
+            }
+        except Exception as err:
+            self._log_info(f"Error with collection file '{fname}': {err}")
+
+    def collect_plugin(self):
         """Collect the data for a plugin."""
         start = time()
         self._collect_copy_specs()
@@ -3113,6 +3179,7 @@ class Plugin():
         self._collect_tailed_files()
         self._collect_cmds()
         self._collect_strings()
+        self._collect_manual()
         fields = (self.name(), time() - start)
         self._log_debug("collected plugin '%s' in %s" % fields)
 
diff --git a/tests/report_tests/plugin_tests/string_collection_tests.py b/tests/report_tests/plugin_tests/collet_manual_tests.py
index 05091a2e..fdcda526 100644
--- a/tests/report_tests/plugin_tests/string_collection_tests.py
+++ b/tests/report_tests/plugin_tests/collet_manual_tests.py
@@ -10,8 +10,8 @@
 from sos_tests import StageOneReportTest
 
 
-class CollectStringTest(StageOneReportTest):
-    """Test to ensure that add_string_as_file() is working for plugins that
+class CollectManualTest(StageOneReportTest):
+    """Test to ensure that collect() is working for plugins that
     directly call it as part of their collections
 
     :avocado: tags=stageone
@@ -30,8 +30,8 @@ class CollectStringTest(StageOneReportTest):
     def test_no_strings_dir(self):
         self.assertFileNotCollected('sos_strings/')
 
-    def test_manifest_strings_correct(self):
+    def test_manifest_collections_correct(self):
         pkgman = self.get_plugin_manifest('unpackaged')
-        self.assertTrue(pkgman['strings']['unpackaged'])
+        self.assertTrue(pkgman['collections']['unpackaged'])
         pyman = self.get_plugin_manifest('python')
-        self.assertTrue(pyman['strings']['digests_json'])
+        self.assertTrue(pyman['collections']['digests'])
diff --git a/tests/unittests/plugin_tests.py b/tests/unittests/plugin_tests.py
index 1d7cc96d..8170a1dd 100644
--- a/tests/unittests/plugin_tests.py
+++ b/tests/unittests/plugin_tests.py
@@ -296,7 +296,7 @@ class PluginTests(unittest.TestCase):
         })
         p.archive = MockArchive()
         p.setup()
-        p.collect()
+        p.collect_plugin()
         self.assertEquals(p.archive.m, {})
 
     def test_postproc_default_on(self):
@@ -358,10 +358,10 @@ class AddCopySpecTests(unittest.TestCase):
         self.mp.sysroot = '/'
         fn = create_file(2)  # create 2MB file, consider a context manager
         self.mp.add_copy_spec(fn, 1)
-        content, fname, _tags = self.mp.copy_strings[0]
-        self.assertTrue("tailed" in fname)
+        fname, _size = self.mp._tail_files_list[0]
+        self.assertTrue(fname == fn)
         self.assertTrue("tmp" in fname)
-        self.assertEquals(1024 * 1024, len(content))
+        self.assertEquals(1024 * 1024, _size)
         os.unlink(fn)
 
     def test_bad_filename(self):
@@ -388,10 +388,9 @@ class AddCopySpecTests(unittest.TestCase):
         create_file(2, dir=tmpdir)
         create_file(2, dir=tmpdir)
         self.mp.add_copy_spec(tmpdir + "/*", 1)
-        self.assertEquals(len(self.mp.copy_strings), 1)
-        content, fname, _tags = self.mp.copy_strings[0]
-        self.assertTrue("tailed" in fname)
-        self.assertEquals(1024 * 1024, len(content))
+        self.assertEquals(len(self.mp._tail_files_list), 1)
+        fname, _size = self.mp._tail_files_list[0]
+        self.assertEquals(1024 * 1024, _size)
         shutil.rmtree(tmpdir)
 
     def test_multiple_files_no_limit(self):
@@ -450,7 +449,7 @@ class RegexSubTests(unittest.TestCase):
     def test_no_replacements(self):
         self.mp.sysroot = '/'
         self.mp.add_copy_spec(j("tail_test.txt"))
-        self.mp.collect()
+        self.mp.collect_plugin()
         replacements = self.mp.do_file_sub(
             j("tail_test.txt"), r"wont_match", "foobar")
         self.assertEquals(0, replacements)
@@ -459,7 +458,7 @@ class RegexSubTests(unittest.TestCase):
         # test uses absolute paths
         self.mp.sysroot = '/'
         self.mp.add_copy_spec(j("tail_test.txt"))
-        self.mp.collect()
+        self.mp.collect_plugin()
         replacements = self.mp.do_file_sub(
             j("tail_test.txt"), r"(tail)", "foobar")
         self.assertEquals(1, replacements)