diff --git a/cmd/knotserver/main.go b/cmd/knotserver/main.go index 6d595985..8789fb33 100644 --- a/cmd/knotserver/main.go +++ b/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 @@ func main() { 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 diff --git a/jetstream/jetstream.go b/jetstream/jetstream.go new file mode 100644 index 00000000..1b8ef7c9 --- /dev/null +++ b/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 +} diff --git a/knotserver/handler.go b/knotserver/handler.go index 6a45ffe8..a961fdcc 100644 --- a/knotserver/handler.go +++ b/knotserver/handler.go @@ -7,6 +7,7 @@ import ( "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 @@ const ( type Handle struct { c *config.Config db *db.DB - jc *JetstreamClient + jc *jetstream.JetstreamClient e *rbac.Enforcer l *slog.Logger @@ -29,7 +30,7 @@ type Handle 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 @@ func Setup(ctx context.Context, c *config.Config, db *db.DB, e *rbac.Enforcer, l db: db, e: e, l: l, + jc: jc, init: make(chan struct{}), } @@ -45,7 +47,7 @@ func Setup(ctx context.Context, c *config.Config, db *db.DB, e *rbac.Enforcer, l 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) } diff --git a/knotserver/jetstream.go b/knotserver/jetstream.go index 4ca30632..440c1ab0 100644 --- a/knotserver/jetstream.go +++ b/knotserver/jetstream.go @@ -8,118 +8,13 @@ import ( "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) pk := db.PublicKey{