Skip to content

Instantly share code, notes, and snippets.

@markuskont
Created February 27, 2019 12:59
Show Gist options
  • Save markuskont/ede5f1188bc4f705087aa412f7c84a01 to your computer and use it in GitHub Desktop.
Save markuskont/ede5f1188bc4f705087aa412f7c84a01 to your computer and use it in GitHub Desktop.
Set CPU affinity and IRQ stuff for moloch-capture with python3 and python-psutil
#!/usr/bin/env python3
import psutil
import subprocess
import re
import sys
import os.path
def get_moloch_capture_parent():
procs = {p.pid: p.info for p in psutil.process_iter(attrs=['pid', 'name', 'username'])}
parent = {k: v for k, v in procs.items() if "moloch-capture" in v["name"]}
parent = list(parent.values())[0]["pid"]
parent = psutil.Process(pid=parent)
return parent
def get_moloch_workers(parent):
workers = parent.threads()
workers = [w.id for w in workers]
workers = [psutil.Process(pid=p) for p in workers]
workers = [{"pid": p.pid, "name": p.name()} for p in workers]
return workers
def get_numa_cores(node):
numa = subprocess.run(['numactl', '--hardware'], stdout=subprocess.PIPE).stdout.decode('utf-8')
numa = numa.split("\n")
numa = [v.split(":")[1].strip() for v in numa if "node {} cpus:".format(NODE) in v][0]
numa = numa.split()
numa = [int(v) for v in numa]
return numa
NODE=1
CAP_IFACE="enp130s0f1"
if __name__ == "__main__":
numa = get_numa_cores(NODE)
intr_thread = numa[0]
cap_thread = numa[1]
worker_threads = numa[2:]
cap_pattern = re.compile("^moloch-(?:capture|simple|af\d+-\d+)$")
pkt_pattern = re.compile("^moloch-pkt\d+$")
parent = get_moloch_capture_parent()
workers = get_moloch_workers(parent)
cap_threads = [t for t in workers if cap_pattern.match(t["name"])]
pkt_threads = [t for t in workers if pkt_pattern.match(t["name"])]
if len(pkt_threads) > len(worker_threads):
print("Too many moloch workers for {} cpu threads".format(len(worker_threads)))
sys.exit(1)
for thread in cap_threads:
subprocess.call(['sudo', 'taskset', '-pc', str(cap_thread), str(thread["pid"])])
for i, thread in enumerate(pkt_threads):
subprocess.call(['sudo', 'taskset', '-pc', str(worker_threads[i]), str(thread["pid"])])
lines = []
with open("/proc/interrupts", "rb") as f:
lines = [l.decode().split(":")[0].lstrip() for l in f if CAP_IFACE in l.decode()]
if len(lines) == 0 or len(lines) > 1:
print("found {} irq for {}, should be 1".format(len(lines), CAP_IFACE))
sys.exit(1)
irq = lines[0]
irq = os.path.join('/proc/irq', str(irq), 'smp_affinity_list')
subprocess.Popen(['/usr/bin/sudo /bin/bash -c \'echo {} > {}\''.format(intr_thread, irq)],
shell=True)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment