@@ -12,22 +12,36 @@ import (
12
12
amqp "github.com/rabbitmq/amqp091-go"
13
13
)
14
14
15
- // RabbitMQBroker implements the Broker interface for RabbitMQ
15
+ // RabbitMQBroker implements the Broker and Poller interfaces for RabbitMQ
16
16
type RabbitMQBroker struct {
17
- connection * amqp.Connection
18
- channel * amqp.Channel
19
- options Options
20
- serializer core.Serializer
21
- queues map [string ]bool // Track declared queues
22
- logger seelog.LoggerInterface
17
+ connection * amqp.Connection
18
+ channel * amqp.Channel
19
+ options Options
20
+ serializer core.Serializer
21
+ declaredQueues map [string ]bool // Track declared queues
22
+ consumerQueues []string // Queues to consume from
23
+ consumerTags map [string ]string // Track consumer tags
24
+ logger seelog.LoggerInterface
23
25
}
24
26
25
27
// NewBroker creates a new RabbitMQ broker
26
28
func NewBroker (options Options , serializer core.Serializer ) * RabbitMQBroker {
27
29
return & RabbitMQBroker {
28
- options : options ,
29
- serializer : serializer ,
30
- queues : make (map [string ]bool ),
30
+ options : options ,
31
+ serializer : serializer ,
32
+ declaredQueues : make (map [string ]bool ),
33
+ consumerTags : make (map [string ]string ),
34
+ }
35
+ }
36
+
37
+ // NewBrokerWithQueues creates a new RabbitMQ broker with consumer queues
38
+ func NewBrokerWithQueues (options Options , serializer core.Serializer , queues []string ) * RabbitMQBroker {
39
+ return & RabbitMQBroker {
40
+ options : options ,
41
+ serializer : serializer ,
42
+ declaredQueues : make (map [string ]bool ),
43
+ consumerQueues : queues ,
44
+ consumerTags : make (map [string ]string ),
31
45
}
32
46
}
33
47
@@ -168,36 +182,14 @@ func (r *RabbitMQBroker) Dequeue(ctx context.Context, queue string) (job.Job, er
168
182
return nil , nil // No message available
169
183
}
170
184
171
- // Create metadata
172
- metadata := job.Metadata {
173
- Queue : queue ,
174
- EnqueuedAt : delivery .Timestamp ,
175
- }
176
-
177
- if delivery .MessageId != "" {
178
- metadata .ID = delivery .MessageId
179
- }
180
-
181
- // Deserialize job
182
- j , err := r .serializer .Deserialize (delivery .Body , metadata )
183
- if err != nil {
184
- // Reject message if we can't deserialize it
185
- if nackErr := delivery .Nack (false , false ); nackErr != nil {
186
- r .logError ("Failed to nack message after deserialization error: %v" , nackErr )
187
- }
185
+ // Convert delivery to job using the shared method
186
+ job := r .convertDeliveryToJob (delivery , queue )
187
+ if job == nil {
188
188
return nil , errors .NewSerializationError (r .serializer .GetFormat (),
189
- fmt .Errorf ("deserialize job: %w" , err ))
189
+ fmt .Errorf ("failed to convert delivery to job" ))
190
190
}
191
191
192
- // Store delivery tag for ACK/NACK
193
- // Always wrap in RMQJob
194
- rmqJob := & RMQJob {
195
- Job : j ,
196
- deliveryTag : delivery .DeliveryTag ,
197
- channel : channel ,
198
- }
199
-
200
- return rmqJob , nil
192
+ return job , nil
201
193
}
202
194
203
195
// Ack acknowledges job completion
@@ -254,7 +246,7 @@ func (r *RabbitMQBroker) CreateQueue(ctx context.Context, name string, options c
254
246
return errors .NewBrokerError ("create_queue" , name , err )
255
247
}
256
248
257
- r .queues [name ] = true
249
+ r .declaredQueues [name ] = true
258
250
return nil
259
251
}
260
252
@@ -270,7 +262,7 @@ func (r *RabbitMQBroker) DeleteQueue(ctx context.Context, name string) error {
270
262
return errors .NewBrokerError ("delete_queue" , name , err )
271
263
}
272
264
273
- delete (r .queues , name )
265
+ delete (r .declaredQueues , name )
274
266
return nil
275
267
}
276
268
@@ -318,7 +310,7 @@ func (r *RabbitMQBroker) ensureQueue(name string) error {
318
310
return err
319
311
}
320
312
321
- if r .queues [name ] {
313
+ if r .declaredQueues [name ] {
322
314
return nil // Already declared
323
315
}
324
316
@@ -335,7 +327,7 @@ func (r *RabbitMQBroker) ensureQueue(name string) error {
335
327
return err
336
328
}
337
329
338
- r .queues [name ] = true
330
+ r .declaredQueues [name ] = true
339
331
return nil
340
332
}
341
333
@@ -346,6 +338,111 @@ func (r *RabbitMQBroker) logError(format string, args ...interface{}) {
346
338
}
347
339
}
348
340
341
+ // Start implements the Poller interface for push-based consumption
342
+ func (r * RabbitMQBroker ) Start (ctx context.Context , jobChan chan <- job.Job ) error {
343
+ r .logger .Infof ("Starting RabbitMQ consumer for queues: %v" , r .consumerQueues )
344
+
345
+ for _ , queue := range r .consumerQueues {
346
+ if err := r .ensureQueue (queue ); err != nil {
347
+ return fmt .Errorf ("failed to ensure queue %s: %w" , queue , err )
348
+ }
349
+
350
+ deliveries , err := r .channel .Consume (
351
+ queue , // queue
352
+ "" , // consumer tag (auto-generated)
353
+ false , // auto-ack
354
+ false , // exclusive
355
+ false , // no-local
356
+ false , // no-wait
357
+ nil , // args
358
+ )
359
+ if err != nil {
360
+ return fmt .Errorf ("failed to start consumer for queue %s: %w" , queue , err )
361
+ }
362
+
363
+ // Store consumer tag for cleanup
364
+ if r .consumerTags == nil {
365
+ r .consumerTags = make (map [string ]string )
366
+ }
367
+
368
+ go r .handleDeliveries (ctx , queue , deliveries , jobChan )
369
+ }
370
+
371
+ // Keep running until context is cancelled
372
+ <- ctx .Done ()
373
+ r .logger .Info ("RabbitMQ consumer stopped" )
374
+ close (jobChan )
375
+ return nil
376
+ }
377
+
378
+ // handleDeliveries processes incoming messages from RabbitMQ
379
+ func (r * RabbitMQBroker ) handleDeliveries (ctx context.Context , queue string , deliveries <- chan amqp.Delivery , jobChan chan <- job.Job ) {
380
+ for {
381
+ select {
382
+ case <- ctx .Done ():
383
+ return
384
+ case delivery , ok := <- deliveries :
385
+ if ! ok {
386
+ r .logger .Warnf ("Delivery channel closed for queue %s" , queue )
387
+ return
388
+ }
389
+
390
+ // Convert delivery to job
391
+ job := r .convertDeliveryToJob (delivery , queue )
392
+ if job != nil {
393
+ select {
394
+ case <- ctx .Done ():
395
+ // Put job back on queue
396
+ if err := delivery .Nack (false , true ); err != nil {
397
+ r .logger .Errorf ("Failed to nack job during shutdown: %v" , err )
398
+ }
399
+ return
400
+ case jobChan <- job :
401
+ r .logger .Debugf ("Job sent to workers: %s" , job .GetClass ())
402
+ }
403
+ }
404
+ }
405
+ }
406
+ }
407
+
408
+ // convertDeliveryToJob converts an AMQP delivery to a Job
409
+ func (r * RabbitMQBroker ) convertDeliveryToJob (delivery amqp.Delivery , queue string ) job.Job {
410
+ // Create metadata
411
+ metadata := job.Metadata {
412
+ Queue : queue ,
413
+ EnqueuedAt : delivery .Timestamp ,
414
+ }
415
+
416
+ if delivery .MessageId != "" {
417
+ metadata .ID = delivery .MessageId
418
+ }
419
+
420
+ // Deserialize job
421
+ j , err := r .serializer .Deserialize (delivery .Body , metadata )
422
+ if err != nil {
423
+ // Reject message if we can't deserialize it
424
+ if nackErr := delivery .Nack (false , false ); nackErr != nil {
425
+ r .logError ("Failed to nack message after deserialization error: %v" , nackErr )
426
+ }
427
+ r .logError ("Failed to deserialize job: %v" , err )
428
+ return nil
429
+ }
430
+
431
+ // Wrap in RMQJob for proper ACK/NACK handling
432
+ rmqJob := & RMQJob {
433
+ Job : j ,
434
+ deliveryTag : delivery .DeliveryTag ,
435
+ channel : r .channel ,
436
+ }
437
+
438
+ return rmqJob
439
+ }
440
+
441
+ // SetConsumerQueues sets the queues that this broker will consume from
442
+ func (r * RabbitMQBroker ) SetConsumerQueues (queues []string ) {
443
+ r .consumerQueues = queues
444
+ }
445
+
349
446
// RMQJob wraps a job with RabbitMQ-specific delivery information
350
447
type RMQJob struct {
351
448
job.Job
0 commit comments