snapshot of https://github.com/centrifugal/centrifugo.git (master, 2026-09-16; blobs > 131072B omitted)
This commit is contained in:
upstream-mirror
354
internal/proxy/subscribe_stream_handler.go
Normal file
354
internal/proxy/subscribe_stream_handler.go
Normal file
@@ -0,0 +1,354 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"io"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/centrifugal/centrifugo/v6/internal/configtypes"
|
||||
"github.com/centrifugal/centrifugo/v6/internal/metrics"
|
||||
"github.com/centrifugal/centrifugo/v6/internal/proxyproto"
|
||||
"github.com/centrifugal/centrifugo/v6/internal/subsource"
|
||||
|
||||
"github.com/centrifugal/centrifuge"
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
"github.com/rs/zerolog/log"
|
||||
)
|
||||
|
||||
// SubscribeStreamHandlerConfig ...
|
||||
type SubscribeStreamHandlerConfig struct {
|
||||
Proxies map[string]*SubscribeStreamProxy
|
||||
}
|
||||
|
||||
// SubscribeStreamHandler ...
|
||||
type SubscribeStreamHandler struct {
|
||||
config SubscribeStreamHandlerConfig
|
||||
|
||||
summary map[string]prometheus.Observer
|
||||
histogram map[string]prometheus.Observer
|
||||
errors map[string]prometheus.Counter
|
||||
inflight map[string]prometheus.Gauge
|
||||
}
|
||||
|
||||
// NewSubscribeStreamHandler ...
|
||||
func NewSubscribeStreamHandler(c SubscribeStreamHandlerConfig) *SubscribeStreamHandler {
|
||||
h := &SubscribeStreamHandler{
|
||||
config: c,
|
||||
}
|
||||
|
||||
summary := map[string]prometheus.Observer{}
|
||||
histogram := map[string]prometheus.Observer{}
|
||||
errCounters := map[string]prometheus.Counter{}
|
||||
inflight := map[string]prometheus.Gauge{}
|
||||
for name := range c.Proxies {
|
||||
summary[name] = metrics.ProxyCallDurationSummary.WithLabelValues("grpc", "subscribe_stream", name)
|
||||
histogram[name] = metrics.ProxyCallDurationHistogram.WithLabelValues("grpc", "subscribe_stream", name)
|
||||
errCounters[name] = metrics.ProxyCallErrorCount.WithLabelValues("grpc", "subscribe_stream", name)
|
||||
inflight[name] = metrics.ProxyCallInflightRequests.WithLabelValues("grpc", "subscribe_stream", name)
|
||||
}
|
||||
h.summary = summary
|
||||
h.histogram = histogram
|
||||
h.errors = errCounters
|
||||
h.inflight = inflight
|
||||
return h
|
||||
}
|
||||
|
||||
// StreamPublishFunc ...
|
||||
type StreamPublishFunc func(data []byte) error
|
||||
|
||||
// SubscribeStreamHandlerFunc ...
|
||||
type SubscribeStreamHandlerFunc func(
|
||||
Client, bool, centrifuge.SubscribeEvent, configtypes.ChannelOptions, PerCallData,
|
||||
) (centrifuge.SubscribeReply, StreamPublishFunc, func(), error)
|
||||
|
||||
// Handle ...
|
||||
func (h *SubscribeStreamHandler) Handle() SubscribeStreamHandlerFunc {
|
||||
return func(
|
||||
client Client, bidi bool, e centrifuge.SubscribeEvent,
|
||||
chOpts configtypes.ChannelOptions, pcd PerCallData,
|
||||
) (centrifuge.SubscribeReply, StreamPublishFunc, func(), error) {
|
||||
started := time.Now()
|
||||
|
||||
var p *SubscribeStreamProxy
|
||||
var summary prometheus.Observer
|
||||
var histogram prometheus.Observer
|
||||
var errCounter prometheus.Counter
|
||||
|
||||
proxyEnabled := chOpts.SubscribeStreamProxyEnabled
|
||||
proxyName := chOpts.SubscribeStreamProxyName
|
||||
if !proxyEnabled {
|
||||
log.Info().Str("channel", e.Channel).Msg("subscribe stream proxy not enabled for a channel")
|
||||
return centrifuge.SubscribeReply{}, nil, nil, centrifuge.ErrorNotAvailable
|
||||
}
|
||||
p = h.config.Proxies[proxyName]
|
||||
summary = h.summary[proxyName]
|
||||
histogram = h.histogram[proxyName]
|
||||
errCounter = h.errors[proxyName]
|
||||
inflight := h.inflight[proxyName]
|
||||
inflight.Inc()
|
||||
defer inflight.Dec()
|
||||
|
||||
req := &proxyproto.SubscribeRequest{
|
||||
Client: client.ID(),
|
||||
Protocol: string(client.Transport().Protocol()),
|
||||
Transport: client.Transport().Name(),
|
||||
Encoding: getEncoding(p.config.BinaryEncoding),
|
||||
|
||||
User: client.UserID(),
|
||||
Channel: e.Channel,
|
||||
Token: e.Token,
|
||||
}
|
||||
if !p.config.BinaryEncoding {
|
||||
req.Data = e.Data
|
||||
} else {
|
||||
req.B64Data = base64.StdEncoding.EncodeToString(e.Data)
|
||||
}
|
||||
if p.config.IncludeConnectionMeta && pcd.Meta != nil {
|
||||
req.Meta = proxyproto.Raw(pcd.Meta)
|
||||
}
|
||||
|
||||
subscriptionReady := make(chan struct{})
|
||||
var subscriptionReadyOnce sync.Once
|
||||
|
||||
subscribeRep, publishFunc, cancelFunc, err := p.SubscribeStream(
|
||||
client.Context(),
|
||||
bidi,
|
||||
req,
|
||||
func(pub *proxyproto.Publication, err error) {
|
||||
subscriptionReadyOnce.Do(func() {
|
||||
select {
|
||||
case <-subscriptionReady:
|
||||
case <-client.Context().Done():
|
||||
return
|
||||
}
|
||||
})
|
||||
select {
|
||||
case <-client.Context().Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
if err != nil {
|
||||
if errors.Is(err, io.EOF) {
|
||||
client.Unsubscribe(e.Channel, centrifuge.Unsubscribe{
|
||||
Code: centrifuge.UnsubscribeCodeServer,
|
||||
Reason: "server unsubscribe",
|
||||
})
|
||||
return
|
||||
}
|
||||
client.Unsubscribe(e.Channel, centrifuge.Unsubscribe{
|
||||
Code: centrifuge.UnsubscribeCodeInsufficient,
|
||||
Reason: "insufficient state",
|
||||
})
|
||||
return
|
||||
}
|
||||
_ = client.WritePublication(e.Channel, ¢rifuge.Publication{
|
||||
Data: pub.Data,
|
||||
Tags: pub.Tags,
|
||||
}, centrifuge.StreamPosition{})
|
||||
},
|
||||
)
|
||||
|
||||
duration := time.Since(started).Seconds()
|
||||
|
||||
if err != nil {
|
||||
select {
|
||||
case <-client.Context().Done():
|
||||
// Client connection already closed.
|
||||
return centrifuge.SubscribeReply{}, nil, nil, centrifuge.DisconnectConnectionClosed
|
||||
default:
|
||||
}
|
||||
summary.Observe(duration)
|
||||
histogram.Observe(duration)
|
||||
errCounter.Inc()
|
||||
log.Error().Err(err).Str("client", client.ID()).Str("channel", e.Channel).Msg("error from subscribe stream proxy")
|
||||
//proxyCallErrorCount.WithLabelValues(proxyName, "subscribe", "internal").Inc()
|
||||
return centrifuge.SubscribeReply{}, nil, nil, err
|
||||
}
|
||||
|
||||
summary.Observe(duration)
|
||||
histogram.Observe(duration)
|
||||
|
||||
if subscribeRep.Disconnect != nil {
|
||||
//proxyCallErrorCount.WithLabelValues(proxyName, "subscribe", "disconnect_"+strconv.FormatUint(uint64(subscribeRep.Disconnect.Code), 10)).Inc()
|
||||
// SubscribeStream returned successfully here, so it has opened the stream
|
||||
// and spawned its reader goroutine. Every rejection/error return below
|
||||
// must cancel it or the stream context and goroutine leak until the
|
||||
// client disconnects.
|
||||
cancelFunc()
|
||||
return centrifuge.SubscribeReply{}, nil, nil, ¢rifuge.Disconnect{
|
||||
Code: subscribeRep.Disconnect.Code,
|
||||
Reason: subscribeRep.Disconnect.Reason,
|
||||
}
|
||||
}
|
||||
if subscribeRep.Error != nil {
|
||||
//proxyCallErrorCount.WithLabelValues(proxyName, "subscribe", "error_"+strconv.FormatUint(uint64(subscribeRep.Error.Code), 10)).Inc()
|
||||
cancelFunc()
|
||||
return centrifuge.SubscribeReply{}, nil, nil, ¢rifuge.Error{
|
||||
Code: subscribeRep.Error.Code,
|
||||
Message: subscribeRep.Error.Message,
|
||||
}
|
||||
}
|
||||
|
||||
presence := chOpts.Presence
|
||||
joinLeave := chOpts.JoinLeave
|
||||
pushJoinLeave := chOpts.ForcePushJoinLeave
|
||||
|
||||
var info []byte
|
||||
var data []byte
|
||||
var expireAt int64
|
||||
|
||||
if subscribeRep.Result != nil {
|
||||
if subscribeRep.Result.B64Info != "" {
|
||||
decodedInfo, err := base64.StdEncoding.DecodeString(subscribeRep.Result.B64Info)
|
||||
if err != nil {
|
||||
log.Error().Err(err).Str("client", client.ID()).Msg("error decoding base64 info")
|
||||
cancelFunc()
|
||||
return centrifuge.SubscribeReply{}, nil, nil, centrifuge.ErrorInternal
|
||||
}
|
||||
info = decodedInfo
|
||||
} else {
|
||||
info = subscribeRep.Result.Info
|
||||
}
|
||||
if subscribeRep.Result.B64Data != "" {
|
||||
decodedData, err := base64.StdEncoding.DecodeString(subscribeRep.Result.B64Data)
|
||||
if err != nil {
|
||||
log.Error().Err(err).Str("client", client.ID()).Msg("error decoding base64 data")
|
||||
cancelFunc()
|
||||
return centrifuge.SubscribeReply{}, nil, nil, centrifuge.ErrorInternal
|
||||
}
|
||||
data = decodedData
|
||||
} else {
|
||||
data = subscribeRep.Result.Data
|
||||
}
|
||||
|
||||
result := subscribeRep.Result
|
||||
|
||||
if result.Override != nil && result.Override.Presence != nil {
|
||||
presence = result.Override.Presence.Value
|
||||
}
|
||||
if result.Override != nil && result.Override.JoinLeave != nil {
|
||||
joinLeave = result.Override.JoinLeave.Value
|
||||
}
|
||||
if result.Override != nil && result.Override.ForcePushJoinLeave != nil {
|
||||
pushJoinLeave = result.Override.ForcePushJoinLeave.Value
|
||||
}
|
||||
|
||||
expireAt = result.ExpireAt
|
||||
}
|
||||
|
||||
return centrifuge.SubscribeReply{
|
||||
Options: centrifuge.SubscribeOptions{
|
||||
ExpireAt: expireAt,
|
||||
ChannelInfo: info,
|
||||
EmitPresence: presence,
|
||||
EmitJoinLeave: joinLeave,
|
||||
PushJoinLeave: pushJoinLeave,
|
||||
EnableRecovery: false, // Not used for subscribe stream proxy.
|
||||
EnablePositioning: false, // Not used for subscribe stream proxy.
|
||||
Data: data,
|
||||
Source: subsource.StreamProxy,
|
||||
HistoryMetaTTL: 0, // Not used for subscribe stream proxy.
|
||||
},
|
||||
ClientSideRefresh: true,
|
||||
SubscriptionReady: subscriptionReady,
|
||||
}, publishFunc, cancelFunc, nil
|
||||
}
|
||||
}
|
||||
|
||||
type OnPublication func(pub *proxyproto.Publication, err error)
|
||||
|
||||
type ChannelStreamReader interface {
|
||||
Recv() (*proxyproto.StreamSubscribeResponse, error)
|
||||
}
|
||||
|
||||
// SubscribeStream ...
|
||||
func (p *SubscribeStreamProxy) SubscribeStream(
|
||||
ctx context.Context,
|
||||
bidi bool,
|
||||
sr *proxyproto.SubscribeRequest,
|
||||
pubFunc OnPublication,
|
||||
) (*proxyproto.SubscribeResponse, StreamPublishFunc, func(), error) {
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
|
||||
var stream ChannelStreamReader
|
||||
|
||||
var publishFunc StreamPublishFunc
|
||||
|
||||
if bidi {
|
||||
bidiStream, err := p.SubscribeBidirectional(ctx)
|
||||
if err != nil {
|
||||
cancel()
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
err = bidiStream.Send(&proxyproto.StreamSubscribeRequest{
|
||||
SubscribeRequest: sr,
|
||||
})
|
||||
if err != nil {
|
||||
cancel()
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
stream = bidiStream
|
||||
publishFunc = func(data []byte) error {
|
||||
return bidiStream.Send(&proxyproto.StreamSubscribeRequest{
|
||||
Publication: &proxyproto.Publication{
|
||||
Data: data,
|
||||
},
|
||||
})
|
||||
}
|
||||
} else {
|
||||
var err error
|
||||
stream, err = p.SubscribeUnidirectional(ctx, sr)
|
||||
if err != nil {
|
||||
cancel()
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
}
|
||||
|
||||
firstMessageReceived := make(chan struct{})
|
||||
|
||||
go func() {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
cancel()
|
||||
return
|
||||
case <-time.After(p.config.Timeout.ToDuration()):
|
||||
cancel()
|
||||
return
|
||||
case <-firstMessageReceived:
|
||||
}
|
||||
}()
|
||||
|
||||
resp, err := stream.Recv()
|
||||
if err != nil {
|
||||
cancel()
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
close(firstMessageReceived)
|
||||
|
||||
go func() {
|
||||
for {
|
||||
pubResp, err := stream.Recv()
|
||||
if err != nil {
|
||||
cancel()
|
||||
pubFunc(nil, err)
|
||||
return
|
||||
}
|
||||
pub := pubResp.GetPublication()
|
||||
if pub != nil {
|
||||
// TODO: better handling of unexpected nil publication.
|
||||
pubFunc(pub, nil)
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
if resp.SubscribeResponse == nil {
|
||||
// The reader goroutine was already spawned above, so cancel before
|
||||
// returning an error or it (and the derived context) leaks.
|
||||
cancel()
|
||||
return nil, nil, nil, errors.New("no subscribe response in first message from stream proxy")
|
||||
}
|
||||
return resp.SubscribeResponse, publishFunc, cancel, nil
|
||||
}
|
||||
Reference in New Issue
Block a user