1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
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
}
|
