diff options
| -rw-r--r-- | main.go | 10 | ||||
| -rw-r--r-- | pubsub.go | 115 | ||||
| -rw-r--r-- | redisDB.go | 7 |
3 files changed, 131 insertions, 1 deletions
@@ -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 +} @@ -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)) |