Add APIs for message broadcast (#9031)
* 🎨 Add API for broadcast * 🎨 Add API `/api/broadcast/getListenerCount` * 🎨 Add API GET `/api/broadcast/channels` * 🎨 Add API POST `/api/broadcast/getChannelInfo` * Update broadcast.go * Update broadcast.go
This commit is contained in:
Yingyi / 颖逸
GitHub
parent
f1e875698a
commit
2f75010785
2 changed files with 231 additions and 0 deletions
226
kernel/api/broadcast.go
Normal file
226
kernel/api/broadcast.go
Normal file
|
|
@ -0,0 +1,226 @@
|
|||
// SiYuan - Refactor your thinking
|
||||
// Copyright (c) 2020-present, b3log.org
|
||||
//
|
||||
// This program is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU Affero General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
//
|
||||
// This program is distributed in the hope that it will be useful,
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU Affero General Public License for more details.
|
||||
//
|
||||
// You should have received a copy of the GNU Affero General Public License
|
||||
// along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
package api
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sync"
|
||||
|
||||
"github.com/88250/gulu"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/olahol/melody"
|
||||
"github.com/siyuan-note/logging"
|
||||
"github.com/siyuan-note/siyuan/kernel/util"
|
||||
)
|
||||
|
||||
type Channel struct {
|
||||
Name string `json:"name"`
|
||||
Count int `json:"count"`
|
||||
}
|
||||
|
||||
var (
|
||||
BroadcastChannels = sync.Map{}
|
||||
)
|
||||
|
||||
/*
|
||||
broadcast create a broadcast channel WebSocket connection
|
||||
|
||||
@param
|
||||
|
||||
query.channel: channel name
|
||||
|
||||
@example
|
||||
|
||||
ws://localhost:6806/ws/broadcast?channel=test
|
||||
*/
|
||||
func broadcast(c *gin.Context) {
|
||||
var (
|
||||
channel string = c.Query("channel")
|
||||
broadcastChannel *melody.Melody
|
||||
)
|
||||
|
||||
if _broadcastChannel, exist := BroadcastChannels.Load(channel); exist {
|
||||
// channel exists, use it
|
||||
broadcastChannel = _broadcastChannel.(*melody.Melody)
|
||||
subscribe(c, broadcastChannel, channel)
|
||||
} else {
|
||||
// channel not found, create a new one
|
||||
broadcastChannel := melody.New()
|
||||
BroadcastChannels.Store(channel, broadcastChannel)
|
||||
subscribe(c, broadcastChannel, channel)
|
||||
|
||||
// broadcast string message to other session
|
||||
broadcastChannel.HandleMessage(func(s *melody.Session, msg []byte) {
|
||||
broadcastChannel.BroadcastOthers(msg, s)
|
||||
})
|
||||
|
||||
// broadcast binary message to other session
|
||||
broadcastChannel.HandleMessageBinary(func(s *melody.Session, msg []byte) {
|
||||
broadcastChannel.BroadcastBinaryOthers(msg, s)
|
||||
})
|
||||
|
||||
// recycling
|
||||
broadcastChannel.HandleClose(func(s *melody.Session, status int, reason string) error {
|
||||
channel := s.Keys["channel"].(string)
|
||||
logging.LogInfof("close broadcast session in channel [%s] with status code %d: %s", channel, status, reason)
|
||||
|
||||
count := broadcastChannel.Len()
|
||||
if count == 0 {
|
||||
BroadcastChannels.Delete(channel)
|
||||
logging.LogInfof("dispose broadcast channel [%s]", channel)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// subscribe creates a new websocket session to a channel
|
||||
func subscribe(c *gin.Context, broadcastChannel *melody.Melody, channel string) {
|
||||
if err := broadcastChannel.HandleRequestWithKeys(
|
||||
c.Writer,
|
||||
c.Request,
|
||||
map[string]interface{}{
|
||||
"channel": channel,
|
||||
},
|
||||
); nil != err {
|
||||
logging.LogErrorf("create broadcast channel failed: %s", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
postMessage send string message to a broadcast channel
|
||||
|
||||
@param
|
||||
|
||||
body.channel: channel name
|
||||
body.message: message payload
|
||||
|
||||
@returns
|
||||
|
||||
body.data.channel.name: channel name
|
||||
body.data.channel.count: indicate how many websocket session received the message
|
||||
*/
|
||||
func postMessage(c *gin.Context) {
|
||||
ret := gulu.Ret.NewResult()
|
||||
defer c.JSON(http.StatusOK, ret)
|
||||
|
||||
arg, ok := util.JsonArg(c, ret)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
channel := arg["channel"].(string)
|
||||
message := arg["message"].(string)
|
||||
|
||||
if _broadcastChannel, ok := BroadcastChannels.Load(channel); !ok {
|
||||
err := fmt.Errorf("broadcast channel [%s] not found", channel)
|
||||
logging.LogWarnf(err.Error())
|
||||
|
||||
ret.Code = -1
|
||||
ret.Msg = err.Error()
|
||||
return
|
||||
} else {
|
||||
var broadcastChannel = _broadcastChannel.(*melody.Melody)
|
||||
if err := broadcastChannel.Broadcast([]byte(message)); nil != err {
|
||||
logging.LogErrorf("broadcast message failed: %s", err)
|
||||
|
||||
ret.Code = -2
|
||||
ret.Msg = err.Error()
|
||||
return
|
||||
}
|
||||
|
||||
count := broadcastChannel.Len()
|
||||
ret.Data = map[string]interface{}{
|
||||
"channel": &Channel{
|
||||
Name: channel,
|
||||
Count: count,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
getChannelInfo gets the information of a broadcast channel
|
||||
|
||||
@param
|
||||
|
||||
body.name: channel name
|
||||
|
||||
@returns
|
||||
|
||||
body.data.channel: the channel information
|
||||
*/
|
||||
func getChannelInfo(c *gin.Context) {
|
||||
ret := gulu.Ret.NewResult()
|
||||
defer c.JSON(http.StatusOK, ret)
|
||||
|
||||
arg, ok := util.JsonArg(c, ret)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
name := arg["name"].(string)
|
||||
|
||||
if _broadcastChannel, ok := BroadcastChannels.Load(name); !ok {
|
||||
err := fmt.Errorf("broadcast channel [%s] not found", name)
|
||||
logging.LogWarnf(err.Error())
|
||||
|
||||
ret.Code = -1
|
||||
ret.Msg = err.Error()
|
||||
return
|
||||
} else {
|
||||
var broadcastChannel = _broadcastChannel.(*melody.Melody)
|
||||
|
||||
count := broadcastChannel.Len()
|
||||
ret.Data = map[string]interface{}{
|
||||
"channel": &Channel{
|
||||
Name: name,
|
||||
Count: count,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
getChannels gets the channel name and lintener number of all broadcast chanel
|
||||
|
||||
@returns
|
||||
|
||||
body.data.channels: {
|
||||
name: channel name
|
||||
count: listener count
|
||||
}[]
|
||||
*/
|
||||
func getChannels(c *gin.Context) {
|
||||
ret := gulu.Ret.NewResult()
|
||||
defer c.JSON(http.StatusOK, ret)
|
||||
|
||||
channels := []*Channel{}
|
||||
BroadcastChannels.Range(func(key, value any) bool {
|
||||
broadcastChannel := value.(*melody.Melody)
|
||||
channels = append(channels, &Channel{
|
||||
Name: key.(string),
|
||||
Count: broadcastChannel.Len(),
|
||||
})
|
||||
return true
|
||||
})
|
||||
ret.Data = map[string]interface{}{
|
||||
"channels": channels,
|
||||
}
|
||||
}
|
||||
|
|
@ -372,4 +372,9 @@ func ServeAPI(ginServer *gin.Engine) {
|
|||
ginServer.Handle("POST", "/api/petal/setPetalEnabled", model.CheckAuth, model.CheckReadonly, setPetalEnabled)
|
||||
|
||||
ginServer.Handle("POST", "/api/network/forwardProxy", model.CheckAuth, model.CheckReadonly, forwardProxy)
|
||||
|
||||
ginServer.Handle("GET", "/ws/broadcast", model.CheckAuth, broadcast)
|
||||
ginServer.Handle("GET", "/api/broadcast/channels", model.CheckAuth, getChannels)
|
||||
ginServer.Handle("POST", "/api/broadcast/postMessage", model.CheckAuth, postMessage)
|
||||
ginServer.Handle("POST", "/api/broadcast/getChannelInfo", model.CheckAuth, getChannelInfo)
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue