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

internal/consuming/azure_service_bus.go

@@ -0,0 +1,419 @@
+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"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
+	"github.com/Azure/azure-sdk-for-go/sdk/messaging/azservicebus"
+)
+
+type AzureServiceBusConsumerConfig = configtypes.AzureServiceBusConsumerConfig
+
+// AzureServiceBusConsumer consumes messages from Azure Service Bus.
+type AzureServiceBusConsumer struct {
+	config     AzureServiceBusConsumerConfig
+	dispatcher Dispatcher
+	client     *azservicebus.Client
+	common     *consumerCommon
+}
+
+// NewAzureServiceBusConsumer creates and initializes a new AzureServiceBusConsumer.
+func NewAzureServiceBusConsumer(
+	cfg AzureServiceBusConsumerConfig,
+	dispatcher Dispatcher,
+	common *consumerCommon,
+) (*AzureServiceBusConsumer, error) {
+	if err := cfg.Validate(); err != nil {
+		return nil, err
+	}
+
+	var client *azservicebus.Client
+	var err error
+	if cfg.UseAzureIdentity {
+		cred, err := azidentity.NewClientSecretCredential(cfg.TenantID, cfg.ClientID, cfg.ClientSecret, nil)
+		if err != nil {
+			return nil, fmt.Errorf("failed to obtain Azure credential: %w", err)
+		}
+		client, err = azservicebus.NewClient(cfg.FullyQualifiedNamespace, cred, nil)
+		if err != nil {
+			return nil, fmt.Errorf("failed to create Azure Service Bus client: %w", err)
+		}
+	} else {
+		client, err = azservicebus.NewClientFromConnectionString(cfg.ConnectionString, nil)
+		if err != nil {
+			return nil, fmt.Errorf("failed to create client from connection string: %w", err)
+		}
+	}
+
+	return &AzureServiceBusConsumer{
+		config:     cfg,
+		dispatcher: dispatcher,
+		client:     client,
+		common:     common,
+	}, nil
+}
+
+// Run starts the consumer until the context is canceled.
+func (c *AzureServiceBusConsumer) Run(ctx context.Context) error {
+	if c.config.UseSessions {
+		return c.runSessionMode(ctx)
+	}
+	return c.runNonSessionMode(ctx)
+}
+
+// runSessionMode consumes messages in session mode using a separate goroutine per queue worker.
+func (c *AzureServiceBusConsumer) runSessionMode(ctx context.Context) error {
+	var wg sync.WaitGroup
+
+	// For each queue, spawn workers.
+	for _, queue := range c.config.Queues {
+		for i := 0; i < c.config.MaxConcurrentCalls; i++ {
+			wg.Add(1)
+			go func(queueName string) {
+				defer wg.Done()
+				for {
+					// Respect context cancellation.
+					if ctx.Err() != nil {
+						return
+					}
+
+					sr, err := c.client.AcceptNextSessionForQueue(ctx, queueName, nil)
+					if err != nil {
+						if errors.Is(err, context.Canceled) {
+							return
+						}
+						c.common.log.Error().Err(err).Msgf("failed to accept session for queue %s", queueName)
+						select {
+						case <-ctx.Done():
+							return
+						case <-time.After(time.Second):
+						}
+						continue
+					}
+
+					// Process the accepted session.
+					c.processSession(ctx, sr, queueName)
+				}
+			}(queue)
+		}
+	}
+
+	wg.Wait()
+	return ctx.Err()
+}
+
+// processSession processes messages for a given session and renews the session lock periodically.
+// The renewal interval is set to 2/3 of the remaining time until the session lock expires.
+// The context passed to message processing is bound to the session lifetime.
+func (c *AzureServiceBusConsumer) processSession(ctx context.Context, sr *azservicebus.SessionReceiver, queueName string) {
+	// Create a session-specific context that will be canceled when the session should end.
+	sessionCtx, cancelSession := context.WithCancel(ctx)
+	defer cancelSession()
+
+	// Derive a context for lock renewal from the session context.
+	renewCtx, cancelRenew := context.WithCancel(sessionCtx)
+	defer cancelRenew()
+
+	// Start a goroutine to renew the session lock at 2/3 of the remaining time until expiration.
+	go func() {
+		for {
+			select {
+			case <-renewCtx.Done():
+				return
+			default:
+			}
+			// Calculate remaining time until lock expiration.
+			lockUntil := sr.LockedUntil() // sr.LockedUntil() returns the time the session is locked until.
+			remaining := time.Until(lockUntil)
+			renewalInterval := time.Duration(float64(remaining) * (2.0 / 3.0))
+			if renewalInterval <= 0 {
+				c.common.log.Error().Msgf("session lock already expired for queue %s", queueName)
+				cancelSession()
+				return
+			}
+			// Wait for the computed renewal interval.
+			select {
+			case <-time.After(renewalInterval):
+				// Attempt to renew the session lock.
+				if err := sr.RenewSessionLock(renewCtx, nil); err != nil {
+					c.common.log.Error().Err(err).Msgf("failed to renew session lock for queue %s", queueName)
+					// On error, cancel the session to abort processing.
+					cancelSession()
+					return
+				} else {
+					c.common.log.Debug().Msgf("renewed session lock for queue %s", queueName)
+				}
+			case <-renewCtx.Done():
+				return
+			}
+		}
+	}()
+
+	// Ensure the session is closed when processing ends.
+	defer func() {
+		if err := sr.Close(sessionCtx); err != nil {
+			c.common.log.Error().Err(err).Msgf("failed to close session for queue %s", queueName)
+		}
+	}()
+
+	// Process messages using the session-bound context.
+	for {
+		if sessionCtx.Err() != nil {
+			return
+		}
+		messages, err := sr.ReceiveMessages(sessionCtx, c.config.MaxReceiveMessages, nil)
+		if err != nil {
+			if errors.Is(err, context.Canceled) {
+				return
+			}
+			c.common.log.Error().Err(err).Msgf("error receiving messages for session on queue %s", queueName)
+			return
+		}
+
+		if logging.Enabled(logging.DebugLevel) {
+			c.common.log.Debug().Str("queue", queueName).Int("num_messages", len(messages)).
+				Msg("received messages from queue")
+		}
+
+		if len(messages) == 0 {
+			// No messages – the session is ended.
+			return
+		}
+		for _, msg := range messages {
+			c.processMessage(sessionCtx, msg, sr)
+		}
+	}
+}
+
+// runNonSessionMode consumes messages without sessions.
+func (c *AzureServiceBusConsumer) runNonSessionMode(ctx context.Context) error {
+	var wg sync.WaitGroup
+
+	// For each queue, spawn receivers.
+	for _, queue := range c.config.Queues {
+		for i := 0; i < c.config.MaxConcurrentCalls; i++ {
+			receiver, err := c.client.NewReceiverForQueue(queue, nil)
+			if err != nil {
+				return fmt.Errorf("failed to create receiver for queue %s: %w", queue, err)
+			}
+			wg.Add(1)
+			go func(r *azservicebus.Receiver, queueName string) {
+				defer wg.Done()
+				defer func() {
+					if err := r.Close(ctx); err != nil {
+						c.common.log.Error().Err(err).Msgf("failed to close receiver for queue %s", queueName)
+					}
+				}()
+				for {
+					if ctx.Err() != nil {
+						return
+					}
+					messages, err := r.ReceiveMessages(ctx, c.config.MaxReceiveMessages, nil)
+					if err != nil {
+						if errors.Is(err, context.Canceled) {
+							return
+						}
+						c.common.log.Error().Err(err).Msgf("error receiving messages from queue %s", queueName)
+						select {
+						case <-ctx.Done():
+							return
+						case <-time.After(time.Second):
+						}
+						continue
+					}
+
+					if logging.Enabled(logging.DebugLevel) {
+						c.common.log.Debug().Str("queue", queueName).Int("num_messages", len(messages)).
+							Msg("received messages from queue")
+					}
+
+					for _, msg := range messages {
+						c.processMessage(ctx, msg, r)
+					}
+				}
+			}(receiver, queue)
+		}
+	}
+
+	wg.Wait()
+	return ctx.Err()
+}
+
+// azureCompleter defines an interface for completing messages.
+type azureCompleter interface {
+	CompleteMessage(context.Context, *azservicebus.ReceivedMessage, *azservicebus.CompleteMessageOptions) error
+	AbandonMessage(context.Context, *azservicebus.ReceivedMessage, *azservicebus.AbandonMessageOptions) error
+}
+
+// processMessage processes a single message with a retry mechanism.
+func (c *AzureServiceBusConsumer) processMessage(ctx context.Context, msg *azservicebus.ReceivedMessage, completer azureCompleter) {
+	var retries int
+	var backoffDuration time.Duration
+	data := msg.Body
+	const maxRetries = 5
+
+	var processErr error
+	for {
+		if c.config.PublicationDataMode.Enabled {
+			processErr = c.processPublicationDataMessage(ctx, msg, data)
+		} else {
+			processErr = c.processCommandMessage(ctx, msg, data)
+		}
+		if processErr == nil {
+			if retries > 0 {
+				c.common.log.Info().Msg("message processed successfully after retries")
+			}
+			break
+		}
+		// Stop retrying if context is canceled or maximum attempts reached.
+		if ctx.Err() != nil || retries >= maxRetries {
+			c.common.log.Error().Err(processErr).Msg("max retries reached or context canceled; abandoning message")
+			break
+		}
+		retries++
+		backoffDuration = getNextBackoffDuration(backoffDuration, retries)
+		c.common.log.Error().Err(processErr).Msgf("error processing message, retrying in %v", backoffDuration)
+		select {
+		case <-time.After(backoffDuration):
+		case <-ctx.Done():
+			return
+		}
+	}
+
+	if processErr != nil {
+		// Processing failed after retries: abandon the message so the broker
+		// redelivers it (and eventually dead-letters at MaxDeliveryCount) instead
+		// of completing it, which settles it as successful and permanently drops it.
+		if err := completer.AbandonMessage(ctx, msg, nil); err != nil {
+			c.common.log.Error().Err(err).Msg("failed to abandon message")
+			metrics.ConsumerErrorsTotal.WithLabelValues(c.common.name).Inc()
+		}
+		return
+	}
+
+	// Complete the message (or log an error on failure).
+	if err := completer.CompleteMessage(ctx, msg, nil); err != nil {
+		c.common.log.Error().Err(err).Msg("failed to complete message")
+		metrics.ConsumerErrorsTotal.WithLabelValues(c.common.name).Inc()
+	} else {
+		metrics.ConsumerProcessedTotal.WithLabelValues(c.common.name).Inc()
+	}
+}
+
+// processPublicationDataMessage handles messages in publication data mode.
+func (c *AzureServiceBusConsumer) processPublicationDataMessage(ctx context.Context, msg *azservicebus.ReceivedMessage, data []byte) error {
+	idempotencyKey, _ := getProperty(msg, c.config.PublicationDataMode.IdempotencyKeyProperty)
+	delta, err := getAzureBoolProperty(msg, c.config.PublicationDataMode.DeltaProperty)
+	if err != nil {
+		c.common.log.Error().Err(err).Msg("error parsing delta property, skipping message")
+		return nil
+	}
+	channelsStr, _ := getProperty(msg, c.config.PublicationDataMode.ChannelsProperty)
+	channels := strings.Split(channelsStr, ",")
+	if len(channels) == 0 || (len(channels) == 1 && channels[0] == "") {
+		c.common.log.Info().Msg("no channels found, skipping message")
+		return nil
+	}
+	tags := getTagsFromProperties(msg, c.config.PublicationDataMode.TagsPropertyPrefix)
+	version, err := getAzureUint64Property(msg, c.config.PublicationDataMode.VersionProperty)
+	if err != nil {
+		c.common.log.Error().Err(err).Msg("error parsing version property, skipping message")
+		return nil
+	}
+	versionEpoch, _ := getProperty(msg, c.config.PublicationDataMode.VersionEpochProperty)
+	return c.dispatcher.DispatchPublication(ctx, channels, api.ConsumedPublication{
+		Data:           data,
+		IdempotencyKey: idempotencyKey,
+		Delta:          delta,
+		Tags:           tags,
+		Version:        version,
+		VersionEpoch:   versionEpoch,
+	})
+}
+
+// processCommandMessage handles messages in command mode.
+func (c *AzureServiceBusConsumer) processCommandMessage(ctx context.Context, msg *azservicebus.ReceivedMessage, data []byte) error {
+	method, _ := getProperty(msg, c.config.MethodProperty)
+	return c.dispatcher.DispatchCommand(ctx, method, data)
+}
+
+func getAzureUint64Property(msg *azservicebus.ReceivedMessage, key string) (uint64, error) {
+	if key == "" {
+		return 0, nil
+	}
+	val, ok := msg.ApplicationProperties[key]
+	if !ok {
+		return 0, nil
+	}
+	strVal, ok := val.(string)
+	if !ok {
+		return 0, fmt.Errorf("property %q is not a string", key)
+	}
+	uintVal, err := strconv.ParseUint(strVal, 10, 64)
+	if err != nil {
+		return 0, fmt.Errorf("error parsing uint64 property %q: %w", key, err)
+	}
+	return uintVal, nil
+}
+
+func getAzureBoolProperty(msg *azservicebus.ReceivedMessage, key string) (bool, error) {
+	if key == "" {
+		return false, nil
+	}
+	val, ok := msg.ApplicationProperties[key]
+	if !ok {
+		return false, nil
+	}
+	strVal, ok := val.(string)
+	if !ok {
+		return false, fmt.Errorf("property %q is not a string", key)
+	}
+	b, err := strconv.ParseBool(strVal)
+	if err != nil {
+		return false, fmt.Errorf("error parsing bool property %q: %w", key, err)
+	}
+	return b, nil
+}
+
+// getProperty extracts a string property from the message's ApplicationProperties.
+func getProperty(msg *azservicebus.ReceivedMessage, key string) (string, bool) {
+	if key == "" {
+		return "", false
+	}
+	if val, ok := msg.ApplicationProperties[key]; ok {
+		if str, ok := val.(string); ok {
+			return str, true
+		}
+	}
+	return "", false
+}
+
+// getTagsFromProperties extracts tag values from message properties using the given prefix.
+func getTagsFromProperties(msg *azservicebus.ReceivedMessage, prefix string) map[string]string {
+	var tags map[string]string
+	if prefix == "" {
+		return tags
+	}
+	for k, v := range msg.ApplicationProperties {
+		if strings.HasPrefix(k, prefix) {
+			if str, ok := v.(string); ok {
+				if tags == nil {
+					tags = make(map[string]string)
+				}
+				tags[strings.TrimPrefix(k, prefix)] = str
+			}
+		}
+	}
+	return tags
+}