commit 9fc61c809487015dc4f357c9ee16ea0f5a0d9a37 · ptr.pet/bsky-repost-likes

+51 -30

server/jetstream.go

···

       35
       35
        
       	optsFn      OptsFn

     

       36
       36
        
       }

     

       37
       37
        
       

     

       38
       38
       -
       func NewStreamManager(logger *slog.Logger, name string, handleEvent HandleEvent, optsFn OptsFn) StreamManager {

     

       39
       39
       -
       	return StreamManager{

     

       38
       38
       +
       func NewStreamManager(logger *slog.Logger, name string, handleEvent HandleEvent, optsFn OptsFn) *StreamManager {

     

       39
       39
       +
       	manager := &StreamManager{

     

       40
       40
        
       		ctx:         context.TODO(),

     

       41
       41
        
       		logger:      logger.With("stream", name),

     

       42
       42
        
       		streamsLock: sync.Mutex{},

     
···

       45
       45
        
       		handleEvent: handleEvent,

     

       46
       46
        
       		optsFn:      optsFn,

     

       47
       47
        
       	}

     

       48
       48
       +
       	go func() {

     

       49
       49
       +
       		for {

     

       50
       50
       +
       			time.Sleep(time.Minute * 2)

     

       51
       51
       +
       			manager.streamsLock.Lock()

     

       52
       52
       +
       			manager.logger.Info("sharding stats", "streamCount", len(manager.streams))

     

       53
       53
       +
       			manager.streamsLock.Unlock()

     

       54
       54
       +
       		}

     

       55
       55
       +
       	}()

     

       56
       56
       +
       	return manager

     

       48
       57
        
       }

     

       49
       58
        
       

     

       50
       59
        
       // doesnt lock streams!!!

     
···

       71
       80
        
       

     

       72
       81
        
       func (manager *StreamManager) updateOpts() {

     

       73
       82
        
       	chunks, userCount := manager.chunkedOpts()

     

       83
       83
       +
       	logger := manager.logger.With("userCount", userCount)

     

       74
       84
        
       	manager.streamsLock.Lock()

     

       75
       75
       -
       	idsSeen := make(map[int]struct{}, 0)

     

       85
       85
       +
       	idsSeen := make(map[int]struct{}, len(manager.streams))

     

       76
       86
        
       	// update existing streams or create new ones

     

       77
       87
        
       	for id, opts := range chunks {

     

       78
       88
        
       		idsSeen[id] = struct{}{}

     
···

       82
       92
        
       				continue

     

       83
       93
        
       			}

     

       84
       94
        
       			if err := stream.inner.SendOptionsUpdate(opts); err != nil {

     

       85
       85
       -
       				manager.logger.Error("couldnt update follow stream opts", "error", err, "streamId", id)

     

       95
       95
       +
       				logger.Error("couldnt update follow stream opts", "error", err, "streamId", id)

     

       86
       96
        
       			}

     

       87
       97
        
       		} else {

     

       88
       98
        
       			manager.startSingle(id, opts)

     
···

       96
       106
        
       		}

     

       97
       107
        
       	}

     

       98
       108
        
       	manager.streamsLock.Unlock()

     

       99
       99
       -
       	manager.logger.Info("updated opts", "userCount", userCount)

     

       109
       109
       +
       	logger.Info("updated opts")

     

       100
       110
        
       }

     

       101
       111
        
       

     

       102
       112
        
       type HandleEvent func(context.Context, *models.Event) error

     
