basic pubsub system

1 files changed, 115 insertions, 0 deletions

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
+}