From 69a5227098b87b048a90ceb1ca7b169c02ba151e Mon Sep 17 00:00:00 2001
From: Tavian Barnes <tavianator@tavianator.com>
Date: Thu, 20 May 2021 13:03:31 -0400
Subject: eval: Raise RLIMIT_NOFILE if possible

This lets us keep more open FDs cached in bftw().  The limit is lowered
before running -exec commands, in case they're incompatible with a high
limit (e.g. due to select()).
---
 ctx.c  |  9 +++++++++
 ctx.h  |  6 ++++++
 eval.c | 48 +++++++++++++++++++++++++++++++++++++-----------
 exec.c | 10 ++++++++++
 4 files changed, 62 insertions(+), 11 deletions(-)

diff --git a/ctx.c b/ctx.c
index 72ed969..b17d148 100644
--- a/ctx.c
+++ b/ctx.c
@@ -93,6 +93,15 @@ struct bfs_ctx *bfs_ctx_new(void) {
 	trie_init(&ctx->files);
 	ctx->nfiles = 0;
 
+	struct rlimit rl;
+	if (getrlimit(RLIMIT_NOFILE, &rl) == 0) {
+		ctx->nofile_soft = rl.rlim_cur;
+		ctx->nofile_hard = rl.rlim_max;
+	} else {
+		ctx->nofile_soft = 1024;
+		ctx->nofile_hard = RLIM_INFINITY;
+	}
+
 	return ctx;
 }
 
diff --git a/ctx.h b/ctx.h
index 35eee76..7af4060 100644
--- a/ctx.h
+++ b/ctx.h
@@ -24,6 +24,7 @@
 #include "bftw.h"
 #include "trie.h"
 #include <stdbool.h>
+#include <sys/resource.h>
 
 /**
  * Various debugging flags.
@@ -117,6 +118,11 @@ struct bfs_ctx {
 	struct trie files;
 	/** The number of files owned by the context. */
 	int nfiles;
+
+	/** The initial RLIMIT_NOFILE soft limit. */
+	rlim_t nofile_soft;
+	/** The initial RLIMIT_NOFILE hard limit. */
+	rlim_t nofile_hard;
 };
 
 /**
diff --git a/eval.c b/eval.c
index ba2a870..c638f10 100644
--- a/eval.c
+++ b/eval.c
@@ -1413,19 +1413,42 @@ done:
 	return state.action;
 }
 
-/**
- * Infer the number of open file descriptors we're allowed to have.
- */
-static int infer_fdlimit(const struct bfs_ctx *ctx) {
-	int ret = 4096;
+/** Compare two rlimit values, accounting for RLIM_INFINITY etc. */
+static int rlim_cmp(rlim_t a, rlim_t b) {
+	// Consider RLIM_{INFINITY,SAVED_{CUR,MAX}} all equally infinite
+	bool a_inf = a == RLIM_INFINITY || a == RLIM_SAVED_CUR || a == RLIM_SAVED_MAX;
+	bool b_inf = b == RLIM_INFINITY || b == RLIM_SAVED_CUR || b == RLIM_SAVED_MAX;
+	if (a_inf || b_inf) {
+		return a_inf - b_inf;
+	}
+
+	return (a > b) - (a < b);
+}
+
+/** Raise RLIMIT_NOFILE if possible, and return the new limit. */
+static int raise_fdlimit(const struct bfs_ctx *ctx) {
+	rlim_t target = 64 << 10;
+	if (rlim_cmp(target, ctx->nofile_hard) > 0) {
+		target = ctx->nofile_hard;
+	}
+
+	int ret = target;
 
-	struct rlimit rl;
-	if (getrlimit(RLIMIT_NOFILE, &rl) == 0) {
-		if (rl.rlim_cur != RLIM_INFINITY) {
-			ret = rl.rlim_cur;
+	if (rlim_cmp(target, ctx->nofile_soft) > 0) {
+		const struct rlimit rl = {
+			.rlim_cur = target,
+			.rlim_max = ctx->nofile_hard,
+		};
+		if (setrlimit(RLIMIT_NOFILE, &rl) != 0) {
+			ret = ctx->nofile_soft;
 		}
 	}
 
+	return ret;
+}
+
+/** Infer the number of file descriptors available to bftw(). */
+static int infer_fdlimit(const struct bfs_ctx *ctx, int limit) {
 	// 3 for std{in,out,err}
 	int nopen = 3 + ctx->nfiles;
 
@@ -1446,7 +1469,7 @@ static int infer_fdlimit(const struct bfs_ctx *ctx) {
 		bfs_closedir(dir);
 	}
 
-	ret -= nopen;
+	int ret = limit - nopen;
 	ret -= ctx->expr->persistent_fds;
 	ret -= ctx->expr->ephemeral_fds;
 
@@ -1512,12 +1535,15 @@ int bfs_eval(const struct bfs_ctx *ctx) {
 		args.seen = &seen;
 	}
 
+	int fdlimit = raise_fdlimit(ctx);
+	fdlimit = infer_fdlimit(ctx, fdlimit);
+
 	struct bftw_args bftw_args = {
 		.paths = ctx->paths,
 		.npaths = darray_length(ctx->paths),
 		.callback = eval_callback,
 		.ptr = &args,
-		.nopenfd = infer_fdlimit(ctx),
+		.nopenfd = fdlimit,
 		.flags = ctx->flags,
 		.strategy = ctx->strategy,
 		.mtab = bfs_ctx_mtab(ctx),
diff --git a/exec.c b/exec.c
index 2856071..bc01808 100644
--- a/exec.c
+++ b/exec.c
@@ -30,6 +30,7 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include <sys/resource.h>
 #include <sys/wait.h>
 #include <unistd.h>
 
@@ -358,6 +359,15 @@ static int bfs_exec_spawn(const struct bfs_exec *execbuf) {
 		goto fail;
 	}
 
+	// Reset RLIMIT_NOFILE, to avoid breaking applications that use select()
+	struct rlimit rl = {
+		.rlim_cur = execbuf->ctx->nofile_soft,
+		.rlim_max = execbuf->ctx->nofile_hard,
+	};
+	if (bfs_spawn_addsetrlimit(&ctx, RLIMIT_NOFILE, &rl) != 0) {
+		goto fail;
+	}
+
 	if (execbuf->wd_fd >= 0) {
 		if (bfs_spawn_addfchdir(&ctx, execbuf->wd_fd) != 0) {
 			goto fail;
-- 
cgit v1.2.3