[CRIU] [PATCH 2/3] p.haul: create base class for all p.haul modules

[Nikita Spiridonov]

Create base class for all p.haul modules (e.g. vz or lxc) and move
final restore logic to this class. This modification is needed to
perform module specific restore.

Signed-off-by: Nikita Spiridonov <nspiridonov at odin.com>
---
 phaul/p_haul_lxc.py     |    3 ++-
 phaul/p_haul_module.py  |   44 ++++++++++++++++++++++++++++++++++++++++++++
 phaul/p_haul_pid.py     |    4 ++--
 phaul/p_haul_service.py |   36 ++----------------------------------
 phaul/p_haul_type.py    |    2 +-
 phaul/p_haul_vz.py      |    3 ++-
 6 files changed, 53 insertions(+), 39 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..7dbc959 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
@@ -15,7 +16,7 @@ lxc_dir = "/var/lib/lxc/"
 lxc_rootfs_dir = "/usr/lib64/lxc/rootfs"
 cg_image_name = "lxccg.img"
 
-class p_haul_type:
+class p_haul_type(p_haul_module.phaul_module):
 	def __init__(self, name):
 		self._ctname = name
 		#
diff --git a/phaul/p_haul_module.py b/phaul/p_haul_module.py
new file mode 100644
index 0000000..34be0fb
--- /dev/null
+++ b/phaul/p_haul_module.py
@@ -0,0 +1,44 @@
+#
+# Base for p.haul modules with generic functionality
+#
+
+import pycriu.rpc
+import criu_req
+
+class phaul_module:
+	"""Basic module with generic functionality"""
+
+	def final_restore(self, img, connection):
+		"""Perform final restore"""
+
+		nroot = self.mount()
+		if nroot:
+			print "Restore root set to %s" % nroot
+
+		req = criu_req.make_restore_req(self, 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
+					#
+					self.prepare_ct(resp.notify.pid)
+				elif resp.notify.script == "network-unlock":
+					self.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
+
+		self.restored(resp.restore.pid)
diff --git a/phaul/p_haul_pid.py b/phaul/p_haul_pid.py
index 5bd2d5d..a17c9bf 100644
--- a/phaul/p_haul_pid.py
+++ b/phaul/p_haul_pid.py
@@ -2,16 +2,16 @@
 # Individual process hauler
 #
 
+import p_haul_module
 import fs_haul_shared
 
 name = "pid"
 
-class p_haul_type:
+class p_haul_type(p_haul_module.phaul_module):
 	def __init__(self, id):
 		self.pid = int(id)
 		self._pidfile = None
 
-
 	#
 	# Initialize itself for source node or destination one
 	#
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..39f9bf5 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
@@ -15,7 +16,7 @@ vz_global_conf = "/etc/vz/vz.conf"
 vz_conf_dir = "/etc/vz/conf/"
 cg_image_name = "ovzcg.img"
 
-class p_haul_type:
+class p_haul_type(p_haul_module.phaul_module):
 	def __init__(self, ctid):
 		self._ctid = ctid
 		self._ct_priv = ""
-- 
1.7.1