1#!/usr/bin/python2 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