context is cancelled, the server will detect it immediately and exit gracefully

test/kafka/docker-compose.yml

@@ -90,11 +90,11 @@ services:
     volumes:
       - seaweedfs-master-data:/data
     healthcheck:
-      test: ["CMD", "wget", "--quiet", "--tries=1", "--spider", "http://localhost:9333/cluster/status"]
+      test: ["CMD-SHELL", "wget --quiet --tries=1 --spider http://localhost:9333/cluster/status || curl -sf http://localhost:9333/cluster/status"]
       interval: 10s
       timeout: 5s
-      retries: 3
-      start_period: 10s
+      retries: 10
+      start_period: 20s
     networks:
       - kafka-test-net
 

weed/mq/kafka/protocol/handler.go

@@ -247,15 +247,29 @@ func (h *Handler) HandleConn(ctx context.Context, conn net.Conn) error {
 			return fmt.Errorf("set read deadline: %w", err)
 		}
 
+		// Check context before reading
+		select {
+		case <-ctx.Done():
+			fmt.Printf("DEBUG: [%s] Context cancelled before reading message header\n", connectionID)
+			return ctx.Err()
+		default:
+		}
+
 		// Read message size (4 bytes)
 		fmt.Printf("DEBUG: [%s] About to read message size header\n", connectionID)
 		var sizeBytes [4]byte
 		if _, err := io.ReadFull(r, sizeBytes[:]); err != nil {
 			if err == io.EOF {
-				fmt.Printf("DEBUG: Client closed connection (clean EOF)\n")
+				fmt.Printf("DEBUG: [%s] Client closed connection (clean EOF)\n", connectionID)
 				return nil // clean disconnect
 			}
 
+			// Check if it's a timeout error
+			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
+				fmt.Printf("DEBUG: [%s] Read timeout due to context cancellation\n", connectionID)
+				return ctx.Err()
+			}
+
 			// Use centralized error classification
 			errorCode := ClassifyNetworkError(err)
 			switch errorCode {