Support synchronous AssertAndFetch for sleep package.

Some synchronization patterns require the ability to simultaneously wake and
sleep a goroutine. For the sleep package, this is the case when a waker must be
asserted when a subsequent fetch is imminent.

Currently, this operation results in significant P churn in the runtime, which
ping-pongs execution between multiple system threads and cores and consumes a
significant amount of host CPU (and because of the context switches, this can
be significant worse with mitigations for side channel vulnerabilities).

The solution is to introduce a dedicated mechanism for a synchronous switch
which does not wake another runtime P (see golang/go#32113). This can be used
by the `AssertAndFetch` API in the sleep package.

The benchmark results for this package are very similiar to raw channel
operations for all cases, with the exception of operations that do not wait.
The primary advantage is more precise control over scheduling. This will be
used in a subsequent change.

```
BenchmarkGoAssertNonWaiting
BenchmarkGoAssertNonWaiting-8                   261364384                4.976 ns/op
BenchmarkGoSingleSelect
BenchmarkGoSingleSelect-8                       20946358                57.77 ns/op
BenchmarkGoMultiSelect
BenchmarkGoMultiSelect-8                         6071697               197.0 ns/op
BenchmarkGoWaitOnSingleSelect
BenchmarkGoWaitOnSingleSelect-8                  4978051               235.4 ns/op
BenchmarkGoWaitOnMultiSelect
BenchmarkGoWaitOnMultiSelect-8                   2309224               520.2 ns/op

BenchmarkSleeperAssertNonWaiting
BenchmarkSleeperAssertNonWaiting-8              447325033                2.657 ns/op
BenchmarkSleeperSingleSelect
BenchmarkSleeperSingleSelect-8                  21488844                55.19 ns/op
BenchmarkSleeperMultiSelect
BenchmarkSleeperMultiSelect-8                   21851674                54.89 ns/op
BenchmarkSleeperWaitOnSingleSelect
BenchmarkSleeperWaitOnSingleSelect-8             2860327               416.4 ns/op
BenchmarkSleeperWaitOnSingleSelectSync
BenchmarkSleeperWaitOnSingleSelectSync-8         2741733               427.1 ns/op
BenchmarkSleeperWaitOnMultiSelect
BenchmarkSleeperWaitOnMultiSelect-8              2867484               418.1 ns/op
BenchmarkSleeperWaitOnMultiSelectSync
BenchmarkSleeperWaitOnMultiSelectSync-8          2789158               427.9 ns/op
```

PiperOrigin-RevId: 415581417

pkg/sleep/sleep_test.go

@@ -17,6 +17,8 @@ package sleep
 import (
 	"math/rand"
 	"runtime"
+	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 )
@@ -332,6 +334,50 @@ func TestDoneFunction(t *testing.T) {
 	}
 }
 
+// TestAssertFetch tests basic assert fetch functionality.
+func TestAssertFetch(t *testing.T) {
+	const sleeperWakers = 100
+	const wakeRequests = 1000
+	const seedAsserts = 10
+
+	ws := make([]Waker, sleeperWakers)
+	ss := make([]Sleeper, sleeperWakers)
+	for i := 0; i < sleeperWakers; i++ {
+		ss[i].AddWaker(&ws[i])
+	}
+	defer func() {
+		for i := 0; i < sleeperWakers; i++ {
+			defer ss[i].Done()
+		}
+	}()
+	var (
+		count int32
+		wg    sync.WaitGroup
+	)
+	for i := 0; i < sleeperWakers; i++ {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			ss[i].Fetch(true /* block */)
+			w := &ws[(i+1)%sleeperWakers]
+			for n := 0; n < wakeRequests; n++ {
+				atomic.AddInt32(&count, 1)
+				ss[i].AssertAndFetch(w)
+			}
+			w.Assert() // Final wake-up.
+		}(i)
+	}
+
+	// Fire the first assertion.
+	ws[0].Assert()
+	wg.Wait()
+
+	// Check what we got.
+	if want := int32(sleeperWakers * wakeRequests); count != want {
+		t.Errorf("unexpected count: got %d, wanted %d", count, want)
+	}
+}
+
 // TestRace tests that multiple wakers can continuously send wake requests to
 // the sleeper.
 func TestRace(t *testing.T) {
@@ -511,6 +557,29 @@ func BenchmarkSleeperWaitOnSingleSelect(b *testing.B) {
 	}
 }
 
