runtime: fix coro interactions with thread-locked goroutines

This change fixes problems with thread-locked goroutines using
newcoro/coroswitch/etc. Currently, the coro paths do not consider
thread-locked goroutines at all and can quickly result in broken
scheduler state or lost/leaked goroutines.

One possible fix to these issues is to fall back on goroutine+channel
semantics, but that turns out to be fairly complicated to implement and
results in significant performance cliffs. More complex thread-lock
state donation tricks also result in some fairly complicated state
tracking that doesn't seem worth it given the use-cases of iter.Pull
(and even then, there will be performance cliffs).

This change implements a much simpler, but more restrictive semantics.
In particular, thread-lock state is tied to the coro at the first call
to newcoro (i.e. iter.Pull). From then on, the invariant is that if the
coro has any thread-lock state *or* a goroutine calling into coroswitch
has any thread-lock state, that the full gamut of thread-lock state must
remain the same as it was when newcoro was called (the full gamut
meaning internal and external lock counts as well as the identity of the
thread that was locked to).

This semantics allows the common cases to be always fast, but comes with
a non-orthogonality caveat. Specifically, when iter.Pull is used in
conjunction with thread-locked goroutines, complex cases (passing next
between goroutines or passing yield between goroutines) are likely to
fail. Simple cases, where any number of iter.Pull iterators are used in
a straightforward way (nested, in series, etc.) from the same
goroutine, will work and will be guaranteed to be fast regardless of
thread-lock state.

This is a compromise for the near-term and we may consider lifting the
restrictions imposed by this CL in the future.

Fixes #65889.
Fixes #65946.

Change-Id: I3fb5791e36a61f5ded50226a229a79d28739b24e
Reviewed-on: https://go-review.googlesource.com/c/go/+/583675
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: David Chase <drchase@google.com>
Reviewed-by: Austin Clements <austin@google.com>

1 files changed, 185 insertions, 0 deletions

diff --git a/src/runtime/testdata/testprogcgo/coro.go b/src/runtime/testdata/testprogcgo/coro.go
new file mode 100644
index 0000000000..e0cb945112
--- /dev/null
+++ b/src/runtime/testdata/testprogcgo/coro.go
@@ -0,0 +1,185 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build goexperiment.rangefunc && !windows
+
+package main
+
+/*
+#include <stdint.h> // for uintptr_t
+
+void go_callback_coro(uintptr_t handle);
+
+static void call_go(uintptr_t handle) {
+	go_callback_coro(handle);
+}
+*/
+import "C"
+
+import (
+	"fmt"
+	"iter"
+	"runtime/cgo"
+)
+
+func init() {
+	register("CoroCgoIterCallback", func() {
+		println("expect: OK")
+		CoroCgo(callerExhaust, iterCallback)
+	})
+	register("CoroCgoIterCallbackYield", func() {
+		println("expect: OS thread locking must match")
+		CoroCgo(callerExhaust, iterCallbackYield)
+	})
+	register("CoroCgoCallback", func() {
+		println("expect: OK")
+		CoroCgo(callerExhaustCallback, iterSimple)
+	})
+	register("CoroCgoCallbackIterNested", func() {
+		println("expect: OK")
+		CoroCgo(callerExhaustCallback, iterNested)
+	})
+	register("CoroCgoCallbackIterCallback", func() {
+		println("expect: OK")
+		CoroCgo(callerExhaustCallback, iterCallback)
+	})
+	register("CoroCgoCallbackIterCallbackYield", func() {
+		println("expect: OS thread locking must match")
+		CoroCgo(callerExhaustCallback, iterCallbackYield)
+	})
+	register("CoroCgoCallbackAfterPull", func() {
+		println("expect: OS thread locking must match")
+		CoroCgo(callerCallbackAfterPull, iterSimple)
+	})
+	register("CoroCgoStopCallback", func() {
+		println("expect: OK")
+		CoroCgo(callerStopCallback, iterSimple)
+	})
+	register("CoroCgoStopCallbackIterNested", func() {
+		println("expect: OK")
+		CoroCgo(callerStopCallback, iterNested)
+	})
+}
+
+var toCall func()
+
+//export go_callback_coro
+func go_callback_coro(handle C.uintptr_t) {
+	h := cgo.Handle(handle)
+	h.Value().(func())()
+	h.Delete()
+}
+
+func callFromC(f func()) {
+	C.call_go(C.uintptr_t(cgo.NewHandle(f)))
+}
+
+func CoroCgo(driver func(iter.Seq[int]) error, seq iter.Seq[int]) {
+	if err := driver(seq); err != nil {
+		println("error:", err.Error())
+		return
+	}
+	println("OK")
+}
+
+func callerExhaust(i iter.Seq[int]) error {
+	next, _ := iter.Pull(i)
+	for {
+		v, ok := next()
+		if !ok {
+			break
+		}
+		if v != 5 {
+			return fmt.Errorf("bad iterator: wanted value %d, got %d", 5, v)
+		}
+	}
+	return nil
+}
+
+func callerExhaustCallback(i iter.Seq[int]) (err error) {
+	callFromC(func() {
+		next, _ := iter.Pull(i)
+		for {
+			v, ok := next()
+			if !ok {
+				break
+			}
+			if v != 5 {
+				err = fmt.Errorf("bad iterator: wanted value %d, got %d", 5, v)
+			}
+		}
+	})
+	return err
+}
+
+func callerStopCallback(i iter.Seq[int]) (err error) {
+	callFromC(func() {
+		next, stop := iter.Pull(i)
+		v, _ := next()
+		stop()
+		if v != 5 {
+			err = fmt.Errorf("bad iterator: wanted value %d, got %d", 5, v)
+		}
+	})
+	return err
+}
+
+func callerCallbackAfterPull(i iter.Seq[int]) (err error) {
+	next, _ := iter.Pull(i)
+	callFromC(func() {
+		for {
+			v, ok := next()
+			if !ok {
+				break
+			}
+			if v != 5 {
+				err = fmt.Errorf("bad iterator: wanted value %d, got %d", 5, v)
+			}
+		}
+	})
+	return err
+}
+
+func iterSimple(yield func(int) bool) {
+	for range 3 {
+		if !yield(5) {
+			return
+		}
+	}
+}
+
+func iterNested(yield func(int) bool) {
+	next, stop := iter.Pull(iterSimple)
+	for {
+		v, ok := next()
+		if ok {
+			if !yield(v) {
+				stop()
+			}
+		} else {
+			return
+		}
+	}
+}
+
+func iterCallback(yield func(int) bool) {
+	for range 3 {
+		callFromC(func() {})
+		if !yield(5) {
+			return
+		}
+	}
+}
+
+func iterCallbackYield(yield func(int) bool) {
+	for range 3 {
+		var ok bool
+		callFromC(func() {
+			ok = yield(5)
+		})
+		if !ok {
+			return
+		}
+	}
+}