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