+// BenchmarkSleeperWaitOnSingleSelectSync is a modification of the similarly
+// named benchmark, except it uses the synchronous AssertAndFetch.
+func BenchmarkSleeperWaitOnSingleSelectSync(b *testing.B) {
+	var (
+		s  Sleeper
+		w  Waker
+		ns Sleeper
+		nw Waker
+	)
+	ns.AddWaker(&nw)
+	s.AddWaker(&w)
+	go func() {
+		ns.Fetch(true)
+		defer w.Assert()
+		for i := 0; i < b.N-1; i++ {
+			ns.AssertAndFetch(&w)
+		}
+	}()
+	for i := 0; i < b.N; i++ {
+		s.AssertAndFetch(&nw)
+	}
+}
+
 // BenchmarkGoWaitOnSingleSelect measures how long it takes to wait on one
 // channel while another goroutine wakes up the sleeper.
 func BenchmarkGoWaitOnSingleSelect(b *testing.B) {
@@ -556,6 +625,34 @@ func BenchmarkSleeperWaitOnMultiSelect(b *testing.B) {
 	}
 }
 
+// BenchmarkSleeperWaitOnMultiSelectSync is a modification of the similarly
+// named benchmark, except it uses the synchronous AssertAndFetch.
+func BenchmarkSleeperWaitOnMultiSelectSync(b *testing.B) {
+	const count = 4
+	var (
+		s  Sleeper
+		ns Sleeper
+		nw Waker
+	)
+	ns.AddWaker(&nw)
+	w := make([]Waker, count)
+	for i := range w {
+		s.AddWaker(&w[i])
+	}
+
+	b.ResetTimer()
+	go func() {
+		ns.Fetch(true)
+		defer w[count-1].Assert()
+		for i := 0; i < b.N-1; i++ {
+			ns.AssertAndFetch(&w[count-1])
+		}
+	}()
+	for i := 0; i < b.N; i++ {
+		s.AssertAndFetch(&nw)
+	}
+}
+
 // BenchmarkGoWaitOnMultiSelect measures how long it takes to wait on 4 channels
 // while another goroutine wakes up the sleeper.
 func BenchmarkGoWaitOnMultiSelect(b *testing.B) {

pkg/sleep/sleep_unsafe.go

@@ -98,10 +98,14 @@ var (
 // returns. These restrictions allow this to be implemented lock-free.
 //
 // This struct is thread-compatible.
+//
+// +stateify savable
 type Sleeper struct {
+	_ sync.NoCopy
+
 	// sharedList is a "stack" of asserted wakers. They atomically add
 	// themselves to the front of this list as they become asserted.
-	sharedList unsafe.Pointer
+	sharedList unsafe.Pointer `state:".(*Waker)"`
 
 	// localList is a list of asserted wakers that is only accessible to the
 	// waiter, and thus doesn't have to be accessed atomically. When
@@ -116,7 +120,17 @@ type Sleeper struct {
 
 	// waitingG holds the G that is sleeping, if any. It is used by wakers
 	// to determine which G, if any, they should wake.
-	waitingG uintptr
+	waitingG uintptr `state:"zero"`
+}
+
+// saveSharedList is invoked by stateify.
+func (s *Sleeper) saveSharedList() *Waker {
+	return (*Waker)(atomic.LoadPointer(&s.sharedList))
+}
+
+// loadSharedList is invoked by stateify.
+func (s *Sleeper) loadSharedList(w *Waker) {
+	atomic.StorePointer(&s.sharedList, unsafe.Pointer(w))
 }
 
 // AddWaker associates the given waker to the sleeper.
@@ -137,7 +151,7 @@ func (s *Sleeper) AddWaker(w *Waker) {
 	for {
 		p := (*Sleeper)(atomic.LoadPointer(&w.s))
 		if p == &assertedSleeper {
-			s.enqueueAssertedWaker(w)
+			s.enqueueAssertedWaker(w, true /* wakep */)
 			return
 		}
 
@@ -148,8 +162,11 @@ func (s *Sleeper) AddWaker(w *Waker) {
 }
 
 // nextWaker returns the next waker in the notification list, blocking if
-// needed.
-func (s *Sleeper) nextWaker(block bool) *Waker {
+// needed. The parameter wakepOrSleep indicates that if the operation does not
+// block, then we will need to explicitly wake a runtime P.
+//
+// Precondition: wakepOrSleep may be true iff block is true.
+func (s *Sleeper) nextWaker(block, wakepOrSleep bool) *Waker {
 	// Attempt to replenish the local list if it's currently empty.
 	if s.localList == nil {
 		for atomic.LoadPointer(&s.sharedList) == nil {
@@ -173,6 +190,10 @@ func (s *Sleeper) nextWaker(block bool) *Waker {
 				break
 			}
 
+			// Since we are sleeping for sure, we no longer
+			// need to wakep once we get a value.
+			wakepOrSleep = false
+
 			// Try to commit the sleep and report it to the
 			// tracer as a select.
 			//
@@ -203,6 +224,11 @@ func (s *Sleeper) nextWaker(block bool) *Waker {
 	w := s.localList
 	s.localList = w.next
 
+	// Do we need to wake a P?
+	if wakepOrSleep {
+		sync.Wakep()
+	}
+
 	return w
 }
 
@@ -218,17 +244,13 @@ func commitSleep(g uintptr, waitingG unsafe.Pointer) bool {
 	return sync.RaceUncheckedAtomicCompareAndSwapUintptr((*uintptr)(waitingG), preparingG, g)
 }
 
-// Fetch fetches the next wake-up notification. If a notification is
-// immediately available, the asserted waker is returned immediately.
-// Otherwise, the behavior depends on the value of 'block': if true, the
-// current goroutine blocks until a notification arrives and returns the
-// asserted waker; if false, nil will be returned.
+// fetch is the backing implementation for Fetch and AssertAndFetch.
 //
-// N.B. This method is *not* thread-safe. Only one goroutine at a time is
-//      allowed to call this method.
-func (s *Sleeper) Fetch(block bool) *Waker {
+// Preconditions are the same as nextWaker.
+//go:nosplit
+func (s *Sleeper) fetch(block, wakepOrSleep bool) *Waker {
 	for {
-		w := s.nextWaker(block)
+		w := s.nextWaker(block, wakepOrSleep)
 		if w == nil {
 			return nil
 		}
@@ -242,6 +264,31 @@ func (s *Sleeper) Fetch(block bool) *Waker {
 	}
 }
 
+// Fetch fetches the next wake-up notification. If a notification is
+// immediately available, the asserted waker is returned immediately.
+// Otherwise, the behavior depends on the value of 'block': if true, the
+// current goroutine blocks until a notification arrives and returns the
+// asserted waker; if false, nil will be returned.
+//
+// N.B. This method is *not* thread-safe. Only one goroutine at a time is
+//      allowed to call this method.
+func (s *Sleeper) Fetch(block bool) *Waker {
+	return s.fetch(block, false /* wakepOrSleep */)
+}
+
+// AssertAndFetch asserts the given waker and fetches the next wake-up notification.
+// Note that this will always be blocking, since there is no value in joining a
+// non-blocking operation.
+//
+// N.B. Like Fetch, this method is *not* thread-safe. This will also yield the current
+//      P to the next goroutine, avoiding associated scheduled overhead.
+//+checkescapes:all
+//go:nosplit
+func (s *Sleeper) AssertAndFetch(n *Waker) *Waker {
+	n.assert(false /* wakep */)
+	return s.fetch(true /* block */, true /* wakepOrSleep*/)
+}
+
 // Done is used to indicate that the caller won't use this Sleeper anymore. It
 // removes the association with all wakers so that they can be safely reused
 // by another sleeper after Done() returns.
@@ -264,7 +311,7 @@ func (s *Sleeper) Done() {
 		// Dequeue exactly one waiter from the list, it may not be
 		// this one but we know this one is in the process. We must
 		// leave it in the asserted state but drop it from our lists.
-		if w := s.nextWaker(true); w != nil {
+		if w := s.nextWaker(true, false); w != nil {
 			prev := &s.allWakers
 			for *prev != w {
 				prev = &((*prev).allWakersNext)
@@ -278,7 +325,7 @@ func (s *Sleeper) Done() {
 
 // enqueueAssertedWaker enqueues an asserted waker to the "ready" circular list
 // of wakers that want to notify the sleeper.
-func (s *Sleeper) enqueueAssertedWaker(w *Waker) {
+func (s *Sleeper) enqueueAssertedWaker(w *Waker, wakep bool) {
 	// Add the new waker to the front of the list.
 	for {
 		v := (*Waker)(atomic.LoadPointer(&s.sharedList))
@@ -298,7 +345,7 @@ func (s *Sleeper) enqueueAssertedWaker(w *Waker) {
 	case 0, preparingG:
 	default:
 		// We managed to get a G. Wake it up.
-		sync.Goready(g, 0)
+		sync.Goready(g, 0, wakep)
 	}
 }
 
@@ -315,6 +362,8 @@ func (s *Sleeper) enqueueAssertedWaker(w *Waker) {
 //
 // Note, it is not safe to copy a Waker as its fields are modified by value
 // (the pointer fields are individually modified with atomic operations).
+//
+// +stateify savable
 type Waker struct {
 	_ sync.NoCopy
 
@@ -327,7 +376,7 @@ type Waker struct {
 	// otherwise -- the waker is not asserted, and is associated with the
 	//     given sleeper. Once it transitions to asserted state, the
 	//     associated sleeper will be woken.
-	s unsafe.Pointer
+	s unsafe.Pointer `state:".(wakerState)"`
 
 	// next is used to form a linked list of asserted wakers in a sleeper.
 	next *Waker
@@ -337,9 +386,32 @@ type Waker struct {
 	allWakersNext *Waker
 }
 
-// Assert moves the waker to an asserted state, if it isn't asserted yet. When
-// asserted, the waker will cause its matching sleeper to wake up.
-func (w *Waker) Assert() {
+type wakerState struct {
+	asserted bool
+	other    *Sleeper
+}
+
+// saveS is invoked by stateify.
+func (w *Waker) saveS() wakerState {
+	s := (*Sleeper)(atomic.LoadPointer(&w.s))
+	if s == &assertedSleeper {
+		return wakerState{asserted: true}
+	}
+	return wakerState{other: s}
+}
+
+// loadS is invoked by stateify.
+func (w *Waker) loadS(ws wakerState) {
+	if ws.asserted {
+		atomic.StorePointer(&w.s, unsafe.Pointer(&assertedSleeper))
+	} else {
+		atomic.StorePointer(&w.s, unsafe.Pointer(ws.other))
+	}
+}
+
+// assert is the implementation for Assert.
+//go:nosplit
+func (w *Waker) assert(wakep bool) {
 	// Nothing to do if the waker is already asserted. This check allows us
 	// to complete this case (already asserted) without any interlocked
 	// operations on x86.
@@ -352,10 +424,16 @@ func (w *Waker) Assert() {
 	case nil:
 	case &assertedSleeper:
 	default:
-		s.enqueueAssertedWaker(w)
+		s.enqueueAssertedWaker(w, wakep)
 	}
 }
 
+// Assert moves the waker to an asserted state, if it isn't asserted yet. When
+// asserted, the waker will cause its matching sleeper to wake up.
+func (w *Waker) Assert() {
+	w.assert(true /* wakep */)
+}
+
 // Clear moves the waker to then non-asserted state and returns whether it was
 // asserted before being cleared.
 //

pkg/sync/BUILD

@@ -22,6 +22,9 @@ go_library(
         "race_amd64.s",
         "race_arm64.s",
         "race_unsafe.go",
+        "runtime_amd64.go",
+        "runtime_amd64.s",
+        "runtime_other.go",
         "runtime_unsafe.go",
         "rwmutex_unsafe.go",
         "seqcount.go",

pkg/sync/runtime_amd64.go

@@ -0,0 +1,29 @@
+// Copyright 2020 The gVisor Authors.
+//
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build amd64 && go1.8 && !go1.19 && !goexperiment.staticlockranking
+// +build amd64,go1.8,!go1.19,!goexperiment.staticlockranking
+
+package sync
+
+import (
+	"sync/atomic"
+)
+
+const supportsWakeSuppression = true
+
+// addrOfSpinning returns the address of runtime.sched.nmspinning.
+func addrOfSpinning() *int32
+
+// nmspinning caches addrOfSpinning.
+var nmspinning = addrOfSpinning()
+
+func preGoReadyWakeSuppression() {
+	atomic.AddInt32(nmspinning, 1)
+}
+
+func postGoReadyWakeSuppression() {
+	atomic.AddInt32(nmspinning, -1)
+}

pkg/sync/runtime_amd64.s

@@ -0,0 +1,25 @@
+// Copyright 2018 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build amd64 && go1.8 && !go1.19 && !goexperiment.staticlockranking
+// +build amd64,go1.8,!go1.19,!goexperiment.staticlockranking
+
+#include "textflag.h"
+
+TEXT ·addrOfSpinning(SB),NOSPLIT,$0-8
+	// The offset specified here is the nmspinning value in sched.
+	LEAQ runtime·sched(SB), AX
+	ADDQ $92, AX
+	MOVQ AX, ret+0(FP)
+	RET

pkg/sync/runtime_other.go

@@ -0,0 +1,14 @@
+// Copyright 2020 The gVisor Authors.
+//
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build !amd64
+// +build !amd64
+
+package sync
+
+const supportsWakeSuppression = false
+
+func preGoReadyWakeSuppression()  {} // Never called.
+func postGoReadyWakeSuppression() {} // Never called.

pkg/sync/runtime_unsafe.go

@@ -35,16 +35,40 @@ func Gopark(unlockf func(uintptr, unsafe.Pointer) bool, lock unsafe.Pointer, rea
 //go:linkname gopark runtime.gopark
 func gopark(unlockf func(uintptr, unsafe.Pointer) bool, lock unsafe.Pointer, reason uint8, traceEv byte, traceskip int)
 
-// Goready is runtime.goready.
+//go:linkname wakep runtime.wakep
+func wakep()
+
+// Wakep is runtime.wakep.
 //
 //go:nosplit
-func Goready(gp uintptr, traceskip int) {
-	goready(gp, traceskip)
+func Wakep() {
+	// This is only supported if we can suppress the wakep called
+	// from  Goready below, which is in certain architectures only.
+	if supportsWakeSuppression {
+		wakep()
+	}
 }
 
 //go:linkname goready runtime.goready
 func goready(gp uintptr, traceskip int)
 
+// Goready is runtime.goready.
+//
+// The additional wakep argument controls whether a new thread will be kicked to
+// execute the P. This should be true in most circumstances. However, if the
+// current thread is about to sleep, then this can be false for efficiency.
+//
+//go:nosplit
+func Goready(gp uintptr, traceskip int, wakep bool) {
+	if supportsWakeSuppression && !wakep {
+		preGoReadyWakeSuppression()
+	}
+	goready(gp, traceskip)
+	if supportsWakeSuppression && !wakep {
+		preGoReadyWakeSuppression()
+	}
+}
+
 // Values for the reason argument to gopark, from Go's src/runtime/runtime2.go.
 const (
 	WaitReasonSelect      uint8 = 9

pkg/syncevent/waiter_unsafe.go

@@ -169,7 +169,7 @@ func (w *Waiter) NotifyPending() {
 	// so. Swap is needed here to ensure that only one call to NotifyPending
 	// calls goready.
 	if g := atomic.SwapUintptr(&w.g, 0); g > preparingG {
-		sync.Goready(g, 0)
+		sync.Goready(g, 0, true /* wakep */)
 	}
 }
 

tools/checklinkname/known.go

@@ -73,6 +73,9 @@ var knownLinknames = map[string]map[string]linknameSignatures{
 		"throw": linknameSignatures{
 			local: "func(s string)",
 		},
+		"wakep": linknameSignatures{
+			local: "func()",
+		},
 	},
 	"sync": map[string]linknameSignatures{
 		"runtime_canSpin": linknameSignatures{