Some final touches on distribution test

Author: 0x19

client/tcp.go

@@ -234,6 +234,11 @@ func (h *tcpEventHandler) OnTraffic(c gnet.Conn) gnet.Action {
 		// Append the chunk to our buffer
 		buf = append(buf, chunk...)
 
+		// Check if we have a complete message in the buffer
+		if isCompleteMessage(buf) {
+			break
+		}
+
 		// If we received less than what would fill a typical buffer,
 		// we've likely received the complete message for now
 		if len(chunk) < 4096 {
@@ -407,6 +412,50 @@ func (h *tcpEventHandler) OnTraffic(c gnet.Conn) gnet.Action {
 	return gnet.None
 }
 
+// isCompleteMessage checks if we have a complete valid message in the buffer
+// by verifying message structure and length requirements
+func isCompleteMessage(data []byte) bool {
+	// Need at least 1 byte for message type
+	if len(data) < 1 {
+		return false
+	}
+
+	// Check if message is a success response (which includes a payload length)
+	msgType := MessageType(data[0])
+	if msgType == MessageType(types.HandlerStatusSuccess.Byte()) {
+		// Need at least 5 bytes for header (1 status + 4 length)
+		if len(data) < 5 {
+			return false
+		}
+
+		// Check if actual data length matches expected length from header
+		expectedLength := binary.BigEndian.Uint32(data[1:5])
+		actualDataLength := len(data) - 5 // subtract header bytes
+
+		// Only complete if we have all the expected data
+		return uint32(actualDataLength) >= expectedLength
+	}
+
+	// For write/read handler messages which have a specific format
+	if msgType == MessageType('W') || msgType == MessageType('R') {
+		// These also use a length field - need at least 5 bytes (1 type + 4 length)
+		if len(data) < 5 {
+			return false
+		}
+
+		// Get the expected message length
+		expectedLength := binary.BigEndian.Uint32(data[1:5])
+		actualDataLength := len(data) - 5
+
+		// Check if we have the complete message
+		return uint32(actualDataLength) >= expectedLength
+	}
+
+	// For other message types where we can't easily determine completeness
+	// by examining headers, use a more conservative approach
+	return false
+}
+
 // OnTick is called periodically
 func (h *tcpEventHandler) OnTick() (time.Duration, gnet.Action) {
 	return time.Second, gnet.None

db/writer.go

@@ -1,6 +1,7 @@
 package db
 
 import (
+	"context"
 	"github.com/erigontech/mdbx-go/mdbx"
 	"github.com/pkg/errors"
 	"go.uber.org/zap"
@@ -89,19 +90,28 @@ func (bw *BatchWriter) BufferWrite(key [32]byte, value []byte) error {
 	// XOR the first and last bytes for slightly better distribution
 	workerID := int(key[0]^key[31]) % bw.workers
 
-	// Use non-blocking send to prevent backpressure
+	// Create a timeout context for the send operation
+	timeoutCtx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
+	defer cancel()
+	
+	// Try to send with timeout
 	select {
 	case bw.workerChannels[workerID] <- WriteRequest{Key: key, Value: value}:
 		// Successfully sent to channel
 		return nil
-	default:
-		// Channel is full, handle gracefully by using a goroutine to avoid blocking
-		// This helps during high-load scenarios to prevent caller blocking
+	case <-timeoutCtx.Done():
+		// Channel send timed out, try in background goroutine
 		go func() {
-			// This will block in the goroutine but not block the caller
-			bw.workerChannels[workerID] <- WriteRequest{Key: key, Value: value}
+			select {
+			case bw.workerChannels[workerID] <- WriteRequest{Key: key, Value: value}:
+				// Successfully sent
+			case <-time.After(5 * time.Second):
+				// If we can't send after a long timeout, log the error
+				zap.L().Error("Failed to queue record for batch writing after extended timeout",
+					zap.Binary("key_prefix", key[:8]))
+			}
 		}()
-		return nil
+		return errors.New("queue operation timed out, write queued in background")
 	}
 }
 

node/db_read_handler.go

@@ -1,8 +1,6 @@
 package node
 
 import (
-	"fmt"
-
 	"github.com/unpackdev/fdb/db"
 	"github.com/unpackdev/fdb/logger"
 	"github.com/unpackdev/fdb/observability"
@@ -122,7 +120,7 @@ func (rh *DbReadHandler) Handle(conn transports.Connection, frame []byte) {
 	// Encode the response to bytes
 	response := dbResp.Encode()
 
-	fmt.Println("SENDING RESPONSE PREFIX", response[0:10], "with status byte:", response[0], "data length:", dbResp.Length)
+	//fmt.Println("SENDING RESPONSE PREFIX", response[0:10], "with status byte:", response[0], "data length:", dbResp.Length)
 
 	// Send the formatted response back to the client
 	conn.Send(response)

node/db_write_handler.go

@@ -1,9 +1,6 @@
 package node
 
 import (
-	"encoding/binary"
-	"fmt"
-
 	"github.com/unpackdev/fdb/db"
 	"github.com/unpackdev/fdb/logger"
 	"github.com/unpackdev/fdb/observability"
@@ -43,15 +40,15 @@ func (wh *DbWriteHandler) ForceFlush() {
 
 // Handle processes the incoming message using the TCPWriteHandler
 func (wh *DbWriteHandler) Handle(conn transports.Connection, frame []byte) {
-	fmt.Println("DID I REACH WRITE DB HANLDER? FRAME LENGTH:", len(frame))
+	//fmt.Println("DID I REACH WRITE DB HANLDER? FRAME LENGTH:", len(frame))
 	// Debug the raw incoming frame
-	fmt.Printf("WRITE HANDLER RAW FRAME (first 20 bytes): %v\n", frame[:min(20, len(frame))])
+	//fmt.Printf("WRITE HANDLER RAW FRAME (first 20 bytes): %v\n", frame[:min(20, len(frame))])
 
 	// Check if first byte is our special marker 0xF0
 	offset := 0
 	if len(frame) > 0 && frame[0] == 0xF0 {
 		offset = 1 // Skip the marker byte
-		fmt.Println("DETECTED MARKER BYTE 0xF0, offset set to", offset)
+		//fmt.Println("DETECTED MARKER BYTE 0xF0, offset set to", offset)
 	}
 
 	// Adjust minimum length check based on whether we have a marker
@@ -80,15 +77,15 @@ func (wh *DbWriteHandler) Handle(conn transports.Connection, frame []byte) {
 	}
 
 	// Debug action byte
-	actionByte := frame[offset]
-	fmt.Printf("ACTION BYTE: 0x%02x\n", actionByte)
+	//actionByte := frame[offset]
+	//fmt.Printf("ACTION BYTE: 0x%02x\n", actionByte)
 
 	// Create a [32]byte key from the frame without using the pool
 	var key [32]byte
 	copy(key[:], frame[offset+1:offset+33]) // Copy directly from frame, accounting for offset
 
 	// Debug the key being extracted
-	fmt.Printf("KEY (first 8 bytes): %v\n", key[:8])
+	//fmt.Printf("KEY (first 8 bytes): %v\n", key[:8])
 
 	// For the message protocol format: [1-byte action][32-byte key][4-byte length][actual data]
 	// We need to skip 4 bytes after the key to get to the actual data
@@ -114,7 +111,7 @@ func (wh *DbWriteHandler) Handle(conn transports.Connection, frame []byte) {
 	// Extract only the actual data, skipping the length field
 	value := frame[valueStart:]
 
-	fmt.Printf("ACTUAL VALUE (first %d bytes): %v\n", min(10, len(value)), value[:min(10, len(value))])
+	//fmt.Printf("ACTUAL VALUE (first %d bytes): %v\n", min(10, len(value)), value[:min(10, len(value))])
 
 	// Add detailed logging for debugging large payload issues
 	wh.logger.Debug("Received write request",
@@ -125,7 +122,7 @@ func (wh *DbWriteHandler) Handle(conn transports.Connection, frame []byte) {
 		zap.Binary("value_prefix", value[:min(10, len(value))]))
 
 	// Log what we're about to write to the database
-	fmt.Printf("ABOUT TO WRITE TO DB - KEY: %v, VALUE (first 20 bytes): %v\n", key[:8], value[:min(20, len(value))])
+	//fmt.Printf("ABOUT TO WRITE TO DB - KEY: %v, VALUE (first 20 bytes): %v\n", key[:8], value[:min(20, len(value))])
 
 	// Buffer the write request with the key as [32]byte
 	err := wh.writer.BufferWrite(key, value)
@@ -144,7 +141,7 @@ func (wh *DbWriteHandler) Handle(conn transports.Connection, frame []byte) {
 
 		// Encode the response to bytes
 		response := dbResp.Encode()
-		fmt.Printf("SENDING ERROR RESPONSE: %v\n", response[:min(20, len(response))])
+		//fmt.Printf("SENDING ERROR RESPONSE: %v\n", response[:min(20, len(response))])
 		conn.Send(response)
 		return
 	}
@@ -174,12 +171,12 @@ func (wh *DbWriteHandler) Handle(conn transports.Connection, frame []byte) {
 	response := dbResp.Encode()
 
 	// Detailed debug of the full response being sent
-	fmt.Println("SENDING WRITE SUCCESS RESPONSE", response[:min(20, len(response))], "with status byte:", response[0])
-	fmt.Printf("FULL RESPONSE DETAILS:\n  - Status: %d\n  - Length bytes: %v (uint32: %d)\n  - Data: %v\n",
-		response[0],
-		response[1:5],
-		binary.BigEndian.Uint32(response[1:5]),
-		response[5:min(25, len(response))])
+	// fmt.Println("SENDING WRITE SUCCESS RESPONSE", response[:min(20, len(response))], "with status byte:", response[0])
+	// fmt.Printf("FULL RESPONSE DETAILS:\n  - Status: %d\n  - Length bytes: %v (uint32: %d)\n  - Data: %v\n",
+	// 	response[0],
+	// 	response[1:5],
+	// 	binary.BigEndian.Uint32(response[1:5]),
+	// 	response[5:min(25, len(response))])
 
 	// Send the formatted response back to the client
 	conn.Send(response)