1#!/usr/bin/python
2# -*- coding: utf-8 -*-
3"""
4Helpers for cgroup testing.
5
6@copyright: 2011 Red Hat Inc.
7@author: Lukas Doktor <ldoktor@redhat.com>
8"""
9import os, logging, subprocess, time, shutil
10from tempfile import mkdtemp
11from autotest_lib.client.bin import utils
12from autotest_lib.client.common_lib import error
13
14
15class Cgroup(object):
16    """
17    Cgroup handling class.
18    """
19    def __init__(self, module, _client):
20        """
21        Constructor
22        @param module: Name of the cgroup module
23        @param _client: Test script pwd + name
24        """
25        self.module = module
26        self._client = _client
27        self.root = None
28
29
30    def initialize(self, modules):
31        """
32        Initializes object for use.
33
34        @param modules: Array of all available cgroup modules.
35        @return: 0 when PASSED.
36        """
37        self.root = modules.get_pwd(self.module)
38        if self.root:
39            return 0
40        else:
41            logging.error("cg.initialize(): Module %s not found", self.module)
42            return -1
43        return 0
44
45
46    def mk_cgroup(self, root=None):
47        """
48        Creates new temporary cgroup
49        @param root: where to create this cgroup (default: self.root)
50        @return: 0 when PASSED
51        """
52        try:
53            if root:
54                pwd = mkdtemp(prefix='cgroup-', dir=root) + '/'
55            else:
56                pwd = mkdtemp(prefix='cgroup-', dir=self.root) + '/'
57        except Exception, inst:
58            logging.error("cg.mk_cgroup(): %s" , inst)
59            return None
60        return pwd
61
62
63    def rm_cgroup(self, pwd, supress=False):
64        """
65        Removes cgroup.
66
67        @param pwd: cgroup directory.
68        @param supress: supress output.
69        @return: 0 when PASSED
70        """
71        try:
72            os.rmdir(pwd)
73        except Exception, inst:
74            if not supress:
75                logging.error("cg.rm_cgroup(): %s" , inst)
76            return -1
77        return 0
78
79
80    def test(self, cmd):
81        """
82        Executes cgroup_client.py with cmd parameter.
83
84        @param cmd: command to be executed
85        @return: subprocess.Popen() process
86        """
87        logging.debug("cg.test(): executing paralel process '%s'", cmd)
88        cmd = self._client + ' ' + cmd
89        process = subprocess.Popen(cmd, shell=True, stdin=subprocess.PIPE,
90                                   stdout=subprocess.PIPE,
91                                   stderr=subprocess.PIPE, close_fds=True)
92        return process
93
94
95    def is_cgroup(self, pid, pwd):
96        """
97        Checks if the 'pid' process is in 'pwd' cgroup
98        @param pid: pid of the process
99        @param pwd: cgroup directory
100        @return: 0 when is 'pwd' member
101        """
102        if open(pwd + '/tasks').readlines().count("%d\n" % pid) > 0:
103            return 0
104        else:
105            return -1
106
107
108    def is_root_cgroup(self, pid):
109        """
110        Checks if the 'pid' process is in root cgroup (WO cgroup)
111        @param pid: pid of the process
112        @return: 0 when is 'root' member
113        """
114        return self.is_cgroup(pid, self.root)
115
116
117    def set_cgroup(self, pid, pwd):
118        """
119        Sets cgroup membership
120        @param pid: pid of the process
121        @param pwd: cgroup directory
122        @return: 0 when PASSED
123        """
124        try:
125            open(pwd+'/tasks', 'w').write(str(pid))
126        except Exception, inst:
127            logging.error("cg.set_cgroup(): %s" , inst)
128            return -1
129        if self.is_cgroup(pid, pwd):
130            logging.error("cg.set_cgroup(): Setting %d pid into %s cgroup "
131                          "failed", pid, pwd)
132            return -1
133        else:
134            return 0
135
136    def set_root_cgroup(self, pid):
137        """
138        Resets the cgroup membership (sets to root)
139        @param pid: pid of the process
140        @return: 0 when PASSED
141        """
142        return self.set_cgroup(pid, self.root)
143
144
145    def get_prop(self, prop, pwd=None, supress=False):
146        """
147        Gets one line of the property value
148        @param prop: property name (file)
149        @param pwd: cgroup directory
150        @param supress: supress the output
151        @return: String value or None when FAILED
152        """
153        tmp = self.get_property(prop, pwd, supress)
154        if tmp:
155            if tmp[0][-1] == '\n':
156                tmp[0] = tmp[0][:-1]
157            return tmp[0]
158        else:
159            return None
160
161
162    def get_property(self, prop, pwd=None, supress=False):
163        """
164        Gets the property value
165        @param prop: property name (file)
166        @param pwd: cgroup directory
167        @param supress: supress the output
168        @return: [] values or None when FAILED
169        """
170        if pwd == None:
171            pwd = self.root
172        try:
173            ret = open(pwd+prop, 'r').readlines()
174        except Exception, inst:
175            ret = None
176            if not supress:
177                logging.error("cg.get_property(): %s" , inst)
178        return ret
179
180
181    def set_prop(self, prop, value, pwd=None, check=True):
182        """
183        Sets the one-line property value concerning the K,M,G postfix
184        @param prop: property name (file)
185        @param value: desired value
186        @param pwd: cgroup directory
187        @param check: check the value after setup
188        @return: 0 when PASSED
189        """
190        _value = value
191        try:
192            value = str(value)
193            if value[-1] == '\n':
194                value = value[:-1]
195            if value[-1] == 'K':
196                value = int(value[:-1]) * 1024
197            elif value[-1] == 'M':
198                value = int(value[:-1]) * 1048576
199            elif value[-1] == 'G':
200                value = int(value[:-1]) * 1073741824
201        except:
202            logging.error("cg.set_prop() fallback into cg.set_property.")
203            value = _value
204        return self.set_property(prop, value, pwd, check)
205
206
207    def set_property(self, prop, value, pwd=None, check=True):
208        """
209        Sets the property value
210        @param prop: property name (file)
211        @param value: desired value
212        @param pwd: cgroup directory
213        @param check: check the value after setup
214        @return: 0 when PASSED
215        """
216        value = str(value)
217        if pwd == None:
218            pwd = self.root
219        try:
220            open(pwd+prop, 'w').write(value)
221        except Exception, inst:
222            logging.error("cg.set_property(): %s" , inst)
223            return -1
224        if check:
225            # Get the first line - '\n'
226            _value = self.get_property(prop, pwd)[0][:-1]
227            if value != _value:
228                logging.error("cg.set_property(): Setting failed: desired = %s,"
229                              " real value = %s", value, _value)
230                return -1
231        return 0
232
233
234    def smoke_test(self):
235        """
236        Smoke test
237        Module independent basic tests
238        """
239        part = 0
240        pwd = self.mk_cgroup()
241        if pwd == None:
242            logging.error("cg.smoke_test[%d]: Can't create cgroup", part)
243            return -1
244
245        part += 1
246        ps = self.test("smoke")
247        if ps == None:
248            logging.error("cg.smoke_test[%d]: Couldn't create process", part)
249            return -1
250
251        part += 1
252        if (ps.poll() != None):
253            logging.error("cg.smoke_test[%d]: Process died unexpectidly", part)
254            return -1
255
256        # New process should be a root member
257        part += 1
258        if self.is_root_cgroup(ps.pid):
259            logging.error("cg.smoke_test[%d]: Process is not a root member",
260                          part)
261            return -1
262
263        # Change the cgroup
264        part += 1
265        if self.set_cgroup(ps.pid, pwd):
266            logging.error("cg.smoke_test[%d]: Could not set cgroup", part)
267            return -1
268
269        # Try to remove used cgroup
270        part += 1
271        if self.rm_cgroup(pwd, supress=True) == 0:
272            logging.error("cg.smoke_test[%d]: Unexpected successful deletion of"
273                          " the used cgroup", part)
274            return -1
275
276        # Return the process into the root cgroup
277        part += 1
278        if self.set_root_cgroup(ps.pid):
279            logging.error("cg.smoke_test[%d]: Could not return the root cgroup "
280                          "membership", part)
281            return -1
282
283        # It should be safe to remove the cgroup now
284        part += 1
285        if self.rm_cgroup(pwd):
286            logging.error("cg.smoke_test[%d]: Can't remove cgroup directory",
287                          part)
288            return -1
289
290        # Finish the process
291        part += 1
292        ps.stdin.write('\n')
293        time.sleep(2)
294        if (ps.poll() == None):
295            logging.error("cg.smoke_test[%d]: Process is not finished", part)
296            return -1
297
298        return 0
299
300
301class CgroupModules(object):
302    """
303    Handles the list of different cgroup filesystems.
304    """
305    def __init__(self):
306        self.modules = []
307        self.modules.append([])
308        self.modules.append([])
309        self.modules.append([])
310        self.mountdir = mkdtemp(prefix='cgroup-') + '/'
311
312
313    def init(self, _modules):
314        """
315        Checks the mounted modules and if necessary mounts them into tmp
316            mountdir.
317        @param _modules: Desired modules.
318        @return: Number of initialized modules.
319        """
320        logging.debug("Desired cgroup modules: %s", _modules)
321        mounts = []
322        fp = open('/proc/mounts', 'r')
323        line = fp.readline().split()
324        while line:
325            if line[2] == 'cgroup':
326                mounts.append(line)
327            line = fp.readline().split()
328        fp.close()
329
330        for module in _modules:
331            # Is it already mounted?
332            i = False
333            for mount in mounts:
334                if mount[3].find(module) != -1:
335                    self.modules[0].append(module)
336                    self.modules[1].append(mount[1] + '/')
337                    self.modules[2].append(False)
338                    i = True
339                    break
340            if not i:
341                # Not yet mounted
342                os.mkdir(self.mountdir + module)
343                cmd = ('mount -t cgroup -o %s %s %s' %
344                       (module, module, self.mountdir + module))
345                try:
346                    utils.run(cmd)
347                    self.modules[0].append(module)
348                    self.modules[1].append(self.mountdir + module)
349                    self.modules[2].append(True)
350                except error.CmdError:
351                    logging.info("Cgroup module '%s' not available", module)
352
353        logging.debug("Initialized cgroup modules: %s", self.modules[0])
354        return len(self.modules[0])
355
356
357    def cleanup(self):
358        """
359        Unmount all cgroups and remove the mountdir.
360        """
361        for i in range(len(self.modules[0])):
362            if self.modules[2][i]:
363                utils.system('umount %s -l' % self.modules[1][i],
364                             ignore_status=True)
365        shutil.rmtree(self.mountdir)
366
367
368    def get_pwd(self, module):
369        """
370        Returns the mount directory of 'module'
371        @param module: desired module (memory, ...)
372        @return: mount directory of 'module' or None
373        """
374        try:
375            i = self.modules[0].index(module)
376        except Exception, inst:
377            logging.error("module %s not found: %s", module, inst)
378            return None
379        return self.modules[1][i]
380