turn pubsub into separate package

1 files changed, 106 insertions, 0 deletions

diff --git a/pkg/pubsub/pubsub.go b/pkg/pubsub/pubsub.go
new file mode 100644
index 0000000..690606c
--- /dev/null
+++ b/pkg/pubsub/pubsub.go
@@ -0,0 +1,106 @@
+package pubsub
+
+import (
+	"net"
+	"redisClone/pkg/core"
+	"sync"
+)
+
+type Hub struct {
+	mu       sync.RWMutex
+	Channels map[string]map[net.Conn]struct{}
+}
+
+func NewHub() *Hub {
+	return &Hub{
+		Channels: make(map[string]map[net.Conn]struct{}),
+	}
+}
+
+func (h *Hub) Subscribe(conn net.Conn, channel string) int {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	if _, exists := h.Channels[channel]; !exists {
+		h.Channels[channel] = make(map[net.Conn]struct{})
+	}
+	h.Channels[channel][conn] = struct{}{}
+
+	subCount := 0
+	for _, subscribers := range h.Channels {
+		if _, ok := subscribers[conn]; ok {
+			subCount++
+		}
+	}
+	return subCount
+}
+
+func (h *Hub) Unsubscribe(conn net.Conn, channels []string) map[string]int {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	if len(channels) == 0 {
+		for channel, subscribers := range h.Channels {
+			if _, ok := subscribers[conn]; ok {
+				channels = append(channels, channel)
+			}
+		}
+	}
+
+	results := make(map[string]int)
+
+	for _, channel := range channels {
+		if subscribers, exists := h.Channels[channel]; exists {
+			delete(subscribers, conn)
+			if len(subscribers) == 0 {
+				delete(h.Channels, channel)
+			}
+		}
+
+		count := 0
+		for _, subs := range h.Channels {
+			if _, ok := subs[conn]; ok {
+				count++
+			}
+		}
+		results[channel] = count
+	}
+	return results
+}
+
+func (h *Hub) Publish(channel, message string) int {
+	h.mu.RLock()
+	subscribers, exists := h.Channels[channel]
+	if !exists || len(subscribers) == 0 {
+		h.mu.RUnlock()
+		return 0
+	}
+
+	conns := make([]net.Conn, 0, len(subscribers))
+	for conn := range subscribers {
+		conns = append(conns, conn)
+	}
+	h.mu.RUnlock()
+
+	payload := core.SerializeRESP([]string{"message", channel, message})
+	count := 0
+	for _, conn := range conns {
+		_, err := conn.Write(payload)
+		if err == nil {
+			count++
+		}
+	}
+	return count
+}
+
+func (h *Hub) HandleDisconnect(conn net.Conn) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	for channel, subscribers := range h.Channels {
+		delete(subscribers, conn)
+		if len(subscribers) == 0 {
+			delete(h.Channels, channel)
+		}
+	}
+}
\ No newline at end of file