From 988d225cf01be202293634cc440f750a7bc53261 Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard@gmail.com>
Date: Thu, 24 Nov 2022 22:09:32 +0100
Subject: [PATCH] unix,win: add uv_get_available_memory() (#3754)

---
 docs/src/misc.rst      |  16 +++-
 include/uv.h           |   1 +
 src/unix/aix.c         |   5 ++
 src/unix/cygwin.c      |   4 +
 src/unix/darwin.c      |   5 ++
 src/unix/freebsd.c     |   5 ++
 src/unix/haiku.c       |   5 ++
 src/unix/hurd.c        |   5 ++
 src/unix/ibmi.c        |   5 ++
 src/unix/linux.c       | 180 ++++++++++++++++++++++++++++++++++-------
 src/unix/netbsd.c      |   5 ++
 src/unix/openbsd.c     |   5 ++
 src/unix/os390.c       |   5 ++
 src/unix/qnx.c         |   5 ++
 src/unix/sunos.c       |   5 ++
 src/win/util.c         |   5 ++
 test/test-get-memory.c |  13 ++-
 17 files changed, 242 insertions(+), 32 deletions(-)

diff --git a/docs/src/misc.rst b/docs/src/misc.rst
index a18040e2..39241df8 100644
--- a/docs/src/misc.rst
+++ b/docs/src/misc.rst
@@ -561,7 +561,7 @@ API
 
 .. c:function:: uint64_t uv_get_constrained_memory(void)
 
-    Gets the amount of memory available to the process (in bytes) based on
+    Gets the total amount of memory available to the process (in bytes) based on
     limits imposed by the OS. If there is no such constraint, or the constraint
     is unknown, `0` is returned. Note that it is not unusual for this value to
     be less than or greater than :c:func:`uv_get_total_memory`.
@@ -572,6 +572,20 @@ API
 
     .. versionadded:: 1.29.0
 
+.. c:function:: uint64_t uv_get_available_memory(void)
+
+    Gets the amount of free memory that is still available to the process (in bytes).
+    This differs from :c:func:`uv_get_free_memory` in that it takes into account any
+    limits imposed by the OS. If there is no such constraint, or the constraint
+    is unknown, the amount returned will be identical to :c:func:`uv_get_free_memory`.
+
+    .. note::
+        This function currently only returns a value that is different from
+        what :c:func:`uv_get_free_memory` reports on Linux, based
+        on cgroups if it is present.
+
+    .. versionadded:: 1.45.0
+
 .. c:function:: uint64_t uv_hrtime(void)
 
     Returns the current high-resolution real time. This is expressed in
diff --git a/include/uv.h b/include/uv.h
index 30e09c61..7189951f 100644
--- a/include/uv.h
+++ b/include/uv.h
@@ -1718,6 +1718,7 @@ UV_EXTERN int uv_chdir(const char* dir);
 UV_EXTERN uint64_t uv_get_free_memory(void);
 UV_EXTERN uint64_t uv_get_total_memory(void);
 UV_EXTERN uint64_t uv_get_constrained_memory(void);
+UV_EXTERN uint64_t uv_get_available_memory(void);
 
 UV_EXTERN uint64_t uv_hrtime(void);
 UV_EXTERN void uv_sleep(unsigned int msec);
diff --git a/src/unix/aix.c b/src/unix/aix.c
index 656869f7..b855282e 100644
--- a/src/unix/aix.c
+++ b/src/unix/aix.c
@@ -391,6 +391,11 @@ uint64_t uv_get_constrained_memory(void) {
 }
 
 
+uint64_t uv_get_available_memory(void) {
+  return uv_get_free_memory();
+}
+
+
 void uv_loadavg(double avg[3]) {
   perfstat_cpu_total_t ps_total;
   int result = perfstat_cpu_total(NULL, &ps_total, sizeof(ps_total), 1);
diff --git a/src/unix/cygwin.c b/src/unix/cygwin.c
index 169958d5..4e541396 100644
--- a/src/unix/cygwin.c
+++ b/src/unix/cygwin.c
@@ -51,3 +51,7 @@ int uv_cpu_info(uv_cpu_info_t** cpu_infos, int* count) {
 uint64_t uv_get_constrained_memory(void) {
   return 0;  /* Memory constraints are unknown. */
 }
+
+uint64_t uv_get_available_memory(void) {
+  return uv_get_free_memory();
+}
diff --git a/src/unix/darwin.c b/src/unix/darwin.c
index 1216770d..90790d70 100644
--- a/src/unix/darwin.c
+++ b/src/unix/darwin.c
@@ -131,6 +131,11 @@ uint64_t uv_get_constrained_memory(void) {
 }
 
 
+uint64_t uv_get_available_memory(void) {
+  return uv_get_free_memory();
+}
+
+
 void uv_loadavg(double avg[3]) {
   struct loadavg info;
   size_t size = sizeof(info);
diff --git a/src/unix/freebsd.c b/src/unix/freebsd.c
index 33a49878..b254b950 100644
--- a/src/unix/freebsd.c
+++ b/src/unix/freebsd.c
@@ -116,6 +116,11 @@ uint64_t uv_get_constrained_memory(void) {
 }
 
 
+uint64_t uv_get_available_memory(void) {
+  return uv_get_free_memory();
+}
+
+
 void uv_loadavg(double avg[3]) {
   struct loadavg info;
   size_t size = sizeof(info);
diff --git a/src/unix/haiku.c b/src/unix/haiku.c
index cf17d836..31284b66 100644
--- a/src/unix/haiku.c
+++ b/src/unix/haiku.c
@@ -84,6 +84,11 @@ uint64_t uv_get_constrained_memory(void) {
 }
 
 
+uint64_t uv_get_available_memory(void) {
+  return uv_get_free_memory();
+}
+
+
 int uv_resident_set_memory(size_t* rss) {
   area_info area;
   ssize_t cookie;
diff --git a/src/unix/hurd.c b/src/unix/hurd.c
index d19ea634..63c87812 100644
--- a/src/unix/hurd.c
+++ b/src/unix/hurd.c
@@ -165,3 +165,8 @@ int uv_cpu_info(uv_cpu_info_t** cpu_infos, int* count) {
 uint64_t uv_get_constrained_memory(void) {
   return 0;  /* Memory constraints are unknown. */
 }
+
+
+uint64_t uv_get_available_memory(void) {
+  return uv_get_free_memory();
+}
diff --git a/src/unix/ibmi.c b/src/unix/ibmi.c
index 8c6ae636..837bba6e 100644
--- a/src/unix/ibmi.c
+++ b/src/unix/ibmi.c
@@ -249,6 +249,11 @@ uint64_t uv_get_constrained_memory(void) {
 }
 
 
+uint64_t uv_get_available_memory(void) {
+  return uv_get_free_memory();
+}
+
+
 void uv_loadavg(double avg[3]) {
   SSTS0200 rcvr;
 
diff --git a/src/unix/linux.c b/src/unix/linux.c
index d2db848b..18dd5029 100644
--- a/src/unix/linux.c
+++ b/src/unix/linux.c
@@ -1343,47 +1343,169 @@ static uint64_t uv__read_uint64(const char* filename) {
 }
 
 
-/* This might return 0 if there was a problem getting the memory limit from
- * cgroups. This is OK because a return value of 0 signifies that the memory
- * limit is unknown.
- */
-static uint64_t uv__get_constrained_memory_fallback(void) {
-  return uv__read_uint64("/sys/fs/cgroup/memory/memory.limit_in_bytes");
+/* Given a buffer with the contents of a cgroup1 /proc/self/cgroups,
+ * finds the location and length of the memory controller mount path.
+ * This disregards the leading / for easy concatenation of paths.
+ * Returns NULL if the memory controller wasn't found. */
+static char* uv__cgroup1_find_memory_controller(char buf[static 1024],
+                                                int* n) {
+  char* p;
+
+  /* Seek to the memory controller line. */
+  p = strchr(buf, ':');
+  while (p != NULL && strncmp(p, ":memory:", 8)) {
+    p = strchr(p, '\n');
+    if (p != NULL)
+      p = strchr(p, ':');
+  }
+
+  if (p != NULL) {
+    /* Determine the length of the mount path. */
+    p = p + strlen(":memory:/");
+    *n = (int) strcspn(p, "\n");
+  }
+
+  return p;
+}
+
+static void uv__get_cgroup1_memory_limits(char buf[static 1024], uint64_t* high,
+                                          uint64_t* max) {
+  char filename[4097];
+  char* p;
+  int n;
+
+  /* Find out where the controller is mounted. */
+  p = uv__cgroup1_find_memory_controller(buf, &n);
+  if (p != NULL) {
+    snprintf(filename, sizeof(filename),
+             "/sys/fs/cgroup/memory/%.*s/memory.soft_limit_in_bytes", n, p);
+    *high = uv__read_uint64(filename);
+
+    snprintf(filename, sizeof(filename),
+             "/sys/fs/cgroup/memory/%.*s/memory.limit_in_bytes", n, p);
+    *max = uv__read_uint64(filename);
+
+    /* If the controller wasn't mounted, the reads above will have failed,
+     * as indicated by uv__read_uint64 returning 0.
+     */
+     if (*high != 0 && *max != 0)
+       return;
+  }
+
+  /* Fall back to the limits of the global memory controller. */
+  *high = uv__read_uint64("/sys/fs/cgroup/memory/memory.soft_limit_in_bytes");
+  *max = uv__read_uint64("/sys/fs/cgroup/memory/memory.limit_in_bytes");
+}
+
+static void uv__get_cgroup2_memory_limits(char buf[static 1024], uint64_t* high,
+                                          uint64_t* max) {
+  char filename[4097];
+  char* p;
+  int n;
+
+  /* Find out where the controller is mounted. */
+  p = buf + strlen("0::/");
+  n = (int) strcspn(p, "\n");
+
+  /* Read the memory limits of the controller. */
+  snprintf(filename, sizeof(filename), "/sys/fs/cgroup/%.*s/memory.max", n, p);
+  *max = uv__read_uint64(filename);
+  snprintf(filename, sizeof(filename), "/sys/fs/cgroup/%.*s/memory.high", n, p);
+  *high = uv__read_uint64(filename);
+}
+
+static uint64_t uv__get_cgroup_constrained_memory(char buf[static 1024]) {
+  uint64_t high;
+  uint64_t max;
+
+  /* In the case of cgroupv2, we'll only have a single entry. */
+  if (strncmp(buf, "0::/", 4))
+    uv__get_cgroup1_memory_limits(buf, &high, &max);
+  else
+    uv__get_cgroup2_memory_limits(buf, &high, &max);
+
+  if (high == 0 || max == 0)
+    return 0;
+
+  return high < max ? high : max;
+}
+
+uint64_t uv_get_constrained_memory(void) {
+  char buf[1024];
+
+  if (uv__slurp("/proc/self/cgroup", buf, sizeof(buf)))
+    return 0;
+
+  return uv__get_cgroup_constrained_memory(buf);
 }
 
 
-uint64_t uv_get_constrained_memory(void) {
+static uint64_t uv__get_cgroup1_current_memory(char buf[static 1024]) {
   char filename[4097];
-  char buf[1024];
-  uint64_t high;
-  uint64_t max;
+  uint64_t current;
   char* p;
+  int n;
+
+  /* Find out where the controller is mounted. */
+  p = uv__cgroup1_find_memory_controller(buf, &n);
+  if (p != NULL) {
+    snprintf(filename, sizeof(filename),
+            "/sys/fs/cgroup/memory/%.*s/memory.usage_in_bytes", n, p);
+    current = uv__read_uint64(filename);
+
+    /* If the controller wasn't mounted, the reads above will have failed,
+     * as indicated by uv__read_uint64 returning 0.
+     */
+    if (current != 0)
+      return current;
+  }
+
+  /* Fall back to the usage of the global memory controller. */
+  return uv__read_uint64("/sys/fs/cgroup/memory/memory.usage_in_bytes");
+}
+
+static uint64_t uv__get_cgroup2_current_memory(char buf[static 1024]) {
+  char filename[4097];
+  char* p;
+  int n;
+
+  /* Find out where the controller is mounted. */
+  p = buf + strlen("0::/");
+  n = (int) strcspn(p, "\n");
+
+  snprintf(filename, sizeof(filename),
+           "/sys/fs/cgroup/%.*s/memory.current", n, p);
+  return uv__read_uint64(filename);
+}
+
+uint64_t uv_get_available_memory(void) {
+  char buf[1024];
+  uint64_t constrained;
+  uint64_t current;
+  uint64_t total;
 
   if (uv__slurp("/proc/self/cgroup", buf, sizeof(buf)))
-    return uv__get_constrained_memory_fallback();
+    return 0;
 
-  if (memcmp(buf, "0::/", 4))
-    return uv__get_constrained_memory_fallback();
+  constrained = uv__get_cgroup_constrained_memory(buf);
+  if (constrained == 0)
+    return uv_get_free_memory();
 
-  p = strchr(buf, '\n');
-  if (p != NULL)
-    *p = '\0';
+  total = uv_get_total_memory();
+  if (constrained > total)
+    return uv_get_free_memory();
 
-  p = buf + 4;
+  /* In the case of cgroupv2, we'll only have a single entry. */
+  if (strncmp(buf, "0::/", 4))
+    current = uv__get_cgroup1_current_memory(buf);
+  else
+    current = uv__get_cgroup2_current_memory(buf);
 
-  snprintf(filename, sizeof(filename), "/sys/fs/cgroup/%s/memory.max", p);
-  max = uv__read_uint64(filename);
+  /* memory usage can be higher than the limit (for short bursts of time) */
+  if (constrained < current)
+    return 0;
 
-  if (max == 0)
-    return uv__get_constrained_memory_fallback();
-
-  snprintf(filename, sizeof(filename), "/sys/fs/cgroup/%s/memory.high", p);
-  high = uv__read_uint64(filename);
-
-  if (high == 0)
-    return uv__get_constrained_memory_fallback();
-
-  return high < max ? high : max;
+  return constrained - current;
 }
 
 
diff --git a/src/unix/netbsd.c b/src/unix/netbsd.c
index a19f5a01..fa21e98e 100644
--- a/src/unix/netbsd.c
+++ b/src/unix/netbsd.c
@@ -131,6 +131,11 @@ uint64_t uv_get_constrained_memory(void) {
 }
 
 
+uint64_t uv_get_available_memory(void) {
+  return uv_get_free_memory();
+}
+
+
 int uv_resident_set_memory(size_t* rss) {
   kvm_t *kd = NULL;
   struct kinfo_proc2 *kinfo = NULL;
diff --git a/src/unix/openbsd.c b/src/unix/openbsd.c
index 3674052f..9c863b6c 100644
--- a/src/unix/openbsd.c
+++ b/src/unix/openbsd.c
@@ -139,6 +139,11 @@ uint64_t uv_get_constrained_memory(void) {
 }
 
 
+uint64_t uv_get_available_memory(void) {
+  return uv_get_free_memory();
+}
+
+
 int uv_resident_set_memory(size_t* rss) {
   struct kinfo_proc kinfo;
   size_t page_size = getpagesize();
diff --git a/src/unix/os390.c b/src/unix/os390.c
index 315b8cd9..3954b2c2 100644
--- a/src/unix/os390.c
+++ b/src/unix/os390.c
@@ -198,6 +198,11 @@ uint64_t uv_get_constrained_memory(void) {
 }
 
 
+uint64_t uv_get_available_memory(void) {
+  return uv_get_free_memory();
+}
+
+
 int uv_resident_set_memory(size_t* rss) {
   char* ascb;
   char* rax;
diff --git a/src/unix/qnx.c b/src/unix/qnx.c
index ca148d34..57ea9dfd 100644
--- a/src/unix/qnx.c
+++ b/src/unix/qnx.c
@@ -88,6 +88,11 @@ uint64_t uv_get_constrained_memory(void) {
 }
 
 
+uint64_t uv_get_available_memory(void) {
+  return uv_get_free_memory();
+}
+
+
 int uv_resident_set_memory(size_t* rss) {
   int fd;
   procfs_asinfo asinfo;
diff --git a/src/unix/sunos.c b/src/unix/sunos.c
index afbc3532..75b6fbad 100644
--- a/src/unix/sunos.c
+++ b/src/unix/sunos.c
@@ -417,6 +417,11 @@ uint64_t uv_get_constrained_memory(void) {
 }
 
 
+uint64_t uv_get_available_memory(void) {
+  return uv_get_free_memory();
+}
+
+
 void uv_loadavg(double avg[3]) {
   (void) getloadavg(avg, 3);
 }
diff --git a/src/win/util.c b/src/win/util.c
index 23dd4136..a3bc1b98 100644
--- a/src/win/util.c
+++ b/src/win/util.c
@@ -362,6 +362,11 @@ uint64_t uv_get_constrained_memory(void) {
 }
 
 
+uint64_t uv_get_available_memory(void) {
+  return uv_get_free_memory();
+}
+
+
 uv_pid_t uv_os_getpid(void) {
   return GetCurrentProcessId();
 }
diff --git a/test/test-get-memory.c b/test/test-get-memory.c
index 4555ba08..9ac42c38 100644
--- a/test/test-get-memory.c
+++ b/test/test-get-memory.c
@@ -26,11 +26,14 @@ TEST_IMPL(get_memory) {
   uint64_t free_mem = uv_get_free_memory();
   uint64_t total_mem = uv_get_total_memory();
   uint64_t constrained_mem = uv_get_constrained_memory();
+  uint64_t available_mem = uv_get_available_memory();
 
-  printf("free_mem=%llu, total_mem=%llu, constrained_mem=%llu\n",
+  printf("free_mem=%llu, total_mem=%llu, constrained_mem=%llu, "
+         "available_mem=%llu\n",
          (unsigned long long) free_mem,
          (unsigned long long) total_mem,
-         (unsigned long long) constrained_mem);
+         (unsigned long long) constrained_mem,
+         (unsigned long long) available_mem);
 
   ASSERT(free_mem > 0);
   ASSERT(total_mem > 0);
@@ -40,5 +43,11 @@ TEST_IMPL(get_memory) {
 #else
   ASSERT(total_mem > free_mem);
 #endif
+  ASSERT_LE(available_mem, total_mem);
+  /* we'd really want to test if available <= free, but that is fragile:
+   * with no limit set, get_available calls and returns get_free; so if
+   * any memory was freed between our calls to get_free and get_available
+   * we would fail such a test test (as observed on CI).
+   */
   return 0;
 }