prober: move to asyncio for probert storage

Signed-off-by: Olivier Gayot <olivier.gayot@canonical.com>

snapcraft.yaml

@@ -215,7 +215,7 @@ parts:
 
     source: https://github.com/canonical/probert.git
     source-type: git
-    source-commit: dacf369e3dedc50018e4b3b86d4d919459da3cc6
+    source-commit: 6c3f5c6b9772ef5eaf95e9be631c7780a280e441
 
     override-build: *pyinstall
 

subiquity/server/controllers/filesystem.py

@@ -30,7 +30,6 @@ from curtin.storage_config import ptable_uuid_to_flag_entry
 import pyudev
 
 from subiquitycore.async_helpers import (
-    run_in_thread,
     schedule_task,
     SingleInstanceTask,
     TaskAlreadyRunningError,
@@ -887,8 +886,7 @@ class FilesystemController(SubiquityController, FilesystemManipulator):
                 probe_types |= {'os'}
             fname = 'probe-data.json'
             key = "ProbeData"
-        storage = await run_in_thread(
-            self.app.prober.get_storage, probe_types)
+        storage = await self.app.prober.get_storage(probe_types)
         # It is possible for the user to submit filesystem config
         # while a probert probe is running. We don't want to overwrite
         # the users config with a blank one if this happens! (See

subiquity/server/controllers/tests/test_filesystem.py

@@ -76,6 +76,7 @@ class TestSubiquityControllerFilesystem(IsolatedAsyncioTestCase):
         self.app.report_start_event = mock.Mock()
         self.app.report_finish_event = mock.Mock()
         self.app.prober = mock.Mock()
+        self.app.prober.get_storage = mock.AsyncMock()
         self.app.block_log_dir = '/inexistent'
         self.app.note_file_for_apport = mock.Mock()
         self.fsc = FilesystemController(app=self.app)
@@ -113,7 +114,7 @@ class TestSubiquityControllerFilesystem(IsolatedAsyncioTestCase):
         self.fsc._configured = False
         self.fsc.locked_probe_data = True
         self.fsc.queued_probe_data = None
-        self.app.prober.get_storage = mock.Mock(return_value={})
+        self.app.prober.get_storage = mock.AsyncMock(return_value={})
         with mock.patch.object(self.fsc.model, 'load_probe_data') as load:
             await self.fsc._probe_once(restricted=True)
         self.assertEqual(self.fsc.queued_probe_data, {})
@@ -125,7 +126,7 @@ class TestSubiquityControllerFilesystem(IsolatedAsyncioTestCase):
         self.fsc._configured = False
         self.fsc.locked_probe_data = False
         self.fsc.queued_probe_data = None
-        self.app.prober.get_storage = mock.Mock(return_value={})
+        self.app.prober.get_storage = mock.AsyncMock(return_value={})
         with mock.patch.object(self.fsc.model, 'load_probe_data') as load:
             await self.fsc._probe_once(restricted=True)
         self.assertIsNone(self.fsc.queued_probe_data, {})
@@ -800,6 +801,7 @@ class TestCoreBootInstallMethods(IsolatedAsyncioTestCase):
         self.app.report_start_event = mock.Mock()
         self.app.report_finish_event = mock.Mock()
         self.app.prober = mock.Mock()
+        self.app.prober.get_storage = mock.AsyncMock()
         self.app.snapdapi = snapdapi.make_api_client(
             AsyncSnapd(get_fake_connection()))
         self.app.dr_cfg = DRConfig()

subiquitycore/prober.py

@@ -13,10 +13,12 @@
 # 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/>.
 
+import asyncio
 import logging
-import time
 import yaml
 
+from subiquitycore.async_helpers import run_in_thread
+
 from probert.network import (
     StoredDataObserver,
     UdevObserver,
@@ -41,13 +43,13 @@ class Prober():
             observer = UdevObserver(receiver)
         return observer, observer.start()
 
-    def get_storage(self, probe_types=None):
+    async def get_storage(self, probe_types=None):
         if self.saved_config is not None:
             flag = 'bpfail-full'
             if probe_types is not None:
                 flag = 'bpfail-restricted'
             if flag in self.debug_flags:
-                time.sleep(2)
+                await asyncio.sleep(2)
                 1/0
             r = self.saved_config['storage'].copy()
             if probe_types is not None and 'defaults' not in probe_types:
@@ -55,5 +57,12 @@ class Prober():
                     if k not in probe_types:
                         r[k] = {}
             return r
+
         from probert.storage import Storage
-        return Storage().probe(probe_types=probe_types)
+
+        # Until probert is completely free of blocking IO, we should continue
+        # running it in a separate thread.
+        def run_probert(probe_types):
+            return asyncio.run(Storage().probe(probe_types=probe_types))
+
+        return await run_in_thread(run_probert, probe_types)

subiquitycore/tests/test_prober.py

@@ -19,9 +19,9 @@ from subiquitycore.prober import Prober
 
 
 class TestProber(SubiTestCase):
-    def test_none_and_defaults_equal(self):
+    async def test_none_and_defaults_equal(self):
         with open('examples/simple.json', 'r') as fp:
             prober = Prober(machine_config=fp, debug_flags=())
-        none_storage = prober.get_storage(probe_types=None)
-        defaults_storage = prober.get_storage(probe_types={'defaults'})
+        none_storage = await prober.get_storage(probe_types=None)
+        defaults_storage = await prober.get_storage(probe_types={'defaults'})
         self.assertEqual(defaults_storage, none_storage)