preoomkiller

#!/usr/bin/env python

"""
Sourced from https://github.com/albertz/helpers/blob/master/cgroup-mem-limit-watcher.py with some modifications

Install (to /usr/local/bin): `$ curl -sSf https://raw.githubusercontent.com/izakp/docker-preoomkiller/master/install-preoomkiller.sh | sh`

Use:
  - Put this in your docker-entrypoint script to run in the background): `exec /usr/local/bin/preoomkiller &`
  - Set ENABLE_PREOOMKILLER=1 in the container's environment to enable
  - Set PREOOMKILLER_MEMORY_USE_FACTOR, PREOOMKILLER_POLL_INTERVAL, PREOOMKILLER_KILL_SIGNAL, PREOOMKILLER_KILL_PID as desired
"""

import sys, time, os, signal

ENABLE_PREOOMKILLER = os.environ.get("ENABLE_PREOOMKILLER")
DEBUG = os.environ.get("PREOOMKILLER_DEBUG")
MEMORY_USE_FACTOR = float(os.environ.get("PREOOMKILLER_MEMORY_USE_FACTOR", "0.95"))
POLL_INTERVAL = int(os.environ.get("PREOOMKILLER_POLL_INTERVAL", "10"))
KILL_SIGNAL = int(os.environ.get("PREOOMKILLER_KILL_SIGNAL", signal.SIGTERM))
PARENT_PID = os.getppid()
KILL_PID = int(os.environ.get("PREOOMKILLER_KILL_PID", PARENT_PID))

EXIT_SIGNALS = [signal.SIGHUP, signal.SIGINT, signal.SIGQUIT, signal.SIGPIPE, signal.SIGTERM]

def on_shudown(signum, frame):
  print("PreOOMKiller: Received signal %s - shutting down." % signum)
  exit(0)

for sig in EXIT_SIGNALS:
  signal.signal(sig, on_shudown)

def byteNumRepr(c):
  if c < 1024: return "%i B" % c
  S = "KMG"
  i = 0
  while i < len(S) - 1:
    if c < 0.8 * 1024 ** (i + 2): break
    i += 1
  f = float(c) / (1024 ** (i + 1))
  return "%.1f %sB" % (f, S[i])

def getRssLimit():
  return int(open("/sys/fs/cgroup/memory/memory.limit_in_bytes").read())

def getTotalRss():
  stats = open("/sys/fs/cgroup/memory/memory.stat").read().splitlines()
  stats = dict([(key,int(value)) for key,value in map(str.split, stats)])
  return stats["rss"]

def getMemoryStats():
  try:
    limit = getRssLimit()
    used = getTotalRss()
    fact = float(used) / limit
  except Exception:
    print("PreOOMKiller: Could not calculate memory use factor - am I running in a container?")
    raise
  return (limit, used, fact)

def printMemoryStats(limit, used, fact):
  print("PreOOMKiller: Mem limit: %s, Current rss: %s, Percent mem used: %s%%" % (byteNumRepr(limit), byteNumRepr(used), round(100.0*fact)))

# Force disable stdout buffering.
sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', 1)

def main():
  if ENABLE_PREOOMKILLER is None:
    print("PreOOMKiller: Not starting preoomkiller as ENABLE_PREOOMKILLER is unset!")
    return

  print("PreOOMKiller: Starting preoomkiller with memory use factor: %s, poll internal: %s seconds" % (MEMORY_USE_FACTOR, POLL_INTERVAL))

  if DEBUG:
    print("PreOOMKiller: ppid: %s" % PARENT_PID)

  while True:
    # If the parent process has died, stop
    if os.getppid() != PARENT_PID:
      print("PreOOMKiller: Parent process has died. Exiting.")
      sys.exit(0)

    limit, used, fact = getMemoryStats()

    if fact >= MEMORY_USE_FACTOR:
      print("PreOOMKiller: Container is over configured memory use factor!")
      printMemoryStats(limit, used, fact)

      # Kill the configured process
      print("PreOOMKiller: Killing process %s with signal %s." % (KILL_PID, KILL_SIGNAL))
      os.kill(KILL_PID, KILL_SIGNAL)

    if DEBUG:
      printMemoryStats(limit, used, fact)

    time.sleep(POLL_INTERVAL)

if __name__ == "__main__":
  main()