stats_manager: Avoid losing data

In the case of multiple StatsManager, we might get data clobbered.
That's particularly bad on long-running tests.
This CL introduces two methods to make sure output files are unique in a
directory.
(1) on init, the user can supply a StatsManagerID (smid) that will be
prepended to all output files.
E.g. if the smid is 'ec' the summary will be stored at ec_summary.txt
(2) should that fail because multiple StatsManagers are using the same
smid supplied (or none were supplied) a simple file rotation is done
where an integer suffix keeps getting incremented until the filename is
unique.
E.g. if summary.txt already exists, then summary1.txt will be created.

This is not threadsafe at all so the user still have to use some caution
to avoid clobbering their own data.

BRANCH=None
BUG=chromium:760267
TEST=unit tests still all pass

Change-Id: I6cc083259362ee20e0242b94ac7cbb1228a06a7a
Signed-off-by: Ruben Rodriguez Buchillon <coconutruben@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/1140029
Reviewed-by: Mengqi Guo <mqg@chromium.org>

2 files changed, 82 insertions, 5 deletions

diff --git a/extra/usb_power/stats_manager.py b/extra/usb_power/stats_manager.py
index cc663ecce0..2607fa0ccf 100644
--- a/extra/usb_power/stats_manager.py
+++ b/extra/usb_power/stats_manager.py
@@ -55,6 +55,9 @@ class StatsManager(object):
   Attributes:
     _data: dict of list of readings for each domain(key)
     _unit: dict of unit for each domain(key)
+    _smid: id supplied to differentiate data output to other StatsManager
+           instances that potentially save to the same directory
+           if smid all output files will be named |smid|_|fname|
     _order: list of formatting order for domains. Domains not listed are
             displayed in sorted order
     _hide_domains: collection of domains to hide when formatting summary string
@@ -67,10 +70,11 @@ class StatsManager(object):
   """
 
   # pylint: disable=W0102
-  def __init__(self, order=[], hide_domains=[]):
+  def __init__(self, smid='', order=[], hide_domains=[]):
     """Initialize infrastructure for data and their statistics."""
     self._data = collections.defaultdict(list)
     self._unit = collections.defaultdict(str)
+    self._smid = smid
     self._order = order
     self._hide_domains = hide_domains
     self._summary = {}
@@ -167,6 +171,47 @@ class StatsManager(object):
     """Getter for summary."""
     return self._summary
 
+  def _MakeUniqueFName(self, fname):
+    """prepend |_smid| to fname & rotate fname to ensure uniqueness.
+
+    Before saving a file through the StatsManager, make sure that the filename
+    is unique, first by prepending the smid if any and otherwise by appending
+    increasing integer suffixes until the filename is unique.
+
+    If |smid| is defined /path/to/example/file.txt becomes
+    /path/to/example/{smid}_file.txt.
+
+    The rotation works by changing /path/to/example/somename.txt to
+    /path/to/example/somename1.txt if the first one already exists on the
+    system.
+
+    Note: this is not thread-safe. While it makes sense to use StatsManager
+    in a threaded data-collection, the data retrieval should happen in a
+    single threaded environment to ensure files don't get potentially clobbered.
+
+    Args:
+      fname: filename to ensure uniqueness.
+
+    Returns:
+      {smid_}fname{tag}.ext
+      the smid portion gets prepended if |smid| is defined
+      the tag portion gets appended if necessary to ensure unique fname
+    """
+    fdir = os.path.dirname(fname)
+    base, ext = os.path.splitext(os.path.basename(fname))
+    if self._smid:
+      base = '%s_%s' % (self._smid, base)
+    unique_fname = os.path.join(fdir, '%s%s' % (base, ext))
+    tag = 0
+    while os.path.exists(unique_fname):
+      old_fname = unique_fname
+      unique_fname = os.path.join(fdir, '%s%d%s' % (base, tag, ext))
+      self._logger.warn('Attempted to store stats information at %s, but file '
+                        'already exists. Attempting to store at %s now.',
+                        old_fname, unique_fname)
+      tag += 1
+    return unique_fname
+
   def SaveSummary(self, directory, fname='summary.txt', prefix=STATS_PREFIX):
     """Save summary to file.
 
@@ -182,7 +227,7 @@ class StatsManager(object):
 
     if not os.path.exists(directory):
       os.makedirs(directory)
