A privacy-first, self-hosted, fully open source personal knowledge management software, written in typescript and golang. (PERSONAL FORK)
0
fork

Configure Feed

Select the types of activity you want to include in your feed.

:recycle: Refactor EventSource broadcast subscribe (#13708)

* :art: Refactor broadcast subscribe using event-bus

* :bug: Fix invalid query param `retry`

* :art: Add event source status log

* :bug: Fix EventSource destroy method

authored by

Yingyi / 颖逸 and committed by
GitHub
ea1081db 866189bc

+45 -85
+44 -84
kernel/api/broadcast.go
··· 23 23 "time" 24 24 25 25 "github.com/88250/gulu" 26 + "github.com/asaskevich/EventBus" 26 27 "github.com/gin-contrib/sse" 27 28 "github.com/gin-gonic/gin" 28 29 "github.com/olahol/melody" ··· 34 35 MessageTypeString MessageType = "string" 35 36 MessageTypeBinary MessageType = "binary" 36 37 MessageTypeClose MessageType = "close" 38 + 39 + EvtBroadcastMessage = "broadcast.message" 37 40 ) 38 41 39 42 var ( 40 43 BroadcastChannels = sync.Map{} // [string (channel-name)] -> *BroadcastChannel 41 - UnifiedSSE = NewEventSourceServer() 42 - messageID = &MessageID{ 44 + UnifiedSSE = &EventSourceServer{ 45 + EventBus: EventBus.New(), 46 + WaitGroup: &sync.WaitGroup{}, 47 + Subscriber: &EventSourceSubscriber{ 48 + lock: &sync.Mutex{}, 49 + count: 0, 50 + }, 51 + } 52 + messageID = &MessageID{ 43 53 lock: &sync.Mutex{}, 44 54 id: 0, 45 55 } ··· 124 134 } 125 135 126 136 func (b *BroadcastChannel) Destroy(force bool) bool { 127 - if force || b.Subscribed() { 137 + if force || !b.Subscribed() { 128 138 b.WebSocket.Close() 129 139 UnifiedSSE.SendEvent(&MessageEvent{ 130 140 Type: MessageTypeClose, ··· 156 166 } 157 167 158 168 type EventSourceServer struct { 159 - Channel MessageEventChannel // message broadcast channel 160 - Open chan MessageEventChannel // SSE connection open channel 161 - Close chan MessageEventChannel // SSE connection close channel 162 - Connections map[MessageEventChannel]bool // SSE connections 163 - 169 + EventBus EventBus.Bus 164 170 WaitGroup *sync.WaitGroup 165 171 Subscriber *EventSourceSubscriber 166 172 } 167 173 168 - // Start starts the SSE server 169 - func (s *EventSourceServer) Start() { 170 - // REF: https://github.com/gin-gonic/examples/blob/master/server-sent-event/main.go 171 - for { 172 - select { 173 - // Add new available client 174 - case channel := <-s.Open: 175 - s.Connections[channel] = true 176 - 177 - // Remove closed client 178 - case channel := <-s.Close: 179 - delete(s.Connections, channel) 180 - close(channel) 181 - 182 - // Broadcast message to client 183 - case event := <-s.Channel: 184 - for connection := range s.Connections { 185 - connection <- event 186 - } 187 - } 188 - } 189 - } 190 - 191 174 // SendEvent sends a message to all subscribers 192 175 func (s *EventSourceServer) SendEvent(event *MessageEvent) bool { 193 176 if event.ID == "" { ··· 198 181 } 199 182 } 200 183 201 - s.Channel <- event 202 - return true 203 - 204 - // select { 205 - // case s.Channel <- event: 206 - // return true 207 - // default: 208 - // logging.LogErrorf("send event failed: %v", event) 209 - // return false 210 - // } 184 + s.EventBus.Publish(EvtBroadcastMessage, event) 185 + return s.EventBus.HasCallback(EvtBroadcastMessage) 211 186 } 212 187 213 188 // Subscribe subscribes to specified broadcast channels 214 - func (s *EventSourceServer) Subscribe(c *gin.Context, messageEventChannel MessageEventChannel, channels ...string) { 215 - defer s.WaitGroup.Done() 216 - s.WaitGroup.Add(1) 217 - 189 + func (s *EventSourceServer) Subscribe(c *gin.Context, retry uint, channels ...string) { 218 190 wg := sync.WaitGroup{} 219 191 wg.Add(len(channels)) 220 192 for _, channel := range channels { ··· 239 211 } 240 212 241 213 c.Writer.Flush() 242 - retry := s.GetRetry(c) 243 - s.Stream(c, messageEventChannel, func(event *MessageEvent, ok bool) bool { 214 + s.Stream(c, func(event *MessageEvent, ok bool) bool { 244 215 if ok { 245 216 if _, exists := channelSet[event.Name]; exists { 246 217 switch event.Type { ··· 288 259 } 289 260 290 261 // SubscribeAll subscribes to all broadcast channels 291 - func (s *EventSourceServer) SubscribeAll(c *gin.Context, messageEventChannel MessageEventChannel) { 292 - defer s.WaitGroup.Done() 293 - s.WaitGroup.Add(1) 294 - 262 + func (s *EventSourceServer) SubscribeAll(c *gin.Context, retry uint) { 295 263 s.Subscriber.updateCount(1) 296 264 297 265 c.Writer.Flush() 298 - retry := s.GetRetry(c) 299 - s.Stream(c, messageEventChannel, func(event *MessageEvent, ok bool) bool { 266 + s.Stream(c, func(event *MessageEvent, ok bool) bool { 300 267 if ok { 301 268 switch event.Type { 302 269 case MessageTypeClose: ··· 324 291 325 292 s.Subscriber.updateCount(-1) 326 293 PruneBroadcastChannels() 327 - 328 294 } 329 295 330 296 // GetRetry gets the retry interval 331 297 // 332 298 // If the retry interval is not specified, it will return 0 333 299 func (s *EventSourceServer) GetRetry(c *gin.Context) uint { 334 - value, err := c.GetQuery("retry") 335 - if !err { 336 - retry, err := strconv.ParseUint(value, 10, 0) 337 - if err == nil { 338 - return uint(retry) 339 - } 300 + value := c.DefaultQuery("retry", "") 301 + retry, err := strconv.ParseUint(value, 10, 0) 302 + if err == nil { 303 + return uint(retry) 340 304 } 341 305 return 0 342 306 } 343 307 344 308 // Stream streams message to client 345 - func (s *EventSourceServer) Stream(c *gin.Context, channel MessageEventChannel, step func(event *MessageEvent, ok bool) bool) bool { 309 + // 310 + // If the client is gone, it will return true 311 + func (s *EventSourceServer) Stream(c *gin.Context, step func(event *MessageEvent, ok bool) bool) bool { 312 + channel := make(MessageEventChannel) 313 + defer close(channel) 314 + 315 + subscriber := func(event *MessageEvent) { 316 + channel <- event 317 + } 318 + s.EventBus.Subscribe(EvtBroadcastMessage, subscriber) 319 + defer s.EventBus.Unsubscribe(EvtBroadcastMessage, subscriber) 320 + 346 321 clientGone := c.Writer.CloseNotify() 347 322 for { 348 323 select { 349 324 case <-clientGone: 325 + logging.LogInfof("event source connection is closed by client") 350 326 return true 351 327 case event, ok := <-channel: 352 328 if step(event, ok) { 353 329 continue 354 330 } 331 + logging.LogInfof("event source connection is closed by server") 355 332 return false 356 333 } 357 334 } ··· 362 339 c.Render(-1, event) 363 340 } 364 341 342 + // Subscribed checks whether the SSE server is subscribed 365 343 func (s *EventSourceServer) Subscribed() bool { 366 344 return s.Subscriber.Count() > 0 367 - } 368 - 369 - func NewEventSourceServer() (server *EventSourceServer) { 370 - server = &EventSourceServer{ 371 - Channel: make(MessageEventChannel, 1024), 372 - Open: make(chan MessageEventChannel, 32), 373 - Close: make(chan MessageEventChannel, 32), 374 - Connections: make(map[MessageEventChannel]bool), 375 - 376 - WaitGroup: &sync.WaitGroup{}, 377 - Subscriber: &EventSourceSubscriber{ 378 - lock: &sync.Mutex{}, 379 - count: 0, 380 - }, 381 - } 382 - go server.Start() 383 - return 384 345 } 385 346 386 347 type ChannelInfo struct { ··· 571 532 c.Writer.Header().Set("Connection", "keep-alive") 572 533 c.Writer.Header().Set("Transfer-Encoding", "chunked") 573 534 574 - messageEventChannel := make(MessageEventChannel) 575 - UnifiedSSE.Open <- messageEventChannel 535 + defer UnifiedSSE.WaitGroup.Done() 536 + UnifiedSSE.WaitGroup.Add(1) 576 537 538 + retry := UnifiedSSE.GetRetry(c) 577 539 channels, ok := c.GetQueryArray("channel") 578 540 if ok { // subscribe specified broadcast channels 579 - UnifiedSSE.Subscribe(c, messageEventChannel, channels...) 541 + UnifiedSSE.Subscribe(c, retry, channels...) 580 542 } else { // subscribe all broadcast channels 581 - UnifiedSSE.SubscribeAll(c, messageEventChannel) 543 + UnifiedSSE.SubscribeAll(c, retry) 582 544 } 583 - 584 - UnifiedSSE.Close <- messageEventChannel 585 545 } 586 546 587 547 // broadcastPublish push multiple binary messages to multiple broadcast channels
+1 -1
kernel/go.mod
··· 18 18 github.com/PuerkitoBio/goquery v1.10.0 19 19 github.com/Xuanwo/go-locale v1.1.2 20 20 github.com/araddon/dateparse v0.0.0-20210429162001-6b43995a97de 21 + github.com/asaskevich/EventBus v0.0.0-20200907212545-49d423059eef 21 22 github.com/common-nighthawk/go-figure v0.0.0-20210622060536-734e95fb86be 22 23 github.com/denisbrodbeck/machineid v1.0.1 23 24 github.com/dgraph-io/ristretto v1.0.0 ··· 91 92 github.com/alecthomas/chroma v0.10.0 // indirect 92 93 github.com/andybalholm/brotli v1.1.1 // indirect 93 94 github.com/andybalholm/cascadia v1.3.2 // indirect 94 - github.com/asaskevich/EventBus v0.0.0-20200907212545-49d423059eef // indirect 95 95 github.com/aws/aws-sdk-go-v2 v1.32.7 // indirect 96 96 github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.7 // indirect 97 97 github.com/aws/aws-sdk-go-v2/config v1.28.7 // indirect