If f calls Do, panic

Change-Id: Ib2674eca814b56db5408009cefc1f8544f5a717c

Author: qiulaidongfeng

src/runtime/mainthread/mainthread.go

@@ -13,25 +13,22 @@
 // yield the main thread temporarily to other [Do] calls by calling [Yield].
 //
 // Each package's initialization functions always run on the main thread,
-// as if by successive calls to Do(init).
+// so call Do in init is unnecessary.
 //
 // For compatibility with earlier versions of Go, if an init function calls [runtime.LockOSThread],
-// then package main's func main also runs on the main thread, as if by Do(main).
+// then package main's func main also runs on the main thread,
+// until the thread is unlocked using [runtime.UnlockOSThread].
 // In this situation, main must explicitly yield the main thread
-// to allow other calls to Do are to proceed.
+// to allow other thread calls to Do are to proceed.
 // See the documentation for [Waiting] for examples.
-package mainthread // imported as "runtime/mainthread"
+package mainthread
 
 import _ "unsafe"
 
 // Do calls f on the main thread.
 // Nothing else runs on the main thread until f returns or calls [Yield].
 // If f calls Do, the nested call panics.
 //
-// Package initialization functions run as if by Do(init).
-// If an init function calls [runtime.LockOSThread], then package main's func main
-// runs as if by Do(main), until the thread is unlocked using [runtime.UnlockOSThread].
-//
 //go:linkname Do runtime.mainThreadDo
 func Do(f func())
 

src/runtime/proc.go

@@ -115,18 +115,57 @@ var modinfo string
 
 var (
 	m0           m
+	mtPtr        atomic.Pointer[mainThreadOnce]
 	g0           g
 	mcache0      *mcache
 	raceprocctx0 uintptr
 	raceFiniLock mutex
 )
 
+func init() {
+	mtPtr.Store(newmainThreadInfo())
+}
+
+type mainThreadOnce struct {
+	//m0once make sure to only send m0wait once if mainthread.Yield is not called.
+	m0once Once
+}
+
+// copy from sync package.
+type Once struct {
+	done atomic.Uint32
+	m    mutex
+}
+
+func (o *Once) Do(f func()) {
+	if o.done.Load() == 0 {
+		o.doSlow(f)
+	}
+}
+
+func (o *Once) doSlow(f func()) {
+	lock(&o.m)
+	defer unlock(&o.m)
+	if o.done.Load() == 0 {
+		defer o.done.Store(1)
+		f()
+	}
+}
+
 var (
+	// m0func pass f from a call to mainthread.Do on non-main thread to the main thread.
 	m0func = make(chan func())
-	waitm0 = make(chan struct{})
-	m0Exec = make(chan struct{}, 1)
+	// m0wait send a signal that the non-main thread is waiting for mainthread.Yield.
+	m0wait = make(chan struct{})
+	// m0exec notifies mainthread.Do when the f passed from Do on the non-main thread
+	// to Yield on the main thread has completed.
+	m0exec = make(chan struct{}, 1)
 )
 
+func newmainThreadInfo() *mainThreadOnce {
+	return new(mainThreadOnce)
+}
+
 // 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
@@ -150,38 +189,68 @@ var runtimeInitTime int64
 var initSigmask sigset
 
 func mainThreadDo(f func()) {
-	gp := getg()
-	if gp.m == &m0 {
+	g := getg()
+	if g.inMainThradDo {
+		panic("runtime: nested call mainthread.Do")
+	}
+	g.inMainThradDo = true
+	if g.m == &m0 {
 		// lock os thread ensure that the main thread always
 		// run only f during a call to f.
 		lockOSThread()
 		defer unlockOSThread()
 		f()
 		return
 	}
-	waitm0 <- struct{}{}
+	mt := mtPtr.Load()
+	mt.m0once.Do(func() {
+		m0wait <- struct{}{}
+	})
 	m0func <- f
-	_ = <-m0Exec
+	_ = <-m0exec
 }
 
 func mainThreadYield() {
 	g := getg()
 	if g.m != &m0 {
 		panic("runtime: call mainthread.Yield must on main thread")
 	}
+	if g.inMainThradDo {
+		panic("runtime: nested call mainthread.Do")
+	}
+	g.inMainThradDo = true
 	// 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()
+	for {
+		select {
+		case f := <-m0func:
+			f()
+			m0exec <- struct{}{}
+		default:
+			// because there is only one main thread,
+			// can use Store directly without considering concurrent call Yield.
+			mtPtr.Store(newmainThreadInfo())
+			for {
+				select {
+				// if there is a new send from m0func before the mtPtr is updated after
+				// all the send from m0func has been received.
+				case f := <-m0func:
+					f()
+					m0exec <- struct{}{}
+				default:
+					return
+				}
+			}
+		}
+	}
 }
 
 func mainThreadWaiting() <-chan struct{} {
-	return waitm0
+	return m0wait
 }
 
 // The main goroutine.

src/runtime/proc_test.go

@@ -1171,7 +1171,7 @@ func TestMainThread(t *testing.T) {
 				runtime.Releasem()
 				if m != runtime.M0 {
 					// don`t use t.Fatal, because it call
-					// Goexit, cause TestMain goroutine exit,
+					// Goexit, cause TestMain goroutine on main thread exit,
 					// and test timeout.
 					t.Fail()
 					t.Log("mainthread.Do.f must on main thread")

src/runtime/runtime2.go

@@ -489,6 +489,8 @@ type g struct {
 
 	coroarg *coro // argument during coroutine transfers
 
+	inMainThradDo bool
+
 	// Per-G tracer state.
 	trace gTraceState
 

src/runtime/testdata/testprog/mainthread.go

@@ -5,6 +5,7 @@
 package main
 
 import (
+	"fmt"
 	"runtime"
 	"runtime/mainthread"
 	"sync"
@@ -29,13 +30,18 @@ func init() {
 		MainThread()
 	})
 	register("MainThread2", func() {
-		println("expect: hello,world")
+		println("expect: hello,runtime: nested call mainthread.Do")
 		MainThread2()
 	})
 }
 
 func MainThread2() {
 	var wg sync.WaitGroup
+	defer func() {
+		if err := recover(); err != nil {
+			fmt.Print(err)
+		}
+	}()
 	runtime.LockOSThread()
 	wg.Add(1)
 	go func() {