[llvm] [lit] Export env vars in script to avoid pruning (PR #105759)

[Keith Smiley via llvm-commits]

https://github.com/keith updated https://github.com/llvm/llvm-project/pull/105759

>From 579bfe7c9467a16189f7695d82483b9d7b14befd Mon Sep 17 00:00:00 2001
From: Keith Smiley <keithbsmiley at gmail.com>
Date: Thu, 22 Aug 2024 17:55:50 -0700
Subject: [PATCH 1/5] [lit] Export env vars in script to avoid pruning

On macOS the dynamic loader prunes dyld specific environment variables
such as `DYLD_INSERT_LIBRARIES`, `DYLD_LIBRARY_PATH`, etc. If these are
set in the lit config it's safe to assume that the user actually wanted
their subprocesses to run with these variables, versus the python
interpreter that gets executed with them before they are pruned. This
change exports all known variables in the shell script instead of
relying on them being passed through.
---
 llvm/utils/lit/lit/TestRunner.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/llvm/utils/lit/lit/TestRunner.py b/llvm/utils/lit/lit/TestRunner.py
index 2d9af9fbbb3634..a5e0bdf41b478b 100644
--- a/llvm/utils/lit/lit/TestRunner.py
+++ b/llvm/utils/lit/lit/TestRunner.py
@@ -1227,6 +1227,11 @@ def executeScript(test, litConfig, tmpBase, commands, cwd):
         if test.config.pipefail:
             f.write(b"set -o pipefail;" if mode == "wb" else "set -o pipefail;")
         f.write(b"set -x;" if mode == "wb" else "set -x;")
+
+        env_str = "\n".join("export {}={};".format(k, shlex.quote(v))
+                            for k, v in test.config.environment.items())
+        f.write(bytes(env_str, "utf-8") if mode == "wb" else env_str)
+
         if sys.version_info > (3, 0) and mode == "wb":
             f.write(bytes("{ " + "; } &&\n{ ".join(commands) + "; }", "utf-8"))
         else:

>From 50d6de5310d298847eb3fa4e87aa9ba201aec4d1 Mon Sep 17 00:00:00 2001
From: Keith Smiley <keithbsmiley at gmail.com>
Date: Fri, 23 Aug 2024 01:02:43 +0000
Subject: [PATCH 2/5] Move set -x later

We don't really care to see the exports I don't think
---
 llvm/utils/lit/lit/TestRunner.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/utils/lit/lit/TestRunner.py b/llvm/utils/lit/lit/TestRunner.py
index a5e0bdf41b478b..21b2475c3c8b33 100644
--- a/llvm/utils/lit/lit/TestRunner.py
+++ b/llvm/utils/lit/lit/TestRunner.py
@@ -1226,11 +1226,11 @@ def executeScript(test, litConfig, tmpBase, commands, cwd):
                     commands[i] += f" && {{ {command}; }}"
         if test.config.pipefail:
             f.write(b"set -o pipefail;" if mode == "wb" else "set -o pipefail;")
-        f.write(b"set -x;" if mode == "wb" else "set -x;")
 
         env_str = "\n".join("export {}={};".format(k, shlex.quote(v))
                             for k, v in test.config.environment.items())
         f.write(bytes(env_str, "utf-8") if mode == "wb" else env_str)
+        f.write(b"set -x;" if mode == "wb" else "set -x;")
 
         if sys.version_info > (3, 0) and mode == "wb":
             f.write(bytes("{ " + "; } &&\n{ ".join(commands) + "; }", "utf-8"))

>From 975b0671a9ac674a5ebfde107bfc888f390a7974 Mon Sep 17 00:00:00 2001
From: Keith Smiley <keithbsmiley at gmail.com>
Date: Fri, 23 Aug 2024 01:11:13 +0000
Subject: [PATCH 3/5] fix python format

---
 llvm/utils/lit/lit/TestRunner.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/llvm/utils/lit/lit/TestRunner.py b/llvm/utils/lit/lit/TestRunner.py
index 21b2475c3c8b33..7079dc46419ca4 100644
--- a/llvm/utils/lit/lit/TestRunner.py
+++ b/llvm/utils/lit/lit/TestRunner.py
@@ -1227,8 +1227,10 @@ def executeScript(test, litConfig, tmpBase, commands, cwd):
         if test.config.pipefail:
             f.write(b"set -o pipefail;" if mode == "wb" else "set -o pipefail;")
 
-        env_str = "\n".join("export {}={};".format(k, shlex.quote(v))
-                            for k, v in test.config.environment.items())
+        env_str = "\n".join(
+            "export {}={};".format(k, shlex.quote(v))
+            for k, v in test.config.environment.items()
+        )
         f.write(bytes(env_str, "utf-8") if mode == "wb" else env_str)
         f.write(b"set -x;" if mode == "wb" else "set -x;")
 

>From b172924c03633b8594916b9c8f676f4ee070aff7 Mon Sep 17 00:00:00 2001
From: Keith Smiley <keithbsmiley at gmail.com>
Date: Fri, 23 Aug 2024 22:10:59 +0000
Subject: [PATCH 4/5] scope to only DYLD env vars

---
 llvm/utils/lit/lit/TestRunner.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/llvm/utils/lit/lit/TestRunner.py b/llvm/utils/lit/lit/TestRunner.py
index 7079dc46419ca4..244038cbf58f51 100644
--- a/llvm/utils/lit/lit/TestRunner.py
+++ b/llvm/utils/lit/lit/TestRunner.py
@@ -1227,9 +1227,13 @@ def executeScript(test, litConfig, tmpBase, commands, cwd):
         if test.config.pipefail:
             f.write(b"set -o pipefail;" if mode == "wb" else "set -o pipefail;")
 
+        # Manually export any DYLD_* variables used by dyld on macOS because
+        # otherwise they are lost when the shell executable is run, before the
+        # lit test is executed.
         env_str = "\n".join(
             "export {}={};".format(k, shlex.quote(v))
             for k, v in test.config.environment.items()
+            if k.startswith("DYLD_")
         )
         f.write(bytes(env_str, "utf-8") if mode == "wb" else env_str)
         f.write(b"set -x;" if mode == "wb" else "set -x;")

>From 0262886ebf92cc7eb54b44a740fbda1529c65dad Mon Sep 17 00:00:00 2001
From: Keith Smiley <keithbsmiley at gmail.com>
Date: Sat, 24 Aug 2024 09:12:02 -0700
Subject: [PATCH 5/5] remove extra whitespace

---
 llvm/utils/lit/lit/TestRunner.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/llvm/utils/lit/lit/TestRunner.py b/llvm/utils/lit/lit/TestRunner.py
index 244038cbf58f51..4dad1412436d93 100644
--- a/llvm/utils/lit/lit/TestRunner.py
+++ b/llvm/utils/lit/lit/TestRunner.py
@@ -1237,7 +1237,6 @@ def executeScript(test, litConfig, tmpBase, commands, cwd):
         )
         f.write(bytes(env_str, "utf-8") if mode == "wb" else env_str)
         f.write(b"set -x;" if mode == "wb" else "set -x;")
-
         if sys.version_info > (3, 0) and mode == "wb":
             f.write(bytes("{ " + "; } &&\n{ ".join(commands) + "; }", "utf-8"))
         else: