-
Notifications
You must be signed in to change notification settings - Fork 7
/
consumer.go
785 lines (707 loc) · 24.3 KB
/
consumer.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
package pgq
import (
"bytes"
"context"
"database/sql"
"encoding/json"
"fmt"
"io"
"log/slog"
"math"
"strings"
"sync"
"time"
"unicode"
"github.com/google/uuid"
"github.com/jackc/pgtype"
"github.com/jmoiron/sqlx"
"github.com/pkg/errors"
"go.opentelemetry.io/otel"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/codes"
"go.opentelemetry.io/otel/metric"
"go.opentelemetry.io/otel/metric/noop"
"go.opentelemetry.io/otel/propagation"
"golang.org/x/sync/semaphore"
"go.dataddo.com/pgq/internal/pg"
"go.dataddo.com/pgq/internal/query"
)
type fatalError struct {
Err error
}
func (e fatalError) Error() string {
return fmt.Sprintf("pgq consumer fatal error: %v", e.Err)
}
func (e fatalError) Unwrap() error { return e.Err }
const (
// MessageProcessed signals that message was processed and shouldn't be processed
// again. If processed with an error, it is expected permanent, and new run would
// result in the same error.
MessageProcessed = true
// MessageNotProcessed signals that message wasn't processed and can be processed
// again. The error interrupting processing is considered temporary.
MessageNotProcessed = false
)
// MessageHandler handles message received from queue. Returning false means
// message wasn't processed correctly and shouldn't be acknowledged. Error
// contains additional information.
//
// Possible combinations:
//
// // | processed | err | description |
// // | --------- | ---------- | ---------------------------------------------------- |
// // | false | <nil> | missing failure info, but the message can be retried |
// // | false | some error | not processed, because of some error, can be retried |
// // | true | <nil> | processed, no error. |
// // | true | some error | processed, ended with error. Don't retry! |
type MessageHandler interface {
HandleMessage(context.Context, *MessageIncoming) (processed bool, err error)
}
// MessageHandlerFunc is MessageHandler implementation by simple function.
type MessageHandlerFunc func(context.Context, *MessageIncoming) (processed bool, err error)
// HandleMessage calls self. It also implements MessageHandler interface.
func (fn MessageHandlerFunc) HandleMessage(ctx context.Context, msg *MessageIncoming) (processed bool, err error) {
return fn(ctx, msg)
}
// consumerConfig contains consumer configuration.
type consumerConfig struct {
// LockDuration is the maximal duration for how long the message remains locked for other consumers.
LockDuration time.Duration
// PollingInterval defines how frequently consumer checks the queue for new messages.
PollingInterval time.Duration
// AckTimeout is the timeout for updating the message status when message is processed.
AckTimeout time.Duration
// MaxParallelMessages sets how many jobs can single consumer process simultaneously.
MaxParallelMessages int
// Metrics define prometheus parameters.
Metrics metric.Meter
// InvalidMessageCallback defines what should happen to messages which are identified as invalid.
// Such messages usually have missing or malformed required fields.
InvalidMessageCallback InvalidMessageCallback
// HistoryLimit means how long in history you want to search for unprocessed messages
// If not set, it will look for message in the whole table.
// You may set this value when using partitioned table to search just in partitions you are interested in
HistoryLimit time.Duration
// MaxConsumeCount is the maximal number of times a message can be consumed before it is ignored.
// This is a safety mechanism to prevent failure infinite loops when a message causes unhandled panic, OOM etc.
MaxConsumeCount uint
// MsgProcessingReserveDuration is the duration for which the message is reserved for handling result state.
MessageProcessingReserveDuration time.Duration
MetadataFilters []MetadataFilter
Logger *slog.Logger
}
var noopLogger = slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.Level(math.MaxInt)}))
var defaultConsumerConfig = consumerConfig{
LockDuration: time.Hour,
PollingInterval: 5 * time.Second,
AckTimeout: 1 * time.Second,
MessageProcessingReserveDuration: 1 * time.Second,
MaxParallelMessages: 1,
InvalidMessageCallback: func(context.Context, InvalidMessage, error) {},
Metrics: noop.Meter{},
MaxConsumeCount: 3,
Logger: noopLogger,
}
// InvalidMessageCallback defines what should happen to messages which are identified as invalid.
// Such messages usually have missing or malformed required fields.
type InvalidMessageCallback func(ctx context.Context, msg InvalidMessage, err error)
// InvalidMessage is definition of invalid message used as argument for InvalidMessageCallback by Consumer.
type InvalidMessage struct {
ID string
Metadata json.RawMessage
Payload json.RawMessage
}
// Consumer is the preconfigured subscriber of the write input messages
type Consumer struct {
db *sqlx.DB
queueName string
cfg consumerConfig
handler MessageHandler
metrics *metrics
sem *semaphore.Weighted
}
// ConsumerOption applies option to consumerConfig.
type ConsumerOption func(c *consumerConfig)
// WithLockDuration sets the maximal duration for how long the message remains
// locked for other consumers.
func WithLockDuration(d time.Duration) ConsumerOption {
return func(c *consumerConfig) {
c.LockDuration = d
}
}
// WithPollingInterval sets how frequently consumer checks the queue for new
// messages.
func WithPollingInterval(d time.Duration) ConsumerOption {
return func(c *consumerConfig) {
c.PollingInterval = d
}
}
// WithAckTimeout sets the timeout for updating the message status when message
// is processed.
func WithAckTimeout(d time.Duration) ConsumerOption {
return func(c *consumerConfig) {
c.AckTimeout = d
}
}
// WithMaxParallelMessages sets how many jobs can single consumer process
// simultaneously.
func WithMaxParallelMessages(n int) ConsumerOption {
return func(c *consumerConfig) {
c.MaxParallelMessages = n
}
}
// WithMetrics sets metrics meter. Default is noop.Meter{}.
func WithMetrics(m metric.Meter) ConsumerOption {
return func(c *consumerConfig) {
c.Metrics = m
}
}
// WithMessageProcessingReserveDuration sets the duration for which the message is reserved for handling result state.
func WithMessageProcessingReserveDuration(d time.Duration) ConsumerOption {
return func(c *consumerConfig) {
c.MessageProcessingReserveDuration = d
}
}
// WithInvalidMessageCallback sets callback for invalid messages.
func WithInvalidMessageCallback(fn InvalidMessageCallback) ConsumerOption {
return func(c *consumerConfig) {
c.InvalidMessageCallback = fn
}
}
// WithHistoryLimit sets how long in history you want to search for unprocessed
// messages (default is no limit). If not set, it will look for message in the
// whole table. You may set this value when using partitioned table to search
// just in partitions you are interested in.
func WithHistoryLimit(d time.Duration) ConsumerOption {
return func(c *consumerConfig) {
c.HistoryLimit = d
}
}
// WithMaxConsumeCount sets the maximal number of times a message can be consumed before it is ignored.
// Unhandled SIGKILL or uncaught panic, OOM error etc. could lead to consumer failure infinite loop.
// Setting this value to greater than 0 will prevent happening this loop.
// Setting this to value 0 disableS this safe mechanism.
func WithMaxConsumeCount(max uint) ConsumerOption {
return func(c *consumerConfig) {
c.MaxConsumeCount = max
}
}
// WithLogger sets logger. Default is no logging.
func WithLogger(logger *slog.Logger) ConsumerOption {
return func(c *consumerConfig) {
c.Logger = logger
}
}
// MetadataFilter is a filter for metadata. Right now support only direct matching of key/value
type (
MetadataOperation string
MetadataFilter struct {
Key string
Operation MetadataOperation
Value string
}
)
const (
OpEqual MetadataOperation = "="
OpNotEqual MetadataOperation = "<>"
)
func WithMetadataFilter(filter *MetadataFilter) ConsumerOption {
return func(c *consumerConfig) {
filters := c.MetadataFilters
if filter == nil {
filters = make([]MetadataFilter, 0, 1)
}
c.MetadataFilters = append(filters, *filter)
}
}
// NewConsumer creates Consumer with proper settings
func NewConsumer(db *sql.DB, queueName string, handler MessageHandler, opts ...ConsumerOption) (*Consumer, error) {
return NewConsumerExt(sqlx.NewDb(db, "pgx"), queueName, handler, opts...)
}
// NewConsumer creates Consumer with proper settings, using sqlx.DB (until refactored to use pgx directly)
func NewConsumerExt(db *sqlx.DB, queueName string, handler MessageHandler, opts ...ConsumerOption) (*Consumer, error) {
config := defaultConsumerConfig
for _, opt := range opts {
opt(&config)
}
metrics, err := prepareProcessMetric(queueName, config.Metrics)
if err != nil {
return nil, errors.Wrap(err, "registering metrics")
}
sem := semaphore.NewWeighted(int64(config.MaxParallelMessages))
return &Consumer{
db: db,
queueName: queueName,
cfg: config,
handler: handler,
metrics: metrics,
sem: sem,
}, nil
}
type metrics struct {
jobsCounter metric.Int64Counter
failedProcessingCounter metric.Int64Counter
}
func prepareProcessMetric(queueName string, meter metric.Meter) (*metrics, error) {
queueName = strings.ReplaceAll(queueName, "/", "_")
// '_total' suffix is added to all counters by default by OpenTelemetry.
jobsCounter, err := meter.Int64Counter(
fmt.Sprintf("pgq_%s_processed_jobs", queueName),
metric.WithDescription("Total number of processed jobs. The label 'resolution' says how the job was handled."),
)
if err != nil {
return nil, errors.WithStack(err)
}
failedProcessingCounter, err := meter.Int64Counter(
fmt.Sprintf("pgq_%s_failed_job_processing", queueName),
metric.WithDescription("Total number of errors during marking a job as processed. Example is a failed job ACK. This metric signals a chance of inconsistencies in the queue."),
)
if err != nil {
return nil, errors.WithStack(err)
}
return &metrics{
jobsCounter: jobsCounter,
failedProcessingCounter: failedProcessingCounter,
}, nil
}
// Run consumes messages until the context isn't cancelled.
func (c *Consumer) Run(ctx context.Context) error {
c.cfg.Logger.InfoContext(ctx, "starting consumption...",
"inputQueue", c.queueName,
)
if err := c.verifyTable(ctx); err != nil {
return errors.Wrap(err, "verifying table")
}
query, err := c.generateQuery()
if err != nil {
return errors.Wrap(err, "generating query")
}
var wg sync.WaitGroup
defer wg.Wait() // wait for handlers to finish
for {
msgs, err := c.consumeMessages(ctx, query)
if err != nil {
if errors.As(err, &fatalError{}) {
return errors.Wrapf(err, "consuming from PostgreSQL queue %s", c.queueName)
}
c.cfg.Logger.InfoContext(ctx, "pgq: consume failed, will retry",
"error", err,
)
continue
}
wg.Add(len(msgs))
for _, msg := range msgs {
go func(msg *MessageIncoming) {
defer wg.Done()
defer c.sem.Release(1)
c.handleMessage(ctx, msg)
}(msg)
}
}
}
func (c *Consumer) verifyTable(ctx context.Context) error {
// --- (1) ----
// Validate the queue mandatory fields
err := ValidateFields(ctx, c.db, c.queueName)
if err != nil {
return errors.Wrap(err, "error validating queue mandatory fields")
}
// --- (2) ----
// Validate the queue mandatory indexes
err = ValidateIndexes(ctx, c.db, c.queueName)
if err != nil {
return errors.Wrap(err, "error validating queue mandatory indexes")
}
return nil
}
func (c *Consumer) generateQuery() (*query.Builder, error) {
qb := query.NewBuilder()
qb.WriteString(`UPDATE `)
qb.WriteString(pg.QuoteIdentifier(c.queueName))
qb.WriteString(` SET locked_until = :locked_until`)
qb.WriteString(`, started_at = CURRENT_TIMESTAMP`)
qb.WriteString(`, consumed_count = consumed_count+1`)
qb.WriteString(` WHERE id IN (`)
{
qb.WriteString(`SELECT id FROM `)
qb.WriteString(pg.QuoteIdentifier(c.queueName))
qb.WriteString(` WHERE`)
if c.cfg.HistoryLimit > 0 {
qb.WriteString(` created_at >= CURRENT_TIMESTAMP - CAST((:history_limit) AS interval) AND`)
qb.WriteString(` created_at < CURRENT_TIMESTAMP AND`)
}
qb.WriteString(` (locked_until IS NULL OR locked_until < CURRENT_TIMESTAMP)`)
if c.cfg.MaxConsumeCount > 0 {
qb.WriteString(` AND consumed_count < :max_consume_count`)
}
for i, filter := range c.cfg.MetadataFilters {
if len(filter.Operation) == 0 {
return nil, fatalError{Err: fmt.Errorf("metadata filter operation is empty")}
}
qb.WriteString(fmt.Sprintf(" AND metadata->>:metadata_key_%d %s :metadata_value_%d", i, filter.Operation, i))
}
qb.WriteString(` AND processed_at IS NULL`)
qb.WriteString(` AND (scheduled_for IS NULL OR scheduled_for < CURRENT_TIMESTAMP)`)
// prioritize scheduled messages and messages with lower consumed_count
qb.WriteString(` ORDER BY scheduled_for ASC NULLS LAST, consumed_count ASC, created_at ASC`)
qb.WriteString(` LIMIT :limit`)
qb.WriteString(` FOR UPDATE SKIP LOCKED`)
}
qb.WriteString(`) RETURNING id, payload, metadata, consumed_count, locked_until`)
return qb, nil
}
func (c *Consumer) handleMessage(ctx context.Context, msg *MessageIncoming) {
ctx, cancel := context.WithDeadline(ctx, msg.Deadline)
defer cancel()
ctxTimeout, cancel := prepareCtxTimeout()
defer cancel()
// TODO configurable Propagator
propagator := otel.GetTextMapPropagator()
carrier := propagation.MapCarrier(msg.Metadata)
ctx = propagator.Extract(ctx, carrier)
ctx, span := otel.Tracer("pgq").Start(ctx, "HandleMessage")
defer span.End()
span.SetAttributes(
attribute.String("messageId", msg.id.String()),
attribute.String("queueName", c.queueName),
)
processed, err := c.handler.HandleMessage(ctx, msg)
if !processed {
reason := "unknown"
if err != nil {
span.RecordError(err)
reason = err.Error()
}
span.SetStatus(codes.Ok, "Message Nacked")
if err := msg.nack(ctxTimeout(c.cfg.AckTimeout), reason); err != nil {
c.cfg.Logger.ErrorContext(ctx, "pgq: nack failed",
"error", err.Error(),
"ackTimeout", c.cfg.AckTimeout,
"reason", reason,
"msg.metadata", msg.Metadata,
)
}
return
}
if err != nil {
span.RecordError(err)
span.SetStatus(codes.Error, "Message Discarded")
discardReason := err.Error()
if err := msg.discard(ctxTimeout(c.cfg.AckTimeout), discardReason); err != nil {
c.cfg.Logger.ErrorContext(ctx, "pgq: discard failed",
"error", err,
"ackTimeout", c.cfg.AckTimeout,
"reason", discardReason,
"msg.metadata", msg.Metadata,
)
}
return
}
span.SetStatus(codes.Ok, "Message Acked")
if err := msg.ack(ctxTimeout(c.cfg.AckTimeout)); err != nil {
c.cfg.Logger.ErrorContext(ctx, "pgq: ack failed",
"error", err,
"ackTimeout", c.cfg.AckTimeout,
"msg.metadata", msg.Metadata,
)
}
}
func prepareCtxTimeout() (func(td time.Duration) context.Context, context.CancelFunc) {
parent, cancel := context.WithCancel(context.Background())
fn := func(td time.Duration) context.Context {
// ctx will be released by parent cancellation
ctx, _ := context.WithTimeout(parent, td) //nolint:govet
return ctx
}
return fn, cancel
}
func (c *Consumer) consumeMessages(ctx context.Context, query *query.Builder) ([]*MessageIncoming, error) {
for {
maxMsg, err := acquireMaxFromSemaphore(ctx, c.sem, int64(c.cfg.MaxParallelMessages))
if err != nil {
return nil, fatalError{Err: errors.WithStack(err)}
}
msgs, err := c.tryConsumeMessages(ctx, query, maxMsg)
if err != nil {
c.sem.Release(maxMsg)
if !errors.Is(err, sql.ErrNoRows) {
return nil, errors.WithStack(err)
}
select {
case <-ctx.Done():
return nil, fatalError{Err: ctx.Err()}
case <-time.After(c.cfg.PollingInterval):
continue
}
}
// release unused resources
c.sem.Release(maxMsg - int64(len(msgs)))
return msgs, nil
}
}
type pgMessage struct {
ID pgtype.UUID
Payload pgtype.JSONB
Metadata pgtype.JSONB
Attempt pgtype.Int4
LockedUntil pgtype.Timestamptz
}
func (c *Consumer) tryConsumeMessages(ctx context.Context, query *query.Builder, limit int64) (_ []*MessageIncoming, err error) {
tx, err := c.db.BeginTxx(ctx, nil)
if err != nil {
// TODO not necessary fatal, network could wiggle.
return nil, fatalError{Err: errors.WithStack(err)}
}
defer func() {
txRollbackErr := tx.Rollback()
if errors.Is(txRollbackErr, sql.ErrTxDone) {
return
}
if txRollbackErr != nil {
c.cfg.Logger.ErrorContext(ctx, "pgq: rollback failed",
"error", txRollbackErr.Error(),
"rollbackReason", err,
)
return
}
}()
lockedUntil := time.Now().Add(c.cfg.LockDuration)
namedParams := map[string]interface{}{
"locked_until": lockedUntil,
"limit": limit,
}
if query.HasParam("history_limit") {
var scanInterval pgtype.Interval
// time.Duration doesn't ever fail
_ = scanInterval.Set(c.cfg.HistoryLimit)
namedParams["history_limit"] = scanInterval
}
if query.HasParam("max_consume_count") {
namedParams["max_consume_count"] = c.cfg.MaxConsumeCount
}
for i, filter := range c.cfg.MetadataFilters {
namedParams[fmt.Sprintf("metadata_key_%d", i)] = filter.Key
namedParams[fmt.Sprintf("metadata_value_%d", i)] = filter.Value
}
queryString, err := query.Build(namedParams)
if err != nil {
return nil, errors.WithStack(err)
}
rows, err := sqlx.NamedQueryContext(ctx, tx, queryString, namedParams)
if err != nil {
if isErrorCode(err, undefinedColumnErrCode) {
return nil, fatalError{Err: err}
}
return nil, errors.WithStack(err)
}
defer rows.Close()
var msgs []*MessageIncoming
for rows.Next() {
msg, err := c.parseRow(ctx, rows)
if err != nil {
return nil, errors.WithStack(err)
}
msgs = append(msgs, msg)
}
if err := rows.Err(); err != nil {
return nil, errors.WithStack(err)
}
if len(msgs) == 0 {
return nil, sql.ErrNoRows
}
if err := tx.Commit(); err != nil {
return nil, errors.Wrap(err, "commit message consumption")
}
return msgs, nil
}
func (c *Consumer) parseRow(ctx context.Context, rows *sqlx.Rows) (*MessageIncoming, error) {
var pgMsg pgMessage
if err := rows.Scan(
&pgMsg.ID,
&pgMsg.Payload,
&pgMsg.Metadata,
&pgMsg.Attempt,
&pgMsg.LockedUntil,
); err != nil {
if isErrorCode(err, undefinedTableErrCode, undefinedColumnErrCode) {
return nil, fatalError{Err: err}
}
return nil, errors.Wrap(err, "retrieving message")
}
msg, err := c.finishParsing(pgMsg)
if err != nil {
c.cfg.Logger.ErrorContext(ctx, "reading message", c.logFields(pgMsg, err)...)
c.discardInvalidMsg(ctx, pgMsg.ID, err)
go c.cfg.InvalidMessageCallback(ctx, InvalidMessage{
ID: uuid.UUID(pgMsg.ID.Bytes).String(),
Metadata: pgMsg.Metadata.Bytes,
Payload: pgMsg.Payload.Bytes,
}, err)
return nil, errors.WithStack(err)
}
return msg, nil
}
func (c *Consumer) logFields(msg pgMessage, err error) []any {
entry := []any{
"msg.id", uuid.UUID(msg.ID.Bytes).String(),
"msg.metadata", json.RawMessage(msg.Metadata.Bytes),
}
if err != nil {
entry = append(entry, []any{
"error", err,
"msg.payload", json.RawMessage(msg.Payload.Bytes),
})
}
return entry
}
func (c *Consumer) discardInvalidMsg(ctx context.Context, id pgtype.UUID, err error) {
ctxTimeout, cancel := prepareCtxTimeout()
defer cancel()
reason := err.Error()
if err := c.discardMessage(c.db, id)(ctxTimeout(c.cfg.AckTimeout), reason); err != nil {
c.cfg.Logger.ErrorContext(ctx, "pgq: discard failed",
"error", err,
"msg.id", id,
"ackTimeout", c.cfg.AckTimeout,
"reason", reason,
)
return
}
}
func (c *Consumer) finishParsing(pgMsg pgMessage) (*MessageIncoming, error) {
msg := &MessageIncoming{
id: uuid.UUID(pgMsg.ID.Bytes),
once: sync.Once{},
ackFn: c.ackMessage(c.db, pgMsg.ID),
nackFn: c.nackMessage(c.db, pgMsg.ID),
discardFn: c.discardMessage(c.db, pgMsg.ID),
}
var err error
msg.Payload, err = parsePayload(pgMsg)
if err != nil {
return msg, errors.Wrap(err, "parsing payload")
}
msg.Metadata, err = parseMetadata(pgMsg)
if err != nil {
return msg, errors.Wrap(err, "parsing metadata")
}
msg.Attempt = int(pgMsg.Attempt.Int)
msg.maxConsumedCount = c.cfg.MaxConsumeCount
msg.Deadline = pgMsg.LockedUntil.Time.Add(-c.cfg.AckTimeout).Add(-c.cfg.MessageProcessingReserveDuration)
return msg, nil
}
func parsePayload(pgMsg pgMessage) (json.RawMessage, error) {
if pgMsg.Payload.Status != pgtype.Present {
return nil, errors.New("missing message payload")
}
if !isJSONObject(pgMsg.Payload.Bytes) {
return nil, errors.New("payload is invalid JSON object")
}
return pgMsg.Payload.Bytes, nil
}
func parseMetadata(pgMsg pgMessage) (map[string]string, error) {
if pgMsg.Metadata.Status != pgtype.Present {
return map[string]string{}, nil
}
var m map[string]string
if err := json.Unmarshal(pgMsg.Metadata.Bytes, &m); err != nil {
if !isJSONObject(pgMsg.Metadata.Bytes) {
return nil, errors.New("metadata is invalid JSON object")
}
return nil, errors.Wrap(err, "parsing metadata")
}
return m, nil
}
func isJSONObject(b json.RawMessage) bool {
if !json.Valid(b) {
return false
}
// remove insignificant characters.
b = bytes.TrimLeftFunc(b, unicode.IsSpace)
return bytes.HasPrefix(b, []byte{'{'})
}
type execer interface {
ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error)
}
func (c *Consumer) ackMessage(exec execer, msgID pgtype.UUID) func(ctx context.Context) error {
query := `UPDATE ` + pg.QuoteIdentifier(c.queueName) + ` SET locked_until = NULL, processed_at = CURRENT_TIMESTAMP WHERE id = $1`
return func(ctx context.Context) error {
if _, err := exec.ExecContext(ctx, query, msgID); err != nil {
c.metrics.failedProcessingCounter.Add(ctx, 1,
metric.WithAttributes(
attribute.String("resolution", "ack"),
attribute.String("queue_name", c.queueName),
),
)
return errors.WithStack(err)
}
c.metrics.jobsCounter.Add(ctx, 1,
metric.WithAttributes(
attribute.String("resolution", "ack"),
attribute.String("queue_name", c.queueName),
),
)
return nil
}
}
func (c *Consumer) nackMessage(exec execer, msgID pgtype.UUID) func(ctx context.Context, reason string) error {
query := `UPDATE ` + pg.QuoteIdentifier(c.queueName) + ` SET locked_until = NULL, error_detail = $2 WHERE id = $1`
return func(ctx context.Context, reason string) error {
if _, err := exec.ExecContext(ctx, query, msgID, reason); err != nil {
c.metrics.failedProcessingCounter.Add(ctx, 1,
metric.WithAttributes(
attribute.String("resolution", "nack"),
attribute.String("queue_name", c.queueName),
),
)
return errors.WithStack(err)
}
c.metrics.jobsCounter.Add(ctx, 1,
metric.WithAttributes(
attribute.String("resolution", "nack"),
attribute.String("queue_name", c.queueName),
),
)
return nil
}
}
func (c *Consumer) discardMessage(exec execer, msgID pgtype.UUID) func(ctx context.Context, reason string) error {
query := `UPDATE ` + pg.QuoteIdentifier(c.queueName) + ` SET locked_until = NULL, processed_at = CURRENT_TIMESTAMP, error_detail = $2 WHERE id = $1`
return func(ctx context.Context, reason string) error {
if _, err := exec.ExecContext(ctx, query, msgID, reason); err != nil {
c.metrics.failedProcessingCounter.Add(ctx, 1,
metric.WithAttributes(
attribute.String("resolution", "discard"),
attribute.String("queue_name", c.queueName),
),
)
return errors.WithStack(err)
}
c.metrics.jobsCounter.Add(ctx, 1,
metric.WithAttributes(
attribute.String("resolution", "discard"),
attribute.String("queue_name", c.queueName),
),
)
return nil
}
}
// acquireMaxFromSemaphore acquires maximum possible weight. It blocks until resources are
// available or ctx is done. On success, returns acquired weight. On failure,
// returns ctx.Err() and leaves the semaphore unchanged.
//
// If ctx is already done, Acquire may still succeed without blocking.
func acquireMaxFromSemaphore(ctx context.Context, w *semaphore.Weighted, size int64) (int64, error) {
for i := size; i > 1; i-- {
if ok := w.TryAcquire(i); ok {
// Same practise like in underlying library.
// Acquired the semaphore after we were canceled. Rather than trying to
// fix up the queue, just pretend we didn't notice the cancellation.
return i, nil
}
}
if err := w.Acquire(ctx, 1); err != nil {
return 0, err
}
return 1, nil
}