[PATCH] Adds cgroup handling library

[Lukas Doktor <ldoktor@redhat.com>]

[new] cgroup_common.py
* library for handling cgroups

Signed-off-by: Lukas Doktor <ldoktor@redhat.com>
---
 client/tests/cgroup/cgroup.py        |    5 +-
 client/tests/cgroup/cgroup_common.py |  327 ++++++++++++++++++++++++++++++++++
 2 files changed, 331 insertions(+), 1 deletions(-)
 create mode 100755 client/tests/cgroup/cgroup_common.py

diff --git a/client/tests/cgroup/cgroup.py b/client/tests/cgroup/cgroup.py
index d043d65..112f012 100755
--- a/client/tests/cgroup/cgroup.py
+++ b/client/tests/cgroup/cgroup.py
@@ -118,6 +118,7 @@ class cgroup(test.test):
         # Fill the memory without cgroup limitation
         # Should pass
         ################################################
+        logging.debug("test_memory: Memfill WO cgroup")
         ps = item.test("memfill %d" % mem)
         ps.stdin.write('\n')
         i = 0
@@ -141,6 +142,7 @@ class cgroup(test.test):
         # memsw: should swap out part of the process and pass
         # WO memsw: should fail (SIGKILL)
         ################################################
+        logging.debug("test_memory: Memfill mem only limit")
         ps = item.test("memfill %d" % mem)
         if item.set_cgroup(ps.pid, pwd):
             logging.error("test_memory: Could not set cgroup")
@@ -187,6 +189,7 @@ class cgroup(test.test):
         # Fill the memory with 1/2 memory+swap limit
         # Should fail
         ################################################
+        logging.debug("test_memory: Memfill mem + swap limit")
         if memsw:
             ps = item.test("memfill %d" % mem)
             if item.set_cgroup(ps.pid, pwd):
@@ -226,11 +229,11 @@ class cgroup(test.test):
             logging.debug("test_memory: Memfill mem+swap cgroup passed")
 
         # cleanup
+        logging.debug("test_memory: Cleanup")
         if item.rm_cgroup(pwd):
             logging.error("test_memory: Can't remove cgroup directory")
             return -1
         os.system("swapon -a")
-        logging.debug("test_memory: Cleanup passed")
 
         logging.info("Leaving 'test_memory': PASSED")
         return 0
