state: Maintain a separate list with leaf nodes for objectDecodeState.

Profiling shows that we spend a LOT of time in decodeState.Load() in firing all
callbacks in the right order. objectDecodeState is organized as a DAG.
Callbacks need to be fired on leaf nodes (where blockedBy==0) and then fired
on internal nodes only when all their dependencies have executed their
callbacks. The earlier approach was to continuously scan a HUGE list of pending
nodes searching for a leaf node, fire callbacks on that, check if any of its
dependents became leaves and then scan the full list again. We were wasting a
lot of time re-evaluating internal nodes repeatedly. For instance, the kernel
object created by checkpointing SGLang server running "gemma-3-27b-pt" had
368,245 objects. The whole pending list was scanned 238,102 times. We spent
2 minutes 10 seconds in just executing callbacks. The total kernel load time
was 2 minutes 30 seconds.

In this change, we track leaf nodes in a separate linked list. We keep
appending to it as and when new nodes become leaves. This way we don't need to
search for leaves. This improved kernel load time to 12 seconds (-92%).

PiperOrigin-RevId: 775402687

Author: ayushr2

pkg/state/BUILD

@@ -19,14 +19,14 @@ go_template_instance(
 )
 
 go_template_instance(
-    name = "complete_list",
-    out = "complete_list.go",
+    name = "ods_list",
+    out = "ods_list.go",
     package = "state",
-    prefix = "complete",
+    prefix = "ods",
     template = "//pkg/ilist:generic_list",
     types = {
-        "Element": "*objectDecodeState",
-        "Linker": "*objectDecodeState",
+        "Element": "*odsListElem",
+        "Linker": "*odsListElem",
     },
 )
 
