schedule: add task(gr3) block/wake function

The current mechenism, sentry uses channel to communicate with go
runtime to block or wake one task.

Add Schedule/Wake function at go runtime, which add two interface,
runtime.Schedule() and runtime.WakeUpG().
Sentry task which need block just uses runtime.Schedule() to block,
the one which need wake one task just uses runtime.WakeUpG().
Sentry task and go runtime scheduler are aware of each other directly
, not by middle lays such as channel.
based on goruntime:
GuhuangLS/go@a01b27f

Signed-off-by: liushi <[email protected]>

Author: GuhuangLS

pkg/sentry/kernel/futex/futex.go

@@ -18,6 +18,8 @@
 package futex
 
 import (
+	"runtime"
+
 	"gvisor.dev/gvisor/pkg/abi/linux"
 	"gvisor.dev/gvisor/pkg/context"
 	"gvisor.dev/gvisor/pkg/hostarch"
@@ -233,6 +235,18 @@ type Waiter struct {
 
 	// tid is the thread ID for the waiter in case this is a PI mutex.
 	tid uint32
+
+	// G is the go routine of the waiter.
+	G uintptr
+
+	// Forever is the select of the waiter: forever or timeout.
+	Forever bool
+
+	// Interrupt informs task/G waked by interrupt.
+	Interrupt bool
+
+	// Woken for test.
+	Woken     bool
 }
 
 // NewWaiter returns a new unqueued Waiter.
@@ -244,7 +258,7 @@ func NewWaiter() *Waiter {
 
 // woken returns true if w has been woken since the last call to WaitPrepare.
 func (w *Waiter) woken() bool {
-	return len(w.C) != 0
+	return len(w.C) != 0 || w.Woken
 }
 
 // bucket holds a list of waiters for a given address hash.
@@ -273,16 +287,23 @@ func (b *bucket) wakeLocked(key *Key, bitmask uint32, n int) int {
 		// Remove from the bucket and wake the waiter.
 		woke := w
 		w = w.Next() // Next iteration.
-		b.wakeWaiterLocked(woke)
+		b.wakeWaiterLocked(woke, key.Kind)
+		woke.Woken = true
 		done++
 	}
 	return done
 }
 
-func (b *bucket) wakeWaiterLocked(w *Waiter) {
+func (b *bucket) wakeWaiterLocked(w *Waiter, kind KeyKind) {
 	// Remove from the bucket and wake the waiter.
 	b.waiters.Remove(w)
-	w.C <- struct{}{}
+
+	if kind == KindPrivate && w.Forever && w.key.Kind == KindPrivate {
+		w.Interrupt = false
+		runtime.WakeG(w.G)
+	} else {
+		w.C <- struct{}{}
+	}
 
 	// NOTE: The above channel write establishes a write barrier according
 	// to the memory model, so nothing may be ordered around it. Since
@@ -396,6 +417,13 @@ func (m *Manager) Fork() *Manager {
 	}
 }
 
+func (m *Manager) KeyMatch(w *Waiter) bool {
+	if w.key.Kind == KindPrivate {
+		return true
+	}
+	return false
+}
+
 // lockBucket returns a locked bucket for the given key.
 func (m *Manager) lockBucket(k *Key) *bucket {
 	var b *bucket
@@ -452,6 +480,9 @@ func (m *Manager) Wake(t Target, addr hostarch.Addr, private bool, bitmask uint3
 	// This function is very hot; avoid defer.
 	k, err := getKey(t, addr, private)
 	if err != nil {
+		if locked {
+			b.mu.Unlock()
+		}
 		return 0, err
 	}
 
@@ -553,22 +584,30 @@ func (m *Manager) WakeOp(t Target, addr1, addr2 hostarch.Addr, private bool, nwa
 // enqueues w to be woken by a send to w.C. If WaitPrepare returns nil, the
 // Waiter must be subsequently removed by calling WaitComplete, whether or not
 // a wakeup is received on w.C.
-func (m *Manager) WaitPrepare(w *Waiter, t Target, addr hostarch.Addr, private bool, val uint32, bitmask uint32) error {
+func (m *Manager) WaitPrepare(w *Waiter, t Target, addr usermem.Addr, private bool, val uint32, bitmask uint32, forever bool) error {
 	k, err := getKey(t, addr, private)
 	if err != nil {
 		return err
 	}
 	// Ownership of k is transferred to w below.
 
 	// Prepare the Waiter before taking the bucket lock.
-	select {
-	case <-w.C:
-	default:
+	if !private || !forever {
+		select {
+		case <-w.C:
+		default:
+		}
 	}
 	w.key = k
 	w.bitmask = bitmask
 
 	b := m.lockBucket(&k)
+	if private {
+		w.Forever = forever
+	} else {
+		w.Forever = false
+	}
+	runtime.ClearGStatus()
 	// This function is very hot; avoid defer.
 
 	// Perform our atomic check.
@@ -578,6 +617,7 @@ func (m *Manager) WaitPrepare(w *Waiter, t Target, addr hostarch.Addr, private b
 		return err
 	}
 
+	w.Woken = false
 	// Add the waiter to the bucket.
 	b.waiters.PushBack(w)
 	w.bucket.Store(b)
@@ -793,6 +833,6 @@ func (m *Manager) unlockPILocked(t Target, addr hostarch.Addr, tid uint32, b *bu
 		return syserror.EINVAL
 	}
 
-	b.wakeWaiterLocked(next)
+	b.wakeWaiterLocked(next, -1)
 	return nil
 }

pkg/sentry/kernel/futex/futex_test.go

@@ -75,7 +75,8 @@ func futexKind(private bool) string {
 
 func newPreparedTestWaiter(t *testing.T, m *Manager, ta Target, addr hostarch.Addr, private bool, val uint32, bitmask uint32) *Waiter {
 	w := NewWaiter()
-	if err := m.WaitPrepare(w, ta, addr, private, val, bitmask); err != nil {
+	w.G = runtime.GetG()
+	if err := m.WaitPrepare(w, ta, addr, private, val, bitmask, true); err != nil {
 		t.Fatalf("WaitPrepare failed: %v", err)
 	}
 	return w
@@ -487,15 +488,16 @@ func (t *testMutex) Lock() {
 
 		// Wait for it to be "not locked".
 		w := NewWaiter()
-		err := t.m.WaitPrepare(w, t.d, t.a, true, testMutexLocked, ^uint32(0))
+		w.G = runtime.GetG()
+		err := t.m.WaitPrepare(w, t.d, t.a, true, testMutexLocked, ^uint32(0), true)
 		if err == unix.EAGAIN {
 			continue
 		}
 		if err != nil {
 			// Should never happen.
 			panic("WaitPrepare returned unexpected error: " + err.Error())
 		}
-		<-w.C
+		runtime.BlockG()
 		t.m.WaitComplete(w, t.d)
 	}
 }

pkg/sentry/kernel/task_block.go

@@ -166,6 +166,14 @@ func (t *Task) block(C <-chan struct{}, timerChan <-chan struct{}) error {
 	}
 }
 
+func (t *Task) TaskStartRegion() *trace.Region {
+	return trace.StartRegion(t.traceContext, blockRegion)
+}
+
+func (t *Task) TaskEndRegion(region *trace.Region) {
+		region.End()
+}
+
 // SleepStart implements context.ChannelSleeper.SleepStart.
 func (t *Task) SleepStart() <-chan struct{} {
 	t.assertTaskGoroutine()
@@ -240,6 +248,11 @@ func (t *Task) interrupt() {
 func (t *Task) interruptSelf() {
 	select {
 	case t.interruptChan <- struct{}{}:
+		w := t.FutexWaiter()
+		if w.Forever && t.Futex().KeyMatch(w) {
+			w.Interrupt = true
+			runtime.WakeG(w.G)
+		}
 	default:
 	}
 	// platform.Context.Interrupt() is unnecessary since a task goroutine

pkg/sentry/syscalls/linux/sys_futex.go

@@ -15,6 +15,7 @@
 package linux
 
 import (
+	"runtime"
 	"time"
 
 	"gvisor.dev/gvisor/pkg/abi/linux"
@@ -53,13 +54,29 @@ func (f *futexWaitRestartBlock) Restart(t *kernel.Task) (uintptr, error) {
 // arguments.
 func futexWaitAbsolute(t *kernel.Task, clockRealtime bool, ts linux.Timespec, forever bool, addr hostarch.Addr, private bool, val, mask uint32) (uintptr, error) {
 	w := t.FutexWaiter()
-	err := t.Futex().WaitPrepare(w, t, addr, private, val, mask)
+	if w.G == 0 {
+		w.G = runtime.GetG()
+	}
+	err := t.Futex().WaitPrepare(w, t, addr, private, val, mask, forever)
 	if err != nil {
+		w.Forever = false
 		return 0, err
 	}
 
 	if forever {
-		err = t.Block(w.C)
+		if private && w.Forever {
+			t.SleepStart()
+			region := t.TaskStartRegion()
+			runtime.BlockG()
+			if w.Interrupt == true {
+				w.Interrupt = false
+				err = syserror.ErrInterrupted
+			}
+			t.TaskEndRegion(region)
+			t.SleepFinish(true)
+		} else {
+			err = t.Block(w.C)
+		}
 	} else if clockRealtime {
 		notifier, tchan := ktime.NewChannelNotifier()
 		timer := ktime.NewTimer(t.Kernel().RealtimeClock(), notifier)
@@ -73,7 +90,10 @@ func futexWaitAbsolute(t *kernel.Task, clockRealtime bool, ts linux.Timespec, fo
 		err = t.BlockWithDeadline(w.C, true, ktime.FromTimespec(ts))
 	}
 
-	t.Futex().WaitComplete(w, t)
+	if !forever || !private || err != nil {
+		t.Futex().WaitComplete(w)
+	}
+	w.Forever = false
 	return 0, syserror.ConvertIntr(err, syserror.ERESTARTSYS)
 }
 
@@ -89,13 +109,33 @@ func futexWaitAbsolute(t *kernel.Task, clockRealtime bool, ts linux.Timespec, fo
 // syscall is restarted with the remaining timeout.
 func futexWaitDuration(t *kernel.Task, duration time.Duration, forever bool, addr hostarch.Addr, private bool, val, mask uint32) (uintptr, error) {
 	w := t.FutexWaiter()
-	err := t.Futex().WaitPrepare(w, t, addr, private, val, mask)
+	if w.G == 0 {
+		w.G = runtime.GetG()
+	}
+	err := t.Futex().WaitPrepare(w, t, addr, private, val, mask, forever)
 	if err != nil {
+		w.Forever = false
 		return 0, err
 	}
 
-	remaining, err := t.BlockWithTimeout(w.C, !forever, duration)
-	t.Futex().WaitComplete(w, t)
+	var remaining time.Duration
+	if private && forever {
+		t.SleepStart()
+		region := t.TaskStartRegion()
+		runtime.BlockG()
+		if w.Interrupt == true {
+			w.Interrupt = false
+			err = syserror.ErrInterrupted
+		}
+		t.TaskEndRegion(region)
+		t.SleepFinish(true)
+	} else {
+		remaining, err = t.BlockWithTimeout(w.C, !forever, duration)
+	}
+	if !forever || !private || err != nil {
+		t.Futex().WaitComplete(w)
+	}
+	w.Forever = false
 	if err == nil {
 		return 0, nil
 	}
@@ -126,6 +166,7 @@ func futexWaitDuration(t *kernel.Task, duration time.Duration, forever bool, add
 
 func futexLockPI(t *kernel.Task, ts linux.Timespec, forever bool, addr hostarch.Addr, private bool) error {
 	w := t.FutexWaiter()
+	w.Forever = false
 	locked, err := t.Futex().LockPI(w, t, addr, uint32(t.ThreadID()), private, false)
 	if err != nil {
 		return err
@@ -154,6 +195,7 @@ func futexLockPI(t *kernel.Task, ts linux.Timespec, forever bool, addr hostarch.
 
 func tryLockPI(t *kernel.Task, addr hostarch.Addr, private bool) error {
 	w := t.FutexWaiter()
+	w.Forever = false
 	locked, err := t.Futex().LockPI(w, t, addr, uint32(t.ThreadID()), private, true)
 	if err != nil {
 		return err