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

internal/consuming/google_pub_sub.go

@@ -0,0 +1,240 @@
+package consuming
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/centrifugal/centrifugo/v6/internal/api"
+	"github.com/centrifugal/centrifugo/v6/internal/configtypes"
+	"github.com/centrifugal/centrifugo/v6/internal/logging"
+	"github.com/centrifugal/centrifugo/v6/internal/metrics"
+
+	"cloud.google.com/go/pubsub/v2"
+	"google.golang.org/api/option"
+)
+
+type GooglePubSubConsumerConfig = configtypes.GooglePubSubConsumerConfig
+
+// pubSubReceiveBackoffReset is the minimum duration a Receive call must run
+// before a subsequent failure is treated as fresh, resetting the retry backoff.
+const pubSubReceiveBackoffReset = 30 * time.Second
+
+// GooglePubSubConsumer represents a Google Pub/Sub consumer.
+type GooglePubSubConsumer struct {
+	config     GooglePubSubConsumerConfig
+	dispatcher Dispatcher
+	client     *pubsub.Client
+	common     *consumerCommon
+}
+
+// NewGooglePubSubConsumer creates a new Google Pub/Sub consumer with the provided auth mechanism.
+func NewGooglePubSubConsumer(
+	config GooglePubSubConsumerConfig, dispatcher Dispatcher, common *consumerCommon,
+) (*GooglePubSubConsumer, error) {
+	if err := config.Validate(); err != nil {
+		return nil, err
+	}
+
+	var clientOpts []option.ClientOption
+	switch strings.ToLower(config.AuthMechanism) {
+	case "", "default":
+		// Use Application Default Credentials.
+	case "service_account":
+		if config.CredentialsFile == "" {
+			return nil, errors.New("credentials_file must be provided for service_account auth")
+		}
+		clientOpts = append(clientOpts, option.WithAuthCredentialsFile(option.ServiceAccount, config.CredentialsFile))
+	default:
+		return nil, fmt.Errorf("unsupported auth mechanism: %s", config.AuthMechanism)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	client, err := pubsub.NewClient(ctx, config.ProjectID, clientOpts...)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create pubsub client: %w", err)
+	}
+
+	consumer := &GooglePubSubConsumer{
+		config:     config,
+		dispatcher: dispatcher,
+		client:     client,
+		common:     common,
+	}
+	return consumer, nil
+}
+
+// Run starts the consumer and blocks until the provided context is canceled or an error occurs.
+func (c *GooglePubSubConsumer) Run(ctx context.Context) error {
+	defer func() {
+		_ = c.client.Close()
+	}()
+
+	// For each subscription in the configuration, spawn a separate goroutine.
+	var wg sync.WaitGroup
+	for _, subID := range c.config.Subscriptions {
+		wg.Add(1)
+		go func(subscriptionID string) {
+			defer wg.Done()
+
+			sub := c.client.Subscriber(subscriptionID)
+			sub.ReceiveSettings.MaxOutstandingMessages = c.config.MaxOutstandingMessages
+			sub.ReceiveSettings.MaxOutstandingBytes = c.config.MaxOutstandingBytes
+			sub.ReceiveSettings.NumGoroutines = 10
+
+			// Receive re-establishes the streaming pull and returns only on ctx
+			// cancellation or an error the SDK deems non-retryable. Such an error is
+			// not necessarily permanent (IAM propagation, quota reset, backend
+			// failover, subscription recreated), and without re-subscribing the
+			// subscription would stay silently dead for the whole process lifetime.
+			// Re-Receive with a capped backoff so a genuinely permanent error can't
+			// spin hot.
+			var backoff time.Duration
+			var retries int
+			for ctx.Err() == nil {
+				start := time.Now()
+				err := sub.Receive(ctx, func(ctx context.Context, msg *pubsub.Message) {
+					if logging.Enabled(logging.DebugLevel) {
+						c.common.log.Debug().Str("subscription", subID).
+							Msg("received message from subscription")
+					}
+					c.dispatchMessage(ctx, msg)
+				})
+				if err == nil || errors.Is(err, context.Canceled) || ctx.Err() != nil {
+					return
+				}
+				// A Receive that ran for a while before failing is a fresh failure,
+				// not an escalating one — reset the backoff ramp.
+				if time.Since(start) > pubSubReceiveBackoffReset {
+					backoff, retries = 0, 0
+				}
+				retries++
+				backoff = getNextBackoffDuration(backoff, retries)
+				metrics.ConsumerErrorsTotal.WithLabelValues(c.common.name).Inc()
+				c.common.log.Error().Err(err).Str("subscription", subscriptionID).
+					Dur("retry_in", backoff).Msg("error receiving messages, will re-subscribe")
+				select {
+				case <-time.After(backoff):
+				case <-ctx.Done():
+					return
+				}
+			}
+		}(subID)
+	}
+
+	wg.Wait()
+	return ctx.Err()
+}
+
+func (c *GooglePubSubConsumer) dispatchMessage(ctx context.Context, msg *pubsub.Message) {
+	c.processSingleMessage(ctx, msg)
+}
+
+func (c *GooglePubSubConsumer) processSingleMessage(ctx context.Context, msg *pubsub.Message) {
+	var err error
+	if c.config.PublicationDataMode.Enabled {
+		err = c.processPublicationDataMessage(ctx, msg)
+	} else {
+		err = c.processCommandMessage(ctx, msg)
+	}
+	if err == nil {
+		msg.Ack()
+		metrics.ConsumerProcessedTotal.WithLabelValues(c.common.name).Inc()
+		return
+	}
+	msg.Nack()
+	metrics.ConsumerErrorsTotal.WithLabelValues(c.common.name).Inc()
+}
+
+// processPublicationDataMessage handles messages in publication data mode.
+func (c *GooglePubSubConsumer) processPublicationDataMessage(ctx context.Context, msg *pubsub.Message) error {
+	data := msg.Data
+	idempotencyKey := getAttributeValue(msg, c.config.PublicationDataMode.IdempotencyKeyAttribute)
+	delta, err := getGoogleBoolAttributeValue(msg, c.config.PublicationDataMode.DeltaAttribute)
+	if err != nil {
+		c.common.log.Error().Err(err).Msg("error parsing delta attribute, skipping message")
+		return nil // Skip message on parsing error.
+	}
+	channelsAttr := getAttributeValue(msg, c.config.PublicationDataMode.ChannelsAttribute)
+	channels := strings.Split(channelsAttr, ",")
+	if len(channels) == 0 || (len(channels) == 1 && channels[0] == "") {
+		c.common.log.Info().Msg("no channels found, skipping message")
+		return nil
+	}
+	tags := publicationTagsFromAttributes(msg, c.config.PublicationDataMode.TagsAttributePrefix)
+	version, err := getGoogleUint64AttributeValue(msg, c.config.PublicationDataMode.VersionAttribute)
+	if err != nil {
+		c.common.log.Error().Err(err).Msg("error parsing version attribute, skipping message")
+		return nil // Skip message on parsing error.
+	}
+	return c.dispatcher.DispatchPublication(ctx, channels, api.ConsumedPublication{
+		Data:           data,
+		IdempotencyKey: idempotencyKey,
+		Delta:          delta,
+		Tags:           tags,
+		Version:        version,
+		VersionEpoch:   getAttributeValue(msg, c.config.PublicationDataMode.VersionEpochAttribute),
+	})
+}
+
+// processCommandMessage handles non-publication messages.
+func (c *GooglePubSubConsumer) processCommandMessage(ctx context.Context, msg *pubsub.Message) error {
+	method := getAttributeValue(msg, c.config.MethodAttribute)
+	return c.dispatcher.DispatchCommand(ctx, method, msg.Data)
+}
+
+func getGoogleUint64AttributeValue(msg *pubsub.Message, key string) (uint64, error) {
+	if key == "" {
+		return 0, nil
+	}
+	val := msg.Attributes[key]
+	if val == "" {
+		return 0, nil
+	}
+	uintVal, err := strconv.ParseUint(val, 10, 64)
+	if err != nil {
+		return 0, fmt.Errorf("error parsing attribute %s: %w", key, err)
+	}
+	return uintVal, nil
+}
+
+func getGoogleBoolAttributeValue(msg *pubsub.Message, key string) (bool, error) {
+	if key == "" {
+		return false, nil
+	}
+	val := msg.Attributes[key]
+	if val == "" {
+		return false, nil
+	}
+	return strconv.ParseBool(val)
+}
+
+// getAttributeValue extracts the value of the given attribute key from the message.
+func getAttributeValue(msg *pubsub.Message, key string) string {
+	if key == "" {
+		return ""
+	}
+	return msg.Attributes[key]
+}
+
+// publicationTagsFromAttributes extracts tag values from message attributes using the given prefix.
+func publicationTagsFromAttributes(msg *pubsub.Message, prefix string) map[string]string {
+	if prefix == "" {
+		return nil
+	}
+	var tags map[string]string
+	for k, v := range msg.Attributes {
+		if strings.HasPrefix(k, prefix) {
+			if tags == nil {
+				tags = make(map[string]string)
+			}
+			tags[strings.TrimPrefix(k, prefix)] = v
+		}
+	}
+	return tags
+}