runtime: conservatively scan extended register state

Conservatively scan the extended register state when GC scans
asynchronously preempted goroutines. This ensures that any pointers
that appear only in vector registers at preemption time are kept alive.

Using vector registers for small memory moves may load pointers into
these registers. If async preemption occurs mid-move, with no write
barrier (e.g., heap-to-stack copies) and the source register clobbered
or source memory modified by a racing goroutine, the pointer may exist
only in the vector register. Without scanning this state, GC could miss
live pointers.

This addresses concerns raised in CL 738261 and enables safe use of
vector registers for operations that may involve pointers.

Change-Id: I5f5ce98d6ed6f7cde34b33da0aea1f880c2fcf41
Reviewed-on: https://go-review.googlesource.com/c/go/+/740681
Reviewed-by: Keith Randall <khr@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Keith Randall <khr@golang.org>
Reviewed-by: Austin Clements <austin@google.com>

3 files changed, 36 insertions, 2 deletions

diff --git a/src/runtime/mgcmark.go b/src/runtime/mgcmark.go
index 714b9a51df..f58b98bd7f 100644
--- a/src/runtime/mgcmark.go
+++ b/src/runtime/mgcmark.go
@@ -967,6 +967,11 @@ func scanstack(gp *g, gcw *gcWork) int64 {
 		scanblock(uintptr(unsafe.Pointer(&gp.sched.ctxt)), goarch.PtrSize, &oneptrmask[0], gcw, &state)
 	}
 
+	// Scan conservatively the extended register state.
+	if gp.asyncSafePoint {
+		xRegScan(gp, gcw, &state)
+	}
+
 	// Scan the stack. Accumulate a list of stack objects.
 	var u unwinder
 	for u.init(gp, 0); u.valid(); u.next() {
diff --git a/src/runtime/preempt_noxreg.go b/src/runtime/preempt_noxreg.go
index 977bf0bcec..ee6f879779 100644
--- a/src/runtime/preempt_noxreg.go
+++ b/src/runtime/preempt_noxreg.go
@@ -25,3 +25,5 @@ func xRegSave(gp *g) {}
 func xRegRestore(gp *g) {}
 
 func (*xRegPerP) free() {}
+
+func xRegScan(gp *g, gcw *gcWork, state *stackScanState) {}
diff --git a/src/runtime/preempt_xreg.go b/src/runtime/preempt_xreg.go
index cc52c5f3c4..8e7c13338c 100644
--- a/src/runtime/preempt_xreg.go
+++ b/src/runtime/preempt_xreg.go
@@ -10,8 +10,8 @@
 // While asynchronous preemption stores general-purpose (GP) registers on the
 // preempted goroutine's own stack, extended register state can be used to save
 // non-GP state off the stack. In particular, this is meant for large vector
-// register files. Currently, we assume this contains only scalar data, though
-// we could change this constraint by conservatively scanning this memory.
+// register files. This memory is conservatively scanned to enable using
+// non-GP registers for operations that may involve pointers.
 //
 // For an architecture to support extended register state, it must provide a Go
 // definition of an xRegState type for storing the state, and its asyncPreempt
@@ -20,6 +20,7 @@
 package runtime
 
 import (
+	"internal/abi"
 	"internal/runtime/sys"
 	"unsafe"
 )
@@ -135,3 +136,29 @@ func (xRegs *xRegPerP) free() {
 		unlock(&xRegAlloc.lock)
 	}
 }
+
+// xRegScan conservatively scans the extended register state.
+//
+// This is supposed to be called only by scanstack when it handles async preemption.
+func xRegScan(gp *g, gcw *gcWork, state *stackScanState) {
+	// Regular async preemption always provides the extended register state.
+	if gp.xRegs.state == nil {
+		var u unwinder
+		for u.init(gp, 0); u.valid(); u.next() {
+			if u.frame.fn.valid() && u.frame.fn.funcID == abi.FuncID_debugCallV2 {
+				return
+			}
+		}
+		println("runtime: gp=", gp, ", goid=", gp.goid)
+		throw("gp.xRegs.state == nil on a scanstack attempt during async preemption")
+	}
+	b := uintptr(unsafe.Pointer(&gp.xRegs.state.regs))
+	n := uintptr(unsafe.Sizeof(gp.xRegs.state.regs))
+	if debugScanConservative {
+		print("begin scan xRegs of goroutine ", gp.goid, " at [", hex(b), ",", hex(b+n), ")\n")
+	}
+	scanConservative(b, n, nil, gcw, state)
+	if debugScanConservative {
+		print("end scan xRegs of goroutine ", gp.goid, "\n")
+	}
+}