Added jetstream monitoring

3 files changed, 149 insertions, 0 deletions

diff --git a/fetchers/cache.go b/fetchers/cache.go
index 11a43f3..4196792 100644
--- a/fetchers/cache.go
+++ b/fetchers/cache.go
@@ -61,6 +61,16 @@ func (c *Cache) Set(uri string, note *notes.Note) {
 	c.client.Set(c.ctx, key, data, c.expiryTime)
 }
 
+func (c *Cache) Delete(uri string) {
+	key := fmt.Sprintf("post:%s", uri)
+	c.client.Del(c.ctx, key)
+}
+
+func (c *Cache) DeleteAuthor(did string) {
+	key := fmt.Sprintf("profile:%s", did)
+	c.client.Del(c.ctx, key)
+}
+
 func (c *Cache) GetAuthor(did string) (*notes.Author, bool) {
 	key := fmt.Sprintf("profile:%s", did)
 	data, err := c.client.Get(c.ctx, key).Bytes()
diff --git a/fetchers/config.go b/fetchers/config.go
index a8aa711..a3a1ccf 100644
--- a/fetchers/config.go
+++ b/fetchers/config.go
@@ -5,6 +5,7 @@ import "log/slog"
 type Config struct {
 	SlingshotURL string
 	PorxieURL    string
+	JetstreamURL string
 	PageSize     int
 	Logger       *slog.Logger
 	Cache        *Cache
diff --git a/fetchers/jetstream.go b/fetchers/jetstream.go
new file mode 100644
index 0000000..78a2550
--- /dev/null
+++ b/fetchers/jetstream.go
@@ -0,0 +1,138 @@
+package fetchers
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"llostre/general"
+	"net/url"
+	"time"
+
+	"github.com/coder/websocket"
+)
+
+type jetstreamEvent struct {
+	DID    string           `json:"did"`
+	Kind   string           `json:"kind"`
+	Commit *jetstreamCommit `json:"commit,omitempty"`
+}
+
+type jetstreamCommit struct {
+	Operation  string          `json:"operation"`
+	Collection string          `json:"collection"`
+	RKey       string          `json:"rkey"`
+	Record     json.RawMessage `json:"record,omitempty"`
+}
+
+var jetstreamCollections = []string{"app.bsky.feed.post", "app.bsky.actor.profile"}
+
+// Meant to be run in its own goroutine; blocks forever, reconnecting on failure.
+// Monitors the jetstream for changes in cached records
+func StartJetstream(conf *Config) {
+	if conf.JetstreamURL == "" {
+		return
+	}
+
+	subscribeURL := buildJetstreamURL(conf.JetstreamURL)
+
+	for {
+		if err := consumeJetstream(conf, subscribeURL); err != nil {
+			conf.Logger.Error("jetstream connection lost, reconnecting", "error", err)
+		}
+		time.Sleep(5 * time.Second)
+	}
+}
+
+func buildJetstreamURL(base string) string {
+	query := url.Values{}
+	for _, collection := range jetstreamCollections {
+		query.Add("wantedCollections", collection)
+	}
+	return fmt.Sprintf("%s?%s", base, query.Encode())
+}
+
+func consumeJetstream(conf *Config, subscribeURL string) error {
+	ctx := context.Background()
+
+	conn, _, err := websocket.Dial(ctx, subscribeURL, nil)
+	if err != nil {
+		return err
+	}
+	defer conn.CloseNow()
+	conn.SetReadLimit(1 << 20) // 1MiB; default 32KiB is too tight for larger posts/profiles
+
+	conf.Logger.Debug("jetstream connected", "url", subscribeURL)
+
+	for {
+		_, data, err := conn.Read(ctx)
+		if err != nil {
+			return err
+		}
+
+		var event jetstreamEvent
+		if err := json.Unmarshal(data, &event); err != nil {
+			conf.Logger.Warn("failed to parse jetstream event", "error", err)
+			continue
+		}
+
+		if event.Kind != "commit" || event.Commit == nil {
+			continue
+		}
+
+		handleJetstreamCommit(conf, event.DID, event.Commit)
+	}
+}
+
+func handleJetstreamCommit(conf *Config, did string, commit *jetstreamCommit) {
+	uri := fmt.Sprintf("at://%s/%s/%s", did, commit.Collection, commit.RKey)
+
+	switch commit.Collection {
+	case "app.bsky.feed.post":
+		handleJetstreamPost(conf, did, uri, commit)
+	case "app.bsky.actor.profile":
+		handleJetstreamProfile(conf, did, commit)
+	}
+}
+
+func handleJetstreamPost(conf *Config, did, uri string, commit *jetstreamCommit) {
+	if commit.Operation == "delete" {
+		if _, found := conf.Cache.Get(uri); found {
+			conf.Cache.Delete(uri)
+		}
+		return
+	}
+
+	// only track posts from authors we already care about
+	author, found := conf.Cache.GetAuthor(did)
+	if !found {
+		return
+	}
+
+	rawRecord := &general.RawRecord{URI: uri, Value: commit.Record}
+
+	// cache the result
+	if _, err := parseAppBskyFeedPostRecord(rawRecord, author, conf); err != nil {
+		conf.Logger.Warn("jetstream: failed to parse post", "uri", uri, "error", err)
+	}
+}
+
+func handleJetstreamProfile(conf *Config, did string, commit *jetstreamCommit) {
+	if commit.Operation == "delete" {
+		if _, found := conf.Cache.GetAuthor(did); found {
+			conf.Cache.DeleteAuthor(did)
+		}
+		return
+	}
+
+	existing, found := conf.Cache.GetAuthor(did)
+	if !found {
+		return
+	}
+
+	rawRecord := &general.RawRecord{URI: fmt.Sprintf("at://%s/app.bsky.actor.profile/self", did), Value: commit.Record}
+
+	// cache the result
+	if _, err := parseAppBskyActorProfile(did, existing.PDS, rawRecord, conf.Cache); err != nil {
+		conf.Logger.Warn("jetstream: failed to parse updated profile", "did", did, "error", err)
+	}
+}