-
Notifications
You must be signed in to change notification settings - Fork 0
/
delayed_queue.go
146 lines (123 loc) · 3.25 KB
/
delayed_queue.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
// Copyright 2021 Zubin. All rights reserved.
package delayedq
import (
"time"
"github.com/marmotedu/log"
"github.com/pkg/errors"
"github.com/streadway/amqp"
)
const (
DEFAULT_PREFETCH_COUNT = 1
DELAYED_EXCHANGE_TYPE = "x-delayed-message"
DELAYED_TYPE = "direct"
)
type DelayedQueue struct {
connection *amqp.Connection
channel *amqp.Channel
url string
exchange string
queue string
key string
prefetchCount int
closeChan chan *amqp.Error
quitChan chan bool
}
func getDefaultTaskQueue() DelayedQueue {
return DelayedQueue{
exchange: "delayed_queue_exchange",
queue: "delayed_queue_queue",
key: "delayed_queue_key",
prefetchCount: DEFAULT_PREFETCH_COUNT,
quitChan: make(chan bool),
}
}
func (dq *DelayedQueue) connect() error {
var err error
dq.connection, err = amqp.Dial(dq.url)
if err != nil {
return errors.Wrap(err, "Failed to connect to RabbitMQ")
}
dq.channel, err = dq.connection.Channel()
if err != nil {
return errors.Wrap(err, "Failed to open a channel")
}
log.Debug("Connect to rabbitMQ established")
dq.closeChan = make(chan *amqp.Error)
dq.connection.NotifyClose(dq.closeChan)
return nil
}
// handleDisconnect handle a disconnection trying to reconnect every 5s.
func (dq *DelayedQueue) handleDisconnect() {
for {
select {
case errChan := <-dq.closeChan:
if errChan != nil {
log.Errorf("RabbitMQ disconnection: %v", errChan)
}
case <-dq.quitChan:
dq.connection.Close()
dq.channel.Close()
dq.connection = nil
dq.channel = nil
log.Debug("RabbitMQ has been shut down...")
dq.quitChan <- true
return
}
log.Info("Trying to reconnect to rabbitMQ...")
time.Sleep(5 * time.Second)
if err := dq.connect(); err != nil {
log.Errorf("RabbitMQ connect error: %v", err)
}
log.Info("Connect to rabbitMQ established")
if err := dq.init(); err != nil {
log.Errorf("RabbitMQ init error: %v", err)
}
}
}
// declare exchange and queue if not exist
func (dq *DelayedQueue) init() (err error) {
args := make(amqp.Table)
args["x-delayed-type"] = DELAYED_TYPE
err = dq.channel.ExchangeDeclare(
dq.exchange, // name
DELAYED_EXCHANGE_TYPE, // type
true, // durable
false, // auto-deleted
false, // internal
false, // noWait
args, // arguments
)
if err != nil {
return errors.Wrap(err, "Failed to declare a exchange")
}
_, err = dq.channel.QueueDeclare(
dq.queue, // name
false, // durable
false, // delete when unused
false, // exclusive
false, // no-wait
nil, // arguments
)
if err != nil {
return errors.Wrap(err, "Failed to declare a queue")
}
err = dq.channel.QueueBind(dq.queue, dq.key, dq.exchange, false, nil)
if err != nil {
return errors.Wrap(err, "Failed to bind queue")
}
err = dq.channel.Qos(
dq.prefetchCount, // prefetch count
0, // prefetch size
false, // global
)
if err != nil {
return errors.Wrap(err, "Failed to set QoS")
}
return nil
}
// Disconnect the channel and connection
func (dq *DelayedQueue) Disconnect() {
dq.quitChan <- true
log.Debug("shutting down rabbitMQ's connection...")
<-dq.quitChan
}