Fix race condition by checking manager stopped state before creating context

Signed-off-by: ZHENYU <[email protected]>

Author: ae86zhizhi

pkg/kvevent/doc.go

@@ -37,14 +37,14 @@
 //
 // The syncIndexerAdapter in pkg/cache might seem like unnecessary complexity,
 // but it serves critical architectural purposes:
-//   1. Prevents circular dependencies between packages
-//   2. Acts as an Anti-Corruption Layer between domains
-//   3. Allows packages to remain independent and testable
-//   4. Handles type conversion between nearly-identical structs
+//  1. Prevents circular dependencies between packages
+//  2. Acts as an Anti-Corruption Layer between domains
+//  3. Allows packages to remain independent and testable
+//  4. Handles type conversion between nearly-identical structs
 //
 // Build Requirements:
 //
 // This package requires ZMQ support. Build with:
-//   go build -tags=zmq
 //
-package kvevent
+//	go build -tags=zmq
+package kvevent

pkg/kvevent/errors.go

@@ -42,4 +42,4 @@ func IsTemporaryError(err error) bool {
 	return errors.Is(err, ErrIndexerNotInitialized) ||
 		errors.Is(err, context.DeadlineExceeded) ||
 		errors.Is(err, context.Canceled)
-}
+}

pkg/kvevent/handler.go

@@ -38,13 +38,7 @@ type eventHandler struct {
 
 // HandleEvent processes KV cache events
 func (h *eventHandler) HandleEvent(event kvcache.KVEvent) error {
-	// Create context with timeout
-	// Use 10s timeout for event processing as sync indexer operations
-	// may involve Redis calls and network I/O under high load
-	ctx, cancel := context.WithTimeout(h.manager.ctx, 10*time.Second)
-	defer cancel()
-
-	// Check if manager is stopped
+	// Check if manager is stopped before using its context
 	h.manager.mu.RLock()
 	stopped := h.manager.stopped
 	h.manager.mu.RUnlock()
@@ -53,6 +47,12 @@ func (h *eventHandler) HandleEvent(event kvcache.KVEvent) error {
 		return ErrManagerStopped
 	}
 
+	// Create context with timeout
+	// Use 10s timeout for event processing as sync indexer operations
+	// may involve Redis calls and network I/O under high load
+	ctx, cancel := context.WithTimeout(h.manager.ctx, 10*time.Second)
+	defer cancel()
+
 	switch e := event.(type) {
 	case *kvcache.BlockStoredEvent:
 		return h.handleBlockStored(ctx, e)
@@ -155,4 +155,4 @@ func tokenIDsToBytes(tokenIDs []int32) []byte {
 		binary.BigEndian.PutUint32(bytes[i*4:], uint32(id))
 	}
 	return bytes
-}
+}

pkg/kvevent/handler_test.go

@@ -83,8 +83,8 @@ func TestTokenIDsToBytes(t *testing.T) {
 			name:     "multiple tokens",
 			tokenIDs: []int32{1, 256, 65535},
 			expected: []byte{
-				0, 0, 0, 1,     // 1
-				0, 0, 1, 0,     // 256
+				0, 0, 0, 1, // 1
+				0, 0, 1, 0, // 256
 				0, 0, 255, 255, // 65535
 			},
 		},
@@ -379,4 +379,4 @@ func TestBinaryEncoding(t *testing.T) {
 			t.Errorf("Byte %d: expected %d, got %d", i, expected[i], bytes[i])
 		}
 	}
-}
+}

pkg/kvevent/interfaces.go

@@ -86,4 +86,4 @@ type BlockRemovedEvent struct {
 	ModelName   string
 	LoraID      int64
 	SourcePod   string
-}
+}

pkg/kvevent/manager.go

@@ -162,6 +162,15 @@ func (m *Manager) OnPodAdd(pod *v1.Pod) {
 		return
 	}
 
+	// Check if manager is stopped before using its context
+	m.mu.RLock()
+	stopped := m.stopped
+	m.mu.RUnlock()
+
+	if stopped {
+		return
+	}
+
 	// Use 5s timeout for pod operations as they involve simple ZMQ ops
 	ctx, cancel := context.WithTimeout(m.ctx, 5*time.Second)
 	defer cancel()
@@ -212,6 +221,15 @@ func (m *Manager) OnPodDelete(pod *v1.Pod) {
 	podKey := utils.GeneratePodKey(pod.Namespace, pod.Name)
 	m.unsubscribeFromPod(podKey)
 
+	// Check if manager is stopped before using its context
+	m.mu.RLock()
+	stopped := m.stopped
+	m.mu.RUnlock()
+
+	if stopped {
+		return
+	}
+
 	// Clean up from sync indexer
 	// Use 5s timeout for cleanup operations as they should be quick
 	ctx, cancel := context.WithTimeout(m.ctx, 5*time.Second)

pkg/kvevent/manager_test.go

@@ -90,4 +90,4 @@ func TestManagerCreation(t *testing.T) {
 
 	// Test Stop (should not panic)
 	manager.Stop()
-}
+}

pkg/kvevent/test_helpers.go

@@ -29,4 +29,4 @@ func NewEventHandlerForTest(manager *Manager, podKey, modelName string, loraID i
 		modelName: modelName,
 		loraID:    loraID,
 	}
-}
+}