X-Git-Url: https://code.grnet.gr/git/snf-image-creator/blobdiff_plain/663f5f807b6685a8af18b539f0802683575103d5..5a380da9c7d19260ec8c033ae8b28b2f1c647d96:/image_creator/rsync.py

diff --git a/image_creator/rsync.py b/image_creator/rsync.py
index 9ae22b6..833d0be 100644
--- a/image_creator/rsync.py
+++ b/image_creator/rsync.py
@@ -1,3 +1,5 @@
+# -*- coding: utf-8 -*-
+#
 # Copyright 2012 GRNET S.A. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or
@@ -31,6 +33,8 @@
 # interpreted as representing official policies, either expressed
 # or implied, of GRNET S.A.
 
+"""This module provides an interface for the rsync utility"""
+
 import subprocess
 import time
 import signal
@@ -82,7 +86,7 @@ class Rsync:
         self._exclude = []
         self._options = ['-v']
 
-    def run(self, src, dest):
+    def run(self, src, dest, slabel='source', dlabel='destination'):
         """Run the actual command"""
         cmd = []
         cmd.append('rsync')
@@ -90,14 +94,15 @@ class Rsync:
         for i in self._exclude:
             cmd.extend(['--exclude', i])
 
-        self._out.output("Calculating total number of host files ...", False)
+        self._out.output("Calculating total number of %s files ..." % slabel,
+                         False)
 
         # If you don't specify a destination, rsync will list the source files.
         dry_run = subprocess.Popen(cmd + [src], shell=False,
                                    stdout=subprocess.PIPE, bufsize=0)
         try:
             total = 0
-            for line in iter(dry_run.stdout.readline, b''):
+            for _ in iter(dry_run.stdout.readline, b''):
                 total += 1
         finally:
             dry_run.communicate()
@@ -106,14 +111,13 @@ class Rsync:
 
         self._out.success("%d" % total)
 
-        progress = self._out.Progress(total,
-                                      "Copying host files into the image")
+        progress = self._out.Progress(total, "Copying files to %s" % dlabel)
         run = subprocess.Popen(cmd + [src, dest], shell=False,
                                stdout=subprocess.PIPE, bufsize=0)
         try:
             t = time.time()
             i = 0
-            for line in iter(run.stdout.readline, b''):
+            for _ in iter(run.stdout.readline, b''):
                 i += 1
                 current = time.time()
                 if current - t > 0.1:
@@ -123,12 +127,19 @@ class Rsync:
             progress.success('done')
 
         finally:
-            run.poll()
-            if run.returncode is None:
-                run.send_signal(signal.SIGHUP)
+            def handler(signum, frame):
+                run.terminate()
+                time.sleep(1)
+                run.poll()
+                if run.returncode is None:
+                    run.kill()
+                run.wait()
+
+            signal.signal(signal.SIGALRM, handler)
+            signal.alarm(2)
             run.communicate()
+            signal.alarm(0)
             if run.returncode != 0:
                 raise FatalError("rsync failed")
 
-
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :