From 0066b862bb3a09f39295abd5d972a53ac8dc1555 Mon Sep 17 00:00:00 2001
From: Peter Bergin <peter@berginkonsult.se>
Date: Wed, 19 Sep 2018 15:12:31 +0200
Subject: [PATCH] rpm/rpmio.c: restrict virtual memory usage if limit set

A solution to avoid OOM situation when the virtual memory is restricted
for a user (ulimit -v). As the lzopen_internal function is run in parallel
one instance per CPU thread the available virtual memory is limited per
CPU thread.

Upstream-Status: Pending [merge of multithreading patches to upstream]

Signed-off-by: Peter Bergin <peter@berginkonsult.se>
Signed-off-by: Ranjitsinh Rathod <ranjitsinh.rathod@kpit.com>
---
 rpmio/rpmio.c | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/rpmio/rpmio.c b/rpmio/rpmio.c
index e051c98..b3c56b6 100644
--- a/rpmio/rpmio.c
+++ b/rpmio/rpmio.c
@@ -845,6 +845,42 @@ static LZFILE *lzopen_internal(const char *mode, int fd, int xz)
 		}
 #endif
 
+		struct rlimit virtual_memory = {RLIM_INFINITY , RLIM_INFINITY};
+		int status = getrlimit(RLIMIT_AS, &virtual_memory);
+		if ((status != -1) && (virtual_memory.rlim_cur != RLIM_INFINITY)) {
+			const uint64_t virtual_memlimit = virtual_memory.rlim_cur;
+			uint32_t threads_max = lzma_cputhreads();
+			const uint64_t virtual_memlimit_per_cpu_thread =
+				virtual_memlimit / ((threads_max == 0) ? 1 : threads_max);
+			rpmlog(RPMLOG_NOTICE, "XZ: virtual memory restricted to %lu and "
+			       "per CPU thread %lu\n", virtual_memlimit, virtual_memlimit_per_cpu_thread);
+			uint64_t memory_usage_virt;
+			/* keep reducing the number of compression threads until memory
+			   usage falls below the limit per CPU thread*/
+			while ((memory_usage_virt = lzma_stream_encoder_mt_memusage(&mt_options)) >
+			       virtual_memlimit_per_cpu_thread) {
+				/* If number of threads goes down to zero or in case of any other error
+				 * lzma_stream_encoder_mt_memusage will return UINT64_MAX. We must check
+				 * for both the cases here to avoid an infinite loop.
+				 * If we get into situation that one thread requires more virtual memory
+				 * than available we set one thread, print error message and try anyway. */
+				if ((--mt_options.threads == 0) || (memory_usage_virt == UINT64_MAX)) {
+					mt_options.threads = 1;
+					rpmlog(RPMLOG_WARNING,
+					       "XZ: Could not adjust number of threads to get below "
+					       "virtual memory limit %lu. usage %lu\n",
+					       virtual_memlimit_per_cpu_thread, memory_usage_virt);
+					break;
+				}
+			}
+			if (threads != (int)mt_options.threads)
+				rpmlog(RPMLOG_NOTICE,
+				       "XZ: Adjusted the number of threads from %d to %d to not "
+				       "exceed the memory usage limit of %lu bytes\n",
+				       threads, mt_options.threads, virtual_memlimit);
+
+		}
+
 		ret = lzma_stream_encoder_mt(&lzfile->strm, &mt_options);
 	    }
 #endif
-- 
2.7.4