[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 09671451af93ebe468971bc3e3b9882d09b7f354 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/4] [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 da7fa86fd39173..7ea9fb9594cc32 100644
--- a/llvm/utils/lit/lit/TestRunner.py
+++ b/llvm/utils/lit/lit/TestRunner.py
@@ -1223,6 +1223,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 f6a256f0607dd2c0565dc52e62c0c8db0bfc9350 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/4] 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 7ea9fb9594cc32..d8972314628222 100644
--- a/llvm/utils/lit/lit/TestRunner.py
+++ b/llvm/utils/lit/lit/TestRunner.py
@@ -1222,11 +1222,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 c6b50d827868e419b9efdc438a7503cca1d88819 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/4] 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 d8972314628222..7fadb0775e3582 100644
--- a/llvm/utils/lit/lit/TestRunner.py
+++ b/llvm/utils/lit/lit/TestRunner.py
@@ -1223,8 +1223,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 268e4b19915fe668cc2cf583a561c0ff5e8a2569 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/4] 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;")