Enhance process pool implementation

[dak.git] / dak / generate_releases.py

diff --git a/dak/generate_releases.py b/dak/generate_releases.py
index 8a371e8a8d3dd5a3ff925d4009950614afbea03d..6c4fed0dcf64c2b84bfb6983cb14556fd858225d 100755 (executable)
--- a/dak/generate_releases.py
+++ b/dak/generate_releases.py
@@ -40,7 +40,6 @@ import bz2
 import apt_pkg
 from tempfile import mkstemp, mkdtemp
 import commands
-from multiprocessing import Pool, TimeoutError
 from sqlalchemy.orm import object_session
 
 from daklib import utils, daklog
@@ -48,10 +47,10 @@ from daklib.regexes import re_gensubrelease, re_includeinrelease
 from daklib.dak_exceptions import *
 from daklib.dbconn import *
 from daklib.config import Config
+from daklib.dakmultiprocessing import DakProcessPool, PROC_STATUS_SUCCESS
 
 ################################################################################
 Logger = None                  #: Our logging object
-results = []                   #: Results of the subprocesses
 
 ################################################################################
 
@@ -74,11 +73,6 @@ SUITE can be a space seperated list, e.g.
 
 ########################################################################
 
-def get_result(arg):
-    global results
-    if arg:
-        results.append(arg)
-
 def sign_release_dir(suite, dirname):
     cnf = Config()
 
@@ -101,7 +95,7 @@ def sign_release_dir(suite, dirname):
 
         # We can only use one key for inline signing so use the first one in
         # the array for consistency
-        firstkey = False
+        firstkey = True
 
         for keyid in suite.signingkeys:
             defkeyid = "--default-key %s" % keyid
@@ -154,7 +148,7 @@ class ReleaseWriter(object):
         suite_suffix = "%s" % (cnf.Find("Dinstall::SuiteSuffix"))
 
         outfile = os.path.join(cnf["Dir::Root"], 'dists', "%s/%s" % (suite.suite_name, suite_suffix), "Release")
-        out = open(outfile, "w")
+        out = open(outfile + ".new", "w")
 
         for key, dbfield in attribs:
             if getattr(suite, dbfield) is not None:
@@ -280,6 +274,7 @@ class ReleaseWriter(object):
                 out.write(" %s %8d %s\n" % (fileinfo[filename][h], fileinfo[filename]['len'], filename))
 
         out.close()
+        os.rename(outfile + '.new', outfile)
 
         sign_release_dir(suite, os.path.dirname(outfile))
 
@@ -289,7 +284,7 @@ class ReleaseWriter(object):
 
 
 def main ():
-    global Logger, results
+    global Logger
 
     cnf = Config()
 
@@ -324,10 +319,8 @@ def main ():
         suites = session.query(Suite).filter(Suite.untouchable == False).all()
 
     broken=[]
-    # For each given suite, run one process
-    results = []
 
-    pool = Pool()
+    pool = DakProcessPool()
 
     for s in suites:
         # Setup a multiprocessing Pool. As many workers as we have CPU cores.
@@ -343,12 +336,10 @@ def main ():
     pool.close()
     pool.join()
 
-    retcode = 0
+    retcode = p.overall_status()
 
-    if len(results) > 0:
-        Logger.log(['Release file generation broken: %s' % (results)])
-        print "Release file generation broken:\n", '\n'.join(results)
-        retcode = 1
+    if retcode > 0:
+        Logger.log(['Release file generation broken: %s' % (p.results)])
 
     Logger.close()
 
@@ -360,13 +351,12 @@ def generate_helper(suite_id):
     '''
     session = DBConn().session()
     suite = Suite.get(suite_id, session)
-    try:
-        rw = ReleaseWriter(suite)
-        rw.generate_release_files()
-    except Exception, e:
-        return str(e)
 
-    return
+    # We allow the process handler to catch and deal with any exceptions
+    rw = ReleaseWriter(suite)
+    rw.generate_release_files()
+
+    return (PROC_STATUS_SUCCESS, 'Release file written for %s' % suite.suite_name)
 
 #######################################################################################