prefix output from worker's subprocesses too

Author: IanButterworth

test/runtests.jl

@@ -14,6 +14,51 @@ include("buildkitetestjson.jl")
 const longrunning_delay = parse(Int, get(ENV, "JULIA_TEST_LONGRUNNING_DELAY", "45")) * 60 # minutes
 const longrunning_interval = parse(Int, get(ENV, "JULIA_TEST_LONGRUNNING_INTERVAL", "15")) * 60 # minutes
 
+# IO wrapper that prefixes each line with a test name
+mutable struct PrefixedIO <: IO
+    io::IO
+    lock::ReentrantLock
+    prefix::Base.RefValue{String}
+    at_line_start::Bool
+end
+PrefixedIO(io::IO, lock::ReentrantLock, prefix::Base.RefValue{String}) = PrefixedIO(io, lock, prefix, true)
+
+function Base.unsafe_write(pio::PrefixedIO, p::Ptr{UInt8}, n::UInt)
+    @lock pio.lock begin
+        written = UInt(0)
+        i = 1
+        while i <= n
+            if pio.at_line_start && !isempty(pio.prefix[])
+                printstyled(pio.io, "  ", pio.prefix[], ": ", color=:light_black)
+                pio.at_line_start = false
+            end
+            # Find next newline
+            j = i
+            while j <= n && unsafe_load(p, j) != UInt8('\n')
+                j += 1
+            end
+            # Write up to and including newline (or end)
+            if j <= n
+                written += Base.unsafe_write(pio.io, p + i - 1, UInt(j - i + 1))
+                pio.at_line_start = true
+                i = j + 1
+            else
+                written += Base.unsafe_write(pio.io, p + i - 1, UInt(j - i))
+                i = j
+            end
+        end
+        return written
+    end
+end
+
+Base.write(pio::PrefixedIO, x::UInt8) = Base.unsafe_write(pio, Ref(x), 1)
+Base.flush(pio::PrefixedIO) = flush(pio.io)
+Base.isopen(pio::PrefixedIO) = isopen(pio.io)
+Base.displaysize(pio::PrefixedIO) = displaysize(pio.io)
+Base.get(pio::PrefixedIO, key, default) = get(pio.io, key, default)
+Base.get(pio::PrefixedIO, key::Symbol, default) = get(pio.io, key, default)
+Base.haskey(pio::PrefixedIO, key) = haskey(pio.io, key)
+
 (; tests, net_on, exit_on_error, use_revise, buildroot, seed) = choosetests(ARGS)
 tests = unique(tests)
 
@@ -163,6 +208,13 @@ cd(@__DIR__) do
         stderr.lock = print_lock
     end
 
+    # Prefix for test output (used by both worker hook and node1 tests)
+    test_output_prefix = Ref("")
+
+    # Create prefixed IO wrappers for node1 tests
+    prefixed_stdout = PrefixedIO(stdout, print_lock, test_output_prefix)
+    prefixed_stderr = PrefixedIO(stderr, print_lock, test_output_prefix)
+
     # Set up hook to display test name with worker output
     Distributed.worker_output_hook[] = (ident, line) -> begin
         wrkr_id = tryparse(Int, ident)
@@ -380,12 +432,20 @@ cd(@__DIR__) do
             isolate = true
             t == "SharedArrays" && (isolate = false)
             before = time()
+            # Set prefix and redirect stdout/stderr to prefixed wrappers
+            test_output_prefix[] = "$t (1)"
             resp, duration = try
-                    r = @invokelatest runtests(t, test_path(t), isolate, seed=seed) # runtests is defined by the include above
-                    r, time() - before
+                    redirect_stdout(prefixed_stdout) do
+                        redirect_stderr(prefixed_stderr) do
+                            r = @invokelatest runtests(t, test_path(t), isolate, seed=seed) # runtests is defined by the include above
+                            r, time() - before
+                        end
+                    end
                 catch e
                     isa(e, InterruptException) && rethrow()
                     Any[CapturedException(e, catch_backtrace())], time() - before
+                finally
+                    test_output_prefix[] = ""
                 end
             if length(resp) == 1
                 print_testworker_errored(t, 1, resp[1])