CQ: add landmines to delete database after unclean shutdown

commit-queue/commit_queue.py

 #!/usr/bin/env python
 # Copyright (c) 2012 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 """Commit queue executable.
 
 Reuse Rietveld and the Chromium Try Server to process and automatically commit
 patches.
 """
 
 import logging
 import logging.handlers
 import optparse
 import os
+import shutil
 import signal
 import sys
+import tempfile
 import time
 
 import find_depot_tools  # pylint: disable=W0611
 import checkout
 import fix_encoding
 import rietveld
 import subprocess2
 
 import async_push
 import cq_alerts
@@ skipping 89 matching lines @@
   logging_rotating_file = logging.handlers.RotatingFileHandler(
       filename=os.path.join(log_directory, 'commit_queue.log'),
       maxBytes= 10*1024*1024,
       backupCount=50)
   logging_rotating_file.setLevel(logging.DEBUG)
   logging_rotating_file.setFormatter(logging.Formatter(
       '%(asctime)s %(levelname)-8s %(module)15s(%(lineno)4d): %(message)s'))
   logging.getLogger().addHandler(logging_rotating_file)
 
 
+class SignalInterrupt(Exception):
+  """Exception that indicates being interrupted by a caught signal."""
+  
+  def __init__(self, signal_set=None, *args, **kwargs):
+    super(SignalInterrupt, self).__init__(*args, **kwargs)
+    self.signal_set = signal_set
+
+
+def SaveDatabaseCopyForDebugging(db_path):
+  """Saves database file for debugging. Returns name of the saved file."""
+  with tempfile.NamedTemporaryFile(
+      dir=os.path.dirname(db_path),
+      prefix='db.debug.',
+      suffix='.json',
+      delete=False) as tmp_file:
+    with open(db_path) as db_file:
+      shutil.copyfileobj(db_file, tmp_file)
+    return tmp_file.name
+
+
 def main():
   parser = optparse.OptionParser(
       description=sys.modules['__main__'].__doc__)
   project_choices = projects.supported_projects()
   parser.add_option('-v', '--verbose', action='store_true')
   parser.add_option(
       '--no-dry-run',
       action='store_false',
       dest='dry_run',
       default=True,
@@ skipping 92 matching lines @@
         # Disable the checkout.
         print 'Using no checkout'
         pc.context.checkout = FakeCheckout()
       else:
         print 'Using read-only checkout'
         pc.context.checkout = checkout.ReadOnlyCheckout(pc.context.checkout)
       # Save pushed events on disk.
       print 'Using read-only chromium-status interface'
       pc.context.status = async_push.AsyncPushStore()
 
+    landmine_path = os.path.join(work_dir,
+                                 pc.context.checkout.project_name + '.landmine')
     db_path = os.path.join(work_dir, pc.context.checkout.project_name + '.json')
     if os.path.isfile(db_path):
-      try:
-        pc.load(db_path)
-      except ValueError:
+      if os.path.isfile(landmine_path):
+        debugging_path = SaveDatabaseCopyForDebugging(db_path)
         os.remove(db_path)
+        logging.warning(('Deleting database because previous shutdown '
+                        'was unclean. The copy of the database is saved '
+                        'as %s.') % debugging_path)
+      else:
+        try:
+          pc.load(db_path)
+        except ValueError as e:
+          debugging_path = SaveDatabaseCopyForDebugging(db_path)
+          os.remove(db_path)
+          logging.warning(('Failed to parse database (%r), deleting it. '
+                          'The copy of the database is saved as %s.') %
+                          (e, debugging_path))
+          raise e
+
+    # Create a file to indicate unclean shutdown.
+    with open(landmine_path, 'w'):
+      pass
 
     sig_handler.installHandlers(
         signal.SIGINT,
         signal.SIGHUP
     )
 
     # Sync every 5 minutes.
     SYNC_DELAY = 5*60
     try:
       if options.query_only:
         pc.look_for_new_pending_commit()
         pc.update_status()
         print(str(pc.queue))
+        os.remove(landmine_path)
         return 0
 
       now = time.time()
       next_loop = now + options.poll_interval
       # First sync is on second loop.
       next_sync = now + options.poll_interval * 2
       while True:
         # In theory, we would gain in performance to parallelize these tasks. In
         # practice I'm not sure it matters.
         pc.look_for_new_pending_commit()
         pc.process_new_pending_commit()
         pc.update_status()
         pc.scan_results()
         if sig_handler.getTriggeredSignals():
-          raise KeyboardInterrupt()
+          raise SignalInterrupt(signal_set=sig_handler.getTriggeredSignals())
         # Save the db at each loop. The db can easily be in the 1mb range so
         # it's slowing down the CQ a tad but it in the 100ms range even for that
         # size.
         pc.save(db_path)
 
         # More than a second to wait and due to sync.
         now = time.time()
         if (next_loop - now) >= 1 and (next_sync - now) <= 0:
           if sys.stdout.isatty():
             sys.stdout.write('Syncing while waiting                \r')
             sys.stdout.flush()
           try:
             pc.context.checkout.prepare(None)
           except subprocess2.CalledProcessError as e:
             # Don't crash, most of the time it's the svn server that is dead.
             # How fun. Send a stack trace to annoy the maintainer.
             errors.send_stack(e)
           next_sync = time.time() + SYNC_DELAY
 
         now = time.time()
         next_loop = max(now, next_loop)
         while True:
           # Abort if any signals are set
           if sig_handler.getTriggeredSignals():
-            raise KeyboardInterrupt()
+            raise SignalInterrupt(signal_set=sig_handler.getTriggeredSignals())
           delay = next_loop - now
           if delay <= 0:
             break
           if sys.stdout.isatty():
             sys.stdout.write('Sleeping for %1.1f seconds          \r' % delay)
             sys.stdout.flush()
           time.sleep(min(delay, 0.1))
           now = time.time()
         if sys.stdout.isatty():
           sys.stdout.write('Running (please do not interrupt)   \r')
           sys.stdout.flush()
         next_loop = time.time() + options.poll_interval
     except:  # Catch all fatal exit conditions.
       logging.exception('CQ loop terminating')
       raise
     finally:
       logging.warning('Saving db...')
       pc.save(db_path)
       pc.close()
       logging.warning('db save successful.')
-  except KeyboardInterrupt as e:
-    print 'Bye bye'
+  except SignalInterrupt:
+    # This is considered a clean shutdown: we only throw this exception
+    # from selected places in the code where the database should be
+    # in a known and consistent state.
+    os.remove(landmine_path)
+
+    print 'Bye bye (SignalInterrupt)'
+    # 23 is an arbitrary value to signal loop.sh that it must stop looping.
+    return 23
+  except KeyboardInterrupt:
+    # This is actually an unclean shutdown. Do not remove the landmine file.
+    # One example of this is user hitting ctrl-c twice at an arbitrary point
+    # inside the CQ loop. There are no guarantees about consistent state
+    # of the database then.
+
+    print 'Bye bye (KeyboardInterrupt - this is considered unclean shutdown)'
     # 23 is an arbitrary value to signal loop.sh that it must stop looping.
     return 23
   except errors.ConfigurationError as e:
     parser.error(str(e))
     return 1
+
+  # CQ generally doesn't exit by itself, but if we ever get here, it looks
+  # like a clean shutdown so remove the landmine file.
+  # TODO(phajdan.jr): Do we ever get here?
+  os.remove(landmine_path)
   return 0
 
 
 if __name__ == '__main__':
   fix_encoding.fix_encoding()
   sys.exit(main())