diff --git a/client/tests/cgroup/cgroup_common.py b/client/tests/cgroup/cgroup_common.py
new file mode 100755
index 0000000..3fd1cf7
--- /dev/null
+++ b/client/tests/cgroup/cgroup_common.py
@@ -0,0 +1,327 @@
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+"""
+Helpers for cgroup testing
+
+@copyright: 2011 Red Hat Inc.
+@author: Lukas Doktor <ldoktor@redhat.com>
+"""
+import os, logging
+import subprocess
+from tempfile import mkdtemp
+import time
+
+class Cgroup:
+    """
+    Cgroup handling class
+    """
+    def __init__(self, module, _client):
+        """
+        Constructor
+        @param module: Name of the cgroup module
+        @param _client: Test script pwd+name
+        """
+        self.module = module
+        self._client = _client
+        self.root = None
+
+
+    def initialize(self, modules):
+        """
+        Inicializes object for use
+        @param modules: array of all available cgroup modules
+        @return: 0 when PASSED
+        """
+        self.root = modules.get_pwd(self.module)
+        if self.root:
+            return 0
+        else:
+            logging.error("cg.initialize(): Module %s not found", self.module)
+            return -1
+        return 0
+
+
+    def mk_cgroup(self, root=None):
+        """
+        Creates new temporary cgroup
+        @param root: where to create this cgroup (default: self.root)
+        @return: 0 when PASSED
+        """
+        try:
+            if root:
+                pwd = mkdtemp(prefix='cgroup-', dir=root) + '/'
+            else:
+                pwd = mkdtemp(prefix='cgroup-', dir=self.root) + '/'
+        except Exception, inst:
+            logging.error("cg.mk_cgroup(): %s" , inst)
+            return None
+        return pwd
+
+
+    def rm_cgroup(self, pwd, supress=False):
+        """
+        Removes cgroup
+        @param pwd: cgroup directory
+        @param supress: supress output
+        @return: 0 when PASSED
+        """
+        try:
+            os.rmdir(pwd)
+        except Exception, inst:
+            if not supress:
+                logging.error("cg.rm_cgroup(): %s" , inst)
+            return -1
+        return 0
+
+
+    def test(self, cmd):
+        """
+        Executes cgroup_client.py with cmd parameter
+        @param cmd: command to be executed
+        @return: subprocess.Popen() process
+        """
+        logging.debug("cg.test(): executing paralel process '%s'" , cmd)
+        process = subprocess.Popen((self._client + ' ' + cmd), shell=True,
+                                stdin=subprocess.PIPE, stdout=subprocess.PIPE,
+                                stderr=subprocess.PIPE, close_fds=True)
+        return process
+
+
+    def is_cgroup(self, pid, pwd):
+        """
+        Checks if the 'pid' process is in 'pwd' cgroup
+        @param pid: pid of the process
+        @param pwd: cgroup directory
+        @return: 0 when is 'pwd' member
+        """
+        if open(pwd+'/tasks').readlines().count("%d\n" % pid) > 0:
+            return 0
+        else:
+            return -1
+
+    def is_root_cgroup(self, pid):
+        """
+        Checks if the 'pid' process is in root cgroup (WO cgroup)
+        @param pid: pid of the process
+        @return: 0 when is 'root' member
+        """
+        return self.is_cgroup(pid, self.root)
+
+    def set_cgroup(self, pid, pwd):
+        """
+        Sets cgroup membership
+        @param pid: pid of the process
+        @param pwd: cgroup directory
+        @return: 0 when PASSED
+        """
+        try:
+            open(pwd+'/tasks', 'w').write(str(pid))
+        except Exception, inst:
+            logging.error("cg.set_cgroup(): %s" , inst)
+            return -1
+        if self.is_cgroup(pid, pwd):
+            logging.error("cg.set_cgroup(): Setting %d pid into %s cgroup "
+                          "failed", pid, pwd)
+            return -1
+        else:
+            return 0
+
+    def set_root_cgroup(self, pid):
+        """
+        Resets the cgroup membership (sets to root)
+        @param pid: pid of the process
+        @return: 0 when PASSED
+        """
+        return self.set_cgroup(pid, self.root)
+
+
+    def get_property(self, prop, pwd=None, supress=False):
+        """
+        Gets the property value
+        @param prop: property name (file)
+        @param pwd: cgroup directory
+        @param supress: supress the output
+        @return: String value or None when FAILED
+        """
+        if pwd == None:
+            pwd = self.root
+        try:
+            ret = open(pwd+prop, 'r').readlines()
+        except Exception, inst:
+            ret = None
+            if not supress:
+                logging.error("cg.get_property(): %s" , inst)
+        return ret
+
+
+    def set_property(self, prop, value, pwd=None, check=True):
+        """
+        Sets the property value
+        @param prop: property name (file)
+        @param value: desired value
+        @param pwd: cgroup directory
+        @param check: check the value after setup
+        @return: 0 when PASSED
+        """
+        value = str(value)
+        if pwd == None:
+            pwd = self.root
+        try:
+            open(pwd+prop, 'w').write(value)
+        except Exception, inst:
+            logging.error("cg.set_property(): %s" , inst)
+            return -1
+        if check:
+            # Get the first line - '\n'
+            _value = self.get_property(prop, pwd)[0][:-1]
+            if value != _value:
+                logging.error("cg.set_property(): Setting failed: desired = %s,"
+                              " real value = %s", value, _value)
+                return -1
+        return 0
+
+
+    def smoke_test(self):
+        """
+        Smoke test
+        Module independent basic tests
+        """
+        part = 0
+        pwd = self.mk_cgroup()
+        if pwd == None:
+            logging.error("cg.smoke_test[%d]: Can't create cgroup", part)
+            return -1
+
+        part += 1
+        ps = self.test("smoke")
+        if ps == None:
+            logging.error("cg.smoke_test[%d]: Couldn't create process", part)
+            return -1
+
+        part += 1
+        if (ps.poll() != None):
+            logging.error("cg.smoke_test[%d]: Process died unexpectidly", part)
+            return -1
+
+        # New process should be a root member
+        part += 1
+        if self.is_root_cgroup(ps.pid):
+            logging.error("cg.smoke_test[%d]: Process is not a root member",
+                          part)
+            return -1
+
+        # Change the cgroup
+        part += 1
+        if self.set_cgroup(ps.pid, pwd):
+            logging.error("cg.smoke_test[%d]: Could not set cgroup", part)
+            return -1
+
+        # Try to remove used cgroup
+        part += 1
+        if self.rm_cgroup(pwd, supress=True) == 0:
+            logging.error("cg.smoke_test[%d]: Unexpected successful deletion of"
+                          " the used cgroup", part)
+            return -1
+
+        # Return the process into the root cgroup
+        part += 1
+        if self.set_root_cgroup(ps.pid):
+            logging.error("cg.smoke_test[%d]: Could not return the root cgroup "
+                          "membership", part)
+            return -1
+
+        # It should be safe to remove the cgroup now
+        part += 1
+        if self.rm_cgroup(pwd):
+            logging.error("cg.smoke_test[%d]: Can't remove cgroup direcotry",
+                          part)
+            return -1
+
+        # Finish the process
+        part += 1
+        ps.stdin.write('\n')
+        time.sleep(2)
+        if (ps.poll() == None):
+            logging.error("cg.smoke_test[%d]: Process is not finished", part)
+            return -1
+
+        return 0
+
+
+class CgroupModules:
+    """
+    Handles the list of different cgroup filesystems
+    """
+    def __init__(self):
+        self.modules = []
+        self.modules.append([])
+        self.modules.append([])
+        self.modules.append([])
+        self.mountdir = mkdtemp(prefix='cgroup-') + '/'
+
+
+    def init(self, _modules):
+        """
+        Checks the mounted modules and if necessarily mounts them into tmp
+            mountdir.
+        @param _modules: desired modules
+        @return: Number of initialized modules
+        """
+        mounts = []
+        fp = open('/proc/mounts', 'r')
+        line = fp.readline().split()
+        while line:
+            if line[2] == 'cgroup':
+                mounts.append(line)
+            line = fp.readline().split()
+        fp.close()
+
+        for module in _modules:
+            # Is it already mounted?
+            i = False
+            for mount in mounts:
+                if mount[3].find(module) != -1:
+                    self.modules[0].append(module)
+                    self.modules[1].append(mount[1]+'/')
+                    self.modules[2].append(False)
+                    i = True
+                    break
+
+            if not i:
+                # Not yet mounted
+                os.mkdir(self.mountdir+module)
+                logging.info('mount -t cgroup -o %s %s %s', module, module,
+                             self.mountdir+module)
+                if (os.system('mount -t cgroup -o %s %s %s'
+                                    % (module, module, self.mountdir+module)) == 0):
+                    self.modules[0].append(module)
+                    self.modules[1].append(self.mountdir+module)
+                    self.modules[2].append(True)
+                else:
+                    logging.error("Module '%s' is not available, mount failed",
+                                    module)
+        return len(self.modules[0])
+
+
+    def cleanup(self):
+        """
+        Unmount all cgroups and remove the mountdir
+        """
+        for i in range(len(self.modules[0])):
+            if self.modules[2][i]:
+                os.system('umount %s -l' % self.modules[1][i])
+        os.system('rm -rf %s' % self.mountdir)
+
+
+    def get_pwd(self, module):
+        """
+        Returns the mount directory of 'module'
+        @param module: desired module (memory, ...)
+        @return: mount directory of 'module' or None
+        """
+        try:
+            i = self.modules[0].index(module)
+        except Exception, inst:
+            logging.error("module %s not found: %s", module, inst)
+            return None
+        return self.modules[1][i]
-- 
1.7.6