-    fname = os.path.join(directory, fname)
+    fname = self._MakeUniqueFName(os.path.join(directory, fname))
     with open(fname, 'w') as f:
       f.write(summary_str)
     return fname
@@ -204,7 +249,7 @@ class StatsManager(object):
       data[domain] = data_entry
     if not os.path.exists(directory):
       os.makedirs(directory)
-    fname = os.path.join(directory, fname)
+    fname = self._MakeUniqueFName(os.path.join(directory, fname))
     with open(fname, 'w') as f:
       json.dump(data, f)
     return fname
@@ -232,8 +277,7 @@ class StatsManager(object):
     for domain, data in self._data.iteritems():
       if not domain.endswith(self._unit[domain]):
         domain = '%s_%s' % (domain, self._unit[domain])
-      fname = '%s.txt' % domain
-      fname = os.path.join(dirname, fname)
+      fname = self._MakeUniqueFName(os.path.join(dirname, '%s.txt' % domain))
       with open(fname, 'w') as f:
         f.write('\n'.join('%.2f' % sample for sample in data) + '\n')
       fnames.append(fname)
diff --git a/extra/usb_power/stats_manager_unittest.py b/extra/usb_power/stats_manager_unittest.py
index 2fb1080f4d..a669ab8e8f 100644
--- a/extra/usb_power/stats_manager_unittest.py
+++ b/extra/usb_power/stats_manager_unittest.py
@@ -143,6 +143,15 @@ class TestStatsManager(unittest.TestCase):
     # Verify nothing gets appended to domain for filename if no unit exists.
     self.assertIn('B.txt', files)
 
+  def test_SaveRawDataSMID(self):
+    """SaveRawData uses the smid when creating output filename."""
+    identifier = 'ec'
+    self.data = stats_manager.StatsManager(smid=identifier)
+    self._populate_dummy_stats()
+    files = self.data.SaveRawData(self.tempdir)
+    for fname in files:
+      self.assertTrue(os.path.basename(fname).startswith(identifier))
+
   def test_SummaryToStringHideDomains(self):
     """Keys indicated in hide_domains are not printed in the summary."""
     data = stats_manager.StatsManager(hide_domains=['A-domain'])
@@ -168,6 +177,14 @@ class TestStatsManager(unittest.TestCase):
     summary_str = data.SummaryToString()
     self.assertRegexpMatches(summary_str, d_b_a_c_regexp)
 
+  def test_MakeUniqueFName(self):
+    data = stats_manager.StatsManager()
+    testfile = os.path.join(self.tempdir, 'testfile.txt')
+    with open(testfile, 'w') as f:
+      f.write('')
+    expected_fname = os.path.join(self.tempdir, 'testfile0.txt')
+    self.assertEqual(expected_fname, data._MakeUniqueFName(testfile))
+
   def test_SaveSummary(self):
     """SaveSummary properly dumps the summary into a file."""
     self._populate_dummy_stats()
@@ -190,6 +207,14 @@ class TestStatsManager(unittest.TestCase):
           '@@   B_mV      3       2.50    0.82       3.50      1.50\n',
           f.readline())
 
+  def test_SaveSummarySMID(self):
+    """SaveSummary uses the smid when creating output filename."""
+    identifier = 'ec'
+    self.data = stats_manager.StatsManager(smid=identifier)
+    self._populate_dummy_stats()
+    fname = os.path.basename(self.data.SaveSummary(self.tempdir))
+    self.assertTrue(fname.startswith(identifier))
+
   def test_SaveSummaryJSON(self):
     """SaveSummaryJSON saves the added data properly in JSON format."""
     self._populate_dummy_stats()
@@ -208,6 +233,14 @@ class TestStatsManager(unittest.TestCase):
       self.assertAlmostEqual(2.5, summary['B']['mean'])
       self.assertEqual('millivolt', summary['B']['unit'])
 
+  def test_SaveSummaryJSONSMID(self):
+    """SaveSummaryJSON uses the smid when creating output filename."""
+    identifier = 'ec'
+    self.data = stats_manager.StatsManager(smid=identifier)
+    self._populate_dummy_stats()
+    fname = os.path.basename(self.data.SaveSummaryJSON(self.tempdir))
+    self.assertTrue(fname.startswith(identifier))
+
   def test_SaveSummaryJSONNoUnit(self):
     """SaveSummaryJSON marks unknown units properly as N/A."""
     self._populate_dummy_stats_no_unit()