basic pubsub system

3 files changed, 131 insertions, 1 deletions

diff --git a/main.go b/main.go
index 84974cd..80b78cb 100644
--- a/main.go
+++ b/main.go
@@ -13,6 +13,7 @@ import (
 const NumShards = 16
 
 var db RedisDB
+var hub Hub
 
 type handlerFunc func(net.Conn, []string)
 
@@ -41,6 +42,9 @@ var commandRegistry = map[string]handlerFunc{
 	"DBSIZE":   handleDbsize,
 	"KEYS":     handleKeys,
 	"LREM":     handleLrem,
+	"SUBSCRIBE":   handleSubscribe,
+	"UNSUBSCRIBE": handleUnsubscribe,
+	"PUBLISH":     handlePublish,
 }
 
 func main() {
@@ -53,6 +57,7 @@ func main() {
     }
 }
 	db = RedisDB{shards: shards}
+	hub = Hub{Channels: make(map[string]map[net.Conn]struct{}),}
 
 	go func() {
 		ticker := time.NewTicker(1 * time.Second)
@@ -88,7 +93,10 @@ func main() {
 }
 
 func handleConnection(conn net.Conn) {
-	defer conn.Close()
+	defer func() {
+		handleDisconnect(conn)
+		conn.Close()
+	}()
 	fmt.Println("Client connected:", conn.RemoteAddr())
 
 	reader := bufio.NewReader(conn)
diff --git a/pubsub.go b/pubsub.go
new file mode 100644
index 0000000..31c3d86
--- /dev/null
+++ b/pubsub.go
@@ -0,0 +1,115 @@
+package main
+
+import (
+	"errors"
+	"net"
+	"sync"
+)
+
+type Hub struct {
+	mu       sync.RWMutex
+	Channels map[string]map[net.Conn]struct{}
+}
+
+func handleSubscribe(conn net.Conn, args []string) {
+	if len(args) < 2 {
+		conn.Write(serializeRESP(errors.New("wrong number of arguments for 'SUBSCRIBE'")))
+		return
+	}
+	channels := args[1:]
+	hub.mu.Lock()
+	defer hub.mu.Unlock()
+	for _, channel := range channels {
+		if _, exists := hub.Channels[channel]; !exists {
+			hub.Channels[channel] = make(map[net.Conn]struct{})
+		}
+		hub.Channels[channel][conn] = struct{}{}
+
+		subCount := 0
+		for _, subscribers := range hub.Channels {
+			if _, ok := subscribers[conn]; ok {
+				subCount++
+			}
+		}
+
+		conn.Write(serializeRESP([]any{"subscribe", channel, subCount}))
+	}
+}
+
+func handleUnsubscribe(conn net.Conn, args []string) {
+	hub.mu.Lock()
+	defer hub.mu.Unlock()
+
+	var channels []string
+	if len(args) < 2 {
+		for channel, subscribers := range hub.Channels {
+			if _, ok := subscribers[conn]; ok {
+				channels = append(channels, channel)
+			}
+		}
+	} else {
+		channels = args[1:]
+	}
+
+	for _, channel := range channels {
+		if subscribers, exists := hub.Channels[channel]; exists {
+			delete(subscribers, conn)
+			if len(subscribers) == 0 {
+				delete(hub.Channels, channel)
+			}
+		}
+
+		subCount := 0
+		for _, subscribers := range hub.Channels {
+			if _, ok := subscribers[conn]; ok {
+				subCount++
+			}
+		}
+
+		conn.Write(serializeRESP([]any{"unsubscribe", channel, subCount}))
+	}
+}
+
+func handlePublish(conn net.Conn, args []string) {
+	if len(args) != 3 {
+		conn.Write(serializeRESP(errors.New("wrong number of arguments for 'PUBLISH'")))
+		return
+	}
+	conn.Write(serializeRESP(hub.Publish(args[1], args[2])))
+}
+
+func handleDisconnect(conn net.Conn) {
+	hub.mu.Lock()
+	defer hub.mu.Unlock()
+	for channel, subscribers := range hub.Channels {
+		delete(subscribers, conn)
+		if len(subscribers) == 0 {
+			delete(hub.Channels, channel)
+		}
+	}
+}
+
+func (hub *Hub) Publish(channel, message string) int {
+	hub.mu.RLock()
+	subscribers, exists := hub.Channels[channel]
+	if !exists || len(subscribers) == 0 {
+		hub.mu.RUnlock()
+		return 0
+	}
+
+	conns := make([]net.Conn, 0, len(subscribers))
+	for conn := range subscribers {
+		conns = append(conns, conn)
+	}
+	hub.mu.RUnlock()
+
+	payload := serializeRESP([]string{"message", channel, message})
+	count := 0
+	for _, conn := range conns {
+		_, err := conn.Write(payload)
+		if err == nil {
+			count++
+		}
+	}
+	return count
+}
diff --git a/redisDB.go b/redisDB.go
index cde14d0..08c044d 100644
--- a/redisDB.go
+++ b/redisDB.go
@@ -35,6 +35,13 @@ func serializeRESP(v any) []byte {
 			result = append(result, serializeRESP(val[i])...)
 		}
 		return result
+	case []any:
+		size := len(val)
+		result := []byte(fmt.Sprintf("*%d\r\n", size))
+		for i := 0; i < size; i++ {
+			result = append(result, serializeRESP(val[i])...)
+		}
+		return result
 	case map[string]string:
 		size := len(val)
 		result := []byte(fmt.Sprintf("*%d\r\n", size*2))