update

consumer.go

@@ -202,7 +202,13 @@ func (c *Consumer) OnClose(_ context.Context, _ net.Conn) error {
 }
 
 func (c *Consumer) OnError(_ context.Context, conn net.Conn, err error) {
-	fmt.Println("Error reading from connection:", err, conn.RemoteAddr())
+	if c.isConnectionClosed(err) {
+		log.Printf("Connection to broker closed for consumer %s at %s", c.id, conn.RemoteAddr())
+		// Trigger reconnection if possible
+		c.reconnectCh <- struct{}{}
+	} else {
+		log.Printf("Error reading from connection: %v", err)
+	}
 }
 
 func (c *Consumer) OnMessage(ctx context.Context, msg *codec.Message, conn net.Conn) error {
@@ -828,6 +834,10 @@ func minInt(a, b int) int {
 	return b
 }
 
+func (c *Consumer) isConnectionClosed(err error) bool {
+	return err == io.EOF || strings.Contains(err.Error(), "connection closed") || strings.Contains(err.Error(), "connection reset")
+}
+
 func isConnectionError(err error) bool {
 	if err == nil {
 		return false

examples/server.go

@@ -13,9 +13,6 @@ func main() {
 		mq.WithCallback(tasks.Callback),
 		mq.WithBrokerURL(":8081"),
 		mq.WithSecurity(true),
-		// mq.WithMonitoring(true),
-		// mq.WithAdminAddr(":8080"),
-		// mq.WithMetricsAddr(":9090"),
 	)
 	InitializeDefaults(b.SecurityManager())
 	// b := mq.NewBroker(mq.WithCallback(tasks.Callback), mq.WithTLS(true, "./certs/server.crt", "./certs/server.key"), mq.WithCAPath("./certs/ca.cert"))

mq.go

@@ -3,6 +3,7 @@ package mq
 import (
 	"context"
 	"fmt"
+	"io"
 	"log"
 	"net"
 	"strings"
@@ -341,6 +342,7 @@ func HeadersWithConsumerIDAndQueue(ctx context.Context, id, queue string) map[st
 
 type QueuedTask struct {
 	Message    *codec.Message
+	Task       *Task
 	RetryCount int
 }
 
@@ -482,6 +484,8 @@ type Broker struct {
 	metricsServer      *MetricsServer
 	authenticatedConns storage.IMap[string, bool]              // authenticated connections
 	taskHeaders        storage.IMap[string, map[string]string] // task headers by task ID
+	pendingTasks       map[string]map[string]*Task             // consumerID -> taskID -> task
+	mu                 sync.RWMutex                            // for pendingTasks
 	isShutdown         int32
 	shutdown           chan struct{}
 	wg                 sync.WaitGroup
@@ -498,6 +502,7 @@ func NewBroker(opts ...Option) *Broker {
 		consumers:    memory.New[string, *consumer](),
 		deadLetter:   memory.New[string, *Queue](),
 		pIDs:         memory.New[string, bool](),
+		pendingTasks: make(map[string]map[string]*Task),
 		opts:         options,
 
 		// Enhanced production features
@@ -611,9 +616,41 @@ func (b *Broker) OnClose(ctx context.Context, conn net.Conn) error {
 
 func (b *Broker) OnError(_ context.Context, conn net.Conn, err error) {
 	if conn != nil {
-		fmt.Println("Error reading from connection:", err, conn.RemoteAddr())
+		if b.isConnectionClosed(err) {
+			log.Printf("Connection closed for consumer at %s", conn.RemoteAddr())
+			// Find and remove the consumer
+			b.consumers.ForEach(func(id string, con *consumer) bool {
+				if con.conn == conn {
+					b.RemoveConsumer(id)
+					b.mu.Lock()
+					if tasks, ok := b.pendingTasks[id]; ok {
+						for _, task := range tasks {
+							// Put back to queue
+							if q, ok := b.queues.Get(task.Topic); ok {
+								select {
+								case q.tasks <- &QueuedTask{Task: task, RetryCount: task.Retries}:
+									log.Printf("Requeued task %s for consumer %s", task.ID, id)
+								default:
+									log.Printf("Failed to requeue task %s, queue full", task.ID)
 								}
 							}
+						}
+						delete(b.pendingTasks, id)
+					}
+					b.mu.Unlock()
+					return false
+				}
+				return true
+			})
+		} else {
+			log.Printf("Error reading from connection: %v", err)
+		}
+	}
+}
+
+func (b *Broker) isConnectionClosed(err error) bool {
+	return err == io.EOF || strings.Contains(err.Error(), "connection closed") || strings.Contains(err.Error(), "connection reset")
+}
 
 func (b *Broker) isAuthenticated(connID string) bool {
 	if b.securityManager == nil {
@@ -722,6 +759,11 @@ func (b *Broker) MessageAck(ctx context.Context, msg *codec.Message) {
 	consumerID, _ := GetConsumerID(ctx)
 	taskID, _ := jsonparser.GetString(msg.Payload, "id")
 	log.Printf("BROKER - MESSAGE_ACK ~> %s on %s for Task %s", consumerID, msg.Queue, taskID)
+	b.mu.Lock()
+	if tasks, ok := b.pendingTasks[consumerID]; ok {
+		delete(tasks, taskID)
+	}
+	b.mu.Unlock()
 }
 
 func (b *Broker) MessageDeny(ctx context.Context, msg *codec.Message) {
@@ -947,7 +989,9 @@ func (b *Broker) receive(ctx context.Context, c net.Conn) (*codec.Message, error
 
 func (b *Broker) broadcastToConsumers(msg *codec.Message) {
 	if queue, ok := b.queues.Get(msg.Queue); ok {
-		task := &QueuedTask{Message: msg, RetryCount: 0}
+		var t Task
+		json.Unmarshal(msg.Payload, &t)
+		task := &QueuedTask{Message: msg, Task: &t, RetryCount: 0}
 		queue.tasks <- task
 	}
 }
@@ -1166,7 +1210,7 @@ func (b *Broker) dispatchWorker(ctx context.Context, queue *Queue) {
 }
 
 func (b *Broker) sendToDLQ(queue *Queue, task *QueuedTask) {
-	id, _ := jsonparser.GetString(task.Message.Payload, "id")
+	id := task.Task.ID
 	if dlq, ok := b.deadLetter.Get(queue.name); ok {
 		log.Printf("Sending task %s to dead-letter queue for %s", id, queue.name)
 		dlq.tasks <- task
@@ -1180,7 +1224,7 @@ func (b *Broker) dispatchTaskToConsumer(ctx context.Context, queue *Queue, task
 	var err error
 
 	// Deduplication: Check if the task has already been processed
-	taskID, _ := jsonparser.GetString(task.Message.Payload, "id")
+	taskID := task.Task.ID
 	if _, exists := b.pIDs.Get(taskID); exists {
 		log.Printf("Task %s already processed, skipping...", taskID)
 		return true
@@ -1193,6 +1237,12 @@ func (b *Broker) dispatchTaskToConsumer(ctx context.Context, queue *Queue, task
 		}
 
 		// Send message asynchronously to avoid blocking
+		b.mu.Lock()
+		if b.pendingTasks[con.id] == nil {
+			b.pendingTasks[con.id] = make(map[string]*Task)
+		}
+		b.pendingTasks[con.id][taskID] = task.Task
+		b.mu.Unlock()
 		go func(consumer *consumer, message *codec.Message) {
 			sendCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 			defer cancel()
@@ -1229,10 +1279,10 @@ func (b *Broker) dispatchTaskToConsumer(ctx context.Context, queue *Queue, task
 	return consumerFound
 }
 
-// Modified backoffRetry: Removed re‑insertion of the task into queue.tasks.
+// Modified backoffRetry: Re-insert the task into queue.tasks after backoff.
 func (b *Broker) backoffRetry(queue *Queue, task *QueuedTask, delay time.Duration) time.Duration {
 	backoffDuration := utils.CalculateJitter(delay, b.opts.jitterPercent)
-	log.Printf("Backing off for %v before retrying task for queue %s", backoffDuration, task.Message.Queue)
+	log.Printf("Backing off for %v before retrying task for queue %s", backoffDuration, task.Task.Topic)
 
 	// Perform backoff sleep in a goroutine to avoid blocking
 	go func() {