snapshot of https://github.com/centrifugal/centrifugo.git (master, 2026-09-16; blobs > 131072B omitted)

internal/survey/survey.go

@@ -0,0 +1,110 @@
+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}
+}