diff --git a/rd-agent/src/side.rs b/rd-agent/src/side.rs index 859a1586..d2881527 100644 --- a/rd-agent/src/side.rs +++ b/rd-agent/src/side.rs @@ -512,6 +512,8 @@ impl Balloon { )?; svc.set_slice(Slice::Sys.name()) + // Make sure memory allocation completed once started + .add_prop("Type".into(), systemd::Prop::String("notify".into())) .add_prop("MemorySwapMax".into(), systemd::Prop::U64(0)) .add_prop( "Slice".into(), diff --git a/rd-agent/src/side/memory-balloon.py b/rd-agent/src/side/memory-balloon.py index c1152192..90aa933c 100755 --- a/rd-agent/src/side/memory-balloon.py +++ b/rd-agent/src/side/memory-balloon.py @@ -2,14 +2,25 @@ # Copyright (c) Facebook, Inc. and its affiliates import mmap +import os import sys import time +import socket if len(sys.argv) < 2: print("Usage: memory-balloon.py BYTES", file=sys.stderr) sys.exit(1) +try: + sd_socket = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) + sd_addr = os.getenv("NOTIFY_SOCKET") + assert sd_addr, "$NOTIFY_SOCKET not available" + sd_socket.connect(sd_addr) +except Exception: + print("Failed to create systemd socket") + raise + nr_pages = int((int(sys.argv[1]) + 4095) / 4096) mm = mmap.mmap(-1, nr_pages * 4096, flags=mmap.MAP_PRIVATE) @@ -21,6 +32,12 @@ print(f"Touched {i * 4096 / (1 << 30):.2f}G") last_at = time.time() +try: + sd_socket.sendall(b"READY=1") +except Exception: + print("Failed to send ready notification to systemd") + raise + print("Allocation done, sleeping...") while True: time.sleep(600)