315 lines
9.0 KiB
Go
315 lines
9.0 KiB
Go
package api
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
|
||
"github.com/centrifugal/centrifugo/v6/internal/apiproto"
|
||
|
||
"github.com/centrifugal/centrifuge"
|
||
"github.com/rs/zerolog/log"
|
||
"go.opentelemetry.io/otel/codes"
|
||
"go.opentelemetry.io/otel/trace"
|
||
)
|
||
|
||
// ConsumingHandlerConfig configures ConsumingHandler.
|
||
type ConsumingHandlerConfig struct {
|
||
UseOpenTelemetry bool
|
||
}
|
||
|
||
// ConsumingHandler is responsible for processing API commands from asynchronous consumers.
|
||
type ConsumingHandler struct {
|
||
node *centrifuge.Node
|
||
config ConsumingHandlerConfig
|
||
api *Executor
|
||
}
|
||
|
||
// NewConsumingHandler creates new ConsumingHandler.
|
||
func NewConsumingHandler(n *centrifuge.Node, apiExecutor *Executor, c ConsumingHandlerConfig) *ConsumingHandler {
|
||
h := &ConsumingHandler{
|
||
node: n,
|
||
config: c,
|
||
api: apiExecutor,
|
||
}
|
||
return h
|
||
}
|
||
|
||
func logNonRetryableConsumingError(err error, method string) {
|
||
log.Error().Err(err).Str("method", method).Msg("non retryable error during consuming, skip message")
|
||
}
|
||
|
||
// ErrShutdown is returned to consumers for messages received after Node shutdown
|
||
// started. Such messages must not be acknowledged: this node is no longer able to
|
||
// deliver them, and depending on the broker a publish may still return success (a
|
||
// memory broker never fails, so without this check a message would be committed
|
||
// while going nowhere). Returning an error makes consumers leave the message
|
||
// unacknowledged, so it's redelivered to another node or after restart.
|
||
//
|
||
// It wraps context.Canceled since consumers treat cancellation as a stop signal
|
||
// and give up immediately instead of retrying with backoff until the shutdown
|
||
// timeout is reached.
|
||
var ErrShutdown = fmt.Errorf("node is shutting down: %w", context.Canceled)
|
||
|
||
type Dispatcher struct {
|
||
handler *ConsumingHandler
|
||
}
|
||
|
||
func NewDispatcher(handler *ConsumingHandler) *Dispatcher {
|
||
return &Dispatcher{
|
||
handler: handler,
|
||
}
|
||
}
|
||
|
||
// shuttingDown returns true once Node shutdown started.
|
||
func (d *Dispatcher) shuttingDown() bool {
|
||
select {
|
||
case <-d.handler.node.NotifyShutdown():
|
||
return true
|
||
default:
|
||
return false
|
||
}
|
||
}
|
||
|
||
func (d *Dispatcher) Publish(ctx context.Context, req *apiproto.PublishRequest) error {
|
||
if d.shuttingDown() {
|
||
return ErrShutdown
|
||
}
|
||
resp := d.handler.api.Publish(ctx, req)
|
||
if d.handler.config.UseOpenTelemetry && resp.Error != nil {
|
||
span := trace.SpanFromContext(ctx)
|
||
span.SetStatus(codes.Error, resp.Error.Error())
|
||
}
|
||
if resp.Error != nil && resp.Error.Code == apiproto.ErrorInternal.Code {
|
||
return resp.Error
|
||
}
|
||
if resp.Error != nil {
|
||
logNonRetryableConsumingError(resp.Error, "publish")
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (d *Dispatcher) Broadcast(ctx context.Context, req *apiproto.BroadcastRequest) error {
|
||
if d.shuttingDown() {
|
||
return ErrShutdown
|
||
}
|
||
resp := d.handler.api.Broadcast(ctx, req)
|
||
if d.handler.config.UseOpenTelemetry && resp.Error != nil {
|
||
span := trace.SpanFromContext(ctx)
|
||
span.SetStatus(codes.Error, resp.Error.Error())
|
||
}
|
||
if resp.Error != nil && resp.Error.Code == apiproto.ErrorInternal.Code {
|
||
return resp.Error
|
||
}
|
||
if resp.Error != nil {
|
||
logNonRetryableConsumingError(resp.Error, "broadcast")
|
||
return nil
|
||
}
|
||
for _, response := range resp.Result.Responses {
|
||
if response.Error != nil && response.Error.Code == apiproto.ErrorInternal.Code {
|
||
// Any internal error in any channel response will result into a retry by a consumer.
|
||
// To prevent duplicate messages publishers may use idempotency keys.
|
||
return response.Error
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
type ConsumedPublication struct {
|
||
Data []byte
|
||
// IdempotencyKey is used to prevent duplicate messages.
|
||
IdempotencyKey string
|
||
// Delta is used to indicate that the message is a delta update.
|
||
Delta bool
|
||
// Tags are used to attach metadata to the message.
|
||
Tags map[string]string
|
||
// Version of publication.
|
||
Version uint64
|
||
// VersionEpoch of publication.
|
||
VersionEpoch string
|
||
}
|
||
|
||
func (d *Dispatcher) DispatchPublication(
|
||
ctx context.Context, channels []string, pub ConsumedPublication,
|
||
) error {
|
||
if d.shuttingDown() {
|
||
return ErrShutdown
|
||
}
|
||
if len(channels) == 0 {
|
||
return nil
|
||
}
|
||
if len(channels) == 1 {
|
||
req := &apiproto.PublishRequest{
|
||
Data: pub.Data,
|
||
Channel: channels[0],
|
||
IdempotencyKey: pub.IdempotencyKey,
|
||
Delta: pub.Delta,
|
||
Tags: pub.Tags,
|
||
Version: pub.Version,
|
||
VersionEpoch: pub.VersionEpoch,
|
||
}
|
||
return d.Publish(ctx, req)
|
||
}
|
||
req := &apiproto.BroadcastRequest{
|
||
Data: pub.Data,
|
||
Channels: channels,
|
||
IdempotencyKey: pub.IdempotencyKey,
|
||
Delta: pub.Delta,
|
||
Tags: pub.Tags,
|
||
Version: pub.Version,
|
||
VersionEpoch: pub.VersionEpoch,
|
||
}
|
||
return d.Broadcast(ctx, req)
|
||
}
|
||
|
||
// Dispatch processes commands received from asynchronous consumers.
|
||
func (d *Dispatcher) dispatchMethodPayload(ctx context.Context, method string, payload []byte) error {
|
||
switch method {
|
||
case "publish":
|
||
_, err := d.handler.handlePublish(ctx, payload)
|
||
if err != nil {
|
||
var apiError *apiproto.Error
|
||
if errors.As(err, &apiError) && apiError.Code == apiproto.ErrorInternal.Code {
|
||
return err
|
||
}
|
||
logNonRetryableConsumingError(err, method)
|
||
return nil
|
||
}
|
||
return nil
|
||
case "broadcast":
|
||
// This one is special as we need to iterate over responses. Ideally we want to use
|
||
// code gen for Dispatch but need to make sure that broadcast logic is preserved.
|
||
res, err := d.handler.handleBroadcast(ctx, payload)
|
||
if err != nil {
|
||
var apiError *apiproto.Error
|
||
if errors.As(err, &apiError) && apiError.Code == apiproto.ErrorInternal.Code {
|
||
return err
|
||
}
|
||
logNonRetryableConsumingError(err, method)
|
||
return nil
|
||
}
|
||
for _, resp := range res.Responses {
|
||
if resp.Error != nil && resp.Error.Code == apiproto.ErrorInternal.Code {
|
||
// Any internal error in any channel response will result into a retry by a consumer.
|
||
// To prevent duplicate messages publishers may use idempotency keys.
|
||
return resp.Error
|
||
}
|
||
}
|
||
return nil
|
||
case "subscribe":
|
||
_, err := d.handler.handleSubscribe(ctx, payload)
|
||
if err != nil {
|
||
var apiError *apiproto.Error
|
||
if errors.As(err, &apiError) && apiError.Code == apiproto.ErrorInternal.Code {
|
||
return err
|
||
}
|
||
logNonRetryableConsumingError(err, method)
|
||
return nil
|
||
}
|
||
return nil
|
||
case "unsubscribe":
|
||
_, err := d.handler.handleUnsubscribe(ctx, payload)
|
||
if err != nil {
|
||
var apiError *apiproto.Error
|
||
if errors.As(err, &apiError) && apiError.Code == apiproto.ErrorInternal.Code {
|
||
return err
|
||
}
|
||
logNonRetryableConsumingError(err, method)
|
||
return nil
|
||
}
|
||
return nil
|
||
case "disconnect":
|
||
_, err := d.handler.handleDisconnect(ctx, payload)
|
||
if err != nil {
|
||
var apiError *apiproto.Error
|
||
if errors.As(err, &apiError) && apiError.Code == apiproto.ErrorInternal.Code {
|
||
return err
|
||
}
|
||
logNonRetryableConsumingError(err, method)
|
||
return nil
|
||
}
|
||
return nil
|
||
case "history_remove":
|
||
_, err := d.handler.handleHistoryRemove(ctx, payload)
|
||
if err != nil {
|
||
var apiError *apiproto.Error
|
||
if errors.As(err, &apiError) && apiError.Code == apiproto.ErrorInternal.Code {
|
||
return err
|
||
}
|
||
logNonRetryableConsumingError(err, method)
|
||
return nil
|
||
}
|
||
return nil
|
||
case "refresh":
|
||
_, err := d.handler.handleRefresh(ctx, payload)
|
||
if err != nil {
|
||
var apiError *apiproto.Error
|
||
if errors.As(err, &apiError) && apiError.Code == apiproto.ErrorInternal.Code {
|
||
return err
|
||
}
|
||
logNonRetryableConsumingError(err, method)
|
||
return nil
|
||
}
|
||
return nil
|
||
default:
|
||
// Skip unsupported.
|
||
log.Info().Msg("skip unsupported API command method")
|
||
return nil
|
||
}
|
||
}
|
||
|
||
var ErrInvalidData = errors.New("invalid data")
|
||
|
||
// JSONRawOrString can decode payload from bytes and from JSON string. This gives
|
||
// us better interoperability. For example, JSONB field is encoded as JSON string in
|
||
// Debezium PostgreSQL connector.
|
||
type JSONRawOrString json.RawMessage
|
||
|
||
func (j *JSONRawOrString) UnmarshalJSON(data []byte) error {
|
||
if len(data) > 0 && data[0] == '"' {
|
||
// Unmarshal as a string, then convert the string to json.RawMessage.
|
||
var str string
|
||
if err := json.Unmarshal(data, &str); err != nil {
|
||
return err
|
||
}
|
||
*j = JSONRawOrString(str)
|
||
} else {
|
||
// Unmarshal directly as json.RawMessage
|
||
*j = data
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// MarshalJSON returns m as the JSON encoding of m.
|
||
func (j JSONRawOrString) MarshalJSON() ([]byte, error) {
|
||
if j == nil {
|
||
return []byte("null"), nil
|
||
}
|
||
return j, nil
|
||
}
|
||
|
||
type MethodWithRequestPayload struct {
|
||
Method string `json:"method"`
|
||
Payload JSONRawOrString `json:"payload"`
|
||
}
|
||
|
||
// DispatchCommand processes commands received from asynchronous consumers.
|
||
func (d *Dispatcher) DispatchCommand(ctx context.Context, method string, payload []byte) error {
|
||
if d.shuttingDown() {
|
||
return ErrShutdown
|
||
}
|
||
if method != "" {
|
||
// If method already known – we can skip decoding into MethodWithRequestPayload.
|
||
return d.dispatchMethodPayload(ctx, method, payload)
|
||
}
|
||
// Otherwise we expect payload to be MethodWithRequestPayload.
|
||
var e MethodWithRequestPayload
|
||
err := json.Unmarshal(payload, &e)
|
||
if err != nil {
|
||
log.Error().Err(err).Msg("skip malformed consumed message")
|
||
return nil
|
||
}
|
||
return d.dispatchMethodPayload(ctx, e.Method, e.Payload)
|
||
}
|