jetstream: init separate package

cmd/knotserver/main.go

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

jetstream/jetstream.go

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

knotserver/handler.go

···
 	"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"
···
 type Handle struct {
 	c  *config.Config
 	db *db.DB
-	jc *JetstreamClient
+	jc *jetstream.JetstreamClient
 	e  *rbac.Enforcer
 	l  *slog.Logger
 
···
 	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{
···
 		db:   db,
 		e:    e,
 		l:    l,
+		jc:   jc,
 		init: make(chan struct{}),
 	}
 
···
 		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)
 	}

knotserver/jetstream.go

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