runtime: lock the main goroutine to the main OS thread during init

We only guarantee that the main goroutine runs on the
main OS thread for initialization.  Programs that wish to
preserve that property for main.main can call runtime.LockOSThread.
This is what programs used to do before we unleashed
goroutines during init, so it is both a simple fix and keeps
existing programs working.

R=iant, r, dave, dvyukov
CC=golang-dev
https://golang.org/cl/5309070

6 files changed, 50 insertions, 48 deletions

diff --git a/src/pkg/runtime/386/asm.s b/src/pkg/runtime/386/asm.s
index 40cef0c1c2..bafef01c7e 100644
--- a/src/pkg/runtime/386/asm.s
+++ b/src/pkg/runtime/386/asm.s
@@ -78,7 +78,7 @@ ok:
 	CALL	runtime·schedinit(SB)
 
 	// create a new goroutine to start program
-	PUSHL	$runtime·mainstart(SB)	// entry
+	PUSHL	$runtime·main(SB)	// entry
 	PUSHL	$0	// arg size
 	CALL	runtime·newproc(SB)
 	POPL	AX
@@ -90,15 +90,6 @@ ok:
 	INT $3
 	RET
 
-TEXT runtime·mainstart(SB),7,$0
-	CALL	main·init(SB)
-	CALL	main·main(SB)
-	PUSHL	$0
-	CALL	runtime·exit(SB)
-	POPL	AX
-	INT $3
-	RET
-
 TEXT runtime·breakpoint(SB),7,$0
 	INT $3
 	RET
diff --git a/src/pkg/runtime/amd64/asm.s b/src/pkg/runtime/amd64/asm.s
index f34ce4a883..866e083106 100644
--- a/src/pkg/runtime/amd64/asm.s
+++ b/src/pkg/runtime/amd64/asm.s
@@ -60,7 +60,7 @@ ok:
 	CALL	runtime·schedinit(SB)
 
 	// create a new goroutine to start program
-	PUSHQ	$runtime·mainstart(SB)		// entry
+	PUSHQ	$runtime·main(SB)		// entry
 	PUSHQ	$0			// arg size
 	CALL	runtime·newproc(SB)
 	POPQ	AX
@@ -72,15 +72,6 @@ ok:
 	CALL	runtime·notok(SB)		// never returns
 	RET
 
-TEXT runtime·mainstart(SB),7,$0
-	CALL	main·init(SB)
-	CALL	main·main(SB)
-	PUSHQ	$0
-	CALL	runtime·exit(SB)
-	POPQ	AX
-	CALL	runtime·notok(SB)
-	RET
-
 TEXT runtime·breakpoint(SB),7,$0
 	BYTE	$0xcc
 	RET
diff --git a/src/pkg/runtime/arm/asm.s b/src/pkg/runtime/arm/asm.s
index a427a465f9..c03422f10c 100644
--- a/src/pkg/runtime/arm/asm.s
+++ b/src/pkg/runtime/arm/asm.s
@@ -43,7 +43,7 @@ TEXT _rt0_arm(SB),7,$-4
 	BL	runtime·schedinit(SB)
 
 	// create a new goroutine to start program
-	MOVW	$runtime·mainstart(SB), R0
+	MOVW	$runtime·main(SB), R0
 	MOVW.W	R0, -4(R13)
 	MOVW	$8, R0
 	MOVW.W	R0, -4(R13)
@@ -60,20 +60,6 @@ TEXT _rt0_arm(SB),7,$-4
 	MOVW	R0, (R1)	// fail hard
 	B	runtime·_dep_dummy(SB)	// Never reached
 
-
-TEXT runtime·mainstart(SB),7,$4
-	BL	main·init(SB)
-	EOR	R0, R0
-	MOVW	R0, 0(R13)
-	BL	main·main(SB)
-	MOVW	$0, R0
-	MOVW	R0, 4(SP)
-	BL	runtime·exit(SB)
-	MOVW	$1234, R0
-	MOVW	$1001, R1
-	MOVW	R0, (R1)	// fail hard
-	RET
-
 // TODO(kaib): remove these once i actually understand how the linker removes symbols
 // pull in dummy dependencies
 TEXT runtime·_dep_dummy(SB),7,$0
