From 226f5145b682026e84d3813d3b1aa771b9124be7 Mon Sep 17 00:00:00 2001
From: Andrei Vagin <avagin@google.com>
Date: Fri, 19 May 2023 00:32:17 -0700
Subject: [PATCH] systrap: preempt long running contexts

A context is preempted if it is running longer than 10ms
and there are other contexts in the queue.

PiperOrigin-RevId: 533376587
---
 pkg/sentry/platform/systrap/shared_context.go | 39 +++++++++++++-
 pkg/sentry/platform/systrap/sysmsg/BUILD      |  2 -
 .../platform/systrap/sysmsg/sysmsg_unsafe.go  | 54 ++-----------------
 3 files changed, 40 insertions(+), 55 deletions(-)

diff --git a/pkg/sentry/platform/systrap/shared_context.go b/pkg/sentry/platform/systrap/shared_context.go
index 4fbaade8f..b2a8faef2 100644
--- a/pkg/sentry/platform/systrap/shared_context.go
+++ b/pkg/sentry/platform/systrap/shared_context.go
@@ -19,8 +19,10 @@ import (
 	"runtime"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/sentry/platform"
 	"gvisor.dev/gvisor/pkg/sentry/platform/systrap/sysmsg"
 	"gvisor.dev/gvisor/pkg/syncevent"
@@ -186,9 +188,42 @@ func (sc *sharedContext) resetAcked() {
 	atomic.StoreUint32(&sc.shared.Acked, ackReset)
 }
 
+const (
+	contextPreemptTimeoutNsec = 10 * 1000 * 1000 // 10ms
+	contextCheckupTimeoutSec  = 5
+	stuckContextTimeout       = 30 * time.Second
+)
+
 func (sc *sharedContext) sleepOnState(state sysmsg.ContextState) {
-	if errno := sc.shared.SleepOnState(state, sc); errno != 0 {
-		panic(fmt.Sprintf("error waiting for state: %v", errno))
+	timeout := unix.Timespec{
+		Sec:  0,
+		Nsec: contextPreemptTimeoutNsec,
+	}
+	sentInterruptOnce := false
+	deadline := time.Now().Add(stuckContextTimeout)
+	for sc.state() == state {
+		errno := sc.shared.SleepOnState(state, &timeout)
+		if errno == 0 {
+			continue
+		}
+		if errno != unix.ETIMEDOUT {
+			panic(fmt.Sprintf("error waiting for state: %v", errno))
+		}
+		if time.Now().After(deadline) {
+			log.Warningf("Systrap task goroutine has been waiting on ThreadContext.State futex too long. ThreadContext: %s", sc.shared)
+		}
+		if sentInterruptOnce {
+			log.Warningf("The context is still running: %s", sc)
+			continue
+		}
+
+		if !sc.isAcked() || sc.subprocess.contextQueue.isEmpty() {
+			continue
+		}
+		sc.NotifyInterrupt()
+		sentInterruptOnce = true
+		timeout.Sec = contextCheckupTimeoutSec
+		timeout.Nsec = 0
 	}
 }
 
diff --git a/pkg/sentry/platform/systrap/sysmsg/BUILD b/pkg/sentry/platform/systrap/sysmsg/BUILD
index 43c583b09..145a752b6 100644
--- a/pkg/sentry/platform/systrap/sysmsg/BUILD
+++ b/pkg/sentry/platform/systrap/sysmsg/BUILD
@@ -131,8 +131,6 @@ go_library(
         "//pkg/abi/linux",
         "//pkg/cpuid",
         "//pkg/hostarch",
-        "//pkg/log",
-        "//pkg/sentry/platform/interrupt",
         "@org_golang_x_sys//unix:go_default_library",
     ],
 )
diff --git a/pkg/sentry/platform/systrap/sysmsg/sysmsg_unsafe.go b/pkg/sentry/platform/systrap/sysmsg/sysmsg_unsafe.go
index 6880f2a91..0261d6e77 100644
--- a/pkg/sentry/platform/systrap/sysmsg/sysmsg_unsafe.go
+++ b/pkg/sentry/platform/systrap/sysmsg/sysmsg_unsafe.go
@@ -20,60 +20,12 @@ import (
 
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/abi/linux"
-	"gvisor.dev/gvisor/pkg/log"
-	"gvisor.dev/gvisor/pkg/sentry/platform/interrupt"
 )
 
-const maxFutexSleepSeconds = 60
-
-// SleepOnState makes the caller sleep on the Msg.State futex.
-func (m *Msg) SleepOnState(curState ThreadState, interruptor interrupt.Receiver) syscall.Errno {
-	futexTimeout := unix.Timespec{
-		Sec:  maxFutexSleepSeconds,
-		Nsec: 0,
-	}
-	sentInterruptOnce := false
-	errno := syscall.Errno(0)
-	for {
-		_, _, errno = unix.Syscall6(unix.SYS_FUTEX, uintptr(unsafe.Pointer(&m.State)),
-			linux.FUTEX_WAIT, uintptr(curState), uintptr(unsafe.Pointer(&futexTimeout)), 0, 0)
-		if errno == unix.ETIMEDOUT {
-			interruptor.NotifyInterrupt()
-			if !sentInterruptOnce {
-				log.Warningf("Systrap task goroutine has been waiting on Msg.State futex too long. Msg: %s", m.String())
-			}
-			sentInterruptOnce = true
-		} else {
-			break
-		}
-	}
-	if errno == unix.EAGAIN || errno == unix.EINTR {
-		errno = 0
-	}
-	return errno
-}
-
 // SleepOnState makes the caller sleep on the ThreadContext.State futex.
-func (c *ThreadContext) SleepOnState(curState ContextState, interruptor interrupt.Receiver) syscall.Errno {
-	futexTimeout := unix.Timespec{
-		Sec:  maxFutexSleepSeconds,
-		Nsec: 0,
-	}
-	sentInterruptOnce := false
-	errno := syscall.Errno(0)
-	for {
-		_, _, errno = unix.Syscall6(unix.SYS_FUTEX, uintptr(unsafe.Pointer(&c.State)),
-			linux.FUTEX_WAIT, uintptr(curState), uintptr(unsafe.Pointer(&futexTimeout)), 0, 0)
-		if errno == unix.ETIMEDOUT {
-			interruptor.NotifyInterrupt()
-			if !sentInterruptOnce {
-				log.Warningf("Systrap task goroutine has been waiting on ThreadContext.State futex too long. ThreadContext: %s", c.String())
-			}
-			sentInterruptOnce = true
-		} else {
-			break
-		}
-	}
+func (c *ThreadContext) SleepOnState(curState ContextState, timeout *unix.Timespec) syscall.Errno {
+	_, _, errno := unix.Syscall6(unix.SYS_FUTEX, uintptr(unsafe.Pointer(&c.State)),
+		linux.FUTEX_WAIT, uintptr(curState), uintptr(unsafe.Pointer(timeout)), 0, 0)
 	if errno == unix.EAGAIN || errno == unix.EINTR {
 		errno = 0
 	}