mirror of
https://github.com/oarkflow/mq.git
synced 2025-10-05 16:06:55 +08:00
576 lines
16 KiB
Go
576 lines
16 KiB
Go
package mq
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/oarkflow/mq/logger"
|
|
)
|
|
|
|
// BrokerEnhancedConfig holds configuration for all enhanced features
|
|
type BrokerEnhancedConfig struct {
|
|
// DLQ Configuration
|
|
DLQStoragePath string
|
|
DLQRetentionPeriod time.Duration
|
|
DLQMaxSize int
|
|
|
|
// WAL Configuration
|
|
WALDirectory string
|
|
WALMaxFileSize int64
|
|
WALSyncInterval time.Duration
|
|
WALFsyncOnWrite bool
|
|
|
|
// Acknowledgment Configuration
|
|
AckTimeout time.Duration
|
|
AckMaxRetries int
|
|
AckCheckInterval time.Duration
|
|
|
|
// Worker Health Configuration
|
|
WorkerHealthTimeout time.Duration
|
|
WorkerCheckInterval time.Duration
|
|
|
|
// Dynamic Scaling Configuration
|
|
MinWorkers int
|
|
MaxWorkers int
|
|
ScaleUpThreshold float64
|
|
ScaleDownThreshold float64
|
|
|
|
// Deduplication Configuration
|
|
DedupWindow time.Duration
|
|
DedupCleanupInterval time.Duration
|
|
DedupPersistent bool
|
|
|
|
// Flow Control Configuration
|
|
FlowControlStrategy FlowControlStrategyType
|
|
FlowControlConfigPath string // Path to flow control config file
|
|
FlowControlEnvPrefix string // Environment variable prefix for flow control
|
|
MaxCredits int64
|
|
MinCredits int64
|
|
CreditRefillRate int64
|
|
CreditRefillInterval time.Duration
|
|
// Token bucket specific
|
|
TokenBucketCapacity int64
|
|
TokenBucketRefillRate int64
|
|
TokenBucketRefillInterval time.Duration
|
|
// Leaky bucket specific
|
|
LeakyBucketCapacity int64
|
|
LeakyBucketLeakInterval time.Duration
|
|
// Credit-based specific
|
|
CreditBasedMaxCredits int64
|
|
CreditBasedRefillRate int64
|
|
CreditBasedRefillInterval time.Duration
|
|
CreditBasedBurstSize int64
|
|
// Rate limiter specific
|
|
RateLimiterRequestsPerSecond int64
|
|
RateLimiterBurstSize int64
|
|
|
|
// Backpressure Configuration
|
|
QueueDepthThreshold int
|
|
MemoryThreshold uint64
|
|
ErrorRateThreshold float64
|
|
|
|
// Snapshot Configuration
|
|
SnapshotDirectory string
|
|
SnapshotInterval time.Duration
|
|
SnapshotRetention time.Duration
|
|
|
|
// Tracing Configuration
|
|
TracingEnabled bool
|
|
TraceRetention time.Duration
|
|
TraceExportInterval time.Duration
|
|
|
|
Logger logger.Logger
|
|
EnableEnhancements bool // Master switch for all enhancements
|
|
}
|
|
|
|
// EnhancedFeatures holds all enhanced feature managers
|
|
type EnhancedFeatures struct {
|
|
ackManager *AckManager
|
|
walLog *WriteAheadLog
|
|
dlqStorage DLQStorage
|
|
dedupManager *DeduplicationManager
|
|
flowController *FlowController
|
|
backpressureMonitor *BackpressureMonitor
|
|
snapshotManager *SnapshotManager
|
|
traceManager *TraceManager
|
|
lifecycleTracker *MessageLifecycleTracker
|
|
config *BrokerEnhancedConfig
|
|
enabled bool
|
|
}
|
|
|
|
// InitializeEnhancements initializes all enhanced features for a broker
|
|
func (b *Broker) InitializeEnhancements(config *BrokerEnhancedConfig) error {
|
|
if config == nil {
|
|
config = DefaultBrokerEnhancedConfig()
|
|
}
|
|
|
|
if !config.EnableEnhancements {
|
|
return nil // Enhancements disabled
|
|
}
|
|
|
|
features := &EnhancedFeatures{
|
|
config: config,
|
|
enabled: true,
|
|
}
|
|
|
|
// Initialize DLQ Storage
|
|
if config.DLQStoragePath != "" {
|
|
dlqStorage, err := NewFileDLQStorage(config.DLQStoragePath, config.Logger)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
features.dlqStorage = dlqStorage
|
|
} else {
|
|
features.dlqStorage = NewInMemoryDLQStorage()
|
|
}
|
|
|
|
// Initialize WAL
|
|
if config.WALDirectory != "" {
|
|
walConfig := WALConfig{
|
|
Directory: config.WALDirectory,
|
|
MaxFileSize: config.WALMaxFileSize,
|
|
SyncInterval: config.WALSyncInterval,
|
|
FsyncOnWrite: config.WALFsyncOnWrite,
|
|
Logger: config.Logger,
|
|
}
|
|
wal, err := NewWriteAheadLog(walConfig)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
features.walLog = wal
|
|
}
|
|
|
|
// Initialize Acknowledgment Manager
|
|
ackConfig := AckManagerConfig{
|
|
AckTimeout: config.AckTimeout,
|
|
MaxRetries: config.AckMaxRetries,
|
|
CheckInterval: config.AckCheckInterval,
|
|
Logger: config.Logger,
|
|
}
|
|
features.ackManager = NewAckManager(ackConfig)
|
|
|
|
// Setup acknowledgment callbacks
|
|
features.ackManager.SetOnReject(func(pm *PendingMessage) {
|
|
// Move to DLQ
|
|
entry := &DLQEntry{
|
|
TaskID: pm.Task.ID,
|
|
QueueName: pm.QueueName,
|
|
OriginalPayload: json.RawMessage(pm.Task.Payload),
|
|
ErrorMessage: "Max retries exceeded",
|
|
FailedAt: time.Now(),
|
|
RetryCount: pm.RetryCount,
|
|
}
|
|
_ = features.dlqStorage.Store(context.Background(), entry)
|
|
})
|
|
|
|
features.ackManager.SetOnRedeliver(func(pm *PendingMessage) {
|
|
// Requeue the message
|
|
if queue, exists := b.queues.Get(pm.QueueName); exists {
|
|
select {
|
|
case queue.tasks <- &QueuedTask{Task: pm.Task}:
|
|
default:
|
|
config.Logger.Warn("Failed to requeue message - queue full",
|
|
logger.Field{Key: "taskID", Value: pm.Task.ID})
|
|
}
|
|
}
|
|
})
|
|
|
|
// Initialize Deduplication Manager
|
|
dedupConfig := DedupConfig{
|
|
Window: config.DedupWindow,
|
|
CleanupInterval: config.DedupCleanupInterval,
|
|
Logger: config.Logger,
|
|
}
|
|
features.dedupManager = NewDeduplicationManager(dedupConfig)
|
|
|
|
// Initialize Flow Controller using factory
|
|
factory := NewFlowControllerFactory()
|
|
|
|
// Try to load configuration from providers
|
|
var flowConfig FlowControlConfig
|
|
var err error
|
|
|
|
// First try file-based configuration
|
|
if config.FlowControlConfigPath != "" {
|
|
fileProvider := NewFileConfigProvider(config.FlowControlConfigPath)
|
|
if loadedConfig, loadErr := fileProvider.GetConfig(); loadErr == nil {
|
|
flowConfig = loadedConfig
|
|
}
|
|
}
|
|
|
|
// If no file config, try environment variables
|
|
if flowConfig.Strategy == "" && config.FlowControlEnvPrefix != "" {
|
|
envProvider := NewEnvConfigProvider(config.FlowControlEnvPrefix)
|
|
if loadedConfig, loadErr := envProvider.GetConfig(); loadErr == nil {
|
|
flowConfig = loadedConfig
|
|
}
|
|
}
|
|
|
|
// If still no config, use broker config defaults based on strategy
|
|
if flowConfig.Strategy == "" {
|
|
flowConfig = FlowControlConfig{
|
|
Strategy: config.FlowControlStrategy,
|
|
Logger: config.Logger,
|
|
}
|
|
|
|
// Set strategy-specific defaults
|
|
switch config.FlowControlStrategy {
|
|
case StrategyTokenBucket:
|
|
flowConfig.MaxCredits = config.TokenBucketCapacity
|
|
flowConfig.RefillRate = config.TokenBucketRefillRate
|
|
flowConfig.RefillInterval = config.TokenBucketRefillInterval
|
|
case StrategyLeakyBucket:
|
|
flowConfig.MaxCredits = config.LeakyBucketCapacity
|
|
flowConfig.RefillInterval = config.LeakyBucketLeakInterval
|
|
case StrategyCreditBased:
|
|
flowConfig.MaxCredits = config.CreditBasedMaxCredits
|
|
flowConfig.RefillRate = config.CreditBasedRefillRate
|
|
flowConfig.RefillInterval = config.CreditBasedRefillInterval
|
|
flowConfig.BurstSize = config.CreditBasedBurstSize
|
|
case StrategyRateLimiter:
|
|
flowConfig.RefillRate = config.RateLimiterRequestsPerSecond
|
|
flowConfig.BurstSize = config.RateLimiterBurstSize
|
|
default:
|
|
// Fallback to token bucket
|
|
flowConfig.Strategy = StrategyTokenBucket
|
|
flowConfig.MaxCredits = config.MaxCredits
|
|
flowConfig.RefillRate = config.CreditRefillRate
|
|
flowConfig.RefillInterval = config.CreditRefillInterval
|
|
}
|
|
}
|
|
|
|
// Ensure logger is set
|
|
flowConfig.Logger = config.Logger
|
|
|
|
// Create flow controller using factory
|
|
features.flowController, err = factory.CreateFlowController(flowConfig)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to create flow controller: %w", err)
|
|
}
|
|
|
|
// Initialize Backpressure Monitor
|
|
backpressureConfig := BackpressureConfig{
|
|
QueueDepthThreshold: config.QueueDepthThreshold,
|
|
MemoryThreshold: config.MemoryThreshold,
|
|
ErrorRateThreshold: config.ErrorRateThreshold,
|
|
Logger: config.Logger,
|
|
}
|
|
features.backpressureMonitor = NewBackpressureMonitor(backpressureConfig)
|
|
|
|
// Initialize Snapshot Manager
|
|
if config.SnapshotDirectory != "" {
|
|
snapshotConfig := SnapshotConfig{
|
|
BaseDir: config.SnapshotDirectory,
|
|
SnapshotInterval: config.SnapshotInterval,
|
|
RetentionPeriod: config.SnapshotRetention,
|
|
Logger: config.Logger,
|
|
}
|
|
snapshotManager, err := NewSnapshotManager(b, snapshotConfig)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
features.snapshotManager = snapshotManager
|
|
}
|
|
|
|
// Initialize Tracing
|
|
if config.TracingEnabled {
|
|
traceConfig := TraceConfig{
|
|
Storage: NewInMemoryTraceStorage(),
|
|
Retention: config.TraceRetention,
|
|
ExportInterval: config.TraceExportInterval,
|
|
Logger: config.Logger,
|
|
}
|
|
features.traceManager = NewTraceManager(traceConfig)
|
|
features.lifecycleTracker = NewMessageLifecycleTracker(features.traceManager, config.Logger)
|
|
}
|
|
|
|
// Store features in broker (we'll need to add this field to Broker struct)
|
|
b.enhanced = features
|
|
|
|
return nil
|
|
}
|
|
|
|
// DefaultBrokerEnhancedConfig returns default configuration
|
|
func DefaultBrokerEnhancedConfig() *BrokerEnhancedConfig {
|
|
return &BrokerEnhancedConfig{
|
|
DLQRetentionPeriod: 7 * 24 * time.Hour,
|
|
DLQMaxSize: 10000,
|
|
WALMaxFileSize: 100 * 1024 * 1024,
|
|
WALSyncInterval: 1 * time.Second,
|
|
WALFsyncOnWrite: false,
|
|
AckTimeout: 30 * time.Second,
|
|
AckMaxRetries: 3,
|
|
AckCheckInterval: 5 * time.Second,
|
|
WorkerHealthTimeout: 30 * time.Second,
|
|
WorkerCheckInterval: 10 * time.Second,
|
|
MinWorkers: 1,
|
|
MaxWorkers: 100,
|
|
ScaleUpThreshold: 0.75,
|
|
ScaleDownThreshold: 0.25,
|
|
DedupWindow: 5 * time.Minute,
|
|
DedupCleanupInterval: 1 * time.Minute,
|
|
// Flow Control defaults (Token Bucket strategy)
|
|
FlowControlStrategy: StrategyTokenBucket,
|
|
FlowControlConfigPath: "",
|
|
FlowControlEnvPrefix: "FLOW_",
|
|
MaxCredits: 1000,
|
|
MinCredits: 100,
|
|
CreditRefillRate: 10,
|
|
CreditRefillInterval: 100 * time.Millisecond,
|
|
TokenBucketCapacity: 1000,
|
|
TokenBucketRefillRate: 100,
|
|
TokenBucketRefillInterval: 100 * time.Millisecond,
|
|
LeakyBucketCapacity: 500,
|
|
LeakyBucketLeakInterval: 200 * time.Millisecond,
|
|
CreditBasedMaxCredits: 1000,
|
|
CreditBasedRefillRate: 100,
|
|
CreditBasedRefillInterval: 200 * time.Millisecond,
|
|
CreditBasedBurstSize: 50,
|
|
RateLimiterRequestsPerSecond: 100,
|
|
RateLimiterBurstSize: 200,
|
|
QueueDepthThreshold: 1000,
|
|
MemoryThreshold: 1 * 1024 * 1024 * 1024, // 1GB
|
|
ErrorRateThreshold: 0.5,
|
|
SnapshotInterval: 5 * time.Minute,
|
|
SnapshotRetention: 24 * time.Hour,
|
|
TracingEnabled: true,
|
|
TraceRetention: 24 * time.Hour,
|
|
TraceExportInterval: 30 * time.Second,
|
|
EnableEnhancements: true,
|
|
}
|
|
}
|
|
|
|
// EnhancedPublish publishes a message with enhanced features
|
|
func (b *Broker) EnhancedPublish(ctx context.Context, task *Task, queueName string) error {
|
|
if b.enhanced == nil || !b.enhanced.enabled {
|
|
// Fall back to regular publish logic
|
|
return b.regularPublish(ctx, task, queueName)
|
|
}
|
|
|
|
// Check for duplicates
|
|
if b.enhanced.dedupManager != nil {
|
|
isDuplicate, err := b.enhanced.dedupManager.CheckDuplicate(ctx, task)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if isDuplicate {
|
|
b.logger.Debug("Duplicate message rejected",
|
|
logger.Field{Key: "taskID", Value: task.ID})
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// Acquire flow control credits
|
|
if b.enhanced.flowController != nil {
|
|
if err := b.enhanced.flowController.AcquireCredit(ctx, 1); err != nil {
|
|
return err
|
|
}
|
|
defer b.enhanced.flowController.ReleaseCredit(1)
|
|
}
|
|
|
|
// Write to WAL
|
|
if b.enhanced.walLog != nil {
|
|
walEntry := &WALEntry{
|
|
EntryType: WALEntryEnqueue,
|
|
TaskID: task.ID,
|
|
QueueName: queueName,
|
|
Payload: json.RawMessage(task.Payload),
|
|
}
|
|
if err := b.enhanced.walLog.WriteEntry(ctx, walEntry); err != nil {
|
|
b.logger.Error("Failed to write WAL entry",
|
|
logger.Field{Key: "error", Value: err})
|
|
}
|
|
}
|
|
|
|
// Start tracing
|
|
if b.enhanced.lifecycleTracker != nil {
|
|
b.enhanced.lifecycleTracker.TrackEnqueue(ctx, task, queueName)
|
|
}
|
|
|
|
// Continue with regular publish
|
|
return b.regularPublish(ctx, task, queueName)
|
|
}
|
|
|
|
// regularPublish is the standard publish logic
|
|
func (b *Broker) regularPublish(ctx context.Context, task *Task, queueName string) error {
|
|
queue, exists := b.queues.Get(queueName)
|
|
if !exists {
|
|
queue = b.NewQueue(queueName)
|
|
}
|
|
|
|
// Enqueue task
|
|
select {
|
|
case queue.tasks <- &QueuedTask{Task: task}:
|
|
// Track for acknowledgment if enhanced features enabled
|
|
if b.enhanced != nil && b.enhanced.ackManager != nil {
|
|
_ = b.enhanced.ackManager.TrackMessage(ctx, task, queueName, "")
|
|
}
|
|
return nil
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
}
|
|
}
|
|
|
|
// RecoverFromWAL recovers broker state from WAL
|
|
func (b *Broker) RecoverFromWAL(ctx context.Context) error {
|
|
if b.enhanced == nil || b.enhanced.walLog == nil {
|
|
return nil
|
|
}
|
|
|
|
b.logger.Info("Starting WAL recovery")
|
|
|
|
return b.enhanced.walLog.Replay(func(entry *WALEntry) error {
|
|
switch entry.EntryType {
|
|
case WALEntryEnqueue:
|
|
var task Task
|
|
if err := json.Unmarshal(entry.Payload, &task); err != nil {
|
|
return err
|
|
}
|
|
|
|
queue, exists := b.queues.Get(entry.QueueName)
|
|
if !exists {
|
|
queue = b.NewQueue(entry.QueueName)
|
|
}
|
|
|
|
select {
|
|
case queue.tasks <- &QueuedTask{Task: &task}:
|
|
default:
|
|
b.logger.Warn("Queue full during recovery",
|
|
logger.Field{Key: "queue", Value: entry.QueueName})
|
|
}
|
|
|
|
case WALEntryComplete, WALEntryFailed:
|
|
// Already processed, no action needed
|
|
}
|
|
|
|
return nil
|
|
})
|
|
}
|
|
|
|
// RecoverFromSnapshot recovers broker state from snapshots
|
|
func (b *Broker) RecoverFromSnapshot(ctx context.Context) error {
|
|
if b.enhanced == nil || b.enhanced.snapshotManager == nil {
|
|
return nil
|
|
}
|
|
|
|
b.logger.Info("Starting snapshot recovery")
|
|
|
|
// Recover all queues
|
|
var recoveredQueues int
|
|
b.queues.ForEach(func(queueName string, _ *Queue) bool {
|
|
if err := b.enhanced.snapshotManager.RestoreFromSnapshot(ctx, queueName); err != nil {
|
|
b.logger.Error("Failed to restore queue from snapshot",
|
|
logger.Field{Key: "queue", Value: queueName},
|
|
logger.Field{Key: "error", Value: err})
|
|
} else {
|
|
recoveredQueues++
|
|
}
|
|
return true
|
|
})
|
|
|
|
b.logger.Info("Snapshot recovery complete",
|
|
logger.Field{Key: "queues", Value: recoveredQueues})
|
|
|
|
return nil
|
|
}
|
|
|
|
// GetEnhancedStats returns comprehensive statistics
|
|
func (b *Broker) GetEnhancedStats() map[string]interface{} {
|
|
stats := make(map[string]interface{})
|
|
|
|
if b.enhanced == nil {
|
|
return stats
|
|
}
|
|
|
|
if b.enhanced.ackManager != nil {
|
|
stats["acknowledgments"] = b.enhanced.ackManager.GetStats()
|
|
}
|
|
|
|
if b.enhanced.walLog != nil {
|
|
stats["wal"] = b.enhanced.walLog.GetStats()
|
|
}
|
|
|
|
if b.enhanced.dedupManager != nil {
|
|
stats["deduplication"] = b.enhanced.dedupManager.GetStats()
|
|
}
|
|
|
|
if b.enhanced.flowController != nil {
|
|
stats["flow_control"] = b.enhanced.flowController.GetStats()
|
|
}
|
|
|
|
if b.enhanced.snapshotManager != nil {
|
|
stats["snapshots"] = b.enhanced.snapshotManager.GetSnapshotStats()
|
|
}
|
|
|
|
if b.enhanced.traceManager != nil {
|
|
stats["tracing"] = b.enhanced.traceManager.GetStats()
|
|
}
|
|
|
|
return stats
|
|
}
|
|
|
|
// ShutdownEnhanced gracefully shuts down all enhanced features
|
|
func (b *Broker) ShutdownEnhanced(ctx context.Context) error {
|
|
if b.enhanced == nil {
|
|
return nil
|
|
}
|
|
|
|
b.logger.Info("Starting enhanced features shutdown")
|
|
|
|
// Shutdown components in order
|
|
if b.enhanced.backpressureMonitor != nil {
|
|
b.enhanced.backpressureMonitor.Shutdown()
|
|
}
|
|
|
|
if b.enhanced.flowController != nil {
|
|
b.enhanced.flowController.Shutdown()
|
|
}
|
|
|
|
if b.enhanced.dedupManager != nil {
|
|
_ = b.enhanced.dedupManager.Shutdown(ctx)
|
|
}
|
|
|
|
if b.enhanced.ackManager != nil {
|
|
_ = b.enhanced.ackManager.Shutdown(ctx)
|
|
}
|
|
|
|
if b.enhanced.snapshotManager != nil {
|
|
_ = b.enhanced.snapshotManager.Shutdown(ctx)
|
|
}
|
|
|
|
if b.enhanced.walLog != nil {
|
|
_ = b.enhanced.walLog.Shutdown(ctx)
|
|
}
|
|
|
|
if b.enhanced.traceManager != nil {
|
|
_ = b.enhanced.traceManager.Shutdown(ctx)
|
|
}
|
|
|
|
if b.enhanced.dlqStorage != nil {
|
|
_ = b.enhanced.dlqStorage.Close()
|
|
}
|
|
|
|
b.logger.Info("Enhanced features shutdown complete")
|
|
|
|
return nil
|
|
}
|
|
|
|
// InitializeEnhancedPool initializes pool with enhanced features
|
|
func (p *Pool) InitializeEnhancedPool(config *BrokerEnhancedConfig) {
|
|
if config == nil || !config.EnableEnhancements {
|
|
return
|
|
}
|
|
|
|
// Add health monitor
|
|
if config.WorkerHealthTimeout > 0 {
|
|
_ = NewWorkerHealthMonitor(p, config.WorkerHealthTimeout, config.WorkerCheckInterval, config.Logger)
|
|
}
|
|
|
|
// Add dynamic scaler
|
|
if config.MinWorkers > 0 && config.MaxWorkers > int(p.numOfWorkers) {
|
|
_ = NewDynamicScaler(p, config.MinWorkers, config.MaxWorkers, config.Logger)
|
|
}
|
|
}
|