···
"github.com/bluesky-social/jetstream/pkg/client"
"github.com/bluesky-social/jetstream/pkg/client/schedulers/sequential"
"github.com/bluesky-social/jetstream/pkg/models"
15
+
func Chunk[T any](slice []T, chunkSize int) [][]T {
17
+
for i := 0; i < len(slice); i += chunkSize {
18
+
end := min(i+chunkSize, len(slice))
19
+
chunks = append(chunks, slice[i:end])
24
+
type Stream struct {
25
+
inner *client.Client
26
+
cancel context.CancelFunc
28
+
type StreamManager struct {
31
+
streamsLock sync.Mutex
32
+
streams map[int]Stream
34
+
handleEvent HandleEvent
38
+
func NewStreamManager(logger *slog.Logger, name string, handleEvent HandleEvent, optsFn OptsFn) StreamManager {
39
+
return StreamManager{
40
+
ctx: context.TODO(),
41
+
logger: logger.With("stream", name),
42
+
streamsLock: sync.Mutex{},
43
+
streams: make(map[int]Stream),
45
+
handleEvent: handleEvent,
50
+
// doesnt lock streams!!!
51
+
func (manager *StreamManager) startSingle(id int, opts models.SubscriberOptionsUpdatePayload) {
52
+
ctx, cancel := context.WithCancel(manager.ctx)
53
+
stream := Stream{inner: nil, cancel: cancel}
54
+
// add to streams and put on wait group
55
+
manager.streams[id] = stream
56
+
go startJetstreamLoop(ctx, manager.logger.With("streamId", id), &stream.inner, fmt.Sprintf("%s_%d", manager.name, id), manager.handleEvent, opts)
59
+
func (manager *StreamManager) chunkedOpts() ([]models.SubscriberOptionsUpdatePayload, int) {
60
+
results := make([]models.SubscriberOptionsUpdatePayload, 0)
61
+
opts := manager.optsFn()
62
+
for _, wantedDidsChunk := range Chunk(opts.WantedDIDs, 9999) {
63
+
results = append(results, models.SubscriberOptionsUpdatePayload{
64
+
WantedCollections: opts.WantedCollections,
65
+
WantedDIDs: wantedDidsChunk,
66
+
MaxMessageSizeBytes: opts.MaxMessageSizeBytes,
69
+
return results, len(opts.WantedDIDs)
72
+
func (manager *StreamManager) updateOpts() {
73
+
chunks, userCount := manager.chunkedOpts()
74
+
manager.streamsLock.Lock()
75
+
idsSeen := make(map[int]struct{}, 0)
76
+
// update existing streams or create new ones
77
+
for id, opts := range chunks {
78
+
idsSeen[id] = struct{}{}
79
+
if len(manager.streams) > id {
80
+
stream := manager.streams[id]
81
+
if stream.inner == nil {
84
+
if err := stream.inner.SendOptionsUpdate(opts); err != nil {
85
+
manager.logger.Error("couldnt update follow stream opts", "error", err, "streamId", id)
88
+
manager.startSingle(id, opts)
91
+
// cancel and delete unused streams
92
+
for k := range manager.streams {
93
+
if _, exists := idsSeen[k]; !exists {
94
+
manager.streams[k].cancel()
95
+
delete(manager.streams, k)
98
+
manager.streamsLock.Unlock()
99
+
manager.logger.Info("updated opts", "userCount", userCount)
type HandleEvent func(context.Context, *models.Event) error
103
+
type OptsFn func() models.SubscriberOptionsUpdatePayload
14
-
func startJetstreamLoop(logger *slog.Logger, outStream **client.Client, name string, handleEvent HandleEvent, optsFn func() models.SubscriberOptionsUpdatePayload) {
105
+
func startJetstreamLoop(ctx context.Context, logger *slog.Logger, outStream **client.Client, name string, handleEvent HandleEvent, opts models.SubscriberOptionsUpdatePayload) {
106
+
backoff := time.Second
16
-
stream, startFn, err := startJetstreamClient(name, optsFn(), handleEvent)
108
+
done := make(chan struct{})
109
+
if ctx.Err() != nil {
112
+
stream, startFn, err := startJetstreamClient(ctx, logger, name, handleEvent)
115
+
logger.Info("starting jetstream client", "collections", opts.WantedCollections, "userCount", len(opts.WantedDIDs))
120
+
// HACK: we need to wait for the websocket connection to start here. so we do
121
+
// need to upstream something to jetstream client
122
+
time.Sleep(time.Second * 2)
123
+
err = stream.SendOptionsUpdate(opts)
22
-
logger.Error("stream failed", "name", name, "error", err)
129
+
logger.Error("stream failed", "error", err, "backoff", backoff)
130
+
time.Sleep(backoff)
131
+
backoff = backoff * 2
133
+
backoff = time.Second
27
-
func startJetstreamClient(name string, opts models.SubscriberOptionsUpdatePayload, handleEvent HandleEvent) (*client.Client, func() error, error) {
28
-
ctx := context.Background()
138
+
func startJetstreamClient(ctx context.Context, logger *slog.Logger, name string, handleEvent HandleEvent) (*client.Client, func() error, error) {
config := client.DefaultClientConfig()
config.WebsocketURL = "wss://jetstream1.us-west.bsky.network/subscribe"
33
-
config.WantedCollections = opts.WantedCollections
34
-
config.WantedDids = opts.WantedDIDs
35
-
config.RequireHello = false
142
+
config.RequireHello = true
scheduler := sequential.NewScheduler(name, logger, handleEvent)
c, err := client.NewClient(config, logger, scheduler)
41
-
logger.Error("failed to create jetstream client", "name", name, "error", err)
148
+
logger.Error("failed to create jetstream client", "error", err)
startFn := func() error {
46
-
logger.Info("starting jetstream client", "name", name, "collections", opts.WantedCollections, "wanted_dids", len(opts.WantedDIDs))
if err := c.ConnectAndRead(ctx, nil); err != nil {
48
-
logger.Error("jetstream client failed", "name", name, "error", err)
154
+
logger.Error("jetstream client failed", "error", err)
ptr.pet