new

Change-Id: I43f2dbd678860136f67e7f5178c122456514ef0c

Author: qiulaidongfeng

src/runtime/crash_test.go

@@ -19,6 +19,7 @@ import (
 	"path/filepath"
 	"regexp"
 	"runtime"
+	"runtime/mainthread"
 	"runtime/trace"
 	"strings"
 	"sync"
@@ -43,7 +44,22 @@ func TestMain(m *testing.M) {
 
 	_, coreErrBefore := os.Stat("core")
 
-	status := m.Run()
+	sch := make(chan int)
+	go func() {
+		sch <- m.Run()
+	}()
+	var status int
+loop:
+	for {
+		select {
+		case s := <-sch:
+			status = s
+			break loop
+		case <-mainthread.Waiting():
+			mainthread.Yield()
+		}
+	}
+
 	for _, file := range toRemove {
 		os.RemoveAll(file)
 	}

src/runtime/export_test.go

@@ -1435,10 +1435,14 @@ func Escape[T any](x T) T {
 }
 
 // Acquirem blocks preemption.
-func Acquirem() {
-	acquirem()
+func Acquirem() *M {
+	return acquirem()
 }
 
+type M = m
+
+var M0 = &m0
+
 func Releasem() {
 	releasem(getg().m)
 }

src/runtime/proc.go

@@ -115,14 +115,18 @@ var modinfo string
 
 var (
 	m0           m
-	m0func       = make(chan func())
-	waitm0       = make(chan struct{})
 	g0           g
 	mcache0      *mcache
 	raceprocctx0 uintptr
 	raceFiniLock mutex
 )
 
+var (
+	m0func = make(chan func())
+	waitm0 = make(chan struct{})
+	m0Exec = make(chan struct{}, 1)
+)
+
 // This slice records the initializing tasks that need to be
 // done to start up the runtime. It is built by the linker.
 var runtime_inittasks []*initTask
@@ -153,13 +157,21 @@ func mainThreadDo(f func()) {
 	}
 	waitm0 <- struct{}{}
 	m0func <- f
+	_ = <-m0Exec
 }
 
 func mainThreadYield() {
 	g := getg()
 	if g.m != &m0 {
 		panic("runtime: call mainthread.Yield must on main thread")
 	}
+	// lock os thread ensure that the main thread always
+	// run only f during a call to f.
+	lockOSThread()
+	defer func() {
+		unlockOSThread()
+		m0Exec <- struct{}{}
+	}()
 	f := <-m0func
 	f()
 }

src/runtime/proc_test.go

@@ -12,6 +12,7 @@ import (
 	"net"
 	"runtime"
 	"runtime/debug"
+	"runtime/mainthread"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -1161,4 +1162,30 @@ func TestBigGOMAXPROCS(t *testing.T) {
 
 func TestMainThread(t *testing.T) {
 	checkCoroTestProgOutput(t, runTestProg(t, "testprog", "MainThread"))
+	t.Run("in mainthread call UnlockOSThread", func(t *testing.T) {
+		mainthread.Do(func() {
+			runtime.UnlockOSThread()
+			for range 1000 {
+				runtime.Gosched()
+				m := runtime.Acquirem()
+				runtime.Releasem()
+				if m != runtime.M0 {
+					// don`t use t.Fatal, because it call
+					// Goexit, cause TestMain goroutine exit,
+					// and test timeout.
+					t.Fail()
+					t.Log("mainthread.Do.f must on main thread")
+					return
+				}
+			}
+			runtime.LockOSThread()
+		})
+	})
+	t.Run("nested call mainthread.Do", func(t *testing.T) {
+		checkCoroTestProgOutput(t, runTestProg(t, "testprog", "MainThread2"))
+	})
+}
+
+func init() {
+	runtime.LockOSThread()
 }

src/runtime/testdata/testprog/mainthread.go

@@ -28,4 +28,26 @@ func init() {
 		println("expect: Ok")
 		MainThread()
 	})
+	register("MainThread2", func() {
+		println("expect: hello,world")
+		MainThread2()
+	})
+}
+
+func MainThread2() {
+	var wg sync.WaitGroup
+	runtime.LockOSThread()
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		mainthread.Do(func() {
+			print("hello,")
+			mainthread.Do(func() {
+				print("world")
+			})
+		})
+	}()
+	<-mainthread.Waiting()
+	mainthread.Yield()
+	wg.Wait()
 }