diff --git a/src/pkg/runtime/debug.go b/src/pkg/runtime/debug.go
index 6370a57d80..124370384c 100644
--- a/src/pkg/runtime/debug.go
+++ b/src/pkg/runtime/debug.go
@@ -10,7 +10,6 @@ func Breakpoint()
 // LockOSThread wires the calling goroutine to its current operating system thread.
 // Until the calling goroutine exits or calls UnlockOSThread, it will always
 // execute in that thread, and no other goroutine can.
-// LockOSThread cannot be used during init functions.
 func LockOSThread()
 
 // UnlockOSThread unwires the calling goroutine from its fixed operating system thread.
diff --git a/src/pkg/runtime/proc.c b/src/pkg/runtime/proc.c
index 3408c06f69..d51e3d3522 100644
--- a/src/pkg/runtime/proc.c
+++ b/src/pkg/runtime/proc.c
@@ -72,6 +72,9 @@ struct Sched {
 	volatile uint32 atomic;	// atomic scheduling word (see below)
 
 	int32 profilehz;	// cpu profiling rate
+	
+	bool init;  // running initialization
+	bool lockmain;  // init called runtime.LockOSThread
 
 	Note	stopped;	// one g can set waitstop and wait here for m's to stop
 };
@@ -171,11 +174,7 @@ setmcpumax(uint32 n)
 //	make & queue new G
 //	call runtime·mstart
 //
-// The new G does:
-//
-//	call main·init_function
-//	call initdone
-//	call main·main
+// The new G calls runtime·main.
 void
 runtime·schedinit(void)
 {
@@ -212,6 +211,32 @@ runtime·schedinit(void)
 	m->nomemprof--;
 }
 
+extern void main·init(void);
+extern void main·main(void);
+
+// The main goroutine.
+void
+runtime·main(void)
+{
+	// Lock the main goroutine onto this, the main OS thread,
+	// during initialization.  Most programs won't care, but a few
+	// do require certain calls to be made by the main thread.
+	// Those can arrange for main.main to run in the main thread
+	// by calling runtime.LockOSThread during initialization
+	// to preserve the lock.
+	runtime·LockOSThread();
+	runtime·sched.init = true;
+	main·init();
+	runtime·sched.init = false;
+	if(!runtime·sched.lockmain)
+		runtime·UnlockOSThread();
+
+	main·main();
+	runtime·exit(0);
+	for(;;)
+		*(int32*)runtime·main = 0;
+}
+
 // Lock the scheduler.
 static void
 schedlock(void)
@@ -1494,13 +1519,6 @@ runtime·Gosched(void)
 	runtime·gosched();
 }
 
-void
-runtime·LockOSThread(void)
-{
-	m->lockedg = g;
-	g->lockedm = m;
-}
-
 // delete when scheduler is stronger
 int32
 runtime·gomaxprocsfunc(int32 n)
@@ -1541,8 +1559,23 @@ runtime·gomaxprocsfunc(int32 n)
 }
 
 void
+runtime·LockOSThread(void)
+{
+	if(m == &runtime·m0 && runtime·sched.init) {
+		runtime·sched.lockmain = true;
+		return;
+	}
+	m->lockedg = g;
+	g->lockedm = m;
+}
+
+void
 runtime·UnlockOSThread(void)
 {
+	if(m == &runtime·m0 && runtime·sched.init) {
+		runtime·sched.lockmain = false;
+		return;
+	}
 	m->lockedg = nil;
 	g->lockedm = nil;
 }
diff --git a/src/pkg/runtime/runtime.h b/src/pkg/runtime/runtime.h
index e3ec197757..e45808f8e0 100644
--- a/src/pkg/runtime/runtime.h
+++ b/src/pkg/runtime/runtime.h
@@ -637,6 +637,8 @@ String	runtime·signame(int32 sig);
 int32	runtime·gomaxprocsfunc(int32 n);
 void	runtime·procyield(uint32);
 void	runtime·osyield(void);
+void	runtime·LockOSThread(void);
+void	runtime·UnlockOSThread(void);
 
 void	runtime·mapassign(MapType*, Hmap*, byte*, byte*);
 void	runtime·mapaccess(MapType*, Hmap*, byte*, byte*, bool*);