@@ -66,12 +66,12 @@ go_library(
     srcs = [
         "addr_range.go",
         "addr_set.go",
-        "complete_list.go",
         "decode.go",
         "decode_unsafe.go",
         "deferred_list.go",
         "encode.go",
         "encode_unsafe.go",
+        "ods_list.go",
         "state.go",
         "state_norace.go",
         "state_race.go",

pkg/state/decode.go

@@ -32,7 +32,7 @@ type internalCallback interface {
 	source() *objectDecodeState
 
 	// callbackRun executes the callback.
-	callbackRun()
+	callbackRun(ds *decodeState)
 }
 
 // userCallback is an implementation of internalCallback.
@@ -44,7 +44,7 @@ func (userCallback) source() *objectDecodeState {
 }
 
 // callbackRun implements internalCallback.callbackRun.
-func (uc userCallback) callbackRun() {
+func (uc userCallback) callbackRun(*decodeState) {
 	uc()
 }
 
@@ -84,7 +84,13 @@ type objectDecodeState struct {
 	// callbacks is a set of callbacks to execute on load.
 	callbacks []internalCallback
 
-	completeEntry
+	pendingEntry odsListElem
+	leafEntry    odsListElem
+}
+
+type odsListElem struct {
+	ods *objectDecodeState
+	odsEntry
 }
 
 // addCallback adds a callback to the objectDecodeState.
@@ -122,8 +128,13 @@ func (ods *objectDecodeState) source() *objectDecodeState {
 }
 
 // callbackRun implements internalCallback.callbackRun.
-func (ods *objectDecodeState) callbackRun() {
+func (ods *objectDecodeState) callbackRun(ds *decodeState) {
 	ods.blockedBy--
+	if ods.blockedBy == 0 {
+		ds.leaves.PushBack(&ods.leafEntry)
+	} else if ods.blockedBy < 0 {
+		Failf("object %d has negative blockedBy: %d", ods.id, ods.blockedBy)
+	}
 }
 
 // decodeState is a graph of objects in the process of being decoded.
@@ -155,7 +166,11 @@ type decodeState struct {
 	deferred map[objectID]wire.Object
 
 	// pending is the set of objects that are not yet complete.
-	pending completeList
+	pending odsList
+
+	// leaves is the set of objects that have no dependencies (blockedBy == 0).
+	// leaves are consumed from the front and appended to the back.
+	leaves odsList
 
 	// stats tracks time data.
 	stats Stats
@@ -185,20 +200,12 @@ func (ds *decodeState) checkComplete(ods *objectDecodeState) bool {
 
 	// Fire all callbacks.
 	for _, ic := range ods.callbacks {
-		ic.callbackRun()
+		ic.callbackRun(ds)
 	}
 
 	// Mark completed.
-	cbs := ods.callbacks
 	ods.callbacks = nil
-	ds.pending.Remove(ods)
-
-	// Recursively check others.
-	for _, ic := range cbs {
-		if other := ic.source(); other != nil && other.blockedBy == 0 {
-			ds.checkComplete(other)
-		}
-	}
+	ds.pending.Remove(&ods.pendingEntry)
 
 	return true // All set.
 }
@@ -223,6 +230,9 @@ func (ds *decodeState) wait(waiter *objectDecodeState, id objectID, callback fun
 
 	// Mark as blocked.
 	waiter.blockedBy++
+	if waiter.blockedBy == 1 {
+		ds.leaves.Remove(&waiter.leafEntry)
+	}
 
 	// No nil can be returned here.
 	other := ds.lookup(id)
@@ -280,6 +290,26 @@ func walkChild(path []wire.Dot, obj reflect.Value) reflect.Value {
 	return obj
 }
 
+func (ds *decodeState) growObjectsByID(id objectID) {
+	if len(ds.objectsByID) < int(id) {
+		ds.objectsByID = append(ds.objectsByID, make([]*objectDecodeState, int(id)-len(ds.objectsByID))...)
+	}
+}
+
+func (ds *decodeState) addObject(id objectID, obj reflect.Value) *objectDecodeState {
+	ods := &objectDecodeState{
+		id:  id,
+		obj: obj,
+	}
+	ods.pendingEntry.ods = ods
+	ods.leafEntry.ods = ods
+	ds.growObjectsByID(id)
+	ds.objectsByID[id-1] = ods
+	ds.pending.PushBack(&ods.pendingEntry)
+	ds.leaves.PushBack(&ods.leafEntry)
+	return ods
+}
+
 // register registers a decode with a type.
 //
 // This type is only used to instantiate a new object if it has not been
@@ -288,11 +318,9 @@ func walkChild(path []wire.Dot, obj reflect.Value) reflect.Value {
 func (ds *decodeState) register(r *wire.Ref, typ reflect.Type) reflect.Value {
 	// Grow the objectsByID slice.
 	id := objectID(r.Root)
-	if len(ds.objectsByID) < int(id) {
-		ds.objectsByID = append(ds.objectsByID, make([]*objectDecodeState, int(id)-len(ds.objectsByID))...)
-	}
 
 	// Does this object already exist?
+	ds.growObjectsByID(id)
 	ods := ds.objectsByID[id-1]
 	if ods != nil {
 		return walkChild(r.Dots, ods.obj)
@@ -303,12 +331,7 @@ func (ds *decodeState) register(r *wire.Ref, typ reflect.Type) reflect.Value {
 		typ = ds.findType(r.Type)
 	}
 	v := reflect.New(typ)
-	ods = &objectDecodeState{
-		id:  id,
-		obj: v.Elem(),
-	}
-	ds.objectsByID[id-1] = ods
-	ds.pending.PushBack(ods)
+	ods = ds.addObject(id, v.Elem())
 
 	// Process any deferred objects & callbacks.
 	if encoded, ok := ds.deferred[id]; ok {
@@ -581,13 +604,8 @@ func (ds *decodeState) Load(obj reflect.Value) {
 		return ds.types.LookupName(id)
 	})
 
-	// Create the root object.
-	rootOds := &objectDecodeState{
-		id:  1,
-		obj: obj,
-	}
-	ds.objectsByID = append(ds.objectsByID, rootOds)
-	ds.pending.PushBack(rootOds)
+	// Add the root object with ID 1.
+	_ = ds.addObject(1, obj)
 
 	// Read the number of objects.
 	numObjects, object, err := ReadHeader(&ds.r)
@@ -603,7 +621,6 @@ func (ds *decodeState) Load(obj reflect.Value) {
 		encoded wire.Object
 		ods     *objectDecodeState
 		id      objectID
-		tid     = typeID(1)
 	)
 	if err := safely(func() {
 		// Decode all objects in the stream.
@@ -616,7 +633,6 @@ func (ds *decodeState) Load(obj reflect.Value) {
 			switch we := encoded.(type) {
 			case *wire.Type:
 				ds.types.Register(we)
-				tid++
 				encoded = nil
 				continue
 			case wire.Uint:
@@ -673,32 +689,22 @@ func (ds *decodeState) Load(obj reflect.Value) {
 	// objects become complete (there is a dependency cycle).
 	//
 	// Note that we iterate backwards here, because there will be a strong
-	// tendendcy for blocking relationships to go from earlier objects to
+	// tendency for blocking relationships to go from earlier objects to
 	// later (deeper) objects in the graph. This will reduce the number of
 	// iterations required to finish all objects.
 	if err := safely(func() {
-		for ds.pending.Back() != nil {
-			thisCycle := false
-			for ods = ds.pending.Back(); ods != nil; {
-				if ds.checkComplete(ods) {
-					thisCycle = true
-					break
-				}
-				ods = ods.Prev()
-			}
-			if !thisCycle {
-				break
-			}
+		for elem := ds.leaves.Front(); elem != nil; elem = elem.Next() {
+			ds.checkComplete(elem.ods)
 		}
 	}); err != nil {
 		Failf("error executing callbacks: %w\nfor object %#v", err, ods.obj.Interface())
 	}
 
 	// Check if we have any remaining dependency cycles. If there are any
 	// objects left in the pending list, then it must be due to a cycle.
-	if ods := ds.pending.Front(); ods != nil {
+	if elem := ds.pending.Front(); elem != nil {
 		// This must be the result of a dependency cycle.
-		cycle := ods.findCycle()
+		cycle := elem.ods.findCycle()
 		var buf bytes.Buffer
 		buf.WriteString("dependency cycle: {")
 		for i, cycleOS := range cycle {