Skip to content

Instantly share code, notes, and snippets.

@quark-zju
Last active August 29, 2015 14:02
Embed
What would you like to do?
Handy memory limit script
#!/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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment