autotest: Prune empty per-request directories

BUG=chromium:955175
TEST=unittests

Change-Id: Ie799f678321b0cd3c1d9beda8c15005b2702906d
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/autotest/+/1808720
Legacy-Commit-Queue: Commit Bot <commit-bot@chromium.org>
Reviewed-by: Aviv Keshet <akeshet@chromium.org>
Reviewed-by: Xinan Lin <linxinan@chromium.org>
Tested-by: Xinan Lin <linxinan@chromium.org>
diff --git a/site_utils/gs_offloader.py b/site_utils/gs_offloader.py
index fbd4979..ec83a50 100755
--- a/site_utils/gs_offloader.py
+++ b/site_utils/gs_offloader.py
@@ -473,6 +473,23 @@
     """
     return TEST_LIST_COLLECTOR in package
 
+
+def _get_swarming_req_dir(path):
+    """
+    Returns the parent directory of |path|, if |path| is a swarming task result.
+
+    @param path: Full path to the result of a task.
+                      e.g. /tmp/results/swarming-44466815c4bc951/1
+
+    @return string of the parent dir or None if not a swarming task.
+    """
+    m_parent = re.match(
+            '(?P<parent_dir>.*/swarming-[0-9a-fA-F]*0)/[1-9a-fA-F]$', path)
+    if m_parent:
+        return m_parent.group('parent_dir')
+    return None
+
+
 def _parse_cts_job_results_file_path(path):
     """Parse CTS file paths an extract required information from them."""
 
@@ -735,6 +752,10 @@
                     correct_results_folder_permission(dir_entry)
             else:
                 self._prune(dir_entry, job_complete_time)
+                swarming_req_dir = _get_swarming_req_dir(dir_entry)
+                if swarming_req_dir:
+                    self._prune_swarming_req_dir(swarming_req_dir)
+
 
     def _try_offload(self, dir_entry, dest_path,
                  stdout_file, stderr_file):
@@ -825,6 +846,20 @@
             # Try again after the permission issue is fixed.
             shutil.rmtree(dir_entry)
 
+    def _prune_swarming_req_dir(self, swarming_req_dir):
+        """Prune swarming request directory, if it is empty.
+
+        @param swarming_req_dir: Directory entry of a swarming request.
+        """
+        try:
+            logging.debug('Pruning swarming request directory %s',
+                          swarming_req_dir)
+            os.rmdir(swarming_req_dir)
+        except OSError as e:
+            # Do nothing and leave this directory to next attempt to remove.
+            logging.debug('Failed to prune swarming request directory %s',
+                          swarming_req_dir)
+
 
 class _OffloadError(Exception):
     """Google Storage offload failed."""
diff --git a/site_utils/gs_offloader_unittest.py b/site_utils/gs_offloader_unittest.py
index 3ca815e..598fb80 100755
--- a/site_utils/gs_offloader_unittest.py
+++ b/site_utils/gs_offloader_unittest.py
@@ -441,7 +441,7 @@
                       `self._resultsroot`.
 
         """
-        os.mkdir(jobdir)
+        os.makedirs(jobdir)
         return _MockJobDirectory(jobdir)
 
 
@@ -554,6 +554,7 @@
         alarm.start()
         self.addCleanup(alarm.stop)
         self.mox.StubOutWithMock(models.test, 'parse_job_keyval')
+        self.should_remove_sarming_req_dir = False
 
 
     def tearDown(self):
@@ -609,6 +610,11 @@
         self.mox.VerifyAll()
         self.assertEqual(not should_succeed,
                          os.path.isdir(self._job.queue_args[0]))
+        swarming_req_dir = gs_offloader._get_swarming_req_dir(
+                self._job.queue_args[0])
+        if swarming_req_dir:
+            self.assertEqual(not self.should_remove_sarming_req_dir,
+                             os.path.exists(swarming_req_dir))
 
 
     def test_offload_success(self):
@@ -627,6 +633,33 @@
         self._run_offload_dir(False, 0)
 
 
+    def test_offload_swarming_req_dir_remove(self):
+        """Test that `offload_dir()` can prune the empty swarming task dir."""
+        should_remove = os.path.join('results', 'swarming-123abc0')
+        self._job = self.make_job(os.path.join(should_remove, '1'))
+        self._mock_offload_dir_calls(['test', '-d'],
+                                     self._job.queue_args)
+
+        os.path.isfile(mox.IgnoreArg()).AndReturn(True)
+        self.should_remove_sarming_req_dir = True
+        self._mock_create_marker_file()
+        self._run_offload_dir(True, 0)
+
+
+    def test_offload_swarming_req_dir_exist(self):
+        """Test that `offload_dir()` keeps the non-empty swarming task dir."""
+        should_not_remove = os.path.join('results', 'swarming-456edf0')
+        self._job = self.make_job(os.path.join(should_not_remove, '1'))
+        self.make_job(os.path.join(should_not_remove, '2'))
+        self._mock_offload_dir_calls(['test', '-d'],
+                                     self._job.queue_args)
+
+        os.path.isfile(mox.IgnoreArg()).AndReturn(True)
+        self.should_remove_sarming_req_dir = False
+        self._mock_create_marker_file()
+        self._run_offload_dir(True, 0)
+
+
     def test_sanitize_dir(self):
         """Test that folder/file name with invalid character can be corrected.
         """