Something went wrong. Try again.
collection of golang services under the Red Dwarf umbrella forked from whey.party/red-dwarf-server
Something went wrong. Try again.
15 kB · 547 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548package backstream
import ( //"bytes" "context" "encoding/json" "errors" "fmt" "log" "net/http" "strings" "sync" "time"
"io" "io/ioutil" "os"
"runtime" "runtime/debug"
"github.com/gorilla/websocket" "github.com/klauspost/compress/zstd"
data "github.com/bluesky-social/indigo/atproto/atdata" atrepo "github.com/bluesky-social/indigo/atproto/repo"
// "github.com/bluesky-social/indigo/repo" "github.com/bluesky-social/indigo/atproto/syntax" "github.com/ipfs/go-cid")
const ( numWorkers = 20)
var DefaultUpgrader = websocket.Upgrader{ CheckOrigin: func(r *http.Request) bool { return true },}
type BackfillHandler struct { Upgrader websocket.Upgrader SessionManager *SessionManager AtpClient *ATProtoClient ZstdDict []byte UseGetRepoMethod bool}
type BackfillParams struct { WantedDIDs []string WantedCollections []string GetRecordFormat bool}
func (h *BackfillHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { compress := (r.URL.Query().Get("compress") == "true") && (h.ZstdDict != nil)
conn, err := h.Upgrader.Upgrade(w, r, nil) if err != nil { log.Printf("Failed to upgrade connection: %v", err) return } defer conn.Close()
if compress { log.Println("Client requested zstd compression. Enabling.") }
params, ticket, err := h.parseQueryParams(r) if err != nil { h.sendError(conn, err.Error()) return }
log.Printf("New connection for ticket: %s. DIDs: %v, Collections: %v, Workers: %d", ticket, params.WantedDIDs, params.WantedCollections, numWorkers)
session := h.SessionManager.GetOrCreate(ticket, params) session.LastAccessed = time.Now()
ctx, cancel := context.WithCancel(r.Context()) defer cancel()
go func() { defer cancel() for { if _, _, err := conn.ReadMessage(); err != nil { if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) { log.Printf("Client disconnected for ticket %s (read error): %v", ticket, err) } break } } }()
var wg sync.WaitGroup jobs := make(chan string, numWorkers) results := make(chan interface{}, 100)
for i := 1; i <= numWorkers; i++ { go h.worker(ctx, i, &wg, jobs, results, session) }
writerDone := make(chan struct{}) if compress { go h.compressedWriter(ctx, cancel, conn, results, writerDone) } else { go h.writer(ctx, cancel, conn, results, writerDone) }
wg.Add(1) go h.producer(ctx, &wg, jobs, session)
wg.Wait() close(results) <-writerDone
log.Printf("Backfill completed for ticket: %s", session.Ticket) h.sendMessage(conn, map[string]string{"status": "complete", "message": "Backfill finished."})}
func (h *BackfillHandler) compressedWriter(ctx context.Context, cancel context.CancelFunc, conn *websocket.Conn, results <-chan interface{}, done chan<- struct{}) { defer close(done)
encoder, err := zstd.NewWriter(nil, zstd.WithEncoderDict(h.ZstdDict)) if err != nil { log.Printf("ERROR: [CompressedWriter] Failed to create zstd encoder with dictionary: %v", err) cancel() return } defer encoder.Close()
for { select { case result, ok := <-results: if !ok { return }
data, err := json.Marshal(result) if err != nil { log.Printf("ERROR: [CompressedWriter] Failed to marshal JSON: %v", err) cancel() return }
compressed := encoder.EncodeAll(data, nil)
if err := conn.WriteMessage(websocket.BinaryMessage, compressed); err != nil { log.Printf("ERROR: [CompressedWriter] Failed to write compressed message: %v", err) cancel() return } case <-ctx.Done(): log.Printf("[CompressedWriter] Context cancelled, stopping.") return } }}
func (h *BackfillHandler) writer(ctx context.Context, cancel context.CancelFunc, conn *websocket.Conn, results <-chan interface{}, done chan<- struct{}) { defer close(done) for { select { case result, ok := <-results: if !ok { return } if err := h.sendMessage(conn, result); err != nil { log.Printf("ERROR: [Writer] Failed to write message, closing connection: %v", err) cancel() return } case <-ctx.Done(): log.Printf("[Writer] Context cancelled, stopping.") return } }}
func (h *BackfillHandler) producer(ctx context.Context, wg *sync.WaitGroup, jobs chan<- string, session *Session) { defer close(jobs) defer wg.Done()
isFullNetwork := len(session.Params.WantedDIDs) == 1 && session.Params.WantedDIDs[0] == "*" isAllCollections := len(session.Params.WantedCollections) == 1 && session.Params.WantedCollections[0] == "*"
if isFullNetwork { if isAllCollections { // --- Case 1: Full Network, All Collections (dids=*&collections=*) --- // We need to list *all* repos from the relay. log.Printf("[Producer] Starting full network scan for all collections.") for { select { case <-ctx.Done(): log.Printf("[Producer] Context cancelled, stopping full repo fetch.") return default: }
log.Printf("[Producer] Fetching all repos with cursor: %s", session.ListReposCursor) repos, nextCursor, err := h.AtpClient.ListRepos(ctx, session.ListReposCursor) if err != nil { log.Printf("ERROR: [Producer] Failed to list all repos: %v", err) return }
for _, repo := range repos { if !session.IsDIDComplete(repo.DID) { wg.Add(1) jobs <- repo.DID } }
session.mu.Lock() session.ListReposCursor = nextCursor session.LastAccessed = time.Now() session.mu.Unlock()
if nextCursor == "" { log.Printf("[Producer] Finished fetching all repos from relay.") break } } } else { // --- Case 2: Full Network, Specific Collections (dids=*&collections=a,b,c) --- // For each specific collection, page through all repos and send DIDs to workers. log.Printf("[Producer] Starting network scan for specific collections: %v", session.Params.WantedCollections) for _, collection := range session.Params.WantedCollections { for { select { case <-ctx.Done(): log.Printf("[Producer] Context cancelled, stopping repo fetch.") return default: }
log.Printf("[Producer] Fetching repos for %s with cursor: %s", collection, session.ListReposCursor) repos, nextCursor, err := h.AtpClient.ListReposByCollection(ctx, collection, session.ListReposCursor) if err != nil { log.Printf("ERROR: [Producer] Failed to list repos for collection %s: %v", collection, err) return }
for _, repo := range repos { if !session.IsDIDComplete(repo.DID) { wg.Add(1) jobs <- repo.DID } }
session.mu.Lock() session.ListReposCursor = nextCursor session.LastAccessed = time.Now() session.mu.Unlock()
if nextCursor == "" { log.Printf("[Producer] Finished fetching all repos for collection %s", collection) break } } } } } else { // --- Case 3: Specific List of DIDs (dids=a,b,c) --- // Send user-provided DIDs to workers. for _, did := range session.Params.WantedDIDs { select { case <-ctx.Done(): log.Printf("[Producer] Context cancelled, stopping DID processing.") return default: if !session.IsDIDComplete(did) { wg.Add(1) jobs <- did } else { log.Printf("[Producer] Skipping already completed DID: %s", did) } } } }}
func (h *BackfillHandler) worker(ctx context.Context, id int, wg *sync.WaitGroup, jobs <-chan string, results chan<- interface{}, session *Session) { for did := range jobs { func(did string) { defer func() { wg.Done()
runtime.GC() debug.FreeOSMemory()
log.Printf("[Worker %d] Cleaned up resources for DID: %s", id, did) }()
select { case <-ctx.Done(): return default: }
log.Printf("[Worker %d] Processing DID: %s", id, did) pdsURL, err := h.AtpClient.ResolveDID(ctx, did) if err != nil { log.Printf("WARN: [Worker %d] Could not resolve DID %s, skipping. Error: %v", id, did, err) return }
if h.UseGetRepoMethod { h.processDIDWithGetRepo(ctx, id, did, pdsURL, results, session) } else { h.processDIDWithListRecords(ctx, id, did, pdsURL, results, session) }
session.MarkDIDComplete(did) log.Printf("[Worker %d] Finished DID: %s", id, did) }(did) }}
func (h *BackfillHandler) processDIDWithGetRepo(ctx context.Context, id int, did, pdsURL string, results chan<- interface{}, session *Session) { log.Printf("[Worker %d] Using streaming getRepo method for %s", id, did) isAllCollections := len(session.Params.WantedCollections) == 1 && session.Params.WantedCollections[0] == "*"
wantedSet := make(map[string]struct{}) if !isAllCollections { for _, coll := range session.Params.WantedCollections { wantedSet[coll] = struct{}{} } }
respBody, err := h.AtpClient.GetRepo(ctx, pdsURL, did) if err != nil { log.Printf("WARN: [Worker %d] Failed to get repo stream for %s: %v", id, did, err) return } defer respBody.Close()
if err := os.MkdirAll("./temp", 0o755); err != nil { panic(err) } tempFile, err := ioutil.TempFile("./temp", "backstream-repo-*.car") if err != nil { log.Printf("ERROR: [Worker %d] Failed to create temp file for %s: %v", id, did, err) return } defer os.Remove(tempFile.Name())
if _, err := io.Copy(tempFile, respBody); err != nil { log.Printf("ERROR: [Worker %d] Failed to write repo to temp file for %s: %v", id, did, err) return }
if err := tempFile.Close(); err != nil { log.Printf("ERROR: [Worker %d] Failed to close temp file for %s: %v", id, did, err) return }
readHandle, err := os.Open(tempFile.Name()) if err != nil { log.Printf("ERROR: [Worker %d] Failed to open temp file for reading %s: %v", id, did, err) return } defer readHandle.Close()
_, r, err := atrepo.LoadRepoFromCAR(ctx, readHandle) if err != nil { log.Printf("WARN: [Worker %d] Failed to read CAR stream for %s from temp file: %v", id, did, err) return }
err = r.MST.Walk(func(k []byte, v cid.Cid) error { select { case <-ctx.Done(): return errors.New("context cancelled during repo walk") default: }
path := string(k) collection, rkey, err := syntax.ParseRepoPath(path) if err != nil { log.Printf("WARN: [Worker %d] Could not parse repo path '%s' for %s, skipping record", id, path, did) return nil }
if !isAllCollections { if _, ok := wantedSet[string(collection)]; !ok { return nil } }
recBytes, _, err := r.GetRecordBytes(ctx, collection, rkey) if err != nil { log.Printf("WARN: [Worker %d] Failed to get record bytes for %s: %v", id, path, err) return nil }
recordVal, err := data.UnmarshalCBOR(recBytes) if err != nil { log.Printf("WARN: [Worker %d] Failed to unmarshal record CBOR for %s: %v", id, path, err) return nil }
record := Record{ URI: fmt.Sprintf("at://%s/%s", did, path), CID: v.String(), Value: recordVal, }
output := h.formatOutput(record, did, string(collection), session.Params.GetRecordFormat) select { case results <- output: case <-ctx.Done(): return errors.New("context cancelled while sending result") }
session.SetListRecordsCursor(did, string(collection), string(rkey)) return nil })
if err != nil && !errors.Is(err, context.Canceled) { log.Printf("WARN: [Worker %d] Error while walking repo for %s: %v", id, did, err) }}
func (h *BackfillHandler) processDIDWithListRecords(ctx context.Context, id int, did, pdsURL string, results chan<- interface{}, session *Session) { log.Printf("[Worker %d] Using listRecords method for %s", id, did) isAllCollections := len(session.Params.WantedCollections) == 1 && session.Params.WantedCollections[0] == "*" var collectionsToProcess []string
if isAllCollections { repoCollections, err := h.AtpClient.DescribeRepo(ctx, pdsURL, did) if err != nil { log.Printf("WARN: [Worker %d] Could not describe repo for %s to find collections, skipping. Error: %v", id, did, err) return } collectionsToProcess = repoCollections log.Printf("[Worker %d] Found %d collections for DID %s", id, len(collectionsToProcess), did) } else { collectionsToProcess = session.Params.WantedCollections }
for _, collection := range collectionsToProcess { cursor := session.GetListRecordsCursor(did, collection) for { select { case <-ctx.Done(): log.Printf("[Worker %d] Context cancelled for DID %s", id, did) return default: }
records, nextCursor, err := h.AtpClient.ListRecords(ctx, pdsURL, did, collection, cursor) if err != nil { if !strings.Contains(err.Error(), "status: 400") { log.Printf("WARN: [Worker %d] Failed to list records for %s/%s, skipping. Error: %v", id, did, collection, err) } break }
for _, record := range records { output := h.formatOutput(record, did, collection, session.Params.GetRecordFormat) select { case results <- output: case <-ctx.Done(): log.Printf("[Worker %d] Context cancelled while sending results for %s", id, did) return } }
session.SetListRecordsCursor(did, collection, nextCursor) cursor = nextCursor if cursor == "" { break } } }}
func (h *BackfillHandler) parseQueryParams(r *http.Request) (BackfillParams, string, error) { query := r.URL.Query() ticket := query.Get("ticket")
wantedDidsStr := query.Get("wantedDids") wantedCollectionsStr := query.Get("wantedCollections")
if wantedCollectionsStr == "" && wantedDidsStr == "" && ticket == "" { ticket = "jetstreamfalse" } else if ticket == "" { ticket = generateTicket() }
if wantedDidsStr == "" { log.Println("Query parameter 'wantedDids' not specified, defaulting to '*' (all repos).") wantedDidsStr = "*" }
if wantedCollectionsStr == "" { log.Println("Query parameter 'wantedCollections' not specified, defaulting to '*' (all collections).") wantedCollectionsStr = "*" }
params := BackfillParams{ WantedDIDs: strings.Split(wantedDidsStr, ","), WantedCollections: strings.Split(wantedCollectionsStr, ","), GetRecordFormat: query.Get("getRecordFormat") == "true", } return params, ticket, nil}
func (h *BackfillHandler) formatOutput(record Record, did, collection string, getRecordFormat bool) interface{} { if getRecordFormat { return GetRecordOutput{ URI: record.URI, CID: record.CID, Value: record.Value, } } uriParts := strings.Split(record.URI, "/") rkey := "" if len(uriParts) == 5 { rkey = uriParts[4] } return JetstreamLikeOutput{ Did: did, Kind: "commit", TimeUS: "1725911162329308", Commit: JetstreamLikeCommit{ Rev: rkey, Operation: "create", Collection: collection, RKey: rkey, Record: record.Value, CID: record.CID, }, }}
func (h *BackfillHandler) sendError(conn *websocket.Conn, message string) { log.Printf("Sending error to client: %s", message) _ = conn.WriteJSON(map[string]string{"error": message})}
func (h *BackfillHandler) sendMessage(conn *websocket.Conn, v interface{}) error { return conn.WriteJSON(v)}