4 files changed, 51 insertions, 42 deletions
diff --git a/pkg/amutex/BUILD b/pkg/amutex/BUILD
index ffc918846..bd3a5cce9 100644
--- a/pkg/amutex/BUILD
+++ b/pkg/amutex/BUILD
@@ -6,7 +6,10 @@ go_library(
     name = "amutex",
     srcs = ["amutex.go"],
     visibility = ["//:sandbox"],
-    deps = ["//pkg/syserror"],
+    deps = [
+        "//pkg/context",
+        "//pkg/syserror",
+    ],
 )
 
 go_test(
diff --git a/pkg/amutex/amutex.go b/pkg/amutex/amutex.go
index a078a31db..d7acc1d9f 100644
--- a/pkg/amutex/amutex.go
+++ b/pkg/amutex/amutex.go
@@ -19,41 +19,17 @@ package amutex
 import (
 	"sync/atomic"
 
+	"gvisor.dev/gvisor/pkg/context"
 	"gvisor.dev/gvisor/pkg/syserror"
 )
 
 // Sleeper must be implemented by users of the abortable mutex to allow for
 // cancellation of waits.
-type Sleeper interface {
-	// SleepStart is called by the AbortableMutex.Lock() function when the
-	// mutex is contended and the goroutine is about to sleep.
-	//
-	// A channel can be returned that causes the sleep to be canceled if
-	// it's readable. If no cancellation is desired, nil can be returned.
-	SleepStart() <-chan struct{}
-
-	// SleepFinish is called by AbortableMutex.Lock() once a contended mutex
-	// is acquired or the wait is aborted.
-	SleepFinish(success bool)
-
-	// Interrupted returns true if the wait is aborted.
-	Interrupted() bool
-}
+type Sleeper = context.ChannelSleeper
 
 // NoopSleeper is a stateless no-op implementation of Sleeper for anonymous
 // embedding in other types that do not support cancelation.
-type NoopSleeper struct{}
-
-// SleepStart implements Sleeper.SleepStart.
-func (NoopSleeper) SleepStart() <-chan struct{} {
-	return nil
-}
-
-// SleepFinish implements Sleeper.SleepFinish.
-func (NoopSleeper) SleepFinish(success bool) {}
-
-// Interrupted implements Sleeper.Interrupted.
-func (NoopSleeper) Interrupted() bool { return false }
+type NoopSleeper = context.Context
 
 // Block blocks until either receiving from ch succeeds (in which case it
 // returns nil) or sleeper is interrupted (in which case it returns
diff --git a/pkg/context/BUILD b/pkg/context/BUILD
index 239f31149..f33e23bf7 100644
--- a/pkg/context/BUILD
+++ b/pkg/context/BUILD
@@ -7,7 +7,6 @@ go_library(
     srcs = ["context.go"],
     visibility = ["//:sandbox"],
     deps = [
-        "//pkg/amutex",
         "//pkg/log",
     ],
 )
diff --git a/pkg/context/context.go b/pkg/context/context.go
index 5319b6d8d..2613bc752 100644
--- a/pkg/context/context.go
+++ b/pkg/context/context.go
@@ -26,7 +26,6 @@ import (
 	"context"
 	"time"
 
-	"gvisor.dev/gvisor/pkg/amutex"
 	"gvisor.dev/gvisor/pkg/log"
 )
 
@@ -68,9 +67,10 @@ func ThreadGroupIDFromContext(ctx Context) (tgid int32, ok bool) {
 // In both cases, values extracted from the Context should be used instead.
 type Context interface {
 	log.Logger
-	amutex.Sleeper
 	context.Context
 
+	ChannelSleeper
+
 	// UninterruptibleSleepStart indicates the beginning of an uninterruptible
 	// sleep state (equivalent to Linux's TASK_UNINTERRUPTIBLE). If deactivate
 	// is true and the Context represents a Task, the Task's AddressSpace is
@@ -85,29 +85,60 @@ type Context interface {
 	UninterruptibleSleepFinish(activate bool)
 }
 
-// NoopSleeper is a noop implementation of amutex.Sleeper and UninterruptibleSleep
-// methods for anonymous embedding in other types that do not implement sleeps.
-type NoopSleeper struct {
-	amutex.NoopSleeper
+// A ChannelSleeper represents a goroutine that may sleep interruptibly, where
+// interruption is indicated by a channel becoming readable.
+type ChannelSleeper interface {
+	// SleepStart is called before going to sleep interruptibly. If SleepStart
+	// returns a non-nil channel and that channel becomes ready for receiving
+	// while the goroutine is sleeping, the goroutine should be woken, and
+	// SleepFinish(false) should be called. Otherwise, SleepFinish(true) should
+	// be called after the goroutine stops sleeping.
+	SleepStart() <-chan struct{}
+
+	// SleepFinish is called after an interruptibly-sleeping goroutine stops
+	// sleeping, as documented by SleepStart.
+	SleepFinish(success bool)
+
+	// Interrupted returns true if the channel returned by SleepStart is
+	// ready for receiving.
+	Interrupted() bool
+}
+
+// NoopSleeper is a noop implementation of ChannelSleeper and
+// Context.UninterruptibleSleep* methods for anonymous embedding in other types
+// that do not implement special behavior around sleeps.
+type NoopSleeper struct{}
+
+// SleepStart implements ChannelSleeper.SleepStart.
+func (NoopSleeper) SleepStart() <-chan struct{} {
+	return nil
+}
+
+// SleepFinish implements ChannelSleeper.SleepFinish.
+func (NoopSleeper) SleepFinish(success bool) {}
+
+// Interrupted implements ChannelSleeper.Interrupted.
+func (NoopSleeper) Interrupted() bool {
+	return false
 }
 
-// UninterruptibleSleepStart does nothing.
-func (NoopSleeper) UninterruptibleSleepStart(bool) {}
+// UninterruptibleSleepStart implements Context.UninterruptibleSleepStart.
+func (NoopSleeper) UninterruptibleSleepStart(deactivate bool) {}
 
-// UninterruptibleSleepFinish does nothing.
-func (NoopSleeper) UninterruptibleSleepFinish(bool) {}
+// UninterruptibleSleepFinish implements Context.UninterruptibleSleepFinish.
+func (NoopSleeper) UninterruptibleSleepFinish(activate bool) {}
 
-// Deadline returns zero values, meaning no deadline.
+// Deadline implements context.Context.Deadline.
 func (NoopSleeper) Deadline() (time.Time, bool) {
 	return time.Time{}, false
 }
 
-// Done returns nil.
+// Done implements context.Context.Done.
 func (NoopSleeper) Done() <-chan struct{} {
 	return nil
 }
 
-// Err returns nil.
+// Err returns context.Context.Err.
 func (NoopSleeper) Err() error {
 	return nil
 }