add in-memory jetstream did filter

appview/db/jetstream.go

···
 }
 
 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
 }

appview/pages/templates/knots.html

···
     <section class="rounded bg-white dark:bg-gray-800 drop-shadow-sm px-6 py-4 mb-6 w-full lg:w-fit">
       <p class="mb-8 dark:text-gray-300">Generate a key to initialize your knot server.</p>
       <form 
-        hx-put="/knots/key"
-        hx-swap="none"
+        hx-post="/knots/key"
         class="max-w-2xl mb-8 space-y-4"
       >
         <input

appview/state/jetstream.go

···
 		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)
 			}
 		}()

appview/state/state.go

···
 	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)
 	}

cmd/jstest/main.go

···
-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")
-	}
-}

cmd/knotserver/main.go

···
 	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)
 	}

jetstream/jetstream.go

···
 	"context"
 	"fmt"
 	"log/slog"
+	"os"
+	"os/signal"
 	"sync"
+	"syscall"
 	"time"
 
 	"github.com/bluesky-social/jetstream/pkg/client"
···
 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
···
 	ident  string
 	l      *slog.Logger
 
+	wantedDids Set[string]
 	db         DB
 	waitForDid bool
 	mu         sync.RWMutex
···
 	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) {
···
 	}
 
 	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 {
···
 
 	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)
 	}()
 
···
 	}
 }
 
+// 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()
···
 		}
 	}
 
-	// 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)
 		}
···
 
 	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
 }

knotserver/db/jetstream.go

···
 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
 }

knotserver/db/pubkeys.go

···
 	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,

knotserver/handler.go

···
 	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)

knotserver/jetstream.go

···
 		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)
···
 		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)

knotserver/routes.go

···
 			return
 		}
 
-		data := make([]map[string]interface{}, 0)
+		data := make([]map[string]any, 0)
 		for _, key := range keys {
 			j := key.JSON()
 			data = append(data, j)
···
 		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)
···
 		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)