Hugging Face's logo

Spaces:
expressapi
/
chroma
Sleeping

App Files Community
chroma
File size: 2,544 Bytes 
287a0bc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
 
package notification

import (
	"context"

	"github.com/apache/pulsar-client-go/pulsar"
	"github.com/chroma/chroma-coordinator/internal/model"
	"github.com/chroma/chroma-coordinator/internal/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(&notificationPb)
		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(&notificationPb)
		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
}