diff --git a/pkg/cpuid/cpuid.go b/pkg/cpuid/cpuid.go
index 413b84930..df5acf67e 100644
--- a/pkg/cpuid/cpuid.go
+++ b/pkg/cpuid/cpuid.go
@@ -38,7 +38,7 @@ import (
 	"gvisor.dev/gvisor/pkg/sync"
 )
 
-// contextID is the package for context.Context.Value keys.
+// contextID is the package for anyContext.Context.Value keys.
 type contextID int
 
 const (
@@ -51,13 +51,13 @@ const (
 	_AT_HWCAP2 = 26
 )
 
-// context represents context.Context.
-type context interface {
+// anyContext represents context.Context.
+type anyContext interface {
 	Value(key any) any
 }
 
 // FromContext returns the FeatureSet from the context, if available.
-func FromContext(ctx context) FeatureSet {
+func FromContext(ctx anyContext) FeatureSet {
 	v := ctx.Value(CtxFeatureSet)
 	if v == nil {
 		return FeatureSet{} // Panics if used.
diff --git a/pkg/cpuid/static_amd64.go b/pkg/cpuid/static_amd64.go
index 09bcf16b2..f21f2e4fb 100644
--- a/pkg/cpuid/static_amd64.go
+++ b/pkg/cpuid/static_amd64.go
@@ -17,6 +17,8 @@
 
 package cpuid
 
+import "context"
+
 // Static is a static CPUID function.
 //
 // +stateify savable
@@ -90,7 +92,7 @@ func (s Static) ToFeatureSet() FeatureSet {
 }
 
 // afterLoad calls normalize.
-func (s Static) afterLoad() {
+func (s Static) afterLoad(context.Context) {
 	s.normalize()
 }
 
diff --git a/pkg/refs/refs_template.go b/pkg/refs/refs_template.go
index 31cda9945..8f4a332cb 100644
--- a/pkg/refs/refs_template.go
+++ b/pkg/refs/refs_template.go
@@ -17,6 +17,7 @@
 package refs_template
 
 import (
+	"context"
 	"fmt"
 
 	"gvisor.dev/gvisor/pkg/atomicbitops"
@@ -156,7 +157,7 @@ func (r *Refs) DecRef(destroy func()) {
 	}
 }
 
-func (r *Refs) afterLoad() {
+func (r *Refs) afterLoad(context.Context) {
 	if r.ReadRefs() > 0 {
 		refs.Register(r)
 	}
diff --git a/pkg/sentry/arch/arch_x86_impl.go b/pkg/sentry/arch/arch_x86_impl.go
index 654168898..ff37cfdb2 100644
--- a/pkg/sentry/arch/arch_x86_impl.go
+++ b/pkg/sentry/arch/arch_x86_impl.go
@@ -19,6 +19,8 @@
 package arch
 
 import (
+	"context"
+
 	"gvisor.dev/gvisor/pkg/sentry/arch/fpu"
 )
 
@@ -35,6 +37,6 @@ type State struct {
 }
 
 // afterLoad is invoked by stateify.
-func (s *State) afterLoad() {
+func (s *State) afterLoad(context.Context) {
 	s.afterLoadFPState()
 }
diff --git a/pkg/sentry/devices/nvproxy/save_restore.go b/pkg/sentry/devices/nvproxy/save_restore.go
index 0b810f4de..1466be6f7 100644
--- a/pkg/sentry/devices/nvproxy/save_restore.go
+++ b/pkg/sentry/devices/nvproxy/save_restore.go
@@ -15,6 +15,7 @@
 package nvproxy
 
 import (
+	goContext "context"
 	"fmt"
 
 	"gvisor.dev/gvisor/pkg/abi/nvgpu"
@@ -30,7 +31,7 @@ func (n *nvproxy) beforeSave() {
 	n.objsLive = nil
 }
 
-func (n *nvproxy) afterLoad() {
+func (n *nvproxy) afterLoad(goContext.Context) {
 	Init()
 	abiCons, ok := abis[n.version]
 	if !ok {
diff --git a/pkg/sentry/fsimpl/devtmpfs/save_restore.go b/pkg/sentry/fsimpl/devtmpfs/save_restore.go
index 28832d850..d6c508c35 100644
--- a/pkg/sentry/fsimpl/devtmpfs/save_restore.go
+++ b/pkg/sentry/fsimpl/devtmpfs/save_restore.go
@@ -14,8 +14,10 @@
 
 package devtmpfs
 
+import "context"
+
 // afterLoad is invoked by stateify.
-func (fst *FilesystemType) afterLoad() {
+func (fst *FilesystemType) afterLoad(context.Context) {
 	if fst.fs != nil {
 		// Ensure that we don't create another filesystem.
 		fst.initOnce.Do(func() {})
diff --git a/pkg/sentry/fsimpl/fuse/save_restore.go b/pkg/sentry/fsimpl/fuse/save_restore.go
index f0b306f80..d5c600d7b 100644
--- a/pkg/sentry/fsimpl/fuse/save_restore.go
+++ b/pkg/sentry/fsimpl/fuse/save_restore.go
@@ -14,6 +14,8 @@
 
 package fuse
 
-func (fRes *futureResponse) afterLoad() {
+import "context"
+
+func (fRes *futureResponse) afterLoad(context.Context) {
 	fRes.ch = make(chan struct{})
 }
diff --git a/pkg/sentry/fsimpl/gofer/save_restore.go b/pkg/sentry/fsimpl/gofer/save_restore.go
index 3bd193df9..c5952d3af 100644
--- a/pkg/sentry/fsimpl/gofer/save_restore.go
+++ b/pkg/sentry/fsimpl/gofer/save_restore.go
@@ -15,6 +15,7 @@
 package gofer
 
 import (
+	goContext "context"
 	"fmt"
 	"io"
 
@@ -128,7 +129,7 @@ func (d *dentry) beforeSave() {
 }
 
 // afterLoad is invoked by stateify.
-func (d *dentry) afterLoad() {
+func (d *dentry) afterLoad(goContext.Context) {
 	d.readFD = atomicbitops.FromInt32(-1)
 	d.writeFD = atomicbitops.FromInt32(-1)
 	d.mmapFD = atomicbitops.FromInt32(-1)
@@ -138,12 +139,12 @@ func (d *dentry) afterLoad() {
 }
 
 // afterLoad is invoked by stateify.
-func (d *directfsDentry) afterLoad() {
+func (d *directfsDentry) afterLoad(goContext.Context) {
 	d.controlFD = -1
 }
 
 // afterLoad is invoked by stateify.
-func (d *dentryPlatformFile) afterLoad() {
+func (d *dentryPlatformFile) afterLoad(goContext.Context) {
 	if d.hostFileMapper.IsInited() {
 		// Ensure that we don't call d.hostFileMapper.Init() again.
 		d.hostFileMapperInitOnce.Do(func() {})
@@ -151,7 +152,7 @@ func (d *dentryPlatformFile) afterLoad() {
 }
 
 // afterLoad is invoked by stateify.
-func (fd *specialFileFD) afterLoad() {
+func (fd *specialFileFD) afterLoad(goContext.Context) {
 	fd.handle.fd = -1
 	if fd.hostFileMapper.IsInited() {
 		// Ensure that we don't call fd.hostFileMapper.Init() again.
diff --git a/pkg/sentry/fsimpl/host/save_restore.go b/pkg/sentry/fsimpl/host/save_restore.go
index 7f46cbb55..0f654bc29 100644
--- a/pkg/sentry/fsimpl/host/save_restore.go
+++ b/pkg/sentry/fsimpl/host/save_restore.go
@@ -15,6 +15,7 @@
 package host
 
 import (
+	"context"
 	"fmt"
 	"io"
 
@@ -57,7 +58,7 @@ func (i *inode) beforeSave() {
 }
 
 // afterLoad is invoked by stateify.
-func (i *inode) afterLoad() {
+func (i *inode) afterLoad(context.Context) {
 	if i.epollable {
 		if err := unix.SetNonblock(i.hostFD, true); err != nil {
 			panic(fmt.Sprintf("host.inode.afterLoad: failed to set host FD %d non-blocking: %v", i.hostFD, err))
diff --git a/pkg/sentry/fsimpl/iouringfs/iouringfs_state.go b/pkg/sentry/fsimpl/iouringfs/iouringfs_state.go
index 94d9d5163..4a57dcb82 100644
--- a/pkg/sentry/fsimpl/iouringfs/iouringfs_state.go
+++ b/pkg/sentry/fsimpl/iouringfs/iouringfs_state.go
@@ -14,6 +14,8 @@
 
 package iouringfs
 
+import "context"
+
 // beforeSave is invoked by stateify.
 func (fd *FileDescription) beforeSave() {
 	if fd.running.Load() != 0 {
@@ -22,7 +24,7 @@ func (fd *FileDescription) beforeSave() {
 }
 
 // afterLoad is invoked by stateify.
-func (fd *FileDescription) afterLoad() {
+func (fd *FileDescription) afterLoad(context.Context) {
 	// Remap shared buffers.
 	fd.remap = true
 	fd.runC = make(chan struct{}, 1)
diff --git a/pkg/sentry/fsimpl/kernfs/save_restore.go b/pkg/sentry/fsimpl/kernfs/save_restore.go
index 9d4130165..9465d2e76 100644
--- a/pkg/sentry/fsimpl/kernfs/save_restore.go
+++ b/pkg/sentry/fsimpl/kernfs/save_restore.go
@@ -15,18 +15,20 @@
 package kernfs
 
 import (
+	"context"
+
 	"gvisor.dev/gvisor/pkg/refs"
 )
 
 // afterLoad is invoked by stateify.
-func (d *Dentry) afterLoad() {
+func (d *Dentry) afterLoad(context.Context) {
 	if d.refs.Load() >= 0 {
 		refs.Register(d)
 	}
 }
 
 // afterLoad is invoked by stateify.
-func (i *inodePlatformFile) afterLoad() {
+func (i *inodePlatformFile) afterLoad(context.Context) {
 	if i.fileMapper.IsInited() {
 		// Ensure that we don't call i.fileMapper.Init() again.
 		i.fileMapperInitOnce.Do(func() {})
diff --git a/pkg/sentry/fsimpl/overlay/save_restore.go b/pkg/sentry/fsimpl/overlay/save_restore.go
index fe78f153a..1c015bb3f 100644
--- a/pkg/sentry/fsimpl/overlay/save_restore.go
+++ b/pkg/sentry/fsimpl/overlay/save_restore.go
@@ -15,10 +15,12 @@
 package overlay
 
 import (
+	"context"
+
 	"gvisor.dev/gvisor/pkg/refs"
 )
 
-func (d *dentry) afterLoad() {
+func (d *dentry) afterLoad(context.Context) {
 	if d.refs.Load() != -1 {
 		refs.Register(d)
 	}
diff --git a/pkg/sentry/fsimpl/tmpfs/save_restore.go b/pkg/sentry/fsimpl/tmpfs/save_restore.go
index 7c99eb358..2692d29f7 100644
--- a/pkg/sentry/fsimpl/tmpfs/save_restore.go
+++ b/pkg/sentry/fsimpl/tmpfs/save_restore.go
@@ -15,6 +15,7 @@
 package tmpfs
 
 import (
+	goContext "context"
 	"fmt"
 
 	"gvisor.dev/gvisor/pkg/context"
@@ -23,7 +24,7 @@ import (
 )
 
 // afterLoad is called by stateify.
-func (fs *filesystem) afterLoad() {
+func (fs *filesystem) afterLoad(goContext.Context) {
 	if !fs.privateMF {
 		fs.mf = fs.mfp.MemoryFile()
 	}
diff --git a/pkg/sentry/fsutil/host_file_mapper_state.go b/pkg/sentry/fsutil/host_file_mapper_state.go
index 576d2a3df..a98b39d6c 100644
--- a/pkg/sentry/fsutil/host_file_mapper_state.go
+++ b/pkg/sentry/fsutil/host_file_mapper_state.go
@@ -14,7 +14,9 @@
 
 package fsutil
 
+import "context"
+
 // afterLoad is invoked by stateify.
-func (f *HostFileMapper) afterLoad() {
+func (f *HostFileMapper) afterLoad(context.Context) {
 	f.mappings = make(map[uint64]mapping)
 }
diff --git a/pkg/sentry/inet/namespace.go b/pkg/sentry/inet/namespace.go
index 4a7971739..d83082366 100644
--- a/pkg/sentry/inet/namespace.go
+++ b/pkg/sentry/inet/namespace.go
@@ -15,6 +15,8 @@
 package inet
 
 import (
+	goContext "context"
+
 	"gvisor.dev/gvisor/pkg/context"
 	"gvisor.dev/gvisor/pkg/sentry/fsimpl/nsfs"
 	"gvisor.dev/gvisor/pkg/sentry/kernel/auth"
@@ -151,7 +153,7 @@ func (n *Namespace) init() {
 }
 
 // afterLoad is invoked by stateify.
-func (n *Namespace) afterLoad() {
+func (n *Namespace) afterLoad(goContext.Context) {
 	n.init()
 }
 
diff --git a/pkg/sentry/kernel/pipe/save_restore.go b/pkg/sentry/kernel/pipe/save_restore.go
index f135827de..af346c104 100644
--- a/pkg/sentry/kernel/pipe/save_restore.go
+++ b/pkg/sentry/kernel/pipe/save_restore.go
@@ -15,11 +15,13 @@
 package pipe
 
 import (
+	"context"
+
 	"gvisor.dev/gvisor/pkg/safemem"
 )
 
 // afterLoad is called by stateify.
-func (p *Pipe) afterLoad() {
+func (p *Pipe) afterLoad(context.Context) {
 	p.bufBlocks[0] = safemem.BlockFromSafeSlice(p.buf)
 	p.bufBlocks[1] = p.bufBlocks[0]
 	p.bufBlockSeq = safemem.BlockSeqFromSlice(p.bufBlocks[:])
diff --git a/pkg/sentry/kernel/task.go b/pkg/sentry/kernel/task.go
index 4b0a131a4..38bfdd0e8 100644
--- a/pkg/sentry/kernel/task.go
+++ b/pkg/sentry/kernel/task.go
@@ -630,7 +630,7 @@ func (t *Task) loadSeccomp(seccompData *taskSeccomp) {
 }
 
 // afterLoad is invoked by stateify.
-func (t *Task) afterLoad() {
+func (t *Task) afterLoad(gocontext.Context) {
 	t.updateInfoLocked()
 	if ts := t.seccomp.Load().(*taskSeccomp); ts != nil {
 		ts.populateCache(t)
diff --git a/pkg/sentry/kernel/timekeeper_state.go b/pkg/sentry/kernel/timekeeper_state.go
index 8e961c832..38ad158f5 100644
--- a/pkg/sentry/kernel/timekeeper_state.go
+++ b/pkg/sentry/kernel/timekeeper_state.go
@@ -15,6 +15,8 @@
 package kernel
 
 import (
+	"context"
+
 	"gvisor.dev/gvisor/pkg/sentry/time"
 )
 
@@ -36,6 +38,6 @@ func (t *Timekeeper) beforeSave() {
 }
 
 // afterLoad is invoked by stateify.
-func (t *Timekeeper) afterLoad() {
+func (t *Timekeeper) afterLoad(context.Context) {
 	t.restored = make(chan struct{})
 }
diff --git a/pkg/sentry/mm/aio_context.go b/pkg/sentry/mm/aio_context.go
index d80ae11d7..f5210b48c 100644
--- a/pkg/sentry/mm/aio_context.go
+++ b/pkg/sentry/mm/aio_context.go
@@ -126,52 +126,52 @@ type AIOContext struct {
 }
 
 // destroy marks the context dead.
-func (ctx *AIOContext) destroy() {
-	ctx.mu.Lock()
-	defer ctx.mu.Unlock()
-	ctx.dead = true
-	ctx.checkForDone()
+func (aio *AIOContext) destroy() {
+	aio.mu.Lock()
+	defer aio.mu.Unlock()
+	aio.dead = true
+	aio.checkForDone()
 }
 
 // Preconditions: ctx.mu must be held by caller.
-func (ctx *AIOContext) checkForDone() {
-	if ctx.dead && ctx.outstanding == 0 {
-		close(ctx.requestReady)
-		ctx.requestReady = nil
+func (aio *AIOContext) checkForDone() {
+	if aio.dead && aio.outstanding == 0 {
+		close(aio.requestReady)
+		aio.requestReady = nil
 	}
 }
 
 // Prepare reserves space for a new request, returning nil if available.
 // Returns EAGAIN if the context is busy and EINVAL if the context is dead.
-func (ctx *AIOContext) Prepare() error {
-	ctx.mu.Lock()
-	defer ctx.mu.Unlock()
-	if ctx.dead {
+func (aio *AIOContext) Prepare() error {
+	aio.mu.Lock()
+	defer aio.mu.Unlock()
+	if aio.dead {
 		// Context died after the caller looked it up.
 		return linuxerr.EINVAL
 	}
-	if ctx.outstanding >= ctx.maxOutstanding {
+	if aio.outstanding >= aio.maxOutstanding {
 		// Context is busy.
 		return linuxerr.EAGAIN
 	}
-	ctx.outstanding++
+	aio.outstanding++
 	return nil
 }
 
 // PopRequest pops a completed request if available, this function does not do
 // any blocking. Returns false if no request is available.
-func (ctx *AIOContext) PopRequest() (any, bool) {
-	ctx.mu.Lock()
-	defer ctx.mu.Unlock()
+func (aio *AIOContext) PopRequest() (any, bool) {
+	aio.mu.Lock()
+	defer aio.mu.Unlock()
 
 	// Is there anything ready?
-	if e := ctx.results.Front(); e != nil {
-		if ctx.outstanding == 0 {
+	if e := aio.results.Front(); e != nil {
+		if aio.outstanding == 0 {
 			panic("AIOContext outstanding is going negative")
 		}
-		ctx.outstanding--
-		ctx.results.Remove(e)
-		ctx.checkForDone()
+		aio.outstanding--
+		aio.results.Remove(e)
+		aio.checkForDone()
 		return e.data, true
 	}
 	return nil, false
@@ -179,17 +179,17 @@ func (ctx *AIOContext) PopRequest() (any, bool) {
 
 // FinishRequest finishes a pending request. It queues up the data
 // and notifies listeners.
-func (ctx *AIOContext) FinishRequest(data any) {
-	ctx.mu.Lock()
-	defer ctx.mu.Unlock()
+func (aio *AIOContext) FinishRequest(data any) {
+	aio.mu.Lock()
+	defer aio.mu.Unlock()
 
 	// Push to the list and notify opportunistically. The channel notify
 	// here is guaranteed to be safe because outstanding must be non-zero.
 	// The requestReady channel is only closed when outstanding reaches zero.
-	ctx.results.PushBack(&ioResult{data: data})
+	aio.results.PushBack(&ioResult{data: data})
 
 	select {
-	case ctx.requestReady <- struct{}{}:
+	case aio.requestReady <- struct{}{}:
 	default:
 	}
 }
@@ -197,46 +197,46 @@ func (ctx *AIOContext) FinishRequest(data any) {
 // WaitChannel returns a channel that is notified when an AIO request is
 // completed. Returns nil if the context is destroyed and there are no more
 // outstanding requests.
-func (ctx *AIOContext) WaitChannel() chan struct{} {
-	ctx.mu.Lock()
-	defer ctx.mu.Unlock()
-	return ctx.requestReady
+func (aio *AIOContext) WaitChannel() chan struct{} {
+	aio.mu.Lock()
+	defer aio.mu.Unlock()
+	return aio.requestReady
 }
 
 // Dead returns true if the context has been destroyed.
-func (ctx *AIOContext) Dead() bool {
-	ctx.mu.Lock()
-	defer ctx.mu.Unlock()
-	return ctx.dead
+func (aio *AIOContext) Dead() bool {
+	aio.mu.Lock()
+	defer aio.mu.Unlock()
+	return aio.dead
 }
 
 // CancelPendingRequest forgets about a request that hasn't yet completed.
-func (ctx *AIOContext) CancelPendingRequest() {
-	ctx.mu.Lock()
-	defer ctx.mu.Unlock()
+func (aio *AIOContext) CancelPendingRequest() {
+	aio.mu.Lock()
+	defer aio.mu.Unlock()
 
-	if ctx.outstanding == 0 {
+	if aio.outstanding == 0 {
 		panic("AIOContext outstanding is going negative")
 	}
-	ctx.outstanding--
-	ctx.checkForDone()
+	aio.outstanding--
+	aio.checkForDone()
 }
 
 // Drain drops all completed requests. Pending requests remain untouched.
-func (ctx *AIOContext) Drain() {
-	ctx.mu.Lock()
-	defer ctx.mu.Unlock()
+func (aio *AIOContext) Drain() {
+	aio.mu.Lock()
+	defer aio.mu.Unlock()
 
-	if ctx.outstanding == 0 {
+	if aio.outstanding == 0 {
 		return
 	}
-	size := uint32(ctx.results.Len())
-	if ctx.outstanding < size {
+	size := uint32(aio.results.Len())
+	if aio.outstanding < size {
 		panic("AIOContext outstanding is going negative")
 	}
-	ctx.outstanding -= size
-	ctx.results.Reset()
-	ctx.checkForDone()
+	aio.outstanding -= size
+	aio.results.Reset()
+	aio.checkForDone()
 }
 
 // aioMappable implements memmap.MappingIdentity and memmap.Mappable for AIO
diff --git a/pkg/sentry/mm/aio_context_state.go b/pkg/sentry/mm/aio_context_state.go
index e8931922f..61ebe8af0 100644
--- a/pkg/sentry/mm/aio_context_state.go
+++ b/pkg/sentry/mm/aio_context_state.go
@@ -14,7 +14,9 @@
 
 package mm
 
+import "context"
+
 // afterLoad is invoked by stateify.
-func (ctx *AIOContext) afterLoad() {
-	ctx.requestReady = make(chan struct{}, 1)
+func (aio *AIOContext) afterLoad(context.Context) {
+	aio.requestReady = make(chan struct{}, 1)
 }
diff --git a/pkg/sentry/mm/save_restore.go b/pkg/sentry/mm/save_restore.go
index 3a5300690..d11ab793e 100644
--- a/pkg/sentry/mm/save_restore.go
+++ b/pkg/sentry/mm/save_restore.go
@@ -15,6 +15,7 @@
 package mm
 
 import (
+	goContext "context"
 	"fmt"
 
 	"gvisor.dev/gvisor/pkg/context"
@@ -52,7 +53,7 @@ func (mm *MemoryManager) beforeSave() {
 }
 
 // afterLoad is invoked by stateify.
-func (mm *MemoryManager) afterLoad() {
+func (mm *MemoryManager) afterLoad(goContext.Context) {
 	mm.mf = mm.mfp.MemoryFile()
 	mm.haveASIO = mm.p.SupportsAddressSpaceIO()
 	for pseg := mm.pmas.FirstSegment(); pseg.Ok(); pseg = pseg.NextSegment() {
diff --git a/pkg/sentry/platform/kvm/context.go b/pkg/sentry/platform/kvm/context.go
index 4c7ba7004..a860fafe8 100644
--- a/pkg/sentry/platform/kvm/context.go
+++ b/pkg/sentry/platform/kvm/context.go
@@ -24,17 +24,17 @@ import (
 	"gvisor.dev/gvisor/pkg/sentry/platform/interrupt"
 )
 
-// context is an implementation of the platform context.
+// platformContext is an implementation of the platform context.
 //
 // This is a thin wrapper around the machine.
-type context struct {
+type platformContext struct {
 	// machine is the parent machine, and is immutable.
 	machine *machine
 
-	// info is the linux.SignalInfo cached for this context.
+	// info is the linux.SignalInfo cached for this platformContext.
 	info linux.SignalInfo
 
-	// interrupt is the interrupt context.
+	// interrupt is the interrupt platformContext.
 	interrupt interrupt.Forwarder
 }
 
@@ -44,8 +44,8 @@ type tryCPUIDError struct{}
 // Error implements error.Error.
 func (tryCPUIDError) Error() string { return "cpuid emulation failed" }
 
-// Switch runs the provided context in the given address space.
-func (c *context) Switch(ctx pkgcontext.Context, mm platform.MemoryManager, ac *arch.Context64, _ int32) (*linux.SignalInfo, hostarch.AccessType, error) {
+// Switch runs the provided platformContext in the given address space.
+func (c *platformContext) Switch(ctx pkgcontext.Context, mm platform.MemoryManager, ac *arch.Context64, _ int32) (*linux.SignalInfo, hostarch.AccessType, error) {
 	as := mm.AddressSpace()
 	localAS := as.(*addressSpace)
 
@@ -114,18 +114,20 @@ restart:
 }
 
 // Interrupt interrupts the running context.
-func (c *context) Interrupt() {
+func (c *platformContext) Interrupt() {
 	c.interrupt.NotifyInterrupt()
 }
 
 // Release implements platform.Context.Release().
-func (c *context) Release() {}
+func (c *platformContext) Release() {}
 
 // FullStateChanged implements platform.Context.FullStateChanged.
-func (c *context) FullStateChanged() {}
+func (c *platformContext) FullStateChanged() {}
 
 // PullFullState implements platform.Context.PullFullState.
-func (c *context) PullFullState(as platform.AddressSpace, ac *arch.Context64) error { return nil }
+func (c *platformContext) PullFullState(as platform.AddressSpace, ac *arch.Context64) error {
+	return nil
+}
 
 // PrepareSleep implements platform.Context.platform.Context.
-func (*context) PrepareSleep() {}
+func (*platformContext) PrepareSleep() {}
diff --git a/pkg/sentry/platform/kvm/kvm.go b/pkg/sentry/platform/kvm/kvm.go
index d7dec72b1..1c1a590a8 100644
--- a/pkg/sentry/platform/kvm/kvm.go
+++ b/pkg/sentry/platform/kvm/kvm.go
@@ -177,7 +177,7 @@ func (k *KVM) NewAddressSpace(any) (platform.AddressSpace, <-chan struct{}, erro
 
 // NewContext returns an interruptible context.
 func (k *KVM) NewContext(pkgcontext.Context) platform.Context {
-	return &context{
+	return &platformContext{
 		machine: k.machine,
 	}
 }
diff --git a/pkg/sentry/platform/systrap/subprocess.go b/pkg/sentry/platform/systrap/subprocess.go
index 965eca2e2..900514d80 100644
--- a/pkg/sentry/platform/systrap/subprocess.go
+++ b/pkg/sentry/platform/systrap/subprocess.go
@@ -133,8 +133,8 @@ type subprocess struct {
 	mu sync.Mutex
 
 	// faultedContexts is the set of contexts for which it's possible that
-	// context.lastFaultSP == this subprocess.
-	faultedContexts map[*context]struct{}
+	// platformContext.lastFaultSP == this subprocess.
+	faultedContexts map[*platformContext]struct{}
 
 	// sysmsgStackPool is a pool of available sysmsg stacks.
 	sysmsgStackPool pool.Pool
@@ -309,7 +309,7 @@ func newSubprocess(create func() (*thread, error), memoryFile *pgalloc.MemoryFil
 	// Ready.
 	sp := &subprocess{
 		requests:          requests,
-		faultedContexts:   make(map[*context]struct{}),
+		faultedContexts:   make(map[*platformContext]struct{}),
 		sysmsgStackPool:   pool.Pool{Start: 0, Limit: maxSystemThreads},
 		threadContextPool: pool.Pool{Start: 0, Limit: maxGuestContexts},
 		memoryFile:        memoryFile,
@@ -726,7 +726,7 @@ func (s *subprocess) decAwakeContexts() {
 // This function returns true on a system call, false on a signal.
 // The second return value is true if a syscall instruction can be replaced on
 // a function call.
-func (s *subprocess) switchToApp(c *context, ac *arch.Context64) (isSyscall bool, shouldPatchSyscall bool, err *platform.ContextError) {
+func (s *subprocess) switchToApp(c *platformContext, ac *arch.Context64) (isSyscall bool, shouldPatchSyscall bool, err *platform.ContextError) {
 	// Reset necessary registers.
 	regs := &ac.StateData().Regs
 	s.resetSysemuRegs(regs)
@@ -957,7 +957,7 @@ func (s *subprocess) Unmap(addr hostarch.Addr, length uint64) {
 	}
 }
 
-func (s *subprocess) PullFullState(c *context, ac *arch.Context64) error {
+func (s *subprocess) PullFullState(c *platformContext, ac *arch.Context64) error {
 	if !c.sharedContext.isActiveInSubprocess(s) {
 		panic("Attempted to PullFullState for context that is not used in subprocess")
 	}
@@ -1140,7 +1140,7 @@ func (s *subprocess) PostFork() {
 // activateContext activates the context in this subprocess.
 // No-op if the context is already active within the subprocess; if not,
 // deactivates it from its last subprocess.
-func (s *subprocess) activateContext(c *context) error {
+func (s *subprocess) activateContext(c *platformContext) error {
 	if !c.sharedContext.isActiveInSubprocess(s) {
 		c.sharedContext.release()
 		c.sharedContext = nil
diff --git a/pkg/sentry/platform/systrap/subprocess_unsafe.go b/pkg/sentry/platform/systrap/subprocess_unsafe.go
index a09d25785..c31695d3a 100644
--- a/pkg/sentry/platform/systrap/subprocess_unsafe.go
+++ b/pkg/sentry/platform/systrap/subprocess_unsafe.go
@@ -87,7 +87,7 @@ func saveFPState(ctx *sharedContext, ac *arch.Context64) {
 
 // restoreFPStateDecoupledContext writes FPState from c to the thread context
 // shared memory region if there is any need to do so.
-func restoreFPState(ctx *sharedContext, c *context, ac *arch.Context64) {
+func restoreFPState(ctx *sharedContext, c *platformContext, ac *arch.Context64) {
 	if !c.needRestoreFPState {
 		return
 	}
diff --git a/pkg/sentry/platform/systrap/sysmsg_thread.go b/pkg/sentry/platform/systrap/sysmsg_thread.go
index fc0b48762..8ea66f2f4 100644
--- a/pkg/sentry/platform/systrap/sysmsg_thread.go
+++ b/pkg/sentry/platform/systrap/sysmsg_thread.go
@@ -44,7 +44,7 @@ type sysmsgThread struct {
 	msg *sysmsg.Msg
 
 	// context is the last context that ran on this thread.
-	context *context
+	context *platformContext
 
 	// stackRange is a sysmsg stack in the memory file.
 	stackRange memmap.FileRange
diff --git a/pkg/sentry/platform/systrap/systrap.go b/pkg/sentry/platform/systrap/systrap.go
index e64a88ad2..40e08732e 100644
--- a/pkg/sentry/platform/systrap/systrap.go
+++ b/pkg/sentry/platform/systrap/systrap.go
@@ -32,8 +32,8 @@
 //   - install seccomp filters to trap user system calls.
 //   - send a fake SIGSEGV to stop the thread in the signal handler.
 //
-// A context is just a collection of temporary variables. Calling Switch on a
-// context does the following:
+// A platformContext is just a collection of temporary variables. Calling Switch on a
+// platformContext does the following:
 //
 //	Set up proper registers and an FPU state on a stub signal frame.
 //	Wake up a stub thread by changing sysmsg->stage and calling FUTEX_WAKE.
@@ -43,7 +43,7 @@
 //
 //	subprocessPool.mu
 //		subprocess.mu
-//			context.mu
+//			platformContext.mu
 //
 // +checkalignedignore
 package systrap
@@ -111,15 +111,15 @@ var (
 	archState sysmsg.ArchState
 )
 
-// context is an implementation of the platform context.
-type context struct {
+// platformContext is an implementation of the platform context.
+type platformContext struct {
 	// signalInfo is the signal info, if and when a signal is received.
 	signalInfo linux.SignalInfo
 
-	// interrupt is the interrupt context.
+	// interrupt is the interrupt platformContext.
 	interrupt interrupt.Forwarder
 
-	// sharedContext is everything related to this context that is resident in
+	// sharedContext is everything related to this platformContext that is resident in
 	// shared memory with the stub thread.
 	// sharedContext is only accessed on the Task goroutine, therefore it is not
 	// mutex protected.
@@ -128,8 +128,8 @@ type context struct {
 	// mu protects the following fields.
 	mu sync.Mutex
 
-	// If lastFaultSP is non-nil, the last context switch was due to a fault
-	// received while executing lastFaultSP. Only context.Switch may set
+	// If lastFaultSP is non-nil, the last platformContext switch was due to a fault
+	// received while executing lastFaultSP. Only platformContext.Switch may set
 	// lastFaultSP to a non-nil value.
 	lastFaultSP *subprocess
 
@@ -151,7 +151,7 @@ type context struct {
 }
 
 // PullFullState implements platform.Context.PullFullState.
-func (c *context) PullFullState(as platform.AddressSpace, ac *arch.Context64) error {
+func (c *platformContext) PullFullState(as platform.AddressSpace, ac *arch.Context64) error {
 	if !c.needToPullFullState {
 		return nil
 	}
@@ -164,13 +164,13 @@ func (c *context) PullFullState(as platform.AddressSpace, ac *arch.Context64) er
 }
 
 // FullStateChanged implements platform.Context.FullStateChanged.
-func (c *context) FullStateChanged() {
+func (c *platformContext) FullStateChanged() {
 	c.needRestoreFPState = true
 	c.needToPullFullState = false
 }
 
-// Switch runs the provided context in the given address space.
-func (c *context) Switch(ctx pkgcontext.Context, mm platform.MemoryManager, ac *arch.Context64, cpu int32) (*linux.SignalInfo, hostarch.AccessType, error) {
+// Switch runs the provided platformContext in the given address space.
+func (c *platformContext) Switch(ctx pkgcontext.Context, mm platform.MemoryManager, ac *arch.Context64, cpu int32) (*linux.SignalInfo, hostarch.AccessType, error) {
 	as := mm.AddressSpace()
 	s := as.(*subprocess)
 	if err := s.activateContext(c); err != nil {
@@ -209,7 +209,7 @@ restart:
 		faultIP = hostarch.Addr(ac.IP())
 	}
 
-	// Update the context to reflect the outcome of this context switch.
+	// Update the platformContext to reflect the outcome of this context switch.
 	c.mu.Lock()
 	lastFaultSP := c.lastFaultSP
 	lastFaultAddr := c.lastFaultAddr
@@ -272,13 +272,13 @@ restart:
 	return &si, at, platform.ErrContextSignal
 }
 
-// Interrupt interrupts the running guest application associated with this context.
-func (c *context) Interrupt() {
+// Interrupt interrupts the running guest application associated with this platformContext.
+func (c *platformContext) Interrupt() {
 	c.interrupt.NotifyInterrupt()
 }
 
-// Release releases all platform resources used by the context.
-func (c *context) Release() {
+// Release releases all platform resources used by the platformContext.
+func (c *platformContext) Release() {
 	if c.sharedContext != nil {
 		c.sharedContext.release()
 		c.sharedContext = nil
@@ -286,7 +286,7 @@ func (c *context) Release() {
 }
 
 // PrepareSleep implements platform.Context.platform.PrepareSleep.
-func (c *context) PrepareSleep() {
+func (c *platformContext) PrepareSleep() {
 	ctx := c.sharedContext
 	if ctx == nil {
 		return
@@ -378,9 +378,9 @@ func (p *Systrap) NewAddressSpace(any) (platform.AddressSpace, <-chan struct{},
 	return as, nil, err
 }
 
-// NewContext returns an interruptible context.
+// NewContext returns an interruptible platformContext.
 func (*Systrap) NewContext(ctx pkgcontext.Context) platform.Context {
-	return &context{
+	return &platformContext{
 		needRestoreFPState:  true,
 		needToPullFullState: false,
 	}
diff --git a/pkg/sentry/socket/netstack/save_restore.go b/pkg/sentry/socket/netstack/save_restore.go
index c7aaf722a..49428fd60 100644
--- a/pkg/sentry/socket/netstack/save_restore.go
+++ b/pkg/sentry/socket/netstack/save_restore.go
@@ -15,11 +15,13 @@
 package netstack
 
 import (
+	"context"
+
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
 )
 
 // afterLoad is invoked by stateify.
-func (s *Stack) afterLoad() {
+func (s *Stack) afterLoad(context.Context) {
 	s.Stack = stack.StackFromEnv // FIXME(b/36201077)
 	if s.Stack == nil {
 		panic("can't restore without netstack/tcpip/stack.Stack")
diff --git a/pkg/sentry/socket/unix/transport/connectioned_state.go b/pkg/sentry/socket/unix/transport/connectioned_state.go
index fea09e1c5..c29c0e50c 100644
--- a/pkg/sentry/socket/unix/transport/connectioned_state.go
+++ b/pkg/sentry/socket/unix/transport/connectioned_state.go
@@ -14,6 +14,8 @@
 
 package transport
 
+import "context"
+
 // saveAcceptedChan is invoked by stateify.
 func (e *connectionedEndpoint) saveAcceptedChan() []*connectionedEndpoint {
 	// If acceptedChan is nil (i.e. we are not listening) then we will save nil.
@@ -60,6 +62,6 @@ func (e *connectionedEndpoint) beforeSave() {
 }
 
 // afterLoad is invoked by stateify.
-func (e *connectionedEndpoint) afterLoad() {
+func (e *connectionedEndpoint) afterLoad(context.Context) {
 	e.ops.InitHandler(e, &stackHandler{}, getSendBufferLimits, getReceiveBufferLimits)
 }
diff --git a/pkg/sentry/socket/unix/transport/connectionless_state.go b/pkg/sentry/socket/unix/transport/connectionless_state.go
index 1bb71baf7..7d4899bc1 100644
--- a/pkg/sentry/socket/unix/transport/connectionless_state.go
+++ b/pkg/sentry/socket/unix/transport/connectionless_state.go
@@ -14,7 +14,9 @@
 
 package transport
 
+import "context"
+
 // afterLoad is invoked by stateify.
-func (e *connectionlessEndpoint) afterLoad() {
+func (e *connectionlessEndpoint) afterLoad(context.Context) {
 	e.ops.InitHandler(e, &stackHandler{}, getSendBufferLimits, getReceiveBufferLimits)
 }
diff --git a/pkg/sentry/socket/unix/transport/save_restore.go b/pkg/sentry/socket/unix/transport/save_restore.go
index 930b00f50..16dcde517 100644
--- a/pkg/sentry/socket/unix/transport/save_restore.go
+++ b/pkg/sentry/socket/unix/transport/save_restore.go
@@ -14,10 +14,13 @@
 
 package transport
 
-import "fmt"
+import (
+	"context"
+	"fmt"
+)
 
 // afterLoad is invoked by stateify.
-func (c *HostConnectedEndpoint) afterLoad() {
+func (c *HostConnectedEndpoint) afterLoad(context.Context) {
 	if err := c.initFromOptions(); err != nil {
 		panic(fmt.Sprintf("initFromOptions failed: %v", err))
 	}
diff --git a/pkg/sentry/vfs/save_restore.go b/pkg/sentry/vfs/save_restore.go
index fe677530f..6732ce929 100644
--- a/pkg/sentry/vfs/save_restore.go
+++ b/pkg/sentry/vfs/save_restore.go
@@ -15,6 +15,7 @@
 package vfs
 
 import (
+	goContext "context"
 	"fmt"
 	"sync/atomic"
 
@@ -122,14 +123,14 @@ func (vfs *VirtualFilesystem) loadMounts(mounts []*Mount) {
 func (mnt *Mount) loadKey(vd VirtualDentry) { mnt.setKey(vd) }
 
 // afterLoad is called by stateify.
-func (mnt *Mount) afterLoad() {
+func (mnt *Mount) afterLoad(goContext.Context) {
 	if mnt.refs.Load() != 0 {
 		refs.Register(mnt)
 	}
 }
 
 // afterLoad is called by stateify.
-func (epi *epollInterest) afterLoad() {
+func (epi *epollInterest) afterLoad(goContext.Context) {
 	// Mark all epollInterests as ready after restore so that the next call to
 	// EpollInstance.ReadEvents() rechecks their readiness.
 	epi.waiter.NotifyEvent(waiter.EventMaskFromLinux(epi.mask))
diff --git a/pkg/state/decode.go b/pkg/state/decode.go
index eb826e60a..5fe76eccd 100644
--- a/pkg/state/decode.go
+++ b/pkg/state/decode.go
@@ -385,7 +385,7 @@ func (ds *decodeState) decodeStruct(ods *objectDecodeState, obj reflect.Value, e
 	if sl, ok := obj.Addr().Interface().(SaverLoader); ok {
 		// Note: may be a registered empty struct which does not
 		// implement the saver/loader interfaces.
-		sl.StateLoad(Source{internal: od})
+		sl.StateLoad(ds.ctx, Source{internal: od})
 	}
 }
 
diff --git a/pkg/state/state.go b/pkg/state/state.go
index 4a9e6eadf..0b62eb9c6 100644
--- a/pkg/state/state.go
+++ b/pkg/state/state.go
@@ -211,7 +211,7 @@ type SaverLoader interface {
 	StateSave(Sink)
 
 	// StateLoad loads the state of the object.
-	StateLoad(Source)
+	StateLoad(context.Context, Source)
 }
 
 // Source is used for Type.StateLoad.
diff --git a/pkg/state/tests/bench.go b/pkg/state/tests/bench.go
index 40869cdfb..d253dbf23 100644
--- a/pkg/state/tests/bench.go
+++ b/pkg/state/tests/bench.go
@@ -14,11 +14,13 @@
 
 package tests
 
+import "context"
+
 // +stateify savable
 type benchStruct struct {
 	B *benchStruct // Must be exported for gob.
 }
 
-func (b *benchStruct) afterLoad() {
+func (b *benchStruct) afterLoad(context.Context) {
 	// Do nothing, just force scheduling.
 }
diff --git a/pkg/state/tests/integer.go b/pkg/state/tests/integer.go
index ca403eed1..cc6ba6061 100644
--- a/pkg/state/tests/integer.go
+++ b/pkg/state/tests/integer.go
@@ -15,6 +15,8 @@
 package tests
 
 import (
+	"context"
+
 	"gvisor.dev/gvisor/pkg/state"
 )
 
@@ -28,7 +30,7 @@ func (t *truncatingUint8) StateSave(m state.Sink) {
 	m.Save(0, &t.save)
 }
 
-func (t *truncatingUint8) StateLoad(m state.Source) {
+func (t *truncatingUint8) StateLoad(_ context.Context, m state.Source) {
 	m.Load(0, &t.load)
 	t.save = uint64(t.load)
 	t.load = 0
@@ -46,7 +48,7 @@ func (t *truncatingUint16) StateSave(m state.Sink) {
 	m.Save(0, &t.save)
 }
 
-func (t *truncatingUint16) StateLoad(m state.Source) {
+func (t *truncatingUint16) StateLoad(_ context.Context, m state.Source) {
 	m.Load(0, &t.load)
 	t.save = uint64(t.load)
 	t.load = 0
@@ -64,7 +66,7 @@ func (t *truncatingUint32) StateSave(m state.Sink) {
 	m.Save(0, &t.save)
 }
 
-func (t *truncatingUint32) StateLoad(m state.Source) {
+func (t *truncatingUint32) StateLoad(_ context.Context, m state.Source) {
 	m.Load(0, &t.load)
 	t.save = uint64(t.load)
 	t.load = 0
@@ -82,7 +84,7 @@ func (t *truncatingInt8) StateSave(m state.Sink) {
 	m.Save(0, &t.save)
 }
 
-func (t *truncatingInt8) StateLoad(m state.Source) {
+func (t *truncatingInt8) StateLoad(_ context.Context, m state.Source) {
 	m.Load(0, &t.load)
 	t.save = int64(t.load)
 	t.load = 0
@@ -100,7 +102,7 @@ func (t *truncatingInt16) StateSave(m state.Sink) {
 	m.Save(0, &t.save)
 }
 
-func (t *truncatingInt16) StateLoad(m state.Source) {
+func (t *truncatingInt16) StateLoad(_ context.Context, m state.Source) {
 	m.Load(0, &t.load)
 	t.save = int64(t.load)
 	t.load = 0
@@ -118,7 +120,7 @@ func (t *truncatingInt32) StateSave(m state.Sink) {
 	m.Save(0, &t.save)
 }
 
-func (t *truncatingInt32) StateLoad(m state.Source) {
+func (t *truncatingInt32) StateLoad(_ context.Context, m state.Source) {
 	m.Load(0, &t.load)
 	t.save = int64(t.load)
 	t.load = 0
@@ -136,7 +138,7 @@ func (t *truncatingFloat32) StateSave(m state.Sink) {
 	m.Save(0, &t.save)
 }
 
-func (t *truncatingFloat32) StateLoad(m state.Source) {
+func (t *truncatingFloat32) StateLoad(_ context.Context, m state.Source) {
 	m.Load(0, &t.load)
 	t.save = float64(t.load)
 	t.load = 0
@@ -154,7 +156,7 @@ func (t *truncatingComplex64) StateSave(m state.Sink) {
 	m.Save(0, &t.save)
 }
 
-func (t *truncatingComplex64) StateLoad(m state.Source) {
+func (t *truncatingComplex64) StateLoad(_ context.Context, m state.Source) {
 	m.Load(0, &t.load)
 	t.save = complex128(t.load)
 	t.load = 0
diff --git a/pkg/state/tests/load.go b/pkg/state/tests/load.go
index 0023948a9..41461c9ca 100644
--- a/pkg/state/tests/load.go
+++ b/pkg/state/tests/load.go
@@ -14,6 +14,8 @@
 
 package tests
 
+import "context"
+
 // +stateify savable
 type genericContainer struct {
 	v any
@@ -24,7 +26,7 @@ type afterLoadStruct struct {
 	v int `state:"nosave"`
 }
 
-func (a *afterLoadStruct) afterLoad() {
+func (a *afterLoadStruct) afterLoad(context.Context) {
 	a.v++
 }
 
@@ -51,7 +53,7 @@ type badCycleStruct struct {
 	b *badCycleStruct `state:"wait"`
 }
 
-func (b *badCycleStruct) afterLoad() {
+func (b *badCycleStruct) afterLoad(context.Context) {
 	if b.b != b {
 		// This is not executable, since AfterLoad requires that the
 		// object and all dependencies are complete. This should cause
diff --git a/pkg/state/tests/register_test.go b/pkg/state/tests/register_test.go
index 2199d6b01..ba61fd0af 100644
--- a/pkg/state/tests/register_test.go
+++ b/pkg/state/tests/register_test.go
@@ -18,6 +18,7 @@
 package tests
 
 import (
+	"context"
 	"testing"
 
 	"gvisor.dev/gvisor/pkg/state"
@@ -44,7 +45,7 @@ type fakerWithSaverLoader struct {
 
 func (f *fakerWithSaverLoader) StateSave(m state.Sink) {}
 
-func (f *fakerWithSaverLoader) StateLoad(m state.Source) {}
+func (f *fakerWithSaverLoader) StateLoad(_ context.Context, m state.Source) {}
 
 // fakerOther calls itself .. uh, itself?
 type fakerOther string
@@ -91,7 +92,7 @@ func (f *fakerOtherSaverLoader) StateFields() []string {
 
 func (f *fakerOtherSaverLoader) StateSave(m state.Sink) {}
 
-func (f *fakerOtherSaverLoader) StateLoad(m state.Source) {}
+func (f *fakerOtherSaverLoader) StateLoad(_ context.Context, m state.Source) {}
 
 func newFakerOtherSaverLoader(name string) *fakerOtherSaverLoader {
 	f := fakerOtherSaverLoader(name)
diff --git a/pkg/tcpip/stack/iptables.go b/pkg/tcpip/stack/iptables.go
index fb0d4a3b0..7295700a2 100644
--- a/pkg/tcpip/stack/iptables.go
+++ b/pkg/tcpip/stack/iptables.go
@@ -15,6 +15,7 @@
 package stack
 
 import (
+	"context"
 	"fmt"
 	"math/rand"
 	"reflect"
@@ -613,7 +614,7 @@ func (it *IPTables) beforeSave() {
 }
 
 // afterLoad is invoked by stateify.
-func (it *IPTables) afterLoad() {
+func (it *IPTables) afterLoad(context.Context) {
 	it.startReaper(reaperDelay)
 }
 
diff --git a/pkg/tcpip/stdclock_state.go b/pkg/tcpip/stdclock_state.go
index 25be1755f..530b46ecf 100644
--- a/pkg/tcpip/stdclock_state.go
+++ b/pkg/tcpip/stdclock_state.go
@@ -14,7 +14,10 @@
 
 package tcpip
 
-import "time"
+import (
+	"context"
+	"time"
+)
 
 // beforeSave is invoked by stateify.
 func (s *stdClock) beforeSave() {
@@ -22,6 +25,6 @@ func (s *stdClock) beforeSave() {
 }
 
 // afterLoad is invoked by stateify.
-func (s *stdClock) afterLoad() {
+func (s *stdClock) afterLoad(context.Context) {
 	s.baseTime = time.Now()
 }
diff --git a/pkg/tcpip/transport/icmp/endpoint_state.go b/pkg/tcpip/transport/icmp/endpoint_state.go
index 54752dd26..6bc6b0e69 100644
--- a/pkg/tcpip/transport/icmp/endpoint_state.go
+++ b/pkg/tcpip/transport/icmp/endpoint_state.go
@@ -15,6 +15,7 @@
 package icmp
 
 import (
+	"context"
 	"fmt"
 	"time"
 
@@ -34,7 +35,7 @@ func (p *icmpPacket) loadReceivedAt(nsec int64) {
 }
 
 // afterLoad is invoked by stateify.
-func (e *endpoint) afterLoad() {
+func (e *endpoint) afterLoad(context.Context) {
 	stack.StackFromEnv.RegisterRestoredEndpoint(e)
 }
 
diff --git a/pkg/tcpip/transport/packet/endpoint_state.go b/pkg/tcpip/transport/packet/endpoint_state.go
index 74203fa83..7543662d0 100644
--- a/pkg/tcpip/transport/packet/endpoint_state.go
+++ b/pkg/tcpip/transport/packet/endpoint_state.go
@@ -15,6 +15,7 @@
 package packet
 
 import (
+	"context"
 	"fmt"
 	"time"
 
@@ -40,7 +41,7 @@ func (ep *endpoint) beforeSave() {
 }
 
 // afterLoad is invoked by stateify.
-func (ep *endpoint) afterLoad() {
+func (ep *endpoint) afterLoad(context.Context) {
 	ep.mu.Lock()
 	defer ep.mu.Unlock()
 
diff --git a/pkg/tcpip/transport/raw/endpoint_state.go b/pkg/tcpip/transport/raw/endpoint_state.go
index 1bda0b8b2..237acf7cd 100644
--- a/pkg/tcpip/transport/raw/endpoint_state.go
+++ b/pkg/tcpip/transport/raw/endpoint_state.go
@@ -15,6 +15,7 @@
 package raw
 
 import (
+	"context"
 	"fmt"
 	"time"
 
@@ -33,7 +34,7 @@ func (p *rawPacket) loadReceivedAt(nsec int64) {
 }
 
 // afterLoad is invoked by stateify.
-func (e *endpoint) afterLoad() {
+func (e *endpoint) afterLoad(context.Context) {
 	stack.StackFromEnv.RegisterRestoredEndpoint(e)
 }
 
diff --git a/pkg/tcpip/transport/tcp/endpoint_state.go b/pkg/tcpip/transport/tcp/endpoint_state.go
index c817b0fb8..281326b10 100644
--- a/pkg/tcpip/transport/tcp/endpoint_state.go
+++ b/pkg/tcpip/transport/tcp/endpoint_state.go
@@ -15,6 +15,7 @@
 package tcp
 
 import (
+	"context"
 	"fmt"
 
 	"gvisor.dev/gvisor/pkg/atomicbitops"
@@ -109,7 +110,7 @@ func (e *endpoint) loadState(epState EndpointState) {
 }
 
 // afterLoad is invoked by stateify.
-func (e *endpoint) afterLoad() {
+func (e *endpoint) afterLoad(context.Context) {
 	// RacyLoad() can be used because we are initializing e.
 	e.origEndpointState = e.state.RacyLoad()
 	// Restore the endpoint to InitialState as it will be moved to
diff --git a/pkg/tcpip/transport/udp/endpoint_state.go b/pkg/tcpip/transport/udp/endpoint_state.go
index 546840b6c..9fe4eb2fa 100644
--- a/pkg/tcpip/transport/udp/endpoint_state.go
+++ b/pkg/tcpip/transport/udp/endpoint_state.go
@@ -15,6 +15,7 @@
 package udp
 
 import (
+	"context"
 	"fmt"
 	"time"
 
@@ -34,7 +35,7 @@ func (p *udpPacket) loadReceivedAt(nsec int64) {
 }
 
 // afterLoad is invoked by stateify.
-func (e *endpoint) afterLoad() {
+func (e *endpoint) afterLoad(context.Context) {
 	stack.StackFromEnv.RegisterRestoredEndpoint(e)
 }
 
diff --git a/tools/go_stateify/main.go b/tools/go_stateify/main.go
index b07b76e06..002dafba5 100644
--- a/tools/go_stateify/main.go
+++ b/tools/go_stateify/main.go
@@ -237,6 +237,7 @@ func main() {
 		once.Do(func() {
 			// Emit the imports.
 			fmt.Fprint(outputFile, "import (\n")
+			fmt.Fprint(outputFile, " \"context\"\n")
 			if *statePkg != "" {
 				fmt.Fprintf(outputFile, "	\"%s\"\n", *statePkg)
 			}
@@ -443,7 +444,7 @@ func main() {
 						methodName: "afterLoad",
 					}]
 					if !hasAfterLoad && generateSaverLoader {
-						fmt.Fprintf(outputFile, "func (%s *%s) afterLoad() {}\n\n", recv, ts.Name.Name)
+						fmt.Fprintf(outputFile, "func (%s *%s) afterLoad(context.Context) {}\n\n", recv, ts.Name.Name)
 					}
 
 					// Generate the load method.
@@ -451,7 +452,7 @@ func main() {
 					// N.B. See the comment above for the save method.
 					if generateSaverLoader {
 						fmt.Fprintf(outputFile, "// +checklocksignore\n")
-						fmt.Fprintf(outputFile, "func (%s *%s) StateLoad(stateSourceObject %sSource) {\n", recv, ts.Name.Name, statePrefix)
+						fmt.Fprintf(outputFile, "func (%s *%s) StateLoad(ctx context.Context, stateSourceObject %sSource) {\n", recv, ts.Name.Name, statePrefix)
 						scanFields(x, scanFunctions{normal: emitLoad, wait: emitLoadWait})
 						scanFields(x, scanFunctions{value: emitLoadValue})
 						if hasAfterLoad {
@@ -459,7 +460,7 @@ func main() {
 							// AfterLoad is called, the object encodes a dependency on
 							// referred objects (i.e. fields). This means that afterLoad
 							// will not be called until the other afterLoads are called.
-							fmt.Fprintf(outputFile, "	stateSourceObject.AfterLoad(%s.afterLoad)\n", recv)
+							fmt.Fprintf(outputFile, "	stateSourceObject.AfterLoad(func () { %s.afterLoad(ctx) })\n", recv)
 						}
 						fmt.Fprintf(outputFile, "}\n\n")
 					}
@@ -504,8 +505,8 @@ func main() {
 							fmt.Fprintf(outputFile, "	(*%s)(%s).StateSave(stateSinkObject)\n", typeName, recv)
 							fmt.Fprintf(outputFile, "}\n\n")
 							fmt.Fprintf(outputFile, "// +checklocksignore\n")
-							fmt.Fprintf(outputFile, "func (%s *%s) StateLoad(stateSourceObject %sSource) {\n", recv, ts.Name.Name, statePrefix)
-							fmt.Fprintf(outputFile, "	(*%s)(%s).StateLoad(stateSourceObject)\n", typeName, recv)
+							fmt.Fprintf(outputFile, "func (%s *%s) StateLoad(ctx context.Context, stateSourceObject %sSource) {\n", recv, ts.Name.Name, statePrefix)
+							fmt.Fprintf(outputFile, "	(*%s)(%s).StateLoad(ctx, stateSourceObject)\n", typeName, recv)
 							fmt.Fprintf(outputFile, "}\n\n")
 						}
 					}