forked from gocelery/gocelery
/
rpc_backend.go
203 lines (178 loc) · 4.9 KB
/
rpc_backend.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
// Copyright (c) 2022 Neil Feng
// This file is part of gocelery which is released under MIT license.
// See file LICENSE for full license details.
package gocelery
import (
"encoding/json"
"fmt"
"github.com/patrickmn/go-cache"
log "github.com/sirupsen/logrus"
"time"
"github.com/streadway/amqp"
)
/* RpcCeleryBackend CeleryBackend for AMQP rpc
The difference between amqpbackend and rpcbackend:
amqpbackend => reply to task_id queue
rpcbackend => reply to celery_backend exchange and route the message to the client
side who have waiting for it through the {oid}_result queue which is binding to
celery_backend exchange.
*/
type RpcCeleryBackend struct {
*AMQPSession
Queue *AMQPQueue
Exchange *AMQPExchange
channel <-chan amqp.Delivery
oid string
task2Reply *cache.Cache
Initialized bool
ExpireDuration time.Duration
}
// NewRpcCeleryBackend creates new RpcCeleryBackend
func NewRpcCeleryBackend(host string) *RpcCeleryBackend {
session, err := NewAMQPSession(host)
if err != nil {
panic(err)
}
return NewRpcCeleryBackendByAMQPSession(session)
}
// NewRpcCeleryBackendByConnAndChannel creates new RpcCeleryBackend by AMQP connection and channel
func NewRpcCeleryBackendByAMQPSession(session *AMQPSession) *RpcCeleryBackend {
backend := &RpcCeleryBackend{
AMQPSession: session,
Queue: &AMQPQueue{
Durable: true,
AutoDelete: true,
},
Exchange: &AMQPExchange{
Name: "celery_backend",
Type: "fanout",
Durable: true,
AutoDelete: true,
},
ExpireDuration: 24 * time.Hour,
task2Reply: cache.New(24*time.Hour, 5*time.Minute),
}
return backend
}
func (b *RpcCeleryBackend) Init(oid string) error {
// autodelete is automatically set to true by python
// (406) PRECONDITION_FAILED - inequivalent arg 'durable' for queue 'bc58c0d895c7421eb7cb2b9bbbd8b36f' in vhost '/': received 'true' but current is 'false'
init := func() error {
args := amqp.Table{"x-expires": int32(b.ExpireDuration.Microseconds())}
b.Queue.Name = fmt.Sprintf("%s_%s", oid, "result")
_, err := b.QueueDeclare(
b.Queue.Name, // name
b.Queue.Durable, // durable
b.Queue.AutoDelete, // autoDelete
false, // exclusive
false, // noWait
args, // args
)
if err != nil {
return err
}
err = b.ExchangeDeclare(
b.Exchange.Name,
b.Exchange.Type,
b.Exchange.Durable,
b.Exchange.AutoDelete,
false,
false,
nil,
)
if err != nil {
return err
}
err = b.QueueBind(b.Queue.Name,
"",
b.Exchange.Name,
false,
nil,
)
if err != nil {
return err
}
return nil
}
b.SetupReconnectHooks(init)
if err := init(); err != nil {
return err
}
b.Initialized = true
return b.startConsume()
}
func (b *RpcCeleryBackend) startConsume() error {
// open channel temporarily
reConsume := func() error {
channel, err := b.Consume(b.Queue.Name, "", false, false, false, false, nil)
if err != nil {
return err
}
b.channel = channel
return nil
}
if err := reConsume(); err != nil {
return err
}
b.SetupReconnectHooks(reConsume)
go func() {
for {
select {
case delivery := <-b.GetConsumerChannel():
deliveryAck(delivery)
var resultMessage ResultMessage
if err := json.Unmarshal(delivery.Body, &resultMessage); err != nil {
log.Warn("Error: unserialize result failed.")
continue
}
// use cache map to ensure there no resource leaked after long-time running.
// 'cause all task_result are broadcast to every client queue.
// So there will be lots of task_result which belongs to other client received.
// We must ensure that those task_result which our client wouldn't care can
// be purged after expiration time.
b.task2Reply.SetDefault(delivery.CorrelationId, &resultMessage)
}
}
}()
return nil
}
// GetResult retrieves result from queue named by oid
func (b *RpcCeleryBackend) GetResult(taskID string) (*ResultMessage, error) {
if b.Initialized == false {
return nil, fmt.Errorf("consumming on an unintialized broker is rejected")
}
// First we lookup the task2Reply Map. If we found what we need, remove it from the map, then return the result.
if result, ok := b.task2Reply.Get(taskID); ok {
b.task2Reply.Delete(taskID)
return result.(*ResultMessage), nil
} else {
return nil, ResultNotAvailableYet
}
}
// SetResult sets result back to result exchange
func (b *RpcCeleryBackend) SetResult(taskID string, result *ResultMessage) error {
result.ID = taskID
resBytes, err := json.Marshal(result)
if err != nil {
return err
}
message := amqp.Publishing{
DeliveryMode: amqp.Persistent,
Timestamp: time.Now(),
ContentType: "application/json",
CorrelationId: taskID,
Body: resBytes,
}
return b.Publish(
b.Exchange.Name,
"",
false,
false,
message,
)
}
func (b *RpcCeleryBackend) GetConsumerChannel() <-chan amqp.Delivery {
b.RWLocker.RLock()
defer b.RWLocker.RUnlock()
return b.channel
}