-
Notifications
You must be signed in to change notification settings - Fork 129
/
publisher.go
212 lines (196 loc) · 6.86 KB
/
publisher.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
package scheduler
import (
"fmt"
"strconv"
"sync"
"time"
"github.com/apache/pulsar-client-go/pulsar"
"github.com/gogo/protobuf/proto"
"github.com/google/uuid"
"github.com/pkg/errors"
"github.com/armadaproject/armada/internal/common/armadacontext"
"github.com/armadaproject/armada/internal/common/eventutil"
"github.com/armadaproject/armada/internal/common/logging"
"github.com/armadaproject/armada/internal/common/schedulers"
"github.com/armadaproject/armada/pkg/armadaevents"
)
const (
// This is half the default pulsar BatchingMaxSize
defaultMaxMessageBatchSize = 64 * 1024
explicitPartitionKey = "armada_pulsar_partition"
)
// Publisher is an interface to be implemented by structs that handle publishing messages to pulsar
type Publisher interface {
// PublishMessages will publish the supplied messages. A LeaderToken is provided and the
// implementor may decide whether to publish based on the status of this token
PublishMessages(ctx *armadacontext.Context, events []*armadaevents.EventSequence, shouldPublish func() bool) error
// PublishMarkers publishes a single marker message for each Pulsar partition. Each marker
// massage contains the supplied group id, which allows all marker messages for a given call
// to be identified. The uint32 returned is the number of messages published
PublishMarkers(ctx *armadacontext.Context, groupId uuid.UUID) (uint32, error)
}
// PulsarPublisher is the default implementation of Publisher
type PulsarPublisher struct {
// Used to send messages to pulsar
producer pulsar.Producer
// Number of partitions on the pulsar topic
numPartitions int
// Timeout after which async messages sends will be considered failed
pulsarSendTimeout time.Duration
// Maximum size (in bytes) of produced pulsar messages.
// This must be below 4MB which is the pulsar message size limit
maxMessageBatchSize uint
}
func NewPulsarPublisher(
pulsarClient pulsar.Client,
producerOptions pulsar.ProducerOptions,
pulsarSendTimeout time.Duration,
) (*PulsarPublisher, error) {
partitions, err := pulsarClient.TopicPartitions(producerOptions.Topic)
if err != nil {
return nil, errors.WithStack(err)
}
producerOptions.MessageRouter = createMessageRouter(producerOptions)
producer, err := pulsarClient.CreateProducer(producerOptions)
if err != nil {
return nil, errors.WithStack(err)
}
maxMessageBatchSize := producerOptions.BatchingMaxSize / 2
if maxMessageBatchSize <= 0 {
maxMessageBatchSize = defaultMaxMessageBatchSize
}
return &PulsarPublisher{
producer: producer,
pulsarSendTimeout: pulsarSendTimeout,
maxMessageBatchSize: maxMessageBatchSize,
numPartitions: len(partitions),
}, nil
}
// PublishMessages publishes all event sequences to pulsar. Event sequences for a given jobset will be combined into
// single event sequences up to maxMessageBatchSize.
func (p *PulsarPublisher) PublishMessages(ctx *armadacontext.Context, events []*armadaevents.EventSequence, shouldPublish func() bool) error {
sequences := eventutil.CompactEventSequences(events)
sequences, err := eventutil.LimitSequencesByteSize(sequences, p.maxMessageBatchSize, true)
if err != nil {
return err
}
msgs := make([]*pulsar.ProducerMessage, len(sequences))
for i, sequence := range sequences {
bytes, err := proto.Marshal(sequence)
if err != nil {
return err
}
msgs[i] = &pulsar.ProducerMessage{
Payload: bytes,
Key: sequences[i].JobSetName,
Properties: map[string]string{
schedulers.PropertyName: schedulers.PulsarSchedulerAttribute,
},
}
}
wg := sync.WaitGroup{}
wg.Add(len(msgs))
// Send messages
if shouldPublish() {
ctx.Debugf("Am leader so will publish")
sendCtx, cancel := armadacontext.WithTimeout(ctx, p.pulsarSendTimeout)
errored := false
for _, msg := range msgs {
p.producer.SendAsync(sendCtx, msg, func(_ pulsar.MessageID, _ *pulsar.ProducerMessage, err error) {
if err != nil {
logging.
WithStacktrace(ctx, err).
Error("error sending message to Pulsar")
errored = true
}
wg.Done()
})
}
wg.Wait()
cancel()
if errored {
return errors.New("One or more messages failed to send to Pulsar")
}
} else {
ctx.Debugf("No longer leader so not publishing")
}
return nil
}
// PublishMarkers sends one pulsar message (containing an armadaevents.PartitionMarker) to each partition
// of the producer's Pulsar topic.
func (p *PulsarPublisher) PublishMarkers(ctx *armadacontext.Context, groupId uuid.UUID) (uint32, error) {
for i := 0; i < p.numPartitions; i++ {
pm := &armadaevents.PartitionMarker{
GroupId: armadaevents.ProtoUuidFromUuid(groupId),
Partition: uint32(i),
}
es := &armadaevents.EventSequence{
Queue: "armada-scheduler",
JobSetName: "armada-scheduler",
Events: []*armadaevents.EventSequence_Event{
{
Created: now(),
Event: &armadaevents.EventSequence_Event_PartitionMarker{
PartitionMarker: pm,
},
},
},
}
bytes, err := proto.Marshal(es)
if err != nil {
return 0, err
}
msg := &pulsar.ProducerMessage{
Properties: map[string]string{
explicitPartitionKey: fmt.Sprintf("%d", i),
schedulers.PropertyName: schedulers.PulsarSchedulerAttribute,
},
Payload: bytes,
}
// use a synchronous send here as the logic is simpler.
// We send relatively few position markers so the performance penalty shouldn't be meaningful
_, err = p.producer.Send(ctx, msg)
if err != nil {
return 0, err
}
}
return uint32(p.numPartitions), nil
}
// createMessageRouter returns a custom Pulsar message router that routes the message to the partition given by the
// explicitPartitionKey msg property. If this property isn't present then it will fall back to the default Pulsar
// message routing logic
func createMessageRouter(options pulsar.ProducerOptions) func(*pulsar.ProducerMessage, pulsar.TopicMetadata) int {
defaultRouter := pulsar.NewDefaultRouter(
JavaStringHash,
options.BatchingMaxMessages,
options.BatchingMaxSize,
options.BatchingMaxPublishDelay,
options.DisableBatching)
return func(msg *pulsar.ProducerMessage, md pulsar.TopicMetadata) int {
explicitPartition, ok := msg.Properties[explicitPartitionKey]
if ok {
partition, err := strconv.ParseInt(explicitPartition, 10, 32)
if err != nil {
panic(errors.Errorf("cannot parse %s as int", explicitPartition))
}
if partition < 0 || uint32(partition) >= md.NumPartitions() {
panic(errors.Errorf("requested partiton %d is not in the range 0-%d", partition, md.NumPartitions()-1))
}
return int(partition)
}
return defaultRouter(msg, md.NumPartitions())
}
}
func now() *time.Time {
t := time.Now()
return &t
}
// JavaStringHash is the default hashing algorithm used by Pulsar
// copied from https://github.com/apache/pulsar-client-go/blob/master/pulsar/internal/hash.go
func JavaStringHash(s string) uint32 {
var h uint32
for i, size := 0, len(s); i < size; i++ {
h = 31*h + uint32(s[i])
}
return h
}