moved dicedb subscription reset logic to cleanup method and added syn…

…chronization for running LivePrefix after subscription state is reset

cli/main.go

@@ -7,6 +7,7 @@ import (
 	"os"
 	"os/exec"
 	"strings"
+	"sync"
 
 	"github.com/c-bata/go-prompt"
 	"github.com/dicedb/dicedb-go"
@@ -21,6 +22,7 @@ type DiceDBClient struct {
 	subCancel  context.CancelFunc
 	addr       string
 	password   string
+	wg         sync.WaitGroup
 }
 
 func Run(host string, port int) {
@@ -57,8 +59,6 @@ func Run(host string, port int) {
 			Key: prompt.ControlC,
 			Fn: func(buf *prompt.Buffer) {
 				if dicedbClient.subscribed {
-					fmt.Println("Exiting watch mode.")
-
 					dicedbClient.handleWatchModeExit()
 				} else {
 					handleExit()
@@ -175,9 +175,10 @@ func (c *DiceDBClient) handleWatchCommand(cmd string, args []string) {
 }
 
 func (c *DiceDBClient) handleWatchModeExit() {
+	fmt.Println("Exiting watch mode.")
+
 	c.subCancel()
-	c.subscribed = false
-	c.subType = ""
+	c.wg.Wait()
 }
 
 func (c *DiceDBClient) handleUnsubscribe() {
@@ -300,13 +301,17 @@ func (c *DiceDBClient) subscribe(channels []string) {
 }
 
 func (c *DiceDBClient) watchCommand(cmd string, args ...interface{}) {
+	c.wg.Add(1)
 	defer func() {
 		c.subscribed = false
 		c.subType = ""
+		if c.watchConn != nil {
+			c.watchConn.Close()
+		}
+		c.wg.Done()
 	}()
 
 	c.watchConn = c.client.WatchConn(c.subCtx)
-	defer c.watchConn.Close()
 
 	// Send the WATCH command
 	firstMsg, err := c.watchConn.Watch(c.subCtx, cmd, args...)