8 files changed, 19 insertions, 23 deletions
diff --git a/pkg/tcpip/transport/tcp/BUILD b/pkg/tcpip/transport/tcp/BUILD
index 6a7153e4d..9ebae6cc7 100644
--- a/pkg/tcpip/transport/tcp/BUILD
+++ b/pkg/tcpip/transport/tcp/BUILD
@@ -1,27 +1,7 @@
 package(licenses = ["notice"])  # Apache 2.0
 
 load("//tools/go_generics:defs.bzl", "go_template_instance")
-load("//tools/go_stateify:defs.bzl", "go_library", "go_stateify", "go_test")
-
-go_stateify(
-    name = "tcp_state",
-    srcs = [
-        "endpoint.go",
-        "endpoint_state.go",
-        "rcv.go",
-        "reno.go",
-        "segment.go",
-        "segment_heap.go",
-        "segment_queue.go",
-        "segment_state.go",
-        "snd.go",
-        "snd_state.go",
-        "tcp_segment_list.go",
-    ],
-    out = "tcp_state.go",
-    imports = ["gvisor.googlesource.com/gvisor/pkg/tcpip/buffer"],
-    package = "tcp",
-)
+load("//tools/go_stateify:defs.bzl", "go_library", "go_test")
 
 go_template_instance(
     name = "tcp_segment_list",
@@ -53,15 +33,14 @@ go_library(
         "snd.go",
         "snd_state.go",
         "tcp_segment_list.go",
-        "tcp_state.go",
         "timer.go",
     ],
     importpath = "gvisor.googlesource.com/gvisor/pkg/tcpip/transport/tcp",
+    imports = ["gvisor.googlesource.com/gvisor/pkg/tcpip/buffer"],
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/rand",
         "//pkg/sleep",
-        "//pkg/state",
         "//pkg/tcpip",
         "//pkg/tcpip/buffer",
         "//pkg/tcpip/header",
diff --git a/pkg/tcpip/transport/tcp/endpoint.go b/pkg/tcpip/transport/tcp/endpoint.go
index 5b8a1e20f..de1883d84 100644
--- a/pkg/tcpip/transport/tcp/endpoint.go
+++ b/pkg/tcpip/transport/tcp/endpoint.go
@@ -54,6 +54,8 @@ const (
 )
 
 // SACKInfo holds TCP SACK related information for a given endpoint.
+//
+// +stateify savable
 type SACKInfo struct {
 	// Blocks is the maximum number of SACK blocks we track
 	// per endpoint.
@@ -69,6 +71,8 @@ type SACKInfo struct {
 // have concurrent goroutines make calls into the endpoint, they are properly
 // synchronized. The protocol implementation, however, runs in a single
 // goroutine.
+//
+// +stateify savable
 type endpoint struct {
 	// workMu is used to arbitrate which goroutine may perform protocol
 	// work. Only the main protocol goroutine is expected to call Lock() on
diff --git a/pkg/tcpip/transport/tcp/rcv.go b/pkg/tcpip/transport/tcp/rcv.go
index b22a00ce1..92ef9c6f7 100644
--- a/pkg/tcpip/transport/tcp/rcv.go
+++ b/pkg/tcpip/transport/tcp/rcv.go
@@ -22,6 +22,8 @@ import (
 
 // receiver holds the state necessary to receive TCP segments and turn them
 // into a stream of bytes.
+//
+// +stateify savable
 type receiver struct {
 	ep *endpoint
 
diff --git a/pkg/tcpip/transport/tcp/reno.go b/pkg/tcpip/transport/tcp/reno.go
index 60f170a27..03ae8d747 100644
--- a/pkg/tcpip/transport/tcp/reno.go
+++ b/pkg/tcpip/transport/tcp/reno.go
@@ -16,6 +16,8 @@ package tcp
 
 // renoState stores the variables related to TCP New Reno congestion
 // control algorithm.
+//
+// +stateify savable
 type renoState struct {
 	s *sender
 }
diff --git a/pkg/tcpip/transport/tcp/segment.go b/pkg/tcpip/transport/tcp/segment.go
index 40928ba2c..8dccea2ba 100644
--- a/pkg/tcpip/transport/tcp/segment.go
+++ b/pkg/tcpip/transport/tcp/segment.go
@@ -36,6 +36,8 @@ const (
 // segment represents a TCP segment. It holds the payload and parsed TCP segment
 // information, and can be added to intrusive lists.
 // segment is mostly immutable, the only field allowed to change is viewToDeliver.
+//
+// +stateify savable
 type segment struct {
 	segmentEntry
 	refCnt int32
diff --git a/pkg/tcpip/transport/tcp/segment_queue.go b/pkg/tcpip/transport/tcp/segment_queue.go
index 2ddcf5f10..6a2d7bc0b 100644
--- a/pkg/tcpip/transport/tcp/segment_queue.go
+++ b/pkg/tcpip/transport/tcp/segment_queue.go
@@ -21,6 +21,8 @@ import (
 )
 
 // segmentQueue is a bounded, thread-safe queue of TCP segments.
+//
+// +stateify savable
 type segmentQueue struct {
 	mu    sync.Mutex  `state:"nosave"`
 	list  segmentList `state:"wait"`
diff --git a/pkg/tcpip/transport/tcp/snd.go b/pkg/tcpip/transport/tcp/snd.go
index e38686e1b..376e81846 100644
--- a/pkg/tcpip/transport/tcp/snd.go
+++ b/pkg/tcpip/transport/tcp/snd.go
@@ -54,6 +54,8 @@ type congestionControl interface {
 }
 
 // sender holds the state necessary to send TCP segments.
+//
+// +stateify savable
 type sender struct {
 	ep *endpoint
 
@@ -133,6 +135,8 @@ type sender struct {
 }
 
 // fastRecovery holds information related to fast recovery from a packet loss.
+//
+// +stateify savable
 type fastRecovery struct {
 	// active whether the endpoint is in fast recovery. The following fields
 	// are only meaningful when active is true.
diff --git a/pkg/tcpip/transport/tcp/snd_state.go b/pkg/tcpip/transport/tcp/snd_state.go
index 33c8867f4..d536839af 100644
--- a/pkg/tcpip/transport/tcp/snd_state.go
+++ b/pkg/tcpip/transport/tcp/snd_state.go
@@ -18,6 +18,7 @@ import (
 	"time"
 )
 
+// +stateify savable
 type unixTime struct {
 	second int64
 	nano   int64