Handy memory limit script

memory-limit.rb

#!/usr/bin/env ruby

# For Linux only
# Limit: 2.5G, change it if needed
GIGA = 1024 ** 3
MEMORY_LIMIT = GIGA * 5 / 2

CG_DIR = "/sys/fs/cgroup/memory/lm.#{MEMORY_LIMIT}"

if ARGV.any?{|s|s.start_with? '-h'}
  puts "Usage: #{$0} process_name [process_name ...]"
  exit
end

def show_memroy_usage
  path = File.join(CG_DIR, 'memory.usage_in_bytes')
  usage = File.read(path).to_i
  puts "Usage: #{
    [usage, MEMORY_LIMIT].map do |v|
      v.to_f / GIGA
    end.map do |v|
      "#{v.round(3)}G"
    end.join(' / ')
  } (#{(usage * 100 / MEMORY_LIMIT)}%)"
end

def system! *args
  result = system *args
  raise "#{args} exec failed" unless result
end

def get_pids(args)
  args.map do |x|
    if x.to_i.to_s == x
      x.to_i
    else
      `pgrep #{x}`.lines.map(&:to_i)
    end
  end.flatten.compact
end

def write_by_root! path, number
  system! "echo #{number.to_i} | sudo tee #{path} >/dev/null"
end

def create_cgroup!
  return if File.directory? CG_DIR
  puts 'Creating cgroup'
  system! 'sudo', 'mkdir', CG_DIR
end

def set_cgroup_option! option, value
  path = File.join(CG_DIR, option)
  current_value = File.read path
  return if current_value.to_i == value.to_i
  puts "Setting #{option} to #{value}"
  write_by_root! path, value
end

def add_tasks! pids
  path = File.join(CG_DIR, 'tasks')
  pids.each do |pid|
    puts "#{pid} (#{File.basename(File.readlink("/proc/#{pid}/exe")) rescue '?'})"
    write_by_root! path, pid
  end
end

# main
pids = get_pids ARGV

create_cgroup!
set_cgroup_option! 'memory.limit_in_bytes', MEMORY_LIMIT
add_tasks! pids
show_memroy_usage