diff --git a/server/configs/config.local.yaml.example b/server/configs/config.local.yaml.example index 258b1fb..76a6330 100644 --- a/server/configs/config.local.yaml.example +++ b/server/configs/config.local.yaml.example @@ -30,10 +30,35 @@ rabbitmq: monitor_result_dlx: monitor.result.dlx monitor_result_dlq: monitor.result.ingest.dlq monitor_result_dlq_routing_key: monitor.result.ingest.dlq + generation_exchange: generation.task + generation_routing_key: generation.task.run + generation_queue: generation.task.run + generation_dlx: generation.task.dlx + generation_dlq: generation.task.run.dlq + generation_dlq_routing_key: generation.task.run.dlq + generation_event_exchange: generation.event + template_assist_exchange: template.assist + template_assist_routing_key: template.assist.run + template_assist_queue: template.assist.run + template_assist_dlx: template.assist.dlx + template_assist_dlq: template.assist.run.dlq + template_assist_dlq_routing_key: template.assist.run.dlq consumer_prefetch: 10 + publish_channel_pool_size: 32 # Prefer environment override when needed: # export RABBITMQ_URL=amqp://geo:geo_dev@localhost:5672/geo +scheduler: + dispatch_interval: 15s + dispatch_timeout: 30s + dispatch_batch_size: 100 + dispatch_concurrency: 4 + generation_queue_backpressure_limit: 1000 + +monitoring_workers: + result_ingest_concurrency: 4 + projection_rebuild_concurrency: 2 + redis: addr: localhost:6379 @@ -104,6 +129,6 @@ retrieval: generation: queue_size: 128 - worker_concurrency: 1 + worker_concurrency: 8 stream_enabled: false article_timeout: 8m diff --git a/server/configs/config.yaml b/server/configs/config.yaml index 39b5c00..91928f2 100644 --- a/server/configs/config.yaml +++ b/server/configs/config.yaml @@ -36,7 +36,32 @@ rabbitmq: monitor_projection_dlx: monitor.projection.dlx monitor_projection_dlq: monitor.projection.rebuild.dlq monitor_projection_dlq_routing_key: monitor.projection.rebuild.dlq + generation_exchange: generation.task + generation_routing_key: generation.task.run + generation_queue: generation.task.run + generation_dlx: generation.task.dlx + generation_dlq: generation.task.run.dlq + generation_dlq_routing_key: generation.task.run.dlq + generation_event_exchange: generation.event + template_assist_exchange: template.assist + template_assist_routing_key: template.assist.run + template_assist_queue: template.assist.run + template_assist_dlx: template.assist.dlx + template_assist_dlq: template.assist.run.dlq + template_assist_dlq_routing_key: template.assist.run.dlq consumer_prefetch: 10 + publish_channel_pool_size: 32 + +scheduler: + dispatch_interval: 15s + dispatch_timeout: 30s + dispatch_batch_size: 100 + dispatch_concurrency: 4 + generation_queue_backpressure_limit: 1000 + +monitoring_workers: + result_ingest_concurrency: 4 + projection_rebuild_concurrency: 2 redis: addr: localhost:6379 @@ -99,7 +124,7 @@ retrieval: generation: queue_size: 128 - worker_concurrency: 1 + worker_concurrency: 8 stream_enabled: false article_timeout: 8m diff --git a/server/go.mod b/server/go.mod index 70e9a01..d4f4e45 100644 --- a/server/go.mod +++ b/server/go.mod @@ -61,6 +61,7 @@ require ( github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect github.com/richardlehane/mscfb v1.0.4 // indirect github.com/richardlehane/msoleps v1.0.3 // indirect + github.com/robfig/cron/v3 v3.0.1 // indirect github.com/rs/xid v1.6.0 // indirect github.com/sagikazarmark/locafero v0.4.0 // indirect github.com/sagikazarmark/slog-shim v0.1.0 // indirect diff --git a/server/go.sum b/server/go.sum index 56f2c6d..1fa9951 100644 --- a/server/go.sum +++ b/server/go.sum @@ -150,6 +150,8 @@ github.com/richardlehane/mscfb v1.0.4/go.mod h1:YzVpcZg9czvAuhk9T+a3avCpcFPMUWm7 github.com/richardlehane/msoleps v1.0.1/go.mod h1:BWev5JBpU9Ko2WAgmZEuiz4/u3ZYTKbjLycmwiWUfWg= github.com/richardlehane/msoleps v1.0.3 h1:aznSZzrwYRl3rLKRT3gUk9am7T/mLNSnJINvN0AQoVM= github.com/richardlehane/msoleps v1.0.3/go.mod h1:BWev5JBpU9Ko2WAgmZEuiz4/u3ZYTKbjLycmwiWUfWg= +github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs= +github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro= github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjRBZyWFQ= github.com/rogpeppe/go-internal v1.10.0/go.mod h1:UQnix2H7Ngw/k4C5ijL5+65zddjncjaFoBhdsK/akog= github.com/rs/xid v1.6.0 h1:fV591PaemRlL6JfRxGDEPl69wICngIQ3shQtzfy2gxU= diff --git a/server/internal/shared/config/config.go b/server/internal/shared/config/config.go index 406c32c..2e57247 100644 --- a/server/internal/shared/config/config.go +++ b/server/internal/shared/config/config.go @@ -14,6 +14,8 @@ type Config struct { Database DatabaseConfig `mapstructure:"database"` MonitoringDatabase DatabaseConfig `mapstructure:"monitoring_database"` RabbitMQ RabbitMQConfig `mapstructure:"rabbitmq"` + Scheduler SchedulerConfig `mapstructure:"scheduler"` + MonitoringWorkers MonitoringConfig `mapstructure:"monitoring_workers"` Redis RedisConfig `mapstructure:"redis"` Qdrant QdrantConfig `mapstructure:"qdrant"` ObjectStorage ObjectStorageConfig `mapstructure:"object_storage"` @@ -65,7 +67,34 @@ type RabbitMQConfig struct { MonitorProjectionDLX string `mapstructure:"monitor_projection_dlx"` MonitorProjectionDLQ string `mapstructure:"monitor_projection_dlq"` MonitorProjectionDLQRouteKey string `mapstructure:"monitor_projection_dlq_routing_key"` + GenerationExchange string `mapstructure:"generation_exchange"` + GenerationRoutingKey string `mapstructure:"generation_routing_key"` + GenerationQueue string `mapstructure:"generation_queue"` + GenerationDLX string `mapstructure:"generation_dlx"` + GenerationDLQ string `mapstructure:"generation_dlq"` + GenerationDLQRouteKey string `mapstructure:"generation_dlq_routing_key"` + GenerationEventExchange string `mapstructure:"generation_event_exchange"` + TemplateAssistExchange string `mapstructure:"template_assist_exchange"` + TemplateAssistRoutingKey string `mapstructure:"template_assist_routing_key"` + TemplateAssistQueue string `mapstructure:"template_assist_queue"` + TemplateAssistDLX string `mapstructure:"template_assist_dlx"` + TemplateAssistDLQ string `mapstructure:"template_assist_dlq"` + TemplateAssistDLQRouteKey string `mapstructure:"template_assist_dlq_routing_key"` ConsumerPrefetch int `mapstructure:"consumer_prefetch"` + PublishChannelPoolSize int `mapstructure:"publish_channel_pool_size"` +} + +type SchedulerConfig struct { + DispatchInterval time.Duration `mapstructure:"dispatch_interval"` + DispatchTimeout time.Duration `mapstructure:"dispatch_timeout"` + DispatchBatchSize int `mapstructure:"dispatch_batch_size"` + DispatchConcurrency int `mapstructure:"dispatch_concurrency"` + GenerationQueueBackpressureLimit int `mapstructure:"generation_queue_backpressure_limit"` +} + +type MonitoringConfig struct { + ResultIngestConcurrency int `mapstructure:"result_ingest_concurrency"` + ProjectionRebuildConcurrency int `mapstructure:"projection_rebuild_concurrency"` } type QdrantConfig struct { @@ -159,6 +188,8 @@ func Load(configPath string) (*Config, error) { applyEnvOverrides(&cfg) normalizeRabbitMQConfig(&cfg.RabbitMQ) + normalizeSchedulerConfig(&cfg.Scheduler) + normalizeMonitoringConfig(&cfg.MonitoringWorkers) return &cfg, nil } @@ -266,6 +297,81 @@ func normalizeRabbitMQConfig(cfg *RabbitMQConfig) { if strings.TrimSpace(cfg.MonitorProjectionDLQRouteKey) == "" { cfg.MonitorProjectionDLQRouteKey = "monitor.projection.rebuild.dlq" } + + if strings.TrimSpace(cfg.GenerationExchange) == "" { + cfg.GenerationExchange = "generation.task" + } + if strings.TrimSpace(cfg.GenerationRoutingKey) == "" { + cfg.GenerationRoutingKey = "generation.task.run" + } + if strings.TrimSpace(cfg.GenerationQueue) == "" { + cfg.GenerationQueue = "generation.task.run" + } + if strings.TrimSpace(cfg.GenerationDLX) == "" { + cfg.GenerationDLX = "generation.task.dlx" + } + if strings.TrimSpace(cfg.GenerationDLQ) == "" { + cfg.GenerationDLQ = "generation.task.run.dlq" + } + if strings.TrimSpace(cfg.GenerationDLQRouteKey) == "" { + cfg.GenerationDLQRouteKey = "generation.task.run.dlq" + } + if strings.TrimSpace(cfg.GenerationEventExchange) == "" { + cfg.GenerationEventExchange = "generation.event" + } + if strings.TrimSpace(cfg.TemplateAssistExchange) == "" { + cfg.TemplateAssistExchange = "template.assist" + } + if strings.TrimSpace(cfg.TemplateAssistRoutingKey) == "" { + cfg.TemplateAssistRoutingKey = "template.assist.run" + } + if strings.TrimSpace(cfg.TemplateAssistQueue) == "" { + cfg.TemplateAssistQueue = "template.assist.run" + } + if strings.TrimSpace(cfg.TemplateAssistDLX) == "" { + cfg.TemplateAssistDLX = "template.assist.dlx" + } + if strings.TrimSpace(cfg.TemplateAssistDLQ) == "" { + cfg.TemplateAssistDLQ = "template.assist.run.dlq" + } + if strings.TrimSpace(cfg.TemplateAssistDLQRouteKey) == "" { + cfg.TemplateAssistDLQRouteKey = "template.assist.run.dlq" + } + if cfg.PublishChannelPoolSize <= 0 { + cfg.PublishChannelPoolSize = 16 + } +} + +func normalizeSchedulerConfig(cfg *SchedulerConfig) { + if cfg == nil { + return + } + + if cfg.DispatchInterval <= 0 { + cfg.DispatchInterval = 15 * time.Second + } + if cfg.DispatchTimeout <= 0 { + cfg.DispatchTimeout = 30 * time.Second + } + if cfg.DispatchBatchSize <= 0 { + cfg.DispatchBatchSize = 20 + } + if cfg.DispatchConcurrency <= 0 { + cfg.DispatchConcurrency = 1 + } +} + +func normalizeMonitoringConfig(cfg *MonitoringConfig) { + if cfg == nil { + return + } + + if cfg.ResultIngestConcurrency <= 0 { + cfg.ResultIngestConcurrency = 1 + } + if cfg.ProjectionRebuildConcurrency <= 0 { + cfg.ProjectionRebuildConcurrency = 1 + } } func lookupFirstNonEmptyEnv(keys ...string) (string, bool) { diff --git a/server/internal/shared/messaging/rabbitmq/client.go b/server/internal/shared/messaging/rabbitmq/client.go index 69c0c15..a5c53fb 100644 --- a/server/internal/shared/messaging/rabbitmq/client.go +++ b/server/internal/shared/messaging/rabbitmq/client.go @@ -12,10 +12,22 @@ import ( ) type Client struct { - cfg config.RabbitMQConfig - mu sync.Mutex - conn *amqp.Connection - closed bool + cfg config.RabbitMQConfig + mu sync.Mutex + conn *amqp.Connection + closed bool + publishPool chan *publishChannel +} + +type publishChannel struct { + conn *amqp.Connection + ch *amqp.Channel +} + +type QueueStats struct { + Name string + Messages int + Consumers int } func New(_ context.Context, cfg config.RabbitMQConfig) (*Client, error) { @@ -23,7 +35,10 @@ func New(_ context.Context, cfg config.RabbitMQConfig) (*Client, error) { return nil, fmt.Errorf("rabbitmq url is required") } - client := &Client{cfg: cfg} + client := &Client{ + cfg: cfg, + publishPool: make(chan *publishChannel, cfg.PublishChannelPoolSize), + } if _, err := client.ensureConnection(); err != nil { return nil, err } @@ -40,8 +55,13 @@ func (c *Client) Close() error { c.closed = true conn := c.conn c.conn = nil + pooledChannels := c.drainPublishChannelsLocked() c.mu.Unlock() + for _, ch := range pooledChannels { + _ = ch.Close() + } + if conn == nil || conn.IsClosed() { return nil } @@ -58,6 +78,10 @@ func (c *Client) Ping(_ context.Context) error { return nil } +func (c *Client) InspectGenerationQueue(ctx context.Context) (QueueStats, error) { + return c.inspectQueue(ctx, c.cfg.GenerationQueue) +} + func (c *Client) PublishMonitorResultIngest(ctx context.Context, body []byte) error { return c.publish(ctx, c.cfg.MonitorResultExchange, c.cfg.MonitorResultRoutingKey, body) } @@ -66,9 +90,21 @@ func (c *Client) PublishMonitorProjectionRebuild(ctx context.Context, body []byt return c.publish(ctx, c.cfg.MonitorProjectionExchange, c.cfg.MonitorProjectionRoutingKey, body) } +func (c *Client) PublishGenerationTask(ctx context.Context, body []byte) error { + return c.publish(ctx, c.cfg.GenerationExchange, c.cfg.GenerationRoutingKey, body) +} + +func (c *Client) PublishGenerationEvent(ctx context.Context, body []byte) error { + return c.publishFanout(ctx, c.cfg.GenerationEventExchange, body) +} + +func (c *Client) PublishTemplateAssistTask(ctx context.Context, body []byte) error { + return c.publish(ctx, c.cfg.TemplateAssistExchange, c.cfg.TemplateAssistRoutingKey, body) +} + func (c *Client) publish(ctx context.Context, exchange, routingKey string, body []byte) error { for attempt := 0; attempt < 2; attempt++ { - conn, ch, err := c.openChannel() + conn, ch, pooled, err := c.acquirePublishChannel() if err != nil { return err } @@ -78,11 +114,12 @@ func (c *Client) publish(ctx context.Context, exchange, routingKey string, body DeliveryMode: amqp.Persistent, Body: body, }) - _ = ch.Close() if publishErr == nil { + c.releasePublishChannel(conn, ch, pooled) return nil } + c.discardPublishChannel(ch) c.invalidateConnection(conn) if attempt == 1 { return fmt.Errorf("publish rabbitmq message: %w", publishErr) @@ -92,6 +129,33 @@ func (c *Client) publish(ctx context.Context, exchange, routingKey string, body return fmt.Errorf("publish rabbitmq message: retry exhausted") } +func (c *Client) publishFanout(ctx context.Context, exchange string, body []byte) error { + for attempt := 0; attempt < 2; attempt++ { + conn, ch, pooled, err := c.acquirePublishChannel() + if err != nil { + return err + } + + publishErr := ch.PublishWithContext(ctx, exchange, "", false, false, amqp.Publishing{ + ContentType: "application/json", + DeliveryMode: amqp.Transient, + Body: body, + }) + if publishErr == nil { + c.releasePublishChannel(conn, ch, pooled) + return nil + } + + c.discardPublishChannel(ch) + c.invalidateConnection(conn) + if attempt == 1 { + return fmt.Errorf("publish rabbitmq fanout message: %w", publishErr) + } + } + + return fmt.Errorf("publish rabbitmq fanout message: retry exhausted") +} + func (c *Client) ConsumeMonitorResultIngest(consumerName string) (<-chan amqp.Delivery, *amqp.Channel, error) { return c.consume(consumerName, c.cfg.MonitorResultQueue) } @@ -100,6 +164,18 @@ func (c *Client) ConsumeMonitorProjectionRebuild(consumerName string) (<-chan am return c.consume(consumerName, c.cfg.MonitorProjectionQueue) } +func (c *Client) ConsumeGenerationTask(consumerName string) (<-chan amqp.Delivery, *amqp.Channel, error) { + return c.consume(consumerName, c.cfg.GenerationQueue) +} + +func (c *Client) ConsumeGenerationEvents(consumerName string) (<-chan amqp.Delivery, *amqp.Channel, error) { + return c.consumeFanout(consumerName, c.cfg.GenerationEventExchange) +} + +func (c *Client) ConsumeTemplateAssistTask(consumerName string) (<-chan amqp.Delivery, *amqp.Channel, error) { + return c.consume(consumerName, c.cfg.TemplateAssistQueue) +} + func (c *Client) consume(consumerName, queue string) (<-chan amqp.Delivery, *amqp.Channel, error) { for attempt := 0; attempt < 2; attempt++ { conn, ch, err := c.openChannel() @@ -142,6 +218,69 @@ func (c *Client) consume(consumerName, queue string) (<-chan amqp.Delivery, *amq return nil, nil, fmt.Errorf("consume rabbitmq queue: retry exhausted") } +func (c *Client) consumeFanout(consumerName, exchange string) (<-chan amqp.Delivery, *amqp.Channel, error) { + for attempt := 0; attempt < 2; attempt++ { + conn, ch, err := c.openChannel() + if err != nil { + return nil, nil, err + } + + queue, err := ch.QueueDeclare( + "", + false, + true, + true, + false, + nil, + ) + if err != nil { + _ = ch.Close() + c.invalidateConnection(conn) + if attempt == 1 { + return nil, nil, fmt.Errorf("declare rabbitmq fanout queue: %w", err) + } + continue + } + + if err := ch.QueueBind( + queue.Name, + "", + exchange, + false, + nil, + ); err != nil { + _ = ch.Close() + c.invalidateConnection(conn) + if attempt == 1 { + return nil, nil, fmt.Errorf("bind rabbitmq fanout queue: %w", err) + } + continue + } + + deliveries, err := ch.Consume( + queue.Name, + consumerName, + false, + true, + false, + false, + nil, + ) + if err != nil { + _ = ch.Close() + c.invalidateConnection(conn) + if attempt == 1 { + return nil, nil, fmt.Errorf("consume rabbitmq fanout queue: %w", err) + } + continue + } + + return deliveries, ch, nil + } + + return nil, nil, fmt.Errorf("consume rabbitmq fanout queue: retry exhausted") +} + func (c *Client) ensureConnection() (*amqp.Connection, error) { if c == nil { return nil, fmt.Errorf("rabbitmq client is nil") @@ -157,6 +296,9 @@ func (c *Client) ensureConnection() (*amqp.Connection, error) { return c.conn, nil } if c.conn != nil { + for _, ch := range c.drainPublishChannelsLocked() { + _ = ch.Close() + } _ = c.conn.Close() c.conn = nil } @@ -195,23 +337,156 @@ func (c *Client) openChannel() (*amqp.Connection, *amqp.Channel, error) { return nil, nil, fmt.Errorf("open rabbitmq channel: retry exhausted") } +func (c *Client) acquirePublishChannel() (*amqp.Connection, *amqp.Channel, bool, error) { + for attempt := 0; attempt < 2; attempt++ { + conn, err := c.ensureConnection() + if err != nil { + return nil, nil, false, err + } + + if ch := c.borrowPublishChannel(conn); ch != nil { + return conn, ch, true, nil + } + + ch, err := conn.Channel() + if err == nil { + return conn, ch, false, nil + } + + c.invalidateConnection(conn) + if attempt == 1 { + return nil, nil, false, fmt.Errorf("open rabbitmq channel: %w", err) + } + } + + return nil, nil, false, fmt.Errorf("open rabbitmq channel: retry exhausted") +} + func (c *Client) invalidateConnection(conn *amqp.Connection) { if c == nil || conn == nil { return } c.mu.Lock() - defer c.mu.Unlock() - if c.conn != conn { + c.mu.Unlock() return } c.conn = nil + pooledChannels := c.drainPublishChannelsLocked() + c.mu.Unlock() + + for _, ch := range pooledChannels { + _ = ch.Close() + } if !conn.IsClosed() { _ = conn.Close() } } +func (c *Client) inspectQueue(ctx context.Context, queue string) (QueueStats, error) { + if err := ctx.Err(); err != nil { + return QueueStats{}, err + } + + for attempt := 0; attempt < 2; attempt++ { + conn, ch, err := c.openChannel() + if err != nil { + return QueueStats{}, err + } + + info, inspectErr := ch.QueueInspect(queue) + _ = ch.Close() + if inspectErr == nil { + return QueueStats{ + Name: info.Name, + Messages: info.Messages, + Consumers: info.Consumers, + }, nil + } + + c.invalidateConnection(conn) + if attempt == 1 { + return QueueStats{}, fmt.Errorf("inspect rabbitmq queue %s: %w", queue, inspectErr) + } + } + + return QueueStats{}, fmt.Errorf("inspect rabbitmq queue %s: retry exhausted", queue) +} + +func (c *Client) borrowPublishChannel(conn *amqp.Connection) *amqp.Channel { + if c == nil || c.publishPool == nil { + return nil + } + + for { + select { + case pooled := <-c.publishPool: + if pooled == nil || pooled.ch == nil { + continue + } + if pooled.conn != conn { + _ = pooled.ch.Close() + continue + } + return pooled.ch + default: + return nil + } + } +} + +func (c *Client) releasePublishChannel(conn *amqp.Connection, ch *amqp.Channel, reusable bool) { + if ch == nil { + return + } + if !reusable || c == nil || c.publishPool == nil || conn == nil || conn.IsClosed() { + _ = ch.Close() + return + } + + c.mu.Lock() + currentConn := c.conn + closed := c.closed + c.mu.Unlock() + + if closed || currentConn != conn { + _ = ch.Close() + return + } + + select { + case c.publishPool <- &publishChannel{conn: conn, ch: ch}: + default: + _ = ch.Close() + } +} + +func (c *Client) discardPublishChannel(ch *amqp.Channel) { + if ch == nil { + return + } + _ = ch.Close() +} + +func (c *Client) drainPublishChannelsLocked() []*amqp.Channel { + if c == nil || c.publishPool == nil { + return nil + } + + drained := make([]*amqp.Channel, 0, len(c.publishPool)) + for { + select { + case pooled := <-c.publishPool: + if pooled != nil && pooled.ch != nil { + drained = append(drained, pooled.ch) + } + default: + return drained + } + } +} + func (c *Client) ensureTopology(conn *amqp.Connection) error { ch, err := conn.Channel() if err != nil { @@ -243,6 +518,42 @@ func (c *Client) ensureTopology(conn *amqp.Connection) error { return err } + if err := c.ensureWorkQueue( + ch, + c.cfg.GenerationExchange, + c.cfg.GenerationRoutingKey, + c.cfg.GenerationQueue, + c.cfg.GenerationDLX, + c.cfg.GenerationDLQ, + c.cfg.GenerationDLQRouteKey, + ); err != nil { + return err + } + + if err := ch.ExchangeDeclare( + c.cfg.GenerationEventExchange, + "fanout", + false, + false, + false, + false, + nil, + ); err != nil { + return fmt.Errorf("declare rabbitmq exchange %s: %w", c.cfg.GenerationEventExchange, err) + } + + if err := c.ensureWorkQueue( + ch, + c.cfg.TemplateAssistExchange, + c.cfg.TemplateAssistRoutingKey, + c.cfg.TemplateAssistQueue, + c.cfg.TemplateAssistDLX, + c.cfg.TemplateAssistDLQ, + c.cfg.TemplateAssistDLQRouteKey, + ); err != nil { + return err + } + return nil }