Merge pull request #299 from facebookexperimental/wait_balloon

Wait for memory balloon allocation

Author: lnyng

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(),

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)