fix: Improve watch command handling after closure

client_wire.go

@@ -17,7 +17,8 @@ type ClientWire struct {
 }
 
 func NewClientWire(maxMsgSize int, host string, port int) (*ClientWire, *wire.WireError) {
-	addr := fmt.Sprintf("%s:%d", host, port)
+	// Use net.JoinHostPort to properly handle both IPv4 and IPv6 addresses
+	addr := net.JoinHostPort(host, fmt.Sprintf("%d", port))
 	conn, err := net.DialTimeout("tcp", addr, 5*time.Second)
 	if err != nil {
 		return nil, &wire.WireError{Kind: wire.NotEstablished, Cause: err}
@@ -43,3 +44,7 @@ func (cw *ClientWire) Receive() (*wire.Result, *wire.WireError) {
 func (cw *ClientWire) Close() {
 	cw.ProtobufTCPWire.Close()
 }
+
+func (cw *ClientWire) IsClosed() bool {
+	return cw.ProtobufTCPWire.IsClosed()
+}

internal/protobuf_wire.go

@@ -4,9 +4,10 @@
 package internal
 
 import (
-	"github.com/dicedb/dicedb-go/wire"
 	"net"
 
+	"github.com/dicedb/dicedb-go/wire"
+
 	"google.golang.org/protobuf/proto"
 )
 
@@ -48,3 +49,7 @@ func (w *ProtobufTCPWire) Receive(dst proto.Message) *wire.WireError {
 func (w *ProtobufTCPWire) Close() {
 	w.tcpWire.Close()
 }
+
+func (w *ProtobufTCPWire) IsClosed() bool {
+	return w.tcpWire.IsClosed()
+}

internal/tcp_wire.go

@@ -5,16 +5,18 @@ package internal
 
 import (
 	"bufio"
+	"context"
 	"encoding/binary"
 	"errors"
 	"fmt"
-	"github.com/dicedb/dicedb-go/wire"
 	"io"
 	"log/slog"
 	"net"
 	"strings"
 	"sync"
 	"time"
+
+	"github.com/dicedb/dicedb-go/wire"
 )
 
 const prefixSize = 4 // bytes
@@ -33,22 +35,27 @@ type TCPWire struct {
 	reader     *bufio.Reader
 	writeMu    sync.Mutex
 	conn       net.Conn
+	ctx        context.Context
+	cancel     context.CancelFunc
 }
 
 func NewTCPWire(maxMsgSize int, conn net.Conn) *TCPWire {
+	ctx, cancel := context.WithCancel(context.Background())
 	return &TCPWire{
 		status:     Open,
 		maxMsgSize: maxMsgSize,
 		conn:       conn,
 		reader:     bufio.NewReader(conn),
+		ctx:        ctx,
+		cancel:     cancel,
 	}
 }
 
 func (w *TCPWire) Send(msg []byte) *wire.WireError {
 	w.writeMu.Lock()
 	defer w.writeMu.Unlock()
 
-	if w.status == Closed {
+	if w.IsClosed() {
 		return &wire.WireError{Kind: wire.Terminated, Cause: errors.New("trying to use closed wire")}
 	}
 
@@ -70,6 +77,11 @@ func (w *TCPWire) Receive() ([]byte, *wire.WireError) {
 		return nil, err
 	}
 
+	if w.IsClosed() {
+		// If the wire is closed, we return an empty buffer to indicate that the read was interrupted
+		return make([]byte, 0), nil
+	}
+
 	if size <= 0 {
 		w.Close()
 		return nil, &wire.WireError{
@@ -95,27 +107,47 @@ func (w *TCPWire) Receive() ([]byte, *wire.WireError) {
 }
 
 func (w *TCPWire) Close() {
-	if w.status == Closed {
+	if w.IsClosed() {
 		return
 	}
 
+	// Cancel the context to stop the read goroutine
+	w.cancel()
 	w.status = Closed
 	err := w.conn.Close()
 	if err != nil {
 		slog.Warn("error closing network connection", "error", err)
-
 		return
 	}
 }
 
+func (w *TCPWire) ReadWithInterruption(buffer []byte, ctx context.Context) (int, error) {
+	done := make(chan struct{})
+	var n int
+	var readErr error
+
+	go func() {
+		n, readErr = io.ReadFull(w.reader, buffer)
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		return n, readErr
+	case <-ctx.Done():
+		<-done
+		return 0, nil
+	}
+}
+
 func (w *TCPWire) readPrefix() (uint32, *wire.WireError) {
 	buffer := make([]byte, prefixSize)
 	delay := 5 * time.Millisecond
 	const maxRetries = 5
 
 	var lastErr error
 	for attempt := 0; attempt < maxRetries; attempt++ {
-		_, err := io.ReadFull(w.reader, buffer)
+		_, err := w.ReadWithInterruption(buffer, w.ctx)
 		if err == nil {
 			return binary.BigEndian.Uint32(buffer), nil
 		}
@@ -166,7 +198,7 @@ func (w *TCPWire) readMessage(size uint32) ([]byte, *wire.WireError) {
 
 	var lastErr error
 	for attempt := 0; attempt < maxRetries; attempt++ {
-		_, err := io.ReadFull(w.reader, buffer)
+		_, err := w.ReadWithInterruption(buffer, w.ctx)
 		if err == nil {
 			return buffer, nil
 		}
@@ -274,3 +306,7 @@ func (w *TCPWire) write(buffer []byte) *wire.WireError {
 func prefix(msgSize int, buffer []byte) {
 	binary.BigEndian.PutUint32(buffer[:prefixSize], uint32(msgSize))
 }
+
+func (w *TCPWire) IsClosed() bool {
+	return w.status == Closed
+}

internal/wire.go

@@ -6,4 +6,5 @@ type Wire interface {
 	Send([]byte) *wire.WireError
 	Receive() ([]byte, *wire.WireError)
 	Close()
+	IsClosed() bool
 }

main.go

@@ -139,7 +139,7 @@ func (c *Client) WatchCh() (<-chan *wire.Result, error) {
 	c.watchRetrier = NewRetrier(5, 5*time.Second)
 	c.watchWire, err = NewClientWire(maxResponseSize, c.host, c.port)
 	if err != nil {
-		return nil, fmt.Errorf("Failed to establish watch connection with server: %w", err)
+		return nil, fmt.Errorf("failed to establish watch connection with server: %w", err)
 	}
 
 	if resp := c.fire(&wire.Command{
@@ -165,15 +165,28 @@ func (c *Client) watch() {
 			break
 		}
 
+		// if the watch connection is closed, break the loop
+		if c.watchWire.IsClosed() {
+			break
+		}
+
 		c.watchCh <- resp
 	}
 }
 
 func (c *Client) Close() {
 	c.mainWire.Close()
+	if c.watchCh != nil {
+		c.CloseWatch()
+	}
+}
+
+func (c *Client) CloseWatch() {
 	if c.watchCh != nil {
 		c.watchWire.Close()
 		close(c.watchCh)
+		// set the watch channel to nil, for next watch to start a new connection
+		c.watchCh = nil
 	}
 }