[CRIU] [PATCH 2/3] p.haul: add file with generic functionality for p.haul modules

[Nikita Spiridonov]

Add file with generic functionality for p.haul modules (e.g. vz or
lxc) and move final restore logic to this file. This modification
is needed to perform module specific restore.

Signed-off-by: Nikita Spiridonov <nspiridonov at odin.com>
---
 phaul/p_haul_lxc.py     |    4 ++++
 phaul/p_haul_module.py  |   41 +++++++++++++++++++++++++++++++++++++++++
 phaul/p_haul_pid.py     |    5 ++++-
 phaul/p_haul_service.py |   36 ++----------------------------------
 phaul/p_haul_type.py    |    2 +-
 phaul/p_haul_vz.py      |    4 ++++
 6 files changed, 56 insertions(+), 36 deletions(-)
 create mode 100644 phaul/p_haul_module.py

diff --git a/phaul/p_haul_lxc.py b/phaul/p_haul_lxc.py
index 13640d2..fbd26d2 100644
--- a/phaul/p_haul_lxc.py
+++ b/phaul/p_haul_lxc.py
@@ -5,6 +5,7 @@
 import os
 import shutil
 import p_haul_cgroup
+import p_haul_module
 import util
 import fs_haul_shared
 import fs_haul_subtree
@@ -122,6 +123,9 @@ class p_haul_type:
 		# Keep this name, we'll need one in prepare_ct()
 		self.cg_img = os.path.join(dir, cg_image_name)
 
+	def final_restore(self, img, connection):
+		p_haul_module.final_restore(self, img, connection)
+
 	#
 	# Create cgroup hierarchy and put root task into it
 	# Hierarchy is unlimited, we will apply config limitations
diff --git a/phaul/p_haul_module.py b/phaul/p_haul_module.py
new file mode 100644
index 0000000..31f7950
--- /dev/null
+++ b/phaul/p_haul_module.py
@@ -0,0 +1,41 @@
+#
+# Generic functionality for p.haul modules
+#
+
+import pycriu.rpc
+import criu_req
+
+def final_restore(htype, img, connection):
+	"""Perform final restore"""
+
+	nroot = htype.mount()
+	if nroot:
+		print "Restore root set to %s" % nroot
+
+	req = criu_req.make_restore_req(htype, img, nroot)
+	resp = connection.send_req(req)
+	while True:
+		if resp.type == pycriu.rpc.NOTIFY:
+			print "\t\tNotify (%s.%d)" % (resp.notify.script, resp.notify.pid)
+			if resp.notify.script == "setup-namespaces":
+				#
+				# At that point we have only one task
+				# living in namespaces and waiting for
+				# us to ACK the notify. Htype might want
+				# to configure namespace (external net
+				# devices) and cgroups
+				#
+				htype.prepare_ct(resp.notify.pid)
+			elif resp.notify.script == "network-unlock":
+				htype.net_unlock()
+			elif resp.notify.script == "network-lock":
+				raise Exception("Locking network on restore?")
+
+			resp = connection.ack_notify()
+			continue
+
+		if not resp.success:
+			raise Exception("Restore failed")
+		break
+
+	htype.restored(resp.restore.pid)
diff --git a/phaul/p_haul_pid.py b/phaul/p_haul_pid.py
index 5bd2d5d..8b2bcc2 100644
--- a/phaul/p_haul_pid.py
+++ b/phaul/p_haul_pid.py
@@ -2,6 +2,7 @@
 # Individual process hauler
 #
 
+import p_haul_module
 import fs_haul_shared
 
 name = "pid"
@@ -11,7 +12,6 @@ class p_haul_type:
 		self.pid = int(id)
 		self._pidfile = None
 
-
 	#
 	# Initialize itself for source node or destination one
 	#
@@ -60,6 +60,9 @@ class p_haul_type:
 	def put_meta_images(self, dir):
 		pass
 
+	def final_restore(self, img, connection):
+		p_haul_module.final_restore(self, img, connection)
+
 	# Things are started to get restored, only the
 	# first task (with pid @pid) is created.
 	def prepare_ct(self, pid):
diff --git a/phaul/p_haul_service.py b/phaul/p_haul_service.py
index bed34c2..ca2d5e1 100644
--- a/phaul/p_haul_service.py
+++ b/phaul/p_haul_service.py
@@ -86,40 +86,8 @@ class phaul_service:
 	def rpc_restore_from_images(self):
 		print "Restoring from images"
 		self.htype.put_meta_images(self.img.image_dir())
-
-		nroot = self.htype.mount()
-		if nroot:
-			print "Restore root set to %s" % nroot
-
-		req = criu_req.make_restore_req(self.htype, self.img, nroot)
-		resp = self.criu.send_req(req)
-		while True:
-			if resp.type == cr_rpc.NOTIFY:
-				print "\t\tNotify (%s.%d)" % (resp.notify.script, resp.notify.pid)
-				if resp.notify.script == "setup-namespaces":
-					#
-					# At that point we have only one task
-					# living in namespaces and waiting for
-					# us to ACK the notify. Htype might want
-					# to configure namespace (external net
-					# devices) and cgroups
-					#
-					self.htype.prepare_ct(resp.notify.pid)
-				elif resp.notify.script == "network-unlock":
-					self.htype.net_unlock()
-				elif resp.notify.script == "network-lock":
-					raise Exception("Locking network on restore?")
-
-				resp = self.criu.ack_notify()
-				continue
-
-			if not resp.success:
-				raise Exception("Restore failed")
-
-			print "Restore succeeded"
-			break
-
-		self.htype.restored(resp.restore.pid)
+		self.htype.final_restore(self.img, self.criu)
+		print "Restore succeeded"
 		self.restored = True
 
 	def rpc_restore_time(self):
diff --git a/phaul/p_haul_type.py b/phaul/p_haul_type.py
index adf654e..1a9e984 100644
--- a/phaul/p_haul_type.py
+++ b/phaul/p_haul_type.py
@@ -20,7 +20,7 @@ def __get(id):
 		return h_type.p_haul_type(id[1])
 	else:
 		print "Unknown type. Try one of", haul_types.keys()
-	 	return None
+		return None
 
 def get_src(id):
 	ht = __get(id)
diff --git a/phaul/p_haul_vz.py b/phaul/p_haul_vz.py
index 5db2f4a..76230d7 100644
--- a/phaul/p_haul_vz.py
+++ b/phaul/p_haul_vz.py
@@ -5,6 +5,7 @@
 import os
 import shlex
 import p_haul_cgroup
+import p_haul_module
 import util
 import fs_haul_shared
 import fs_haul_subtree
@@ -139,6 +140,9 @@ class p_haul_type:
 		# Keep this name, we'll need one in prepare_ct()
 		self.cg_img = os.path.join(path, cg_image_name)
 
+	def final_restore(self, img, connection):
+		p_haul_module.final_restore(self, img, connection)
+
 	def prepare_ct(self, pid):
 		"""Create cgroup hierarchy and put root task into it.
 
-- 
1.7.1