diff --git a/jetstream/jetstream.go b/jetstream/jetstream.go --- a/jetstream/jetstream.go +++ b/jetstream/jetstream.go @@ -4,7 +4,10 @@ "context" "fmt" "log/slog" + "os" + "os/signal" "sync" + "syscall" "time" "github.com/bluesky-social/jetstream/pkg/client" @@ -16,8 +19,9 @@ type DB interface { GetLastTimeUs() (int64, error) SaveLastTimeUs(int64) error - UpdateLastTimeUs(int64) error } + +type Set[T comparable] map[T]struct{} type JetstreamClient struct { cfg *client.ClientConfig @@ -25,6 +29,7 @@ ident string l *slog.Logger + wantedDids Set[string] db DB waitForDid bool mu sync.RWMutex @@ -37,25 +42,24 @@ if did == "" { return } + j.mu.Lock() - j.cfg.WantedDids = append(j.cfg.WantedDids, did) + j.wantedDids[did] = struct{}{} j.mu.Unlock() } -func (j *JetstreamClient) UpdateDids(dids []string) { - j.mu.Lock() - for _, did := range dids { - if did != "" { - j.cfg.WantedDids = append(j.cfg.WantedDids, did) +type processor func(context.Context, *models.Event) error + +func (j *JetstreamClient) withDidFilter(processFunc processor) processor { + // since this closure references j.WantedDids; it should auto-update + // existing instances of the closure when j.WantedDids is mutated + return func(ctx context.Context, evt *models.Event) error { + if _, ok := j.wantedDids[evt.Did]; ok { + return processFunc(ctx, evt) + } else { + return nil } } - j.mu.Unlock() - - j.cancelMu.Lock() - if j.cancel != nil { - j.cancel() - } - j.cancelMu.Unlock() } func NewJetstreamClient(endpoint, ident string, collections []string, cfg *client.ClientConfig, logger *slog.Logger, db DB, waitForDid bool) (*JetstreamClient, error) { @@ -66,23 +70,24 @@ } return &JetstreamClient{ - cfg: cfg, - ident: ident, - db: db, - l: logger, + cfg: cfg, + ident: ident, + db: db, + l: logger, + wantedDids: make(map[string]struct{}), // This will make the goroutine in StartJetstream wait until - // cfg.WantedDids has been populated, typically using UpdateDids. + // j.wantedDids has been populated, typically using addDids. waitForDid: waitForDid, }, nil } // StartJetstream starts the jetstream client and processes events using the provided processFunc. -// The caller is responsible for saving the last time_us to the database (just use your db.SaveLastTimeUs). +// The caller is responsible for saving the last time_us to the database (just use your db.UpdateLastTimeUs). func (j *JetstreamClient) StartJetstream(ctx context.Context, processFunc func(context.Context, *models.Event) error) error { logger := j.l - sched := sequential.NewScheduler(j.ident, logger, processFunc) + sched := sequential.NewScheduler(j.ident, logger, j.withDidFilter(processFunc)) client, err := client.NewClient(j.cfg, log.New("jetstream"), sched) if err != nil { @@ -92,11 +97,15 @@ go func() { if j.waitForDid { - for len(j.cfg.WantedDids) == 0 { + for len(j.wantedDids) == 0 { time.Sleep(time.Second) } } logger.Info("done waiting for did") + + go j.periodicLastTimeSave(ctx) + j.saveIfKilled(ctx) + j.connectAndRead(ctx) }() @@ -130,6 +139,21 @@ } } +// save cursor periodically +func (j *JetstreamClient) periodicLastTimeSave(ctx context.Context) { + ticker := time.NewTicker(time.Minute) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + j.db.SaveLastTimeUs(time.Now().UnixMicro()) + } + } +} + func (j *JetstreamClient) getLastTimeUs(ctx context.Context) *int64 { l := log.FromContext(ctx) lastTimeUs, err := j.db.GetLastTimeUs() @@ -142,11 +166,11 @@ } } - // If last time is older than a week, start from now + // If last time is older than 2 days, start from now if time.Now().UnixMicro()-lastTimeUs > 2*24*60*60*1000*1000 { lastTimeUs = time.Now().UnixMicro() l.Warn("last time us is older than 2 days; discarding that and starting from now") - err = j.db.UpdateLastTimeUs(lastTimeUs) + err = j.db.SaveLastTimeUs(lastTimeUs) if err != nil { l.Error("failed to save last time us", "error", err) } @@ -154,4 +178,42 @@ l.Info("found last time_us", "time_us", lastTimeUs) return &lastTimeUs +} + +func (j *JetstreamClient) saveIfKilled(ctx context.Context) context.Context { + ctxWithCancel, cancel := context.WithCancel(ctx) + + sigChan := make(chan os.Signal, 1) + + signal.Notify(sigChan, + syscall.SIGINT, + syscall.SIGTERM, + syscall.SIGQUIT, + syscall.SIGHUP, + syscall.SIGKILL, + syscall.SIGSTOP, + ) + + go func() { + sig := <-sigChan + j.l.Info("Received signal, initiating graceful shutdown", "signal", sig) + + lastTimeUs := time.Now().UnixMicro() + if err := j.db.SaveLastTimeUs(lastTimeUs); err != nil { + j.l.Error("Failed to save last time during shutdown", "error", err) + } + j.l.Info("Saved lastTimeUs before shutdown", "lastTimeUs", lastTimeUs) + + j.cancelMu.Lock() + if j.cancel != nil { + j.cancel() + } + j.cancelMu.Unlock() + + cancel() + + os.Exit(0) + }() + + return ctxWithCancel } diff --git a/knotserver/handler.go b/knotserver/handler.go --- a/knotserver/handler.go +++ b/knotserver/handler.go @@ -63,7 +63,9 @@ if len(dids) > 0 { h.knotInitialized = true close(h.init) - // h.jc.UpdateDids(dids) + for _, d := range dids { + h.jc.AddDid(d) + } } r.Get("/", h.Index) diff --git a/knotserver/jetstream.go b/knotserver/jetstream.go --- a/knotserver/jetstream.go +++ b/knotserver/jetstream.go @@ -53,6 +53,7 @@ l.Error("failed to add did", "error", err) return fmt.Errorf("failed to add did: %w", err) } + h.jc.AddDid(did) if err := h.fetchAndAddKeys(ctx, did); err != nil { return fmt.Errorf("failed to fetch and add keys: %w", err) @@ -115,10 +116,9 @@ eventTime := event.TimeUS lastTimeUs := eventTime + 1 fmt.Println("lastTimeUs", lastTimeUs) - if err := h.db.UpdateLastTimeUs(lastTimeUs); err != nil { + if err := h.db.SaveLastTimeUs(lastTimeUs); err != nil { err = fmt.Errorf("(deferred) failed to save last time us: %w", err) } - // h.jc.UpdateDids([]string{did}) }() raw := json.RawMessage(event.Commit.Record) diff --git a/knotserver/routes.go b/knotserver/routes.go --- a/knotserver/routes.go +++ b/knotserver/routes.go @@ -448,7 +448,7 @@ return } - data := make([]map[string]interface{}, 0) + data := make([]map[string]any, 0) for _, key := range keys { j := key.JSON() data = append(data, j) @@ -684,8 +684,8 @@ writeError(w, err.Error(), http.StatusInternalServerError) return } - h.jc.AddDid(did) + if err := h.e.AddMember(ThisServer, did); err != nil { l.Error("adding member", "error", err.Error()) writeError(w, err.Error(), http.StatusInternalServerError) @@ -768,8 +768,8 @@ writeError(w, err.Error(), http.StatusInternalServerError) return } + h.jc.AddDid(data.Did) - // h.jc.UpdateDids([]string{data.Did}) if err := h.e.AddOwner(ThisServer, data.Did); err != nil { l.Error("adding owner", "error", err.Error()) writeError(w, err.Error(), http.StatusInternalServerError) diff --git a/appview/db/jetstream.go b/appview/db/jetstream.go --- a/appview/db/jetstream.go +++ b/appview/db/jetstream.go @@ -5,21 +5,17 @@ } func (db DbWrapper) SaveLastTimeUs(lastTimeUs int64) error { - _, err := db.Exec(`insert into _jetstream (last_time_us) values (?)`, lastTimeUs) + _, err := db.Exec(` + insert into _jetstream (id, last_time_us) + values (1, ?) + on conflict(id) do update set last_time_us = excluded.last_time_us + `, lastTimeUs) return err -} - -func (db DbWrapper) UpdateLastTimeUs(lastTimeUs int64) error { - _, err := db.Exec(`update _jetstream set last_time_us = ? where rowid = 1`, lastTimeUs) - if err != nil { - return err - } - return nil } func (db DbWrapper) GetLastTimeUs() (int64, error) { var lastTimeUs int64 - row := db.QueryRow(`select last_time_us from _jetstream`) + row := db.QueryRow(`select last_time_us from _jetstream where id = 1;`) err := row.Scan(&lastTimeUs) return lastTimeUs, err } diff --git a/appview/state/jetstream.go b/appview/state/jetstream.go --- a/appview/state/jetstream.go +++ b/appview/state/jetstream.go @@ -20,7 +20,7 @@ defer func() { eventTime := e.TimeUS lastTimeUs := eventTime + 1 - if err := d.UpdateLastTimeUs(lastTimeUs); err != nil { + if err := d.SaveLastTimeUs(lastTimeUs); err != nil { err = fmt.Errorf("(deferred) failed to save last time us: %w", err) } }() diff --git a/appview/state/state.go b/appview/state/state.go --- a/appview/state/state.go +++ b/appview/state/state.go @@ -60,7 +60,15 @@ resolver := appview.NewResolver() wrapper := db.DbWrapper{d} - jc, err := jetstream.NewJetstreamClient(config.JetstreamEndpoint, "appview", []string{tangled.GraphFollowNSID}, nil, slog.Default(), wrapper, false) + jc, err := jetstream.NewJetstreamClient( + config.JetstreamEndpoint, + "appview", + []string{tangled.GraphFollowNSID}, + nil, + slog.Default(), + wrapper, + false, + ) if err != nil { return nil, fmt.Errorf("failed to create jetstream client: %w", err) } diff --git a/cmd/jstest/main.go b/cmd/jstest/main.go deleted file mode 100644 --- a/cmd/jstest/main.go +++ /dev/null @@ -1,150 +0,0 @@ -package main - -import ( - "context" - "flag" - "log/slog" - "os" - "os/signal" - "strings" - "syscall" - "time" - - "github.com/bluesky-social/jetstream/pkg/client" - "github.com/bluesky-social/jetstream/pkg/models" - "tangled.sh/tangled.sh/core/jetstream" -) - -// Simple in-memory implementation of DB interface -type MemoryDB struct { - lastTimeUs int64 -} - -func (m *MemoryDB) GetLastTimeUs() (int64, error) { - if m.lastTimeUs == 0 { - return time.Now().UnixMicro(), nil - } - return m.lastTimeUs, nil -} - -func (m *MemoryDB) SaveLastTimeUs(ts int64) error { - m.lastTimeUs = ts - return nil -} - -func (m *MemoryDB) UpdateLastTimeUs(ts int64) error { - m.lastTimeUs = ts - return nil -} - -func main() { - // Setup logger - logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{ - Level: slog.LevelInfo, - })) - - // Create in-memory DB - db := &MemoryDB{} - - // Get query URL from flag - var queryURL string - flag.StringVar(&queryURL, "query-url", "", "Jetstream query URL containing DIDs") - flag.Parse() - - if queryURL == "" { - logger.Error("No query URL provided, use --query-url flag") - os.Exit(1) - } - - // Extract wantedDids parameters - didParams := strings.Split(queryURL, "&wantedDids=") - dids := make([]string, 0, len(didParams)-1) - for i, param := range didParams { - if i == 0 { - // Skip the first part (the base URL with cursor) - continue - } - dids = append(dids, param) - } - - // Extract collections - collections := []string{"sh.tangled.publicKey", "sh.tangled.knot.member"} - - // Create client configuration - cfg := client.DefaultClientConfig() - cfg.WebsocketURL = "wss://jetstream2.us-west.bsky.network/subscribe" - cfg.WantedCollections = collections - - // Create jetstream client - jsClient, err := jetstream.NewJetstreamClient( - cfg.WebsocketURL, - "tangled-jetstream", - collections, - cfg, - logger, - db, - false, - ) - if err != nil { - logger.Error("Failed to create jetstream client", "error", err) - os.Exit(1) - } - - // Update DIDs - jsClient.UpdateDids(dids) - - // Create a context that will be canceled on SIGINT or SIGTERM - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - // Setup signal handling with a buffered channel - sigCh := make(chan os.Signal, 1) - signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM) - - // Process function for events - processFunc := func(ctx context.Context, event *models.Event) error { - // Log the event details - logger.Info("Received event", - "collection", event.Commit.Collection, - "did", event.Did, - "rkey", event.Commit.RKey, - "action", event.Kind, - "time_us", event.TimeUS, - ) - - // Save the last time_us - if err := db.UpdateLastTimeUs(event.TimeUS); err != nil { - logger.Error("Failed to update last time_us", "error", err) - } - - return nil - } - - // Start jetstream - if err := jsClient.StartJetstream(ctx, processFunc); err != nil { - logger.Error("Failed to start jetstream", "error", err) - os.Exit(1) - } - - // Wait for signal instead of context.Done() - sig := <-sigCh - logger.Info("Received signal, shutting down", "signal", sig) - cancel() // Cancel context after receiving signal - - // Shutdown gracefully with a timeout - shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second) - defer shutdownCancel() - - done := make(chan struct{}) - go func() { - jsClient.Shutdown() - close(done) - }() - - select { - case <-done: - logger.Info("Jetstream client shut down gracefully") - case <-shutdownCtx.Done(): - logger.Warn("Shutdown timed out, forcing exit") - } -} diff --git a/cmd/knotserver/main.go b/cmd/knotserver/main.go --- a/cmd/knotserver/main.go +++ b/cmd/knotserver/main.go @@ -49,7 +49,7 @@ jc, err := jetstream.NewJetstreamClient(c.Server.JetstreamEndpoint, "knotserver", []string{ tangled.PublicKeyNSID, tangled.KnotMemberNSID, - }, nil, l, db, false) + }, nil, l, db, true) if err != nil { l.Error("failed to setup jetstream", "error", err) } diff --git a/knotserver/db/jetstream.go b/knotserver/db/jetstream.go --- a/knotserver/db/jetstream.go +++ b/knotserver/db/jetstream.go @@ -1,21 +1,17 @@ package db func (d *DB) SaveLastTimeUs(lastTimeUs int64) error { - _, err := d.db.Exec(`insert into _jetstream (last_time_us) values (?)`, lastTimeUs) + _, err := d.db.Exec(` + insert into _jetstream (id, last_time_us) + values (1, ?) + on conflict(id) do update set last_time_us = excluded.last_time_us + `, lastTimeUs) return err -} - -func (d *DB) UpdateLastTimeUs(lastTimeUs int64) error { - _, err := d.db.Exec(`update _jetstream set last_time_us = ? where rowid = 1`, lastTimeUs) - if err != nil { - return err - } - return nil } func (d *DB) GetLastTimeUs() (int64, error) { var lastTimeUs int64 - row := d.db.QueryRow(`select last_time_us from _jetstream`) + row := d.db.QueryRow(`select last_time_us from _jetstream where id = 1;`) err := row.Scan(&lastTimeUs) return lastTimeUs, err } diff --git a/knotserver/db/pubkeys.go b/knotserver/db/pubkeys.go --- a/knotserver/db/pubkeys.go +++ b/knotserver/db/pubkeys.go @@ -44,8 +44,8 @@ return err } -func (pk *PublicKey) JSON() map[string]interface{} { - return map[string]interface{}{ +func (pk *PublicKey) JSON() map[string]any { + return map[string]any{ "did": pk.Did, "key": pk.Key, "created": pk.Created, diff --git a/appview/pages/templates/knots.html b/appview/pages/templates/knots.html --- a/appview/pages/templates/knots.html +++ b/appview/pages/templates/knots.html @@ -8,8 +8,7 @@

Generate a key to initialize your knot server.