1 files changed, 29 insertions, 3 deletions
diff --git a/src/runtime/export_test.go b/src/runtime/export_test.go
index 929bb35db6..f2fa11dc98 100644
--- a/src/runtime/export_test.go
+++ b/src/runtime/export_test.go
@@ -983,9 +983,35 @@ func MapHashCheck(m interface{}, k interface{}) (uintptr, uintptr) {
 	return x, y
 }
 
-func MSpanCountAlloc(bits []byte) int {
-	s := (*mspan)(mheap_.spanalloc.alloc())
+// mspan wrapper for testing.
+//go:notinheap
+type MSpan mspan
+
+// Allocate an mspan for testing.
+func AllocMSpan() *MSpan {
+	var s *mspan
+	systemstack(func() {
+		lock(&mheap_.lock)
+		s = (*mspan)(mheap_.spanalloc.alloc())
+		unlock(&mheap_.lock)
+	})
+	return (*MSpan)(s)
+}
+
+// Free an allocated mspan.
+func FreeMSpan(s *MSpan) {
+	systemstack(func() {
+		lock(&mheap_.lock)
+		mheap_.spanalloc.free(unsafe.Pointer(s))
+		unlock(&mheap_.lock)
+	})
+}
+
+func MSpanCountAlloc(ms *MSpan, bits []byte) int {
+	s := (*mspan)(ms)
 	s.nelems = uintptr(len(bits) * 8)
 	s.gcmarkBits = (*gcBits)(unsafe.Pointer(&bits[0]))
-	return s.countAlloc()
+	result := s.countAlloc()
+	s.gcmarkBits = nil
+	return result
 }