jetrelay ehhh

.gitignore

···
-cmd/aturilist/badger_data
+cmd/aturilist/badger_data
+cmd/backstream/temp

cmd/aturilist/main.go

···
 	"github.com/dgraph-io/badger/v4"
 	"github.com/gin-gonic/gin"
 
-	// Restored your specific imports
 	"tangled.org/whey.party/red-dwarf-server/auth"
 	"tangled.org/whey.party/red-dwarf-server/microcosm"
 	"tangled.org/whey.party/red-dwarf-server/microcosm/slingshot"
···
 	db     *badger.DB
 	logger *slog.Logger
 
-	// Locks for specific operations if needed, though Badger is thread-safe
 	backfillTracker map[string]*sync.WaitGroup
 	backfillMutex   sync.Mutex
 }
···
 
 	initURLs(*prod)
 
-	// 1. Initialize DB
 	db, err := badger.Open(badger.DefaultOptions(*dbPath))
 	if err != nil {
 		logger.Error("Failed to open BadgerDB", "error", err)
···
 		logger: logger,
 	}
 
-	// 2. Initialize Auth
 	auther, err := auth.NewAuth(
 		100_000,
 		time.Hour*12,
 		5,
-		serviceWebDID, //+"#bsky_appview",
+		serviceWebDID,
 	)
 	if err != nil {
 		log.Fatalf("Failed to create Auth: %v", err)
 	}
 
-	// 3. Initialize Clients
 	ctx := context.Background()
 	sl := slingshot.NewSlingshot(SLINGSHOT_URL)
 
-	// 4. Initialize Jetstream
 	config := client.DefaultClientConfig()
 	config.WebsocketURL = JETSTREAM_URL
 	config.Compress = true
···
 		return
 	}
 
-	// Connect with cursor (5 minutes ago)
 	cursor := time.Now().Add(-5 * time.Minute).UnixMicro()
 
 	go func() {
 		logger.Info("Connecting to Jetstream...")
-		/*
-			If you resume a jetstream firehose from a cursor, everything works fine until you catch up to real time.
-			At that point, the connection drops. If you connect without a cursor (going straight to realtime), it keeps working.
-		*/
 		for {
 			if err := c.ConnectAndRead(ctx, &cursor); err != nil {
 				logger.Error("jetstream connection disconnected", "error", err)
···
 
 			select {
 			case <-ctx.Done():
-				return // Context cancelled, exit loop
+				return
 			default:
 				logger.Info("Reconnecting to Jetstream in 5 seconds...", "cursor", cursor)
 				time.Sleep(5 * time.Second)
···
 		}
 	}()
 
-	// 5. Initialize Router
 	router := gin.New()
 	router.Use(auther.AuthenticateGinRequestViaJWT)
 
···
 
 	router.GET("/xrpc/app.reddwarf.aturilist.countRecords", srv.handleCountRecords)
 
-	// heavily rate limited because can be used for spam.
 	router.POST("/xrpc/app.reddwarf.aturilist.indexRecord", func(c *gin.Context) {
 		srv.handleIndexRecord(c, sl)
 	})
···
 
 	router.Run(":7155")
 }
-
-// --- Jetstream Handler ---
 
 type JetstreamHandler struct {
 	srv *Server
···
 func (h *JetstreamHandler) HandleEvent(ctx context.Context, event *models.Event) error {
 	if event != nil {
 		if event.Commit != nil {
-			// Identify Delete operation
 			isDelete := event.Commit.Operation == models.CommitOperationDelete
 
-			// Process
 			h.srv.processRecord(event.Did, event.Commit.Collection, event.Commit.RKey, isDelete)
 
 		}
···
 	return nil
 }
 
-// --- DB Helpers ---
-
 func makeKey(repo, collection, rkey string) []byte {
 	return []byte(fmt.Sprintf("%s|%s|%s", repo, collection, rkey))
 }
···
 	return parts[0], parts[1], parts[2], nil
 }
 
-// processRecord handles the DB write/delete.
-// isDelete=true removes the key. isDelete=false sets the key.
 func (s *Server) processRecord(repo, collection, rkey string, isDelete bool) {
 	key := makeKey(repo, collection, rkey)
 
···
 		if isDelete {
 			return txn.Delete(key)
 		}
-		// On create/update, store current timestamp.
-		// You can store more data (Cid, etc) here if needed later.
 		return txn.Set(key, []byte(time.Now().Format(time.RFC3339)))
 	})
 
···
 		s.logger.Error("Failed to update DB", "repo", repo, "rkey", rkey, "err", err)
 	}
 }
