python/pakfire/cgroup.py

   1 #!/usr/bin/python
   2
   3 import os
   4 import shutil
   5 import signal
   6 import time
   7
   8 import logging
   9 log = logging.getLogger("pakfire.cgroups")
  10
  11 CGROUP_PATH_CANDIDATES = (
  12         "/sys/fs/cgroup/systemd/system",
  13         "/sys/fs/cgroup",
  14 )
  15
  16 def find_cgroup_path():
  17         """
  18                 This function tries to find the right place
  19                 where to put the cgroups.
  20         """
  21         for path in CGROUP_PATH_CANDIDATES:
  22                 check_path = os.path.join(path, "tasks")
  23                 if not os.path.exists(check_path):
  24                         continue
  25
  26                 return path
  27
  28 CGROUP_PATH = find_cgroup_path()
  29
  30 def supported():
  31         """
  32                 Returns True or False depending on
  33                 whether cgroups are supported or not.
  34         """
  35         if CGROUP_PATH is None:
  36                 return False
  37
  38         return True
  39
  40 class CGroup(object):
  41         def __init__(self, name):
  42                 assert supported(), "cgroups are not supported by this kernel"
  43
  44                 self.name = name
  45                 self.path = os.path.join(CGROUP_PATH, name)
  46                 self.path = os.path.abspath(self.path)
  47
  48                 # The parent cgroup.
  49                 self._parent = None
  50
  51                 # Initialize the cgroup.
  52                 self.create()
  53
  54                 log.debug("cgroup '%s' has been successfully initialized." % self.name)
  55
  56         def __repr__(self):
  57                 return "<%s %s>" % (self.__class__.__name__, self.name)
  58
  59         def __cmp__(self, other):
  60                 return cmp(self.path, other.path)
  61
  62         def create(self):
  63                 """
  64                         Creates the filesystem structure for
  65                         the cgroup.
  66                 """
  67                 if os.path.exists(self.path):
  68                         return
  69
  70                 log.debug("cgroup '%s' has been created." % self.name)
  71                 os.makedirs(self.path)
  72
  73         def attach(self):
  74                 """
  75                         Attaches this task to the cgroup.
  76                 """
  77                 pid = os.getpid()
  78                 self.attach_task(pid)
  79
  80         def destroy(self):
  81                 """
  82                         Deletes the cgroup.
  83
  84                         All running tasks will be migrated to the parent cgroup.
  85                 """
  86                 # Don't delete the root cgroup.
  87                 if self == self.root:
  88                         return
  89
  90                 # Move all tasks to the parent.
  91                 self.migrate(self.parent)
  92
  93                 # Just make sure the statement above worked.
  94                 assert self.is_empty(recursive=True), "cgroup must be empty to be destroyed"
  95                 assert not self.processes
  96
  97                 # Remove the file tree.
  98                 try:
  99                         os.rmdir(self.path)
 100                 except OSError, e:
 101                         # Ignore "Device or resource busy".
 102                         if e.errno == 16:
 103                                 return
 104
 105                         raise
 106
 107         def _read(self, file):
 108                 """
 109                         Reads the contect of file in the cgroup directory.
 110                 """
 111                 file = os.path.join(self.path, file)
 112
 113                 with open(file) as f:
 114                         return f.read()
 115
 116         def _read_pids(self, file):
 117                 """
 118                         Reads file and interprets the lines as a sorted list.
 119                 """
 120                 _pids = self._read(file)
 121
 122                 pids = []
 123
 124                 for pid in _pids.splitlines():
 125                         try:
 126                                 pid = int(pid)
 127                         except ValueError:
 128                                 continue
 129
 130                         if pid in pids:
 131                                 continue
 132
 133                         pids.append(pid)
 134
 135                 return sorted(pids)
 136
 137         def _write(self, file, what):
 138                 """
 139                         Writes what to file in the cgroup directory.
 140                 """
 141                 file = os.path.join(self.path, file)
 142
 143                 f = open(file, "w")
 144                 f.write("%s" % what)
 145                 f.close()
 146
 147         @property
 148         def root(self):
 149                 if self.parent:
 150                         return self.parent.root
 151
 152                 return self
 153
 154         @property
 155         def parent(self):
 156                 # Cannot go above CGROUP_PATH.
 157                 if self.path == CGROUP_PATH:
 158                         return
 159
 160                 if self._parent is None:
 161                         parent_name = os.path.dirname(self.name)
 162                         self._parent = CGroup(parent_name)
 163
 164                 return self._parent
 165
 166         @property
 167         def subgroups(self):
 168                 subgroups = []
 169
 170                 for name in os.listdir(self.path):
 171                         path = os.path.join(self.path, name)
 172                         if not os.path.isdir(path):
 173                                 continue
 174
 175                         name = os.path.join(self.name, name)
 176                         group = CGroup(name)
 177
 178                         subgroups.append(group)
 179
 180                 return subgroups
 181
 182         def is_empty(self, recursive=False):
 183                 """
 184                         Returns True if the cgroup is empty.
 185
 186                         Otherwise returns False.
 187                 """
 188                 if self.tasks:
 189                         return False
 190
 191                 if recursive:
 192                         for subgroup in self.subgroups:
 193                                 if subgroup.is_empty(recursive=recursive):
 194                                         continue
 195
 196                                 return False
 197
 198                 return True
 199
 200         @property
 201         def tasks(self):
 202                 """
 203                         Returns a list of pids of all tasks
 204                         in this process group.
 205                 """
 206                 return self._read_pids("tasks")
 207
 208         @property
 209         def processes(self):
 210                 """
 211                         Returns a list of pids of all processes
 212                         that are currently running within the cgroup.
 213                 """
 214                 return self._read_pids("cgroup.procs")
 215
 216         def attach_task(self, pid):
 217                 """
 218                         Attaches the task with the given PID to
 219                         the cgroup.
 220                 """
 221                 self._write("tasks", pid)
 222
 223         def migrate_task(self, other, pid):
 224                 """
 225                         Migrates a single task to another cgroup.
 226                 """
 227                 other.attach_task(pid)
 228
 229         def migrate(self, other):
 230                 if self.is_empty(recursive=True):
 231                         return
 232
 233                 log.info("Migrating all tasks from '%s' to '%s'." \
 234                         % (self.name, other.name))
 235
 236                 while True:
 237                         # Migrate all tasks to the new cgroup.
 238                         for task in self.tasks:
 239                                 self.migrate_task(other, task)
 240
 241                         # Also do that for all subgroups.
 242                         for subgroup in self.subgroups:
 243                                 subgroup.migrate(other)
 244
 245                         if self.is_empty():
 246                                 break
 247
 248         def kill(self, sig=signal.SIGTERM, recursive=True):
 249                 killed_processes = []
 250
 251                 mypid = os.getpid()
 252
 253                 while True:
 254                         for proc in self.processes:
 255                                 # Don't kill myself.
 256                                 if proc == mypid:
 257                                         continue
 258
 259                                 # Skip all processes that have already been killed.
 260                                 if proc in killed_processes:
 261                                         continue
 262
 263                                 # If we haven't killed the process yet, we kill it.
 264                                 log.debug("Sending signal %s to process %s..." % (sig, proc))
 265
 266                                 try:
 267                                         os.kill(proc, sig)
 268                                 except OSError, e:
 269                                         raise
 270
 271                                 # Save all killed processes to a list.
 272                                 killed_processes.append(proc)
 273
 274                         else:
 275                                 # If no processes are left to be killed, we end the loop.
 276                                 break
 277
 278                 # Nothing more to do if not in recursive mode.
 279                 if not recursive:
 280                         return
 281
 282                 # Kill all processes in subgroups as well.
 283                 for subgroup in self.subgroups:
 284                         subgroup.kill(sig=sig, recursive=recursive)
 285
 286         def kill_and_wait(self):
 287                 # Safely kill all processes in the cgroup.
 288                 # This first sends SIGTERM and then checks 8 times
 289                 # after 200ms whether the group is empty. If not,
 290                 # everything what's still in there gets SIGKILL
 291                 # and it is five more times checked if everything
 292                 # went away.
 293
 294                 sig = None
 295                 for i in range(15):
 296                         if i == 0:
 297                                 sig = signal.SIGTERM
 298                         elif i == 9:
 299                                 sig = signal.SIGKILL
 300                         else:
 301                                 sig = None
 302
 303                         # If no signal is given and there are no processes
 304                         # left, our job is done and we can exit.
 305                         if not self.processes:
 306                                 break
 307
 308                         if sig:
 309                                 # Send sig to all processes in the cgroup.
 310                                 log.info("Sending signal %s to all processes in '%s'." % (sig, self.name))
 311                                 self.kill(sig=sig, recursive=True)
 312
 313                         # Sleep for 200ms.
 314                         time.sleep(0.2)
 315
 316                 return self.is_empty()