···

       105
       115
        
       func startJetstreamLoop(ctx context.Context, logger *slog.Logger, outStream **client.Client, name string, handleEvent HandleEvent, opts models.SubscriberOptionsUpdatePayload) {

     

       106
       116
        
       	backoff := time.Second

     

       107
       117
        
       	for {

     

       108
       108
       -
       		done := make(chan struct{})

     

       109
       109
       -
       		if ctx.Err() != nil {

     

       110
       110
       -
       			break

     

       111
       111
       -
       		}

     

       112
       112
       -
       		stream, startFn, err := startJetstreamClient(ctx, logger, name, handleEvent)

     

       113
       113
       -
       		*outStream = stream

     

       114
       114
       -
       		if startFn != nil {

     

       115
       115
       -
       			logger.Info("starting jetstream client", "collections", opts.WantedCollections, "userCount", len(opts.WantedDIDs))

     

       116
       116
       -
       			go func() {

     

       117
       117
       -
       				err = startFn()

     

       118
       118
       -
       				done <- struct{}{}

     

       119
       119
       -
       			}()

     

       120
       120
       -
       			// HACK: we need to wait for the websocket connection to start here. so we do

     

       121
       121
       -
       			// need to upstream something to jetstream client

     

       122
       122
       -
       			time.Sleep(time.Second * 2)

     

       123
       123
       -
       			err = stream.SendOptionsUpdate(opts)

     

       124
       124
       -
       			if err == nil {

     

       125
       125
       -
       				<-done

     

       118
       118
       +
       		select {

     

       119
       119
       +
       		case <-ctx.Done():

     

       120
       120
       +
       			return

     

       121
       121
       +
       		default:

     

       122
       122
       +
       			streamDone := make(chan struct{})

     

       123
       123
       +
       			stream, startFn, err := startJetstreamClient(ctx, logger, name, handleEvent)

     

       124
       124
       +
       			*outStream = stream

     

       125
       125
       +
       			if startFn != nil {

     

       126
       126
       +
       				logger.Info("starting jetstream client", "collections", opts.WantedCollections, "userCount", len(opts.WantedDIDs))

     

       127
       127
       +
       				go func() {

     

       128
       128
       +
       					err = startFn()

     

       129
       129
       +
       					streamDone <- struct{}{}

     

       130
       130
       +
       				}()

     

       131
       131
       +
       				// HACK: we need to wait for the websocket connection to start here. so we do

     

       132
       132
       +
       				// need to upstream something to jetstream client

     

       133
       133
       +
       				time.Sleep(time.Second * 2)

     

       134
       134
       +
       				go func() {

     

       135
       135
       +
       					// HACK: also silly because it panics if the connection isnt established yet (why????????????)

     

       136
       136
       +
       					defer func() {

     

       137
       137
       +
       						panic := recover()

     

       138
       138
       +
       						if panic != nil {

     

       139
       139
       +
       							err = fmt.Errorf("%s", panic)

     

       140
       140
       +
       						}

     

       141
       141
       +
       					}()

     

       142
       142
       +
       					err = stream.SendOptionsUpdate(opts)

     

       143
       143
       +
       				}()

     

       144
       144
       +
       				if err == nil {

     

       145
       145
       +
       					<-streamDone

     

       146
       146
       +
       				}

     

       126
       147
        
       			}

     

       127
       127
       -
       		}

     

       128
       128
       -
       		if err != nil {

     

       129
       129
       -
       			logger.Error("stream failed", "error", err, "backoff", backoff)

     

       130
       130
       -
       			time.Sleep(backoff)

     

       131
       131
       -
       			backoff = backoff * 2

     

       132
       132
       -
       		} else {

     

       133
       133
       -
       			backoff = time.Second

     

       148
       148
       +
       			if err != nil {

     

       149
       149
       +
       				logger.Error("stream failed", "error", err, "backoff", backoff)

     

       150
       150
       +
       				time.Sleep(backoff)

     

       151
       151
       +
       				backoff = backoff * 2

     

       152
       152
       +
       			} else {

     

       153
       153
       +
       				backoff = time.Second

     

       154
       154
       +
       			}

     

       134
       155
        
       		}

     

       135
       156
        
       	}

     

       136
       157
        
       }

+2 -2

server/main.go

···

       67
       67
        
       	subscribers = hashmap.New[string, *SubscriberData]()

     

       68
       68
        
       	actorData   = hashmap.New[syntax.DID, *ActorData]()

     

       69
       69
        
       

     

       70
       70
       -
       	likeStreams   StreamManager

     

       71
       71
       -
       	followStreams StreamManager

     

       70
       70
       +
       	likeStreams   *StreamManager

     

       71
       71
       +
       	followStreams *StreamManager

     

       72
       72
        
       

     

       73
       73
        
       	upgrader = websocket.Upgrader{

     

       74
       74
        
       		CheckOrigin: func(r *http.Request) bool {