111 lines
2.9 KiB
Go
111 lines
2.9 KiB
Go
package survey
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"github.com/centrifugal/centrifugo/v6/internal/apiproto"
|
|
|
|
"github.com/centrifugal/centrifuge"
|
|
"github.com/gobwas/glob"
|
|
"google.golang.org/protobuf/proto"
|
|
)
|
|
|
|
// Handler can handle survey.
|
|
type Handler func(node *centrifuge.Node, data []byte) centrifuge.SurveyReply
|
|
|
|
type Caller struct {
|
|
node *centrifuge.Node
|
|
handlers map[string]Handler
|
|
}
|
|
|
|
func NewCaller(node *centrifuge.Node) *Caller {
|
|
c := &Caller{
|
|
node: node,
|
|
handlers: map[string]Handler{
|
|
"channels": respondChannelsSurvey,
|
|
},
|
|
}
|
|
c.node.OnSurvey(func(event centrifuge.SurveyEvent, cb centrifuge.SurveyCallback) {
|
|
h, ok := c.handlers[event.Op]
|
|
if !ok {
|
|
cb(centrifuge.SurveyReply{Code: MethodNotFound})
|
|
return
|
|
}
|
|
cb(h(c.node, event.Data))
|
|
})
|
|
return c
|
|
}
|
|
|
|
const (
|
|
InternalError uint32 = 1
|
|
InvalidRequest uint32 = 2
|
|
MethodNotFound uint32 = 3
|
|
)
|
|
|
|
func (c *Caller) Channels(ctx context.Context, cmd *apiproto.ChannelsRequest) (map[string]*apiproto.ChannelInfo, error) {
|
|
return surveyChannels(ctx, c.node, cmd)
|
|
}
|
|
|
|
func surveyChannels(ctx context.Context, node *centrifuge.Node, cmd *apiproto.ChannelsRequest) (map[string]*apiproto.ChannelInfo, error) {
|
|
req, _ := proto.Marshal(cmd)
|
|
results, err := node.Survey(ctx, "channels", req, "")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
channels := map[string]*apiproto.ChannelInfo{}
|
|
for nodeID, result := range results {
|
|
if result.Code > 0 {
|
|
return nil, fmt.Errorf("non-zero code from node %s: %d", nodeID, result.Code)
|
|
}
|
|
var nodeChannels apiproto.ChannelsResult
|
|
err := proto.Unmarshal(result.Data, &nodeChannels)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("error unmarshalling data from node %s: %w", nodeID, err)
|
|
}
|
|
for ch, chInfo := range nodeChannels.Channels {
|
|
info, ok := channels[ch]
|
|
if !ok {
|
|
channels[ch] = &apiproto.ChannelInfo{NumClients: chInfo.NumClients}
|
|
} else {
|
|
info.NumClients += chInfo.NumClients
|
|
channels[ch] = info
|
|
}
|
|
}
|
|
}
|
|
return channels, nil
|
|
}
|
|
|
|
func respondChannelsSurvey(node *centrifuge.Node, params []byte) centrifuge.SurveyReply {
|
|
var req apiproto.ChannelsRequest
|
|
err := proto.Unmarshal(params, &req)
|
|
if err != nil {
|
|
return centrifuge.SurveyReply{Code: InvalidRequest}
|
|
}
|
|
var g *glob.Pattern
|
|
if req.Pattern != "" {
|
|
var err error
|
|
g, err = glob.Compile(req.Pattern)
|
|
if err != nil {
|
|
return centrifuge.SurveyReply{Code: InvalidRequest}
|
|
}
|
|
}
|
|
channels := node.Hub().Channels()
|
|
channelsMap := make(map[string]*apiproto.ChannelInfo, len(channels))
|
|
for _, ch := range channels {
|
|
if g != nil && !g.Match(ch) {
|
|
continue
|
|
}
|
|
if numSubscribers := node.Hub().NumSubscribers(ch); numSubscribers > 0 {
|
|
channelsMap[ch] = &apiproto.ChannelInfo{
|
|
NumClients: uint32(numSubscribers),
|
|
}
|
|
}
|
|
}
|
|
data, err := proto.Marshal(&apiproto.ChannelsResult{Channels: channelsMap})
|
|
if err != nil {
|
|
return centrifuge.SurveyReply{Code: InternalError}
|
|
}
|
|
return centrifuge.SurveyReply{Data: data}
|
|
}
|