[CRIU] [PATCH 3/3] p.haul: implement Virtuozzo specific container restore

[Nikita Spiridonov]

It is needed to exec vzctl restore instead of usual criu restore
request to successfully restore Virtuozzo containers. Implement
Virtuozzo specific restore.

Signed-off-by: Nikita Spiridonov <nspiridonov at odin.com>
---
 phaul/criu_api.py  |    5 ++++-
 phaul/criu_req.py  |    4 ++--
 phaul/p_haul_vz.py |   31 +++++++++++++++++++++++++++++++
 3 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/phaul/criu_api.py b/phaul/criu_api.py
index 1b7781d..2946e6d 100644
--- a/phaul/criu_api.py
+++ b/phaul/criu_api.py
@@ -54,7 +54,7 @@ class criu_conn:
 
 	def send_req(self, req):
 		req.opts.log_level = self.verb
-		req.opts.log_file = "criu_%s.%d.log" % (criu_req.get_name(req), self._iter)
+		req.opts.log_file = self.get_log_name(req.type)
 		self._cs.send(req.SerializeToString())
 		self._iter += 1
 		self._last_req = req.type
@@ -69,6 +69,9 @@ class criu_conn:
 
 		return self._recv_resp()
 
+	def get_log_name(self, req_type):
+		return "criu_%s.%d.log" % (criu_req.get_name(req_type), self._iter)
+
 #
 # Helper to read CRIU-generated statistics
 #
diff --git a/phaul/criu_req.py b/phaul/criu_req.py
index fc03093..42225aa 100644
--- a/phaul/criu_req.py
+++ b/phaul/criu_req.py
@@ -13,9 +13,9 @@ _names = {
 	pycriu.rpc.CPUINFO_CHECK: "cpuinfo-check",
 }
 
-def get_name(req):
+def get_name(req_type):
 	"""Return printable request name"""
-	return _names[req.type]
+	return _names[req_type]
 
 def _make_req(typ, htype):
 	"""Prepare generic criu request"""
diff --git a/phaul/p_haul_vz.py b/phaul/p_haul_vz.py
index 39f9bf5..027b063 100644
--- a/phaul/p_haul_vz.py
+++ b/phaul/p_haul_vz.py
@@ -3,6 +3,7 @@
 #
 
 import os
+import subprocess
 import shlex
 import p_haul_cgroup
 import p_haul_module
@@ -15,6 +16,7 @@ name = "vz"
 vz_global_conf = "/etc/vz/vz.conf"
 vz_conf_dir = "/etc/vz/conf/"
 cg_image_name = "ovzcg.img"
+vzctl_bin = "vzctl"
 
 class p_haul_type(p_haul_module.phaul_module):
 	def __init__(self, ctid):
@@ -140,6 +142,35 @@ class p_haul_type(p_haul_module.phaul_module):
 		# Keep this name, we'll need one in prepare_ct()
 		self.cg_img = os.path.join(path, cg_image_name)
 
+	def __setup_restore_extra_args(self, path, img, connection):
+		"""Create temporary file with extra arguments for criu restore"""
+		extra_args = [
+			"VE_WORK_DIR={0}\n".format(img.work_dir()),
+			"VE_RESTORE_LOG_PATH={0}\n".format(
+				connection.get_log_name(pycriu.rpc.RESTORE))]
+		with open(path, "w") as f:
+			f.writelines(extra_args)
+
+	def __remove_restore_extra_args(self, path):
+		"""Remove temporary file with extra arguments for criu restore"""
+		if os.path.isfile(path):
+			os.remove(path)
+
+	def final_restore(self, img, connection):
+		"""Perform Virtuozzo-specific final restore"""
+		try:
+			# Setup restore extra arguments
+			args_path = os.path.join(img.image_dir(), "restore-extra-args")
+			self.__setup_restore_extra_args(args_path, img, connection)
+			# Run vzctl restore
+			proc = subprocess.Popen([vzctl_bin, "restore", self._ctid,
+				"--dumpfile", img.image_dir()])
+			if proc.wait() != 0:
+				raise Exception("Restore failed ({0})".format(proc.returncode))
+		finally:
+			# Remove restore extra arguments
+			self.__remove_restore_extra_args(args_path)
+
 	def prepare_ct(self, pid):
 		"""Create cgroup hierarchy and put root task into it.
 
-- 
1.7.1