355 lines
10 KiB
Go
355 lines
10 KiB
Go
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
|
|
}
|