jetstream: init separate package
Anirudh Oppiliappan 2 weeks ago 4 files (+152, -109)
MODIFIED
cmd/knotserver/main.go
MODIFIED
cmd/knotserver/main.go
@@ -4,6 +4,8 @@ import ("context""net/http"+ "github.com/sotangled/tangled/api/tangled"+ "github.com/sotangled/tangled/jetstream""github.com/sotangled/tangled/knotserver""github.com/sotangled/tangled/knotserver/config""github.com/sotangled/tangled/knotserver/db"@@ -40,7 +42,15 @@ l.Error("failed to setup rbac enforcer", "error", err)return}- mux, err := knotserver.Setup(ctx, c, db, e, l)+ jc, err := jetstream.NewJetstreamClient("knotserver", []string{+ tangled.PublicKeyNSID,+ tangled.KnotMemberNSID,+ }, nil, db)+ if err != nil {+ l.Error("failed to setup jetstream", "error", err)+ }++ mux, err := knotserver.Setup(ctx, c, db, e, jc, l)if err != nil {l.Error("failed to setup server", "error", err)return
ADDED
jetstream/jetstream.go
ADDED
jetstream/jetstream.go
@@ -0,0 +1,136 @@+package jetstream++import (+ "context"+ "fmt"+ "sync"+ "time"++ "github.com/bluesky-social/jetstream/pkg/client"+ "github.com/bluesky-social/jetstream/pkg/client/schedulers/sequential"+ "github.com/bluesky-social/jetstream/pkg/models"+ "github.com/sotangled/tangled/log"+)++type DB interface {+ GetLastTimeUs() (int64, error)+ SaveLastTimeUs(int64) error+}++type JetstreamClient struct {+ cfg *client.ClientConfig+ client *client.Client+ ident string++ db DB+ reconnectCh chan struct{}+ mu sync.RWMutex+}++func (j *JetstreamClient) AddDid(did string) {+ j.mu.Lock()+ j.cfg.WantedDids = append(j.cfg.WantedDids, did)+ j.mu.Unlock()+ j.reconnectCh <- struct{}{}+}++func (j *JetstreamClient) UpdateDids(dids []string) {+ j.mu.Lock()+ j.cfg.WantedDids = dids+ j.mu.Unlock()+ j.reconnectCh <- struct{}{}+}++func NewJetstreamClient(ident string, collections []string, cfg *client.ClientConfig, db DB) (*JetstreamClient, error) {+ if cfg == nil {+ cfg = client.DefaultClientConfig()+ cfg.WebsocketURL = "wss://jetstream1.us-west.bsky.network/subscribe"+ cfg.WantedCollections = collections+ }++ return &JetstreamClient{+ cfg: cfg,+ ident: ident,+ db: db,+ reconnectCh: make(chan struct{}, 1),+ }, nil+}++func (j *JetstreamClient) StartJetstream(ctx context.Context, processFunc func(context.Context, *models.Event) error) error {+ logger := log.FromContext(ctx)++ pf := func(ctx context.Context, e *models.Event) error {+ err := processFunc(ctx, e)+ if err != nil {+ return err+ }++ if err := j.db.SaveLastTimeUs(e.TimeUS); err != nil {+ return err+ }++ return nil+ }++ sched := sequential.NewScheduler(j.ident, logger, pf)++ client, err := client.NewClient(j.cfg, log.New("jetstream"), sched)+ if err != nil {+ return fmt.Errorf("failed to create jetstream client: %w", err)+ }+ j.client = client++ go func() {+ lastTimeUs := j.getLastTimeUs(ctx)+ for len(j.cfg.WantedDids) == 0 {+ time.Sleep(time.Second)+ }+ j.connectAndRead(ctx, &lastTimeUs)+ }()++ return nil+}++func (j *JetstreamClient) connectAndRead(ctx context.Context, cursor *int64) {+ l := log.FromContext(ctx)+ for {+ select {+ case <-j.reconnectCh:+ l.Info("(re)connecting jetstream client")+ j.client.Scheduler.Shutdown()+ if err := j.client.ConnectAndRead(ctx, cursor); err != nil {+ l.Error("error reading jetstream", "error", err)+ }+ default:+ if err := j.client.ConnectAndRead(ctx, cursor); err != nil {+ l.Error("error reading jetstream", "error", err)+ }+ }+ }+}++func (j *JetstreamClient) getLastTimeUs(ctx context.Context) int64 {+ l := log.FromContext(ctx)+ lastTimeUs, err := j.db.GetLastTimeUs()+ if err != nil {+ l.Warn("couldn't get last time us, starting from now", "error", err)+ lastTimeUs = time.Now().UnixMicro()+ err = j.db.SaveLastTimeUs(lastTimeUs)+ if err != nil {+ l.Error("failed to save last time us")+ }+ }++ // If last time is older than a week, start from now+ if time.Now().UnixMicro()-lastTimeUs > 7*24*60*60*1000*1000 {+ lastTimeUs = time.Now().UnixMicro()+ l.Warn("last time us is older than a week. discarding that and starting from now")+ err = j.db.SaveLastTimeUs(lastTimeUs)+ if err != nil {+ l.Error("failed to save last time us")+ }+ }++ l.Info("found last time_us", "time_us", lastTimeUs)+ return lastTimeUs+}
MODIFIED
knotserver/handler.go
MODIFIED
knotserver/handler.go
@@ -7,6 +7,7 @@ "log/slog""net/http""github.com/go-chi/chi/v5"+ "github.com/sotangled/tangled/jetstream""github.com/sotangled/tangled/knotserver/config""github.com/sotangled/tangled/knotserver/db""github.com/sotangled/tangled/rbac"@@ -19,7 +20,7 @@type Handle struct {c *config.Configdb *db.DB- jc *JetstreamClient+ jc *jetstream.JetstreamCliente *rbac.Enforcerl *slog.Logger@@ -29,7 +30,7 @@ init chan struct{}knotInitialized bool}-func Setup(ctx context.Context, c *config.Config, db *db.DB, e *rbac.Enforcer, l *slog.Logger) (http.Handler, error) {+func Setup(ctx context.Context, c *config.Config, db *db.DB, e *rbac.Enforcer, jc *jetstream.JetstreamClient, l *slog.Logger) (http.Handler, error) {r := chi.NewRouter()h := Handle{@@ -37,6 +38,7 @@ c: c,db: db,e: e,l: l,+ jc: jc,init: make(chan struct{}),}@@ -45,7 +47,7 @@ if err != nil {return nil, fmt.Errorf("failed to setup enforcer: %w", err)}- err = h.StartJetstream(ctx)+ err = h.jc.StartJetstream(ctx, h.processMessages)if err != nil {return nil, fmt.Errorf("failed to start jetstream: %w", err)}
MODIFIED
knotserver/jetstream.go
MODIFIED
knotserver/jetstream.go
@@ -8,117 +8,12 @@ "io""net/http""net/url""strings"- "sync"- "time"- "github.com/bluesky-social/jetstream/pkg/client"- "github.com/bluesky-social/jetstream/pkg/client/schedulers/sequential""github.com/bluesky-social/jetstream/pkg/models""github.com/sotangled/tangled/api/tangled""github.com/sotangled/tangled/knotserver/db""github.com/sotangled/tangled/log")--type JetstreamClient struct {- cfg *client.ClientConfig- client *client.Client- reconnectCh chan struct{}- mu sync.RWMutex-}--func (h *Handle) StartJetstream(ctx context.Context) error {- l := h.l- ctx = log.IntoContext(ctx, l)- collections := []string{tangled.PublicKeyNSID, tangled.KnotMemberNSID}- dids := []string{}-- cfg := client.DefaultClientConfig()- cfg.WebsocketURL = "wss://jetstream1.us-west.bsky.network/subscribe"- cfg.WantedCollections = collections- cfg.WantedDids = dids-- sched := sequential.NewScheduler("knotserver", l, h.processMessages)-- client, err := client.NewClient(cfg, l, sched)- if err != nil {- l.Error("failed to create jetstream client", "error", err)- }-- jc := &JetstreamClient{- cfg: cfg,- client: client,- reconnectCh: make(chan struct{}, 1),- }-- h.jc = jc-- go func() {- lastTimeUs := h.getLastTimeUs(ctx)- for len(h.jc.cfg.WantedDids) == 0 {- time.Sleep(time.Second)- }- h.connectAndRead(ctx, &lastTimeUs)- }()- return nil-}--func (h *Handle) connectAndRead(ctx context.Context, cursor *int64) {- l := log.FromContext(ctx)- for {- select {- case <-h.jc.reconnectCh:- l.Info("(re)connecting jetstream client")- h.jc.client.Scheduler.Shutdown()- if err := h.jc.client.ConnectAndRead(ctx, cursor); err != nil {- l.Error("error reading jetstream", "error", err)- }- default:- if err := h.jc.client.ConnectAndRead(ctx, cursor); err != nil {- l.Error("error reading jetstream", "error", err)- }- }- }-}--func (j *JetstreamClient) AddDid(did string) {- j.mu.Lock()- j.cfg.WantedDids = append(j.cfg.WantedDids, did)- j.mu.Unlock()- j.reconnectCh <- struct{}{}-}--func (j *JetstreamClient) UpdateDids(dids []string) {- j.mu.Lock()- j.cfg.WantedDids = dids- j.mu.Unlock()- j.reconnectCh <- struct{}{}-}--func (h *Handle) getLastTimeUs(ctx context.Context) int64 {- l := log.FromContext(ctx)- lastTimeUs, err := h.db.GetLastTimeUs()- if err != nil {- l.Warn("couldn't get last time us, starting from now", "error", err)- lastTimeUs = time.Now().UnixMicro()- err = h.db.SaveLastTimeUs(lastTimeUs)- if err != nil {- l.Error("failed to save last time us")- }- }-- // If last time is older than a week, start from now- if time.Now().UnixMicro()-lastTimeUs > 7*24*60*60*1000*1000 {- lastTimeUs = time.Now().UnixMicro()- l.Warn("last time us is older than a week. discarding that and starting from now")- err = h.db.SaveLastTimeUs(lastTimeUs)- if err != nil {- l.Error("failed to save last time us")- }- }-- l.Info("found last time_us", "time_us", lastTimeUs)- return lastTimeUs-}func (h *Handle) processPublicKey(ctx context.Context, did string, record tangled.PublicKey) error {l := log.FromContext(ctx)