/
notifier.go
94 lines (82 loc) · 2.47 KB
/
notifier.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
package notification
import (
"context"
"github.com/apache/pulsar-client-go/pulsar"
"github.com/chroma-core/chroma/go/pkg/model"
"github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb"
"github.com/pingcap/log"
"go.uber.org/zap"
"google.golang.org/protobuf/proto"
)
type Notifier interface {
Notify(ctx context.Context, notifications []model.Notification) error
}
type PulsarNotifier struct {
producer pulsar.Producer
}
var _ Notifier = &PulsarNotifier{}
func NewPulsarNotifier(producer pulsar.Producer) *PulsarNotifier {
return &PulsarNotifier{
producer: producer,
}
}
func (p *PulsarNotifier) Notify(ctx context.Context, notifications []model.Notification) error {
for _, notification := range notifications {
notificationPb := coordinatorpb.Notification{
CollectionId: notification.CollectionID,
Type: notification.Type,
Status: notification.Status,
}
payload, err := proto.Marshal(¬ificationPb)
if err != nil {
log.Error("Failed to marshal notification", zap.Error(err))
return err
}
message := &pulsar.ProducerMessage{
Key: notification.CollectionID,
Payload: payload,
}
// Since the number of notifications is small, we can send them synchronously
// for now. This is easy to reason about hte order of notifications.
//
// As follow up optimizations, we can send them asynchronously in batches and
// track failed messages.
_, err = p.producer.Send(ctx, message)
if err != nil {
log.Error("Failed to send message", zap.Error(err))
return err
}
log.Info("Published message", zap.Any("message", message))
}
return nil
}
type MemoryNotifier struct {
queue []pulsar.ProducerMessage
}
var _ Notifier = &MemoryNotifier{}
func NewMemoryNotifier() *MemoryNotifier {
return &MemoryNotifier{
queue: make([]pulsar.ProducerMessage, 0),
}
}
func (m *MemoryNotifier) Notify(ctx context.Context, notifications []model.Notification) error {
for _, notification := range notifications {
notificationPb := coordinatorpb.Notification{
CollectionId: notification.CollectionID,
Type: notification.Type,
Status: notification.Status,
}
payload, err := proto.Marshal(¬ificationPb)
if err != nil {
log.Error("Failed to marshal notification", zap.Error(err))
return err
}
message := pulsar.ProducerMessage{
Key: notification.CollectionID,
Payload: payload,
}
m.queue = append(m.queue, message)
log.Info("Published message", zap.Any("message", message))
}
return nil
}