-
-// --- HTTP Handlers ---
 
 func (s *Server) handleListRecords(c *gin.Context) {
 	repo := c.Query("repo")
 	collection := c.Query("collection")
 	cursor := c.Query("cursor")
-	reverse := c.Query("reverse") == "true" // 1. Check param
+	reverse := c.Query("reverse") == "true"
 	limit := 50
 
 	if repo == "" || collection == "" {
···
 		return
 	}
 
-	// Base prefix: "repo|collection|"
 	prefixStr := fmt.Sprintf("%s|%s|", repo, collection)
 	prefix := []byte(prefixStr)
 
···
 	var lastRkey string
 
 	err := s.db.View(func(txn *badger.Txn) error {
-		// 2. Configure Iterator Options
 		opts := badger.DefaultIteratorOptions
 		opts.PrefetchValues = false
-		opts.Reverse = reverse // Set reverse mode
+		opts.Reverse = reverse
 
 		it := txn.NewIterator(opts)
 		defer it.Close()
 
-		// 3. Determine Start Key
 		var startKey []byte
 		if cursor != "" {
-			// If cursor exists, we seek to it regardless of direction
 			startKey = makeKey(repo, collection, cursor)
 		} else {
 			if reverse {
-				// REVERSE START: "repo|collection|" + 0xFF
-				// This seeks to the theoretical end of this prefix range
 				startKey = append([]byte(prefixStr), 0xFF)
 			} else {
-				// FORWARD START: "repo|collection|"
 				startKey = prefix
 			}
 		}
 
-		// 4. Seek and Iterate
 		it.Seek(startKey)
 
-		// Handle Cursor Pagination Skip
-		// If we provided a cursor, we likely landed exactly ON that cursor.
-		// We want the record *after* (or *before* in reverse) the cursor.
 		if cursor != "" && it.Valid() {
-			// Badger's Seek moves to key >= seek_key (even in reverse mode logic varies slightly,
-			// but practically we check if we landed on the exact cursor).
 			if string(it.Item().Key()) == string(startKey) {
-				it.Next() // Skip the cursor itself
+				it.Next()
 			}
 		}
 
-		// Iterate as long as the key still starts with our prefix
 		for ; it.ValidForPrefix(prefix); it.Next() {
 			if len(aturis) >= limit {
 				break
···
 		"count":  len(aturis),
 	}
 
-	// Only return cursor if we hit the limit, allowing the client to request the next page
 	if lastRkey != "" && len(aturis) == limit {
 		resp["cursor"] = lastRkey
 	}
···
 	})
 }
 
-// handleIndexRecord now takes the Slingshot client specifically
 func (s *Server) handleIndexRecord(c *gin.Context, sl *microcosm.MicrocosmClient) {
-	//authedUserDid := c.GetString("user_did")
-	// Support JSON body preferentially, fallback to Query/Form
 	var req struct {
 		Collection string `json:"collection"`
 		Repo       string `json:"repo"`
···
 		return
 	}
 
-	// Verify existence using Slingshot/Agnostic
 	recordResponse, err := agnostic.RepoGetRecord(c.Request.Context(), sl, "", req.Collection, req.Repo, req.RKey)
 	if err != nil {
-		// Does not exist remotely -> Delete locally
 		s.processRecord(req.Repo, req.Collection, req.RKey, true)
 
-		// You might want to return 200 even if deleted, to confirm "indexing done"
 		c.Status(200)
 		return
 	}
 
-	// Exists remotely -> Parse and Insert locally
 	uri := recordResponse.Uri
 	aturi, err := syntax.ParseATURI(uri)
 	if err != nil {

cmd/jetrelay/main.go

···
 package main
 
 import (
-	"flag"
+	"context"
+	"encoding/json"
 	"fmt"
+	"io"
+	"log"
+	"net/http"
+	"sort"
+	"sync"
+	"time"
+
+	"github.com/gorilla/websocket"
+	"github.com/klauspost/compress/zstd"
 )
 
-type multiFlag []string
+const (
+	ServerPort     = ":3878"
+	DictionaryURL  = "https://raw.githubusercontent.com/bluesky-social/jetstream/main/pkg/models/zstd_dictionary"
+	BufferSize     = 100000
+	ReconnectDelay = 5 * time.Second
+)
 
-func (m *multiFlag) String() string {
-	return fmt.Sprint(*m)
+var SourceJetstreams = []string{
+	"ws://localhost:6008/subscribe", // local jetstream
+	"ws://localhost:3877/subscribe", // local backstream
+}
+
+type Event struct {
+	Kind   string          `json:"kind"`
+	TimeUS int64           `json:"time_us"`
+	Commit json.RawMessage `json:"commit,omitempty"`
+}
+
+type BufferedEvent struct {
+	RelayTimeUS int64
+	RawJSON     []byte
 }
 
-func (m *multiFlag) Set(value string) error {
-	*m = append(*m, value)
+type History struct {
+	events []BufferedEvent
+	mu     sync.RWMutex
+}
+
+func (h *History) Add(jsonBytes []byte, relayTime int64) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	h.events = append(h.events, BufferedEvent{
+		RelayTimeUS: relayTime,
+		RawJSON:     jsonBytes,
+	})
+
+	if len(h.events) > BufferSize {
+		h.events = h.events[len(h.events)-BufferSize:]
+	}
+}
+
+func (h *History) GetSince(cursor int64) []BufferedEvent {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+
+	idx := sort.Search(len(h.events), func(i int) bool {
+		return h.events[i].RelayTimeUS > cursor
+	})
+
+	if idx < len(h.events) {
+		result := make([]BufferedEvent, len(h.events)-idx)
+		copy(result, h.events[idx:])
+		return result
+	}
 	return nil
 }
 
+var (
+	history  = &History{events: make([]BufferedEvent, 0, BufferSize)}
+	zstdDict []byte
+	hub      *Hub
+	upgrader = websocket.Upgrader{
+		CheckOrigin: func(r *http.Request) bool { return true },
+	}
+)
+
 func main() {
-	var js multiFlag
-	flag.Var(&js, "j", "jetstream instances 'write multiple to input more than one'")
+	log.Println("Initializing Relay...")
+
+	var err error
+	zstdDict, err = downloadDictionary()
+	if err != nil {
+		log.Fatalf("Failed to load dictionary: %v", err)
+	}
+
+	hub = newHub()
+	go hub.run()
+
+	ctx := context.Background()
+	for i, url := range SourceJetstreams {
+		go runUpstreamConsumer(ctx, i, url)
+	}
+
+	http.HandleFunc("/subscribe", serveWs)
+	log.Printf("🔥 Relay Active on %s", ServerPort)
+	if err := http.ListenAndServe(ServerPort, nil); err != nil {
+		log.Fatal(err)
+	}
+}
+
+func runUpstreamConsumer(ctx context.Context, id int, baseURL string) {
+	var lastSeenCursor int64 = 0
+
+	for {
+		connectURL := baseURL
+		if lastSeenCursor > 0 {
+			connectURL = fmt.Sprintf("%s?cursor=%d", baseURL, lastSeenCursor)
+			log.Printf("[Input %d] Reconnecting with cursor: %d", id, lastSeenCursor)
+		} else {
+			log.Printf("[Input %d] Connecting fresh...", id)
+		}
+
+		conn, _, err := websocket.DefaultDialer.Dial(connectURL, nil)
+		if err != nil {
+			log.Printf("[Input %d] Connect failed: %v. Retrying...", id, err)
+			time.Sleep(ReconnectDelay)
+			continue
+		}
+
+		log.Printf("[Input %d] Connected.", id)
+
+		for {
+			_, msg, err := conn.ReadMessage()
+			if err != nil {
+				log.Printf("[Input %d] Read error: %v", id, err)
+				break
+			}
+
+			var genericEvent map[string]interface{}
+			if err := json.Unmarshal(msg, &genericEvent); err != nil {
+				continue
+			}
+
+			if t, ok := genericEvent["time_us"].(float64); ok {
+				lastSeenCursor = int64(t)
+			}
+
+			nowUS := time.Now().UnixMicro()
+			genericEvent["time_us"] = nowUS
+
+			finalBytes, err := json.Marshal(genericEvent)
+			if err != nil {
+				continue
+			}
+
+			history.Add(finalBytes, nowUS)
+
+			hub.broadcast <- BufferedEvent{RelayTimeUS: nowUS, RawJSON: finalBytes}
+		}
+		conn.Close()
+		time.Sleep(ReconnectDelay)
+	}
+}
+
+func serveWs(w http.ResponseWriter, r *http.Request) {
+	conn, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		return
+	}
+
+	compress := r.URL.Query().Get("compress") == "true"
+
+	var clientCursor int64 = 0
+	cursorStr := r.URL.Query().Get("cursor")
+	if cursorStr != "" {
+		fmt.Sscanf(cursorStr, "%d", &clientCursor)
+	}
+
+	client := &Client{
+		hub:        hub,
+		conn:       conn,
+		send:       make(chan BufferedEvent, 2048),
+		compress:   compress,
+		lastSentUS: 0,
+	}
+
+	if compress {
+		enc, _ := zstd.NewWriter(nil, zstd.WithEncoderDict(zstdDict))
+		client.encoder = enc
+	}
+
+	client.hub.register <- client
+
+	go client.writePump()
+
+	if clientCursor > 0 {
+		log.Printf("Client requested replay from %d", clientCursor)
+		missedEvents := history.GetSince(clientCursor)
+		for _, evt := range missedEvents {
+			client.send <- evt
+		}
+	}
+
+	go client.readPump()
+}
+
+type Client struct {
+	hub        *Hub
+	conn       *websocket.Conn
+	send       chan BufferedEvent
+	compress   bool
+	encoder    *zstd.Encoder
+	lastSentUS int64
+}
+
+type Hub struct {
+	clients    map[*Client]bool
+	broadcast  chan BufferedEvent
+	register   chan *Client
+	unregister chan *Client
+	mu         sync.RWMutex
+}
+
+func newHub() *Hub {
+	return &Hub{
+		clients:    make(map[*Client]bool),
+		broadcast:  make(chan BufferedEvent, 10000),
+		register:   make(chan *Client),
+		unregister: make(chan *Client),
+	}
+}
+
+func (h *Hub) run() {
+	for {
+		select {
+		case client := <-h.register:
+			h.mu.Lock()
+			h.clients[client] = true
+			h.mu.Unlock()
+
+		case client := <-h.unregister:
+			h.mu.Lock()
+			if _, ok := h.clients[client]; ok {
+				delete(h.clients, client)
+				close(client.send)
+				if client.encoder != nil {
+					client.encoder.Close()
+				}
+			}
+			h.mu.Unlock()
+
+		case msg := <-h.broadcast:
+			h.mu.RLock()
+			for client := range h.clients {
+				select {
+				case client.send <- msg:
+				default:
+					go func(c *Client) {
+						h.unregister <- c
+						c.conn.Close()
+					}(client)
+				}
+			}
+			h.mu.RUnlock()
+		}
+	}
+}
+
+func (c *Client) writePump() {
+	defer c.conn.Close()
+
+	for msg := range c.send {
+		if msg.RelayTimeUS <= c.lastSentUS {
+			continue
+		}
+
+		c.lastSentUS = msg.RelayTimeUS
+
+		if c.compress {
+			compressed := c.encoder.EncodeAll(msg.RawJSON, nil)
+			if err := c.conn.WriteMessage(websocket.BinaryMessage, compressed); err != nil {
+				return
+			}
+		} else {
+			if err := c.conn.WriteMessage(websocket.TextMessage, msg.RawJSON); err != nil {
+				return
+			}
+		}
+	}
+}
 
-	flag.Parse()
+func (c *Client) readPump() {
+	defer func() {
+		c.hub.unregister <- c
+		c.conn.Close()
+	}()
+	for {
+		if _, _, err := c.conn.ReadMessage(); err != nil {
+			break
+		}
+	}
+}
 
-	fmt.Println(js) // prints: [hi hello what]
+func downloadDictionary() ([]byte, error) {
+	resp, err := http.Get(DictionaryURL)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+	return io.ReadAll(resp.Body)
 }