diff --git a/pkg/buffer/buffer_state.go b/pkg/buffer/buffer_state.go
index 8b8e15ea2..d57dfa022 100644
--- a/pkg/buffer/buffer_state.go
+++ b/pkg/buffer/buffer_state.go
@@ -14,12 +14,16 @@
 
 package buffer
 
+import (
+	"context"
+)
+
 // saveData is invoked by stateify.
 func (b *Buffer) saveData() []byte {
 	return b.Flatten()
 }
 
 // loadData is invoked by stateify.
-func (b *Buffer) loadData(data []byte) {
+func (b *Buffer) loadData(_ context.Context, data []byte) {
 	*b = MakeWithData(data)
 }
diff --git a/pkg/cpuid/cpuid_amd64.go b/pkg/cpuid/cpuid_amd64.go
index 8118077bd..178428d44 100644
--- a/pkg/cpuid/cpuid_amd64.go
+++ b/pkg/cpuid/cpuid_amd64.go
@@ -18,6 +18,7 @@
 package cpuid
 
 import (
+	"context"
 	"fmt"
 	"io"
 )
@@ -56,7 +57,7 @@ func (fs *FeatureSet) saveFunction() Static {
 }
 
 // loadFunction saves the function as a static query.
-func (fs *FeatureSet) loadFunction(s Static) {
+func (fs *FeatureSet) loadFunction(_ context.Context, s Static) {
 	fs.Function = s
 }
 
diff --git a/pkg/segment/set_state.go b/pkg/segment/set_state.go
index f7cc10105..bd8b2f289 100644
--- a/pkg/segment/set_state.go
+++ b/pkg/segment/set_state.go
@@ -14,6 +14,10 @@
 
 package segment
 
+import (
+	"context"
+)
+
 func (s *Set) saveRoot() []FlatSegment {
 	fs := s.ExportSlice()
 	// The state package saves data in slice capacity beyond slice length; save
@@ -22,7 +26,7 @@ func (s *Set) saveRoot() []FlatSegment {
 	return fs
 }
 
-func (s *Set) loadRoot(fs []FlatSegment) {
+func (s *Set) loadRoot(_ context.Context, fs []FlatSegment) {
 	if err := s.ImportSlice(fs); err != nil {
 		panic(err)
 	}
diff --git a/pkg/sentry/fsimpl/erofs/save_restore.go b/pkg/sentry/fsimpl/erofs/save_restore.go
index 399acf8e8..110b9b459 100644
--- a/pkg/sentry/fsimpl/erofs/save_restore.go
+++ b/pkg/sentry/fsimpl/erofs/save_restore.go
@@ -15,6 +15,7 @@
 package erofs
 
 import (
+	goContext "context"
 	"fmt"
 	"os"
 
@@ -62,6 +63,6 @@ func (d *dentry) saveParent() *dentry {
 }
 
 // loadParent is called by stateify.
-func (d *dentry) loadParent(parent *dentry) {
+func (d *dentry) loadParent(_ goContext.Context, parent *dentry) {
 	d.parent.Store(parent)
 }
diff --git a/pkg/sentry/fsimpl/fuse/connection.go b/pkg/sentry/fsimpl/fuse/connection.go
index 3f5a2987d..f83438c89 100644
--- a/pkg/sentry/fsimpl/fuse/connection.go
+++ b/pkg/sentry/fsimpl/fuse/connection.go
@@ -15,6 +15,7 @@
 package fuse
 
 import (
+	goContext "context"
 	"sync"
 
 	"gvisor.dev/gvisor/pkg/abi/linux"
@@ -193,7 +194,7 @@ func (conn *connection) saveInitializedChan() bool {
 	}
 }
 
-func (conn *connection) loadInitializedChan(closed bool) {
+func (conn *connection) loadInitializedChan(_ goContext.Context, closed bool) {
 	conn.initializedChan = make(chan struct{}, 1)
 	if closed {
 		close(conn.initializedChan)
diff --git a/pkg/sentry/fsimpl/fuse/dev_state.go b/pkg/sentry/fsimpl/fuse/dev_state.go
index 0c79566f9..ba0487acf 100644
--- a/pkg/sentry/fsimpl/fuse/dev_state.go
+++ b/pkg/sentry/fsimpl/fuse/dev_state.go
@@ -14,10 +14,14 @@
 
 package fuse
 
+import (
+	"context"
+)
+
 func (fd *DeviceFD) saveFullQueueCh() int {
 	return cap(fd.fullQueueCh)
 }
 
-func (fd *DeviceFD) loadFullQueueCh(capacity int) {
+func (fd *DeviceFD) loadFullQueueCh(_ context.Context, capacity int) {
 	fd.fullQueueCh = make(chan struct{}, capacity)
 }
diff --git a/pkg/sentry/fsimpl/gofer/save_restore.go b/pkg/sentry/fsimpl/gofer/save_restore.go
index c5952d3af..80dca6845 100644
--- a/pkg/sentry/fsimpl/gofer/save_restore.go
+++ b/pkg/sentry/fsimpl/gofer/save_restore.go
@@ -166,7 +166,7 @@ func (d *dentry) saveParent() *dentry {
 }
 
 // loadParent is called by stateify.
-func (d *dentry) loadParent(parent *dentry) {
+func (d *dentry) loadParent(_ goContext.Context, parent *dentry) {
 	d.parent.Store(parent)
 }
 
diff --git a/pkg/sentry/fsimpl/kernfs/save_restore.go b/pkg/sentry/fsimpl/kernfs/save_restore.go
index 9465d2e76..3d9e2beef 100644
--- a/pkg/sentry/fsimpl/kernfs/save_restore.go
+++ b/pkg/sentry/fsimpl/kernfs/save_restore.go
@@ -41,6 +41,6 @@ func (d *Dentry) saveParent() *Dentry {
 }
 
 // loadParent is called by stateify.
-func (d *Dentry) loadParent(parent *Dentry) {
+func (d *Dentry) loadParent(_ context.Context, parent *Dentry) {
 	d.parent.Store(parent)
 }
diff --git a/pkg/sentry/fsimpl/overlay/save_restore.go b/pkg/sentry/fsimpl/overlay/save_restore.go
index 1c015bb3f..d3e2c5ada 100644
--- a/pkg/sentry/fsimpl/overlay/save_restore.go
+++ b/pkg/sentry/fsimpl/overlay/save_restore.go
@@ -32,6 +32,6 @@ func (d *dentry) saveParent() *dentry {
 }
 
 // loadParent is called by stateify.
-func (d *dentry) loadParent(parent *dentry) {
+func (d *dentry) loadParent(_ context.Context, parent *dentry) {
 	d.parent.Store(parent)
 }
diff --git a/pkg/sentry/fsimpl/tmpfs/save_restore.go b/pkg/sentry/fsimpl/tmpfs/save_restore.go
index 2692d29f7..668bbb3b6 100644
--- a/pkg/sentry/fsimpl/tmpfs/save_restore.go
+++ b/pkg/sentry/fsimpl/tmpfs/save_restore.go
@@ -36,7 +36,7 @@ func (d *dentry) saveParent() *dentry {
 }
 
 // loadParent is called by stateify.
-func (d *dentry) loadParent(parent *dentry) {
+func (d *dentry) loadParent(_ goContext.Context, parent *dentry) {
 	d.parent.Store(parent)
 }
 
diff --git a/pkg/sentry/kernel/fd_table.go b/pkg/sentry/kernel/fd_table.go
index 2d29d94b3..d818f3697 100644
--- a/pkg/sentry/kernel/fd_table.go
+++ b/pkg/sentry/kernel/fd_table.go
@@ -15,6 +15,7 @@
 package kernel
 
 import (
+	goContext "context"
 	"fmt"
 	"math"
 	"strings"
@@ -101,7 +102,7 @@ func (f *FDTable) saveDescriptorTable() map[int32]descriptor {
 	return m
 }
 
-func (f *FDTable) loadDescriptorTable(m map[int32]descriptor) {
+func (f *FDTable) loadDescriptorTable(_ goContext.Context, m map[int32]descriptor) {
 	ctx := context.Background()
 	f.initNoLeakCheck() // Initialize table.
 	f.fdBitmap = bitmap.New(uint32(math.MaxUint16))
diff --git a/pkg/sentry/kernel/kernel_state.go b/pkg/sentry/kernel/kernel_state.go
index 2f8acd37e..65247f3a7 100644
--- a/pkg/sentry/kernel/kernel_state.go
+++ b/pkg/sentry/kernel/kernel_state.go
@@ -15,6 +15,8 @@
 package kernel
 
 import (
+	"context"
+
 	"gvisor.dev/gvisor/pkg/tcpip"
 )
 
@@ -24,7 +26,7 @@ func (k *Kernel) saveDanglingEndpoints() []tcpip.Endpoint {
 }
 
 // loadDanglingEndpoints is invoked by stateify.
-func (k *Kernel) loadDanglingEndpoints(es []tcpip.Endpoint) {
+func (k *Kernel) loadDanglingEndpoints(_ context.Context, es []tcpip.Endpoint) {
 	for _, e := range es {
 		tcpip.AddDanglingEndpoint(e)
 	}
diff --git a/pkg/sentry/kernel/pending_signals_state.go b/pkg/sentry/kernel/pending_signals_state.go
index e77f1a254..b3e9cf104 100644
--- a/pkg/sentry/kernel/pending_signals_state.go
+++ b/pkg/sentry/kernel/pending_signals_state.go
@@ -14,7 +14,11 @@
 
 package kernel
 
-import "gvisor.dev/gvisor/pkg/abi/linux"
+import (
+	"context"
+
+	"gvisor.dev/gvisor/pkg/abi/linux"
+)
 
 // +stateify savable
 type savedPendingSignal struct {
@@ -37,7 +41,7 @@ func (p *pendingSignals) saveSignals() []savedPendingSignal {
 }
 
 // loadSignals is invoked by stateify.
-func (p *pendingSignals) loadSignals(pending []savedPendingSignal) {
+func (p *pendingSignals) loadSignals(_ context.Context, pending []savedPendingSignal) {
 	for _, sps := range pending {
 		p.enqueue(sps.si, sps.timer)
 	}
diff --git a/pkg/sentry/kernel/syscalls_state.go b/pkg/sentry/kernel/syscalls_state.go
index 0b17a562e..9b646f320 100644
--- a/pkg/sentry/kernel/syscalls_state.go
+++ b/pkg/sentry/kernel/syscalls_state.go
@@ -15,6 +15,7 @@
 package kernel
 
 import (
+	"context"
 	"fmt"
 
 	"gvisor.dev/gvisor/pkg/abi"
@@ -38,7 +39,7 @@ func (image *TaskImage) saveSt() syscallTableInfo {
 }
 
 // loadSt loads the SyscallTable.
-func (image *TaskImage) loadSt(sti syscallTableInfo) {
+func (image *TaskImage) loadSt(_ context.Context, sti syscallTableInfo) {
 	st, ok := LookupSyscallTable(sti.OS, sti.Arch)
 	if !ok {
 		panic(fmt.Sprintf("syscall table not found for OS %v, Arch %v", sti.OS, sti.Arch))
diff --git a/pkg/sentry/kernel/task.go b/pkg/sentry/kernel/task.go
index 08fb4e910..6bb5808fe 100644
--- a/pkg/sentry/kernel/task.go
+++ b/pkg/sentry/kernel/task.go
@@ -615,7 +615,7 @@ func (t *Task) savePtraceTracer() *Task {
 	return t.ptraceTracer.Load()
 }
 
-func (t *Task) loadPtraceTracer(tracer *Task) {
+func (t *Task) loadPtraceTracer(_ gocontext.Context, tracer *Task) {
 	t.ptraceTracer.Store(tracer)
 }
 
@@ -623,7 +623,7 @@ func (t *Task) saveSeccomp() *taskSeccomp {
 	return t.seccomp.Load()
 }
 
-func (t *Task) loadSeccomp(seccompData *taskSeccomp) {
+func (t *Task) loadSeccomp(_ gocontext.Context, seccompData *taskSeccomp) {
 	t.seccomp.Store(seccompData)
 }
 
diff --git a/pkg/sentry/kernel/thread_group.go b/pkg/sentry/kernel/thread_group.go
index 8896b9ac3..4e5dd4f6c 100644
--- a/pkg/sentry/kernel/thread_group.go
+++ b/pkg/sentry/kernel/thread_group.go
@@ -15,6 +15,7 @@
 package kernel
 
 import (
+	goContext "context"
 	"sync/atomic"
 
 	"gvisor.dev/gvisor/pkg/abi/linux"
@@ -293,7 +294,7 @@ func (tg *ThreadGroup) saveOldRSeqCritical() *OldRSeqCriticalRegion {
 }
 
 // loadOldRSeqCritical is invoked by stateify.
-func (tg *ThreadGroup) loadOldRSeqCritical(r *OldRSeqCriticalRegion) {
+func (tg *ThreadGroup) loadOldRSeqCritical(_ goContext.Context, r *OldRSeqCriticalRegion) {
 	tg.oldRSeqCritical.Store(r)
 }
 
diff --git a/pkg/sentry/loader/vdso_state.go b/pkg/sentry/loader/vdso_state.go
index db378e90a..949afd14c 100644
--- a/pkg/sentry/loader/vdso_state.go
+++ b/pkg/sentry/loader/vdso_state.go
@@ -15,6 +15,7 @@
 package loader
 
 import (
+	"context"
 	"debug/elf"
 )
 
@@ -40,7 +41,7 @@ func (v *VDSO) savePhdrs() []elfProgHeader {
 }
 
 // loadPhdrs is invoked by stateify.
-func (v *VDSO) loadPhdrs(s []elfProgHeader) {
+func (v *VDSO) loadPhdrs(_ context.Context, s []elfProgHeader) {
 	v.phdrs = make([]elf.ProgHeader, 0, len(s))
 	for _, h := range s {
 		v.phdrs = append(v.phdrs, elf.ProgHeader(h))
diff --git a/pkg/sentry/mm/save_restore.go b/pkg/sentry/mm/save_restore.go
index d11ab793e..0391d6aa5 100644
--- a/pkg/sentry/mm/save_restore.go
+++ b/pkg/sentry/mm/save_restore.go
@@ -113,7 +113,7 @@ func (v *vma) saveRealPerms() int {
 	return b
 }
 
-func (v *vma) loadRealPerms(b int) {
+func (v *vma) loadRealPerms(_ goContext.Context, b int) {
 	if b&vmaRealPermsRead > 0 {
 		v.realPerms.Read = true
 	}
diff --git a/pkg/sentry/socket/netstack/netstack_state.go b/pkg/sentry/socket/netstack/netstack_state.go
index d995cc679..57d9dc310 100644
--- a/pkg/sentry/socket/netstack/netstack_state.go
+++ b/pkg/sentry/socket/netstack/netstack_state.go
@@ -15,6 +15,7 @@
 package netstack
 
 import (
+	"context"
 	"time"
 )
 
@@ -24,7 +25,7 @@ func (s *sock) saveTimestamp() int64 {
 	return s.timestamp.UnixNano()
 }
 
-func (s *sock) loadTimestamp(nsec int64) {
+func (s *sock) loadTimestamp(_ context.Context, nsec int64) {
 	s.readMu.Lock()
 	defer s.readMu.Unlock()
 	s.timestamp = time.Unix(0, nsec)
diff --git a/pkg/sentry/socket/socket_state.go b/pkg/sentry/socket/socket_state.go
index 32e12b238..229f2528c 100644
--- a/pkg/sentry/socket/socket_state.go
+++ b/pkg/sentry/socket/socket_state.go
@@ -15,6 +15,7 @@
 package socket
 
 import (
+	"context"
 	"time"
 )
 
@@ -22,6 +23,6 @@ func (i *IPControlMessages) saveTimestamp() int64 {
 	return i.Timestamp.UnixNano()
 }
 
-func (i *IPControlMessages) loadTimestamp(nsec int64) {
+func (i *IPControlMessages) loadTimestamp(_ context.Context, nsec int64) {
 	i.Timestamp = time.Unix(0, nsec)
 }
diff --git a/pkg/sentry/socket/unix/transport/connectioned_state.go b/pkg/sentry/socket/unix/transport/connectioned_state.go
index c29c0e50c..a2223aa24 100644
--- a/pkg/sentry/socket/unix/transport/connectioned_state.go
+++ b/pkg/sentry/socket/unix/transport/connectioned_state.go
@@ -42,7 +42,7 @@ func (e *connectionedEndpoint) saveAcceptedChan() []*connectionedEndpoint {
 }
 
 // loadAcceptedChan is invoked by stateify.
-func (e *connectionedEndpoint) loadAcceptedChan(acceptedSlice []*connectionedEndpoint) {
+func (e *connectionedEndpoint) loadAcceptedChan(_ context.Context, acceptedSlice []*connectionedEndpoint) {
 	// If acceptedSlice is nil, then acceptedChan should also be nil.
 	if acceptedSlice != nil {
 		// Otherwise, create a new channel with the same capacity as acceptedSlice.
diff --git a/pkg/sentry/vfs/file_description_impl_util.go b/pkg/sentry/vfs/file_description_impl_util.go
index 499d9fed3..1aedc9ca8 100644
--- a/pkg/sentry/vfs/file_description_impl_util.go
+++ b/pkg/sentry/vfs/file_description_impl_util.go
@@ -16,6 +16,7 @@ package vfs
 
 import (
 	"bytes"
+	goContext "context"
 	"io"
 	"math"
 
@@ -289,7 +290,7 @@ func (fd *DynamicBytesFileDescriptionImpl) saveBuf() []byte {
 	return fd.buf.Bytes()
 }
 
-func (fd *DynamicBytesFileDescriptionImpl) loadBuf(p []byte) {
+func (fd *DynamicBytesFileDescriptionImpl) loadBuf(_ goContext.Context, p []byte) {
 	fd.buf.Write(p)
 }
 
diff --git a/pkg/sentry/vfs/mount.go b/pkg/sentry/vfs/mount.go
index 854ff516a..8790ffd2c 100644
--- a/pkg/sentry/vfs/mount.go
+++ b/pkg/sentry/vfs/mount.go
@@ -893,7 +893,7 @@ func (vfs *VirtualFilesystem) disconnectLocked(mnt *Mount) VirtualDentry {
 		delete(mnt.ns.mountpoints, vd.dentry)
 	}
 	vfs.mounts.removeSeqed(mnt)
-	mnt.loadKey(VirtualDentry{}) // Clear mnt.key.
+	mnt.setKey(VirtualDentry{}) // Clear mnt.key.
 	vfsmpmounts := vfs.mountpoints[vd.dentry]
 	delete(vfsmpmounts, mnt)
 	if len(vfsmpmounts) == 0 {
diff --git a/pkg/sentry/vfs/mount_test.go b/pkg/sentry/vfs/mount_test.go
index 060dd8c7c..b05961b71 100644
--- a/pkg/sentry/vfs/mount_test.go
+++ b/pkg/sentry/vfs/mount_test.go
@@ -80,7 +80,7 @@ const enableComparativeBenchmarks = false
 
 func newBenchMount() *Mount {
 	mount := &Mount{}
-	mount.loadKey(VirtualDentry{&Mount{}, &Dentry{}})
+	mount.setKey(VirtualDentry{&Mount{}, &Dentry{}})
 	return mount
 }
 
diff --git a/pkg/sentry/vfs/save_restore.go b/pkg/sentry/vfs/save_restore.go
index 6732ce929..8a5880079 100644
--- a/pkg/sentry/vfs/save_restore.go
+++ b/pkg/sentry/vfs/save_restore.go
@@ -109,7 +109,7 @@ func (vfs *VirtualFilesystem) saveMounts() []*Mount {
 func (mnt *Mount) saveKey() VirtualDentry { return mnt.getKey() }
 
 // loadMounts is called by stateify.
-func (vfs *VirtualFilesystem) loadMounts(mounts []*Mount) {
+func (vfs *VirtualFilesystem) loadMounts(_ goContext.Context, mounts []*Mount) {
 	if mounts == nil {
 		return
 	}
@@ -120,7 +120,7 @@ func (vfs *VirtualFilesystem) loadMounts(mounts []*Mount) {
 }
 
 // loadKey is called by stateify.
-func (mnt *Mount) loadKey(vd VirtualDentry) { mnt.setKey(vd) }
+func (mnt *Mount) loadKey(_ goContext.Context, vd VirtualDentry) { mnt.setKey(vd) }
 
 // afterLoad is called by stateify.
 func (mnt *Mount) afterLoad(goContext.Context) {
diff --git a/pkg/sleep/sleep_unsafe.go b/pkg/sleep/sleep_unsafe.go
index ca7ee6576..eab682dc0 100644
--- a/pkg/sleep/sleep_unsafe.go
+++ b/pkg/sleep/sleep_unsafe.go
@@ -68,6 +68,7 @@
 package sleep
 
 import (
+	"context"
 	"sync/atomic"
 	"unsafe"
 
@@ -129,7 +130,7 @@ func (s *Sleeper) saveSharedList() *Waker {
 }
 
 // loadSharedList is invoked by stateify.
-func (s *Sleeper) loadSharedList(w *Waker) {
+func (s *Sleeper) loadSharedList(_ context.Context, w *Waker) {
 	atomic.StorePointer(&s.sharedList, unsafe.Pointer(w))
 }
 
@@ -408,7 +409,7 @@ func (w *Waker) saveS() wakerState {
 }
 
 // loadS is invoked by stateify.
-func (w *Waker) loadS(ws wakerState) {
+func (w *Waker) loadS(_ context.Context, ws wakerState) {
 	if ws.asserted {
 		atomic.StorePointer(&w.s, unsafe.Pointer(&assertedSleeper))
 	} else {
diff --git a/pkg/state/tests/load.go b/pkg/state/tests/load.go
index 41461c9ca..846314622 100644
--- a/pkg/state/tests/load.go
+++ b/pkg/state/tests/load.go
@@ -39,7 +39,7 @@ func (v *valueLoadStruct) saveV() int64 {
 	return int64(v.v) // Save as int64.
 }
 
-func (v *valueLoadStruct) loadV(value int64) {
+func (v *valueLoadStruct) loadV(_ context.Context, value int64) {
 	v.v = int(value) // Load as int.
 }
 
diff --git a/pkg/state/tests/struct.go b/pkg/state/tests/struct.go
index 012bf5782..28b86287d 100644
--- a/pkg/state/tests/struct.go
+++ b/pkg/state/tests/struct.go
@@ -14,6 +14,10 @@
 
 package tests
 
+import (
+	"context"
+)
+
 type unregisteredEmptyStruct struct{}
 
 // typeOnlyEmptyStruct just implements the state.Type interface.
@@ -82,7 +86,7 @@ func (ifv *innerFieldValue) saveV() *savedFieldValue {
 	return &savedFieldValue{ifv.v}
 }
 
-func (ifv *innerFieldValue) loadV(sfv *savedFieldValue) {
+func (ifv *innerFieldValue) loadV(_ context.Context, sfv *savedFieldValue) {
 	ifv.v = sfv.v
 }
 
diff --git a/pkg/sync/atomicptr/generic_atomicptr_unsafe.go b/pkg/sync/atomicptr/generic_atomicptr_unsafe.go
index fba08b74b..a57db410d 100644
--- a/pkg/sync/atomicptr/generic_atomicptr_unsafe.go
+++ b/pkg/sync/atomicptr/generic_atomicptr_unsafe.go
@@ -8,6 +8,7 @@
 package seqatomic
 
 import (
+	"context"
 	"sync/atomic"
 	"unsafe"
 )
@@ -31,7 +32,7 @@ func (p *AtomicPtr) savePtr() *Value {
 	return p.Load()
 }
 
-func (p *AtomicPtr) loadPtr(v *Value) {
+func (p *AtomicPtr) loadPtr(_ context.Context, v *Value) {
 	p.Store(v)
 }
 
diff --git a/pkg/tcpip/tcpip_state.go b/pkg/tcpip/tcpip_state.go
index 5181e355a..0603ff049 100644
--- a/pkg/tcpip/tcpip_state.go
+++ b/pkg/tcpip/tcpip_state.go
@@ -15,6 +15,7 @@
 package tcpip
 
 import (
+	"context"
 	"time"
 )
 
@@ -22,6 +23,6 @@ func (c *ReceivableControlMessages) saveTimestamp() int64 {
 	return c.Timestamp.UnixNano()
 }
 
-func (c *ReceivableControlMessages) loadTimestamp(nsec int64) {
+func (c *ReceivableControlMessages) loadTimestamp(_ context.Context, nsec int64) {
 	c.Timestamp = time.Unix(0, nsec)
 }
diff --git a/pkg/tcpip/transport/icmp/endpoint_state.go b/pkg/tcpip/transport/icmp/endpoint_state.go
index a58cd5d0e..d7ce14230 100644
--- a/pkg/tcpip/transport/icmp/endpoint_state.go
+++ b/pkg/tcpip/transport/icmp/endpoint_state.go
@@ -30,7 +30,7 @@ func (p *icmpPacket) saveReceivedAt() int64 {
 }
 
 // loadReceivedAt is invoked by stateify.
-func (p *icmpPacket) loadReceivedAt(nsec int64) {
+func (p *icmpPacket) loadReceivedAt(_ context.Context, nsec int64) {
 	p.receivedAt = time.Unix(0, nsec)
 }
 
diff --git a/pkg/tcpip/transport/packet/endpoint_state.go b/pkg/tcpip/transport/packet/endpoint_state.go
index 228f09b7a..78e4a27d3 100644
--- a/pkg/tcpip/transport/packet/endpoint_state.go
+++ b/pkg/tcpip/transport/packet/endpoint_state.go
@@ -29,7 +29,7 @@ func (p *packet) saveReceivedAt() int64 {
 }
 
 // loadReceivedAt is invoked by stateify.
-func (p *packet) loadReceivedAt(nsec int64) {
+func (p *packet) loadReceivedAt(_ context.Context, nsec int64) {
 	p.receivedAt = time.Unix(0, nsec)
 }
 
diff --git a/pkg/tcpip/transport/raw/endpoint_state.go b/pkg/tcpip/transport/raw/endpoint_state.go
index 7d20de0a5..427a9d8e8 100644
--- a/pkg/tcpip/transport/raw/endpoint_state.go
+++ b/pkg/tcpip/transport/raw/endpoint_state.go
@@ -29,7 +29,7 @@ func (p *rawPacket) saveReceivedAt() int64 {
 }
 
 // loadReceivedAt is invoked by stateify.
-func (p *rawPacket) loadReceivedAt(nsec int64) {
+func (p *rawPacket) loadReceivedAt(_ context.Context, nsec int64) {
 	p.receivedAt = time.Unix(0, nsec)
 }
 
diff --git a/pkg/tcpip/transport/tcp/endpoint_state.go b/pkg/tcpip/transport/tcp/endpoint_state.go
index fea5539de..7a6bc8d7f 100644
--- a/pkg/tcpip/transport/tcp/endpoint_state.go
+++ b/pkg/tcpip/transport/tcp/endpoint_state.go
@@ -69,7 +69,7 @@ func (a *acceptQueue) saveEndpoints() []*endpoint {
 }
 
 // loadEndpoints is invoked by stateify.
-func (a *acceptQueue) loadEndpoints(acceptedEndpoints []*endpoint) {
+func (a *acceptQueue) loadEndpoints(_ context.Context, acceptedEndpoints []*endpoint) {
 	for _, ep := range acceptedEndpoints {
 		a.endpoints.PushBack(ep)
 	}
@@ -90,7 +90,7 @@ var connectingLoading sync.WaitGroup
 // Bound endpoint loading happens last.
 
 // loadState is invoked by stateify.
-func (e *endpoint) loadState(epState EndpointState) {
+func (e *endpoint) loadState(_ context.Context, epState EndpointState) {
 	// This is to ensure that the loading wait groups include all applicable
 	// endpoints before any asynchronous calls to the Wait() methods.
 	// For restore purposes we treat TimeWait like a connected endpoint.
diff --git a/pkg/tcpip/transport/tcp/segment_state.go b/pkg/tcpip/transport/tcp/segment_state.go
index 57bbd69ff..76ab56294 100644
--- a/pkg/tcpip/transport/tcp/segment_state.go
+++ b/pkg/tcpip/transport/tcp/segment_state.go
@@ -14,6 +14,10 @@
 
 package tcp
 
+import (
+	"context"
+)
+
 // saveOptions is invoked by stateify.
 func (s *segment) saveOptions() []byte {
 	// We cannot save s.options directly as it may point to s.data's trimmed
@@ -23,7 +27,7 @@ func (s *segment) saveOptions() []byte {
 }
 
 // loadOptions is invoked by stateify.
-func (s *segment) loadOptions(options []byte) {
+func (s *segment) loadOptions(_ context.Context, options []byte) {
 	// NOTE: We cannot point s.options back into s.data's trimmed tail. But
 	// it is OK as they do not need to aliased. Plus, options is already
 	// allocated so there is no cost here.
diff --git a/pkg/tcpip/transport/udp/endpoint_state.go b/pkg/tcpip/transport/udp/endpoint_state.go
index 27d0d8476..46b5afc86 100644
--- a/pkg/tcpip/transport/udp/endpoint_state.go
+++ b/pkg/tcpip/transport/udp/endpoint_state.go
@@ -30,7 +30,7 @@ func (p *udpPacket) saveReceivedAt() int64 {
 }
 
 // loadReceivedAt is invoked by stateify.
-func (p *udpPacket) loadReceivedAt(nsec int64) {
+func (p *udpPacket) loadReceivedAt(_ context.Context, nsec int64) {
 	p.receivedAt = time.Unix(0, nsec)
 }
 
diff --git a/tools/go_stateify/main.go b/tools/go_stateify/main.go
index 002dafba5..a9d96d1e5 100644
--- a/tools/go_stateify/main.go
+++ b/tools/go_stateify/main.go
@@ -370,7 +370,7 @@ func main() {
 						emitField(name)
 					}
 					emitLoadValue := func(name, typName string) {
-						fmt.Fprintf(outputFile, "	stateSourceObject.LoadValue(%d, new(%s), func(y any) { %s.load%s(y.(%s)) })\n", fields[name], typName, recv, camelCased(name), typName)
+						fmt.Fprintf(outputFile, "	stateSourceObject.LoadValue(%d, new(%s), func(y any) { %s.load%s(ctx, y.(%s)) })\n", fields[name], typName, recv, camelCased(name), typName)
 					}
 					emitLoad := func(name string) {
 						fmt.Fprintf(outputFile, "	stateSourceObject.Load(%d, &%s.%s)\n", fields[name], recv, name)