From 4c97a940cab8eeb97631d6ba9834a37b512fc423 Mon Sep 17 00:00:00 2001
From: Jesse Gross <jesse@kernel.org>
Date: Wed, 27 May 2026 13:07:33 -0700
Subject: [PATCH] mlxthread: preserve the original stack when worker work
 panics

Work that panics on the locked MLX worker goroutine was recovered and
re-raised on the caller, so the printed trace pointed at the re-panic
site in this package rather than the code that actually panicked.

Capture the worker stack at recovery and carry it through a value that
implements error, so the runtime prints the original location in the
fatal trace.
---
 x/internal/mlxthread/thread.go      | 35 ++++++++++++++++++++---------
 x/internal/mlxthread/thread_test.go | 13 +++++++++--
 2 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/x/internal/mlxthread/thread.go b/x/internal/mlxthread/thread.go
index 5c041931..3eea2287 100644
--- a/x/internal/mlxthread/thread.go
+++ b/x/internal/mlxthread/thread.go
@@ -3,7 +3,9 @@ package mlxthread
 import (
 	"context"
 	"errors"
+	"fmt"
 	"runtime"
+	"runtime/debug"
 	"sync/atomic"
 )
 
@@ -24,8 +26,21 @@ type job struct {
 }
 
 type result struct {
-	err        error
-	panicValue any
+	err   error
+	panic *panicError
+}
+
+// panicError carries a value recovered from the worker goroutine together with
+// the stack captured at recovery, before the original stack unwinds. Because it
+// implements error, re-panicking with it makes the runtime print the original
+// worker location in the fatal trace instead of this package's re-panic site.
+type panicError struct {
+	value any
+	stack []byte
+}
+
+func (p *panicError) Error() string {
+	return fmt.Sprintf("%v\n\nmlx worker stack:\n%s", p.value, p.stack)
 }
 
 // Start creates a long-lived worker goroutine locked to one OS thread.
@@ -40,8 +55,8 @@ func Start(name string, init func() error) (*Thread, error) {
 	go t.loop(init, initResult)
 
 	res := <-initResult
-	if res.panicValue != nil {
-		panic(res.panicValue)
+	if res.panic != nil {
+		panic(res.panic)
 	}
 	if res.err != nil {
 		return nil, res.err
@@ -60,8 +75,8 @@ func (t *Thread) Do(ctx context.Context, fn func() error) error {
 	if err != nil {
 		return err
 	}
-	if res.panicValue != nil {
-		panic(res.panicValue)
+	if res.panic != nil {
+		panic(res.panic)
 	}
 	return res.err
 }
@@ -101,8 +116,8 @@ func (t *Thread) Stop(ctx context.Context, cleanup func()) error {
 		}
 		return err
 	}
-	if res.panicValue != nil {
-		panic(res.panicValue)
+	if res.panic != nil {
+		panic(res.panic)
 	}
 	if res.err != nil {
 		return res.err
@@ -123,7 +138,7 @@ func (t *Thread) loop(init func() error, initResult chan<- result) {
 
 	res := run(init)
 	initResult <- res
-	if res.err != nil || res.panicValue != nil {
+	if res.err != nil || res.panic != nil {
 		close(t.done)
 		return
 	}
@@ -166,7 +181,7 @@ func (t *Thread) enqueue(ctx context.Context, fn func() error, stop, allowStoppi
 func run(fn func() error) (res result) {
 	defer func() {
 		if v := recover(); v != nil {
-			res.panicValue = v
+			res.panic = &panicError{value: v, stack: debug.Stack()}
 		}
 	}()
 	if fn != nil {
diff --git a/x/internal/mlxthread/thread_test.go b/x/internal/mlxthread/thread_test.go
index 9d511351..d15ea599 100644
--- a/x/internal/mlxthread/thread_test.go
+++ b/x/internal/mlxthread/thread_test.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"reflect"
 	"runtime"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -42,8 +43,16 @@ func TestDoPropagatesPanicToCaller(t *testing.T) {
 	defer thread.Stop(context.Background(), nil)
 
 	defer func() {
-		if got := recover(); got != "boom" {
-			t.Fatalf("got panic %v, want boom", got)
+		got := recover()
+		pe, ok := got.(*panicError)
+		if !ok {
+			t.Fatalf("got panic %T (%v), want *panicError", got, got)
+		}
+		if pe.value != "boom" {
+			t.Fatalf("got panic value %v, want boom", pe.value)
+		}
+		if !strings.Contains(string(pe.stack), "TestDoPropagatesPanicToCaller") {
+			t.Fatalf("captured stack missing original frame:\n%s", pe.stack)
 		}
 	}()