From 8e96174e1ff6f6118f0ef54407de92b2375acdf2 Mon Sep 17 00:00:00 2001
From: Yatao Li <yatao.li@live.com>
Date: Sun, 23 Jan 2022 03:35:27 +0800
Subject: [PATCH] scan memory threshold

---
 src/cli.c   |  6 ++++++
 src/cli.h   |  1 +
 src/ctx.h   |  1 +
 src/main.c  |  2 ++
 src/tpool.c | 36 ++++++++++++++++++++++++++++++++++++
 5 files changed, 46 insertions(+)

diff --git a/src/cli.c b/src/cli.c
index fc2b762..afd0163 100644
--- a/src/cli.c
+++ b/src/cli.c
@@ -19,6 +19,8 @@
 
 #define DEFAULT_MAX_MEM_BUFFER 2000
 
+#define DEFAULT_THROTTLE_MEMORY_THRESHOLD 0
+
 const char *TESS_DATAPATHS[] = {
         "/usr/share/tessdata/",
         "/usr/share/tesseract-ocr/tessdata/",
@@ -255,6 +257,10 @@ int scan_args_validate(scan_args_t *args, int argc, const char **argv) {
         args->max_memory_buffer = DEFAULT_MAX_MEM_BUFFER;
     }
 
+    if (args->throttle_memory_threshold <= 0) {
+        args->throttle_memory_threshold = DEFAULT_THROTTLE_MEMORY_THRESHOLD;
+    }
+
     if (args->list_path != NULL) {
         if (strcmp(args->list_path, "-") == 0) {
             args->list_file = stdin;
diff --git a/src/cli.h b/src/cli.h
index a32d40a..1cd92ba 100644
--- a/src/cli.h
+++ b/src/cli.h
@@ -10,6 +10,7 @@ typedef struct scan_args {
     int size;
     int content_size;
     int threads;
+    int throttle_memory_threshold;
     char *incremental;
     char *output;
     char *rewrite_url;
diff --git a/src/ctx.h b/src/ctx.h
index fb55998..6d6b224 100644
--- a/src/ctx.h
+++ b/src/ctx.h
@@ -35,6 +35,7 @@ typedef struct {
     int threads;
     int depth;
     int calculate_checksums;
+    size_t mem_limit;
 
     size_t stat_tn_size;
     size_t stat_index_size;
diff --git a/src/main.c b/src/main.c
index a1b567f..669037f 100644
--- a/src/main.c
+++ b/src/main.c
@@ -253,6 +253,7 @@ void initialize_scan_context(scan_args_t *args) {
 
     ScanCtx.threads = args->threads;
     ScanCtx.depth = args->depth;
+    ScanCtx.mem_limit = args->throttle_memory_threshold * 1024 * 1024;
 
     strncpy(ScanCtx.index.path, args->output, sizeof(ScanCtx.index.path));
     strncpy(ScanCtx.index.desc.name, args->name, sizeof(ScanCtx.index.desc.name));
@@ -586,6 +587,7 @@ int main(int argc, const char *argv[]) {
 
             OPT_GROUP("Scan options"),
             OPT_INTEGER('t', "threads", &common_threads, "Number of threads. DEFAULT=1"),
+            OPT_STRING(0, "mem-throttle", &scan_args->throttle_memory_threshold, "Total memory threshold in MB for scan throttling. DEFAULT=0"),
             OPT_FLOAT('q', "quality", &scan_args->quality,
                       "Thumbnail quality, on a scale of 1.0 to 31.0, 1.0 being the best. DEFAULT=3"),
             OPT_INTEGER(0, "size", &scan_args->size,
diff --git a/src/tpool.c b/src/tpool.c
index e0de5e8..be2fff5 100644
--- a/src/tpool.c
+++ b/src/tpool.c
@@ -114,6 +114,35 @@ int tpool_add_work(tpool_t *pool, thread_func_t func, void *arg) {
     return 1;
 }
 
+/**
+ * see: https://github.com/htop-dev/htop/blob/f782f821f7f8081cb43bbad1c37f32830a260a81/linux/LinuxProcessList.c
+ */
+__always_inline
+static size_t _get_total_mem_mb() {
+    FILE* statmfile = fopen("/proc/self/statm", "r");
+    if (!statmfile)
+      return 0;
+
+    long int dummy, dummy2, dummy3, dummy4, dummy5, dummy6;
+    long int m_resident;
+
+    int r = fscanf(statmfile, "%ld %ld %ld %ld %ld %ld %ld",
+        &dummy, /* m_virt */
+        &m_resident,
+        &dummy2, /* m_share */
+        &dummy3, /* m_trs */
+        &dummy4, /* unused since Linux 2.6; always 0 */
+        &dummy5, /* m_drs */
+        &dummy6); /* unused since Linux 2.6; always 0 */
+    fclose(statmfile);
+
+    if (r == 7) {
+        return m_resident * 4 / 1024; // XXX assume 4KB pages.
+    } else {
+        return 0;
+    }
+}
+
 /**
  * Thread worker function
  */
@@ -142,6 +171,13 @@ static void *tpool_worker(void *arg) {
                 break;
             }
 
+            while(ScanCtx.mem_limit > 0 && _get_total_mem_mb() >= ScanCtx.mem_limit) {
+                if (pool->stop) {
+                    break;
+                }
+                usleep(10000);
+            }
+
             work->func(work->arg);
             if (pool->free_arg) {
                 free(work->arg);