Server connection management; add bulk string reading optimization, improve connection limit checks, and refine idle connection handling

Author: l00pss

protocol.go

@@ -118,10 +118,22 @@ func (c *Connection) readBulkString(sizeBytes []byte) (RedisValue, error) {
 	}
 
 	// Read the bulk data plus CRLF
-	data := make([]byte, size+2)
-	_, err = io.ReadFull(c.reader, data)
-	if err != nil {
-		return RedisValue{}, err
+	var data []byte
+	if size > 1024*1024 {
+		r := io.LimitReader(c.reader, int64(size+2))
+		data, err = io.ReadAll(r)
+		if err != nil {
+			return RedisValue{}, err
+		}
+		if len(data) != size+2 {
+			return RedisValue{}, fmt.Errorf("unexpected EOF reading bulk string")
+		}
+	} else {
+		data = make([]byte, size+2)
+		_, err = io.ReadFull(c.reader, data)
+		if err != nil {
+			return RedisValue{}, err
+		}
 	}
 
 	return RedisValue{Type: BulkString, Bulk: data[:size]}, nil

server.go

@@ -119,19 +119,24 @@ func (s *Server) Serve() error {
 		}
 
 		s.wg.Add(1)
-		go func(netConn net.Conn) {
-			defer s.wg.Done()
 
-			// Check connection limit after Accept to prevent TOCTOU race
-			if s.MaxConnections > 0 && s.connCount.Add(1) > int64(s.MaxConnections) {
+		if s.MaxConnections > 0 {
+			if s.connCount.Add(1) > int64(s.MaxConnections) {
 				s.connCount.Add(-1)
-				netConn.Close()
-				s.Logger.Warn("Connection limit reached, rejecting connection from %s", netConn.RemoteAddr())
-				return
+				s.wg.Done()
+				conn.Close()
+				s.Logger.Warn("Connection limit reached, rejecting connection from %s", conn.RemoteAddr())
+				continue
 			}
+		} else {
+			s.connCount.Add(1)
+		}
+
+		go func(netConn net.Conn) {
+			defer s.wg.Done()
+			defer s.connCount.Add(-1)
 
 			s.handleConnectionInternal(netConn)
-			s.connCount.Add(-1)
 		}(conn)
 	}
 }
@@ -151,13 +156,18 @@ func (s *Server) Shutdown(ctx context.Context) error {
 
 	// Close all active connections
 	s.mu.RLock()
+	conns := make([]*Connection, 0, len(s.activeConns))
 	for conn := range s.activeConns {
+		conns = append(conns, conn)
+	}
+	s.mu.RUnlock()
+
+	for _, conn := range conns {
 		err := conn.Close()
 		if err != nil {
 			return err
 		}
 	}
-	s.mu.RUnlock()
 
 	// Run shutdown hooks
 	s.mu.Lock()
@@ -252,9 +262,9 @@ func (s *Server) handleConnectionInternal(netConn net.Conn) {
 
 		s.Logger.Debug("Command from %s: %s %v", netConn.RemoteAddr(), cmd.Name, cmd.Args)
 
-		s.setConnectionActive(conn)
-
+		conn.setState(StateProcessing)
 		response := s.handleCommand(conn, cmd)
+		conn.setState(StateActive)
 
 		if s.WriteTimeout > 0 {
 			err := netConn.SetWriteDeadline(time.Now().Add(s.WriteTimeout))
@@ -377,14 +387,14 @@ func (s *Server) checkIdleConnections() {
 
 		currentState := ConnState(conn.state.Load())
 
-		if currentState == StateActive && lastUsed.Before(idleThreshold) {
+		if (currentState == StateActive || currentState == StateIdle) && lastUsed.Before(idleThreshold) {
 			idleConns = append(idleConns, conn)
 		}
 	}
 
 	for _, conn := range idleConns {
-		conn.setState(StateIdle)
-		s.Logger.Debug("Connection %s marked as idle", conn.RemoteAddr())
+		s.Logger.Info("Closing idle connection %s", conn.RemoteAddr())
+		conn.Close()
 	}
 }
 

types.go

@@ -139,6 +139,7 @@ const (
 	StateActive
 	StateIdle
 	StateClosed
+	StateProcessing
 )
 
 type RedisValue struct {