diff --git a/cmd/collectiondir/Dockerfile b/cmd/collectiondir/Dockerfile new file mode 100644 index 00000000..3ab8e247 --- /dev/null +++ b/cmd/collectiondir/Dockerfile @@ -0,0 +1,43 @@ +FROM golang:1.23-bullseye AS build-env + +ENV DEBIAN_FRONTEND=noninteractive +ENV TZ=Etc/UTC +ENV GODEBUG="netdns=go" +ENV GOOS="linux" +ENV GOARCH="amd64" +ENV CGO_ENABLED="1" + +WORKDIR /usr/src/collectiondir + +COPY . . + +RUN go mod download && \ + go mod verify + +RUN go build \ + -v \ + -trimpath \ + -tags timetzdata \ + -o /collectiondir-bin \ + ./cmd/collectiondir + +FROM debian:bullseye-slim + +ENV DEBIAN_FRONTEND="noninteractive" +ENV TZ=Etc/UTC +ENV GODEBUG="netdns=go" + +RUN apt-get update && apt-get install --yes \ + dumb-init \ + ca-certificates \ + runit + +WORKDIR /collectiondir +COPY --from=build-env /collectiondir-bin /usr/bin/collectiondir + +ENTRYPOINT ["/usr/bin/dumb-init", "--"] +CMD ["/usr/bin/collectiondir"] + +LABEL org.opencontainers.image.source=https://github.com/bluesky-social/indigo +LABEL org.opencontainers.image.description="collectiondir " +LABEL org.opencontainers.image.licenses=MIT diff --git a/cmd/collectiondir/README.md b/cmd/collectiondir/README.md new file mode 100644 index 00000000..5136bdaf --- /dev/null +++ b/cmd/collectiondir/README.md @@ -0,0 +1,56 @@ +# Collection Directory + +Maintain a directory of which repos use which collections of records. + +e.g. "app.bsky.feed.post" is used by did:alice did:bob + +Firehose consumer and crawler of PDS via listRepos and describeRepo. + +The primary query is: + +``` +/v1/getDidsForCollection?collection={}&cursor={} +``` + +It returns JSON: + +```json +{"dids":["did:A", "..."], +"cursor":"opaque text"} +``` + +query parameter `collection` may be repeated up to 10 times. They must always be sent in the same order or the cursor will break. + +If multiple collections are specified, the result stream is not guaranteed to be de-duplicated on Did and Dids may be repeated. +(A merge window is used so that the service is _likely_ to not send duplicate Dids.) + + +### Analytics queries + +``` +/v1/listCollections?c={}&cursor={}&limit={50<=limit<=1000} +``` + +`listCollections` returns JSON with a map of collection name to approximate number of dids implementing it. +With no `c` parameter it returns all known collections with cursor paging. +With up to 20 repeated `c` paramaters it returns only those collections (no paging). +It may be the cached result of a computation, up to several minutes out of date. +```json +{"collections":{"app.bsky.feed.post": 123456789, "some collection": 42}, +"cursor":"opaque text"} +``` + + +## Design + +### Schema + +The primary database is (collection, seen time int64 milliseconds, did) + +This allows for efficient cursor fetching of more dids for a collection. + +e.g. A new service starts consuming the firehose for events it wants in collection `com.newservice.data.thing`, +it then calls the collection directory for a list of repos which may have already created data in this collection, +and does `getRepo` calls to those repo's PDSes to get prior data. +By the time it is done paging forward through the collection directory results and getting those repos, +it will have backfilled data and new data it has collected live off the firehose. \ No newline at end of file diff --git a/cmd/collectiondir/collectiondir.go b/cmd/collectiondir/collectiondir.go new file mode 100644 index 00000000..fac2f0dc --- /dev/null +++ b/cmd/collectiondir/collectiondir.go @@ -0,0 +1,151 @@ +package main + +import ( + "compress/gzip" + "encoding/csv" + "encoding/json" + "errors" + "fmt" + "github.com/carlmjohnson/versioninfo" + "github.com/urfave/cli/v2" + "io" + "log/slog" + "os" + "strings" +) + +func main() { + app := cli.App{ + Name: "collectiondir", + Usage: "collection directory service", + Version: versioninfo.Short(), + Flags: []cli.Flag{ + &cli.BoolFlag{ + Name: "verbose", + }, + }, + Commands: []*cli.Command{ + serveCmd, + crawlCmd, + buildCmd, + statsCmd, + }, + } + err := app.Run(os.Args) + if err != nil { + fmt.Fprintf(os.Stderr, "%s\n", err.Error()) + os.Exit(1) + } +} + +var statsCmd = &cli.Command{ + Name: "stats", + Flags: []cli.Flag{ + &cli.StringFlag{ + Name: "pebble", + Usage: "path to store pebble db", + Required: true, + }, + }, + Action: func(cctx *cli.Context) error { + logLevel := slog.LevelInfo + if cctx.Bool("verbose") { + logLevel = slog.LevelDebug + } + log := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: logLevel})) + slog.SetDefault(log) + pebblePath := cctx.String("pebble") + var db PebbleCollectionDirectory + db.log = log + err := db.Open(pebblePath) + if err != nil { + return err + } + defer db.Close() + + stats, err := db.GetCollectionStats() + if err != nil { + return err + } + blob, err := json.MarshalIndent(stats, "", " ") + os.Stdout.Write(blob) + os.Stdout.Write([]byte{'\n'}) + return nil + }, +} + +var buildCmd = &cli.Command{ + Name: "build", + Usage: "collect csv into a database", + Flags: []cli.Flag{ + &cli.StringFlag{ + Name: "csv", + Required: true, + }, + &cli.StringFlag{ + Name: "pebble", + Usage: "path to store pebble db", + Required: true, + }, + }, + Action: func(cctx *cli.Context) error { + logLevel := slog.LevelInfo + if cctx.Bool("verbose") { + logLevel = slog.LevelDebug + } + log := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: logLevel})) + slog.SetDefault(log) + pebblePath := cctx.String("pebble") + var db PebbleCollectionDirectory + db.log = log + err := db.Open(pebblePath) + if err != nil { + return err + } + defer db.Close() + csvPath := cctx.String("csv") + var fin io.Reader + if csvPath == "-" { + fin = os.Stdin + } else if strings.HasSuffix(csvPath, ".gz") { + osin, err := os.Open(csvPath) + if err != nil { + return fmt.Errorf("%s: could not open csv, %w", csvPath, err) + } + defer osin.Close() + gzin, err := gzip.NewReader(osin) + if err != nil { + return fmt.Errorf("%s: could not open csv, %w", csvPath, err) + } + defer gzin.Close() + fin = gzin + } else { + osin, err := os.Open(csvPath) + if err != nil { + return fmt.Errorf("%s: could not open csv, %w", csvPath, err) + } + defer osin.Close() + fin = osin + } + reader := csv.NewReader(fin) + rowcount := 0 + results := make(chan DidCollection, 100) + go db.SetFromResults(results) + for { + row, err := reader.Read() + if errors.Is(err, io.EOF) { + break + } + did := row[0] + collection := row[1] + results <- DidCollection{ + Did: did, + Collection: collection, + } + rowcount++ + } + close(results) + log.Debug("read csv", "rows", rowcount) + return nil + }, +} diff --git a/cmd/collectiondir/crawl.go b/cmd/collectiondir/crawl.go new file mode 100644 index 00000000..a86f54e9 --- /dev/null +++ b/cmd/collectiondir/crawl.go @@ -0,0 +1,152 @@ +package main + +import ( + "context" + "encoding/csv" + "fmt" + "golang.org/x/time/rate" + "io" + "log/slog" + "net/http" + "net/url" + "os" + + "github.com/urfave/cli/v2" + + "github.com/bluesky-social/indigo/api/atproto" + "github.com/bluesky-social/indigo/xrpc" +) + +type DidCollection struct { + Did string `json:"d"` + Collection string `json:"c"` +} + +func DidCollectionsToCsv(out io.Writer, sources <-chan DidCollection) { + writer := csv.NewWriter(out) + defer writer.Flush() + var row [2]string + for dc := range sources { + row[0] = dc.Did + row[1] = dc.Collection + writer.Write(row[:]) + } +} + +var crawlCmd = &cli.Command{ + Name: "crawl", + Usage: "crawl a PDS", + Flags: []cli.Flag{ + &cli.StringFlag{ + Name: "host", + Usage: "hostname or URL of PDS", + }, + &cli.StringFlag{ + Name: "csv-out", + Usage: "path for output or - for stdout", + }, + &cli.Float64Flag{ + Name: "qps", + Usage: "queries per second to do vs target PDS", + Value: 50, // large PDS: 500_000 repos, 10_000 seconds, ~3 hours + }, + &cli.StringFlag{ + Name: "ratelimit-header", + Usage: "secret for friend PDSes", + EnvVars: []string{"BSKY_SOCIAL_RATE_LIMIT_SKIP", "RATE_LIMIT_HEADER"}, + }, + }, + Action: func(cctx *cli.Context) error { + log := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelInfo})) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + hostname := cctx.String("host") + hosturl, err := url.Parse(hostname) + if err != nil { + hosturl = new(url.URL) + hosturl.Scheme = "https" + hosturl.Host = hostname + } + httpClient := http.Client{} + rpcClient := xrpc.Client{ + Host: hosturl.String(), + Client: &httpClient, + } + if cctx.IsSet("ratelimit-header") { + rpcClient.Headers = map[string]string{ + "x-ratelimit-bypass": cctx.String("ratelimit-header"), + } + } + log.Info("will crawl", "url", rpcClient.Host) + csvOutPath := cctx.String("csv-out") + var fout io.Writer = os.Stdout + if csvOutPath != "" { + if csvOutPath == "-" { + fout = os.Stdout + } else { + fout, err = os.Create(csvOutPath) + if err != nil { + return fmt.Errorf("%s: could not open for writing: %w", csvOutPath, err) + } + } + } + qps := cctx.Float64("qps") + results := make(chan DidCollection, 100) + defer close(results) + go DidCollectionsToCsv(fout, results) + crawler := Crawler{ + Ctx: ctx, + RpcClient: &rpcClient, + QPS: qps, + Results: results, + Log: log, + } + err = crawler.CrawlPDSRepoCollections() + log.Info("done") + + return err + }, +} + +type Crawler struct { + Ctx context.Context + RpcClient *xrpc.Client + QPS float64 + Results chan<- DidCollection + Log *slog.Logger +} + +// CrawlPDSRepoCollections +// write results to chan +// does _not_ close chan +// (allow multiple threads of PDS queries running to one output chan, e.g. feeding into SetFromResults() ) +func (cr *Crawler) CrawlPDSRepoCollections() error { + var cursor string + limiter := rate.NewLimiter(rate.Limit(cr.QPS), 1) + for { + limiter.Wait(cr.Ctx) + repos, err := atproto.SyncListRepos(cr.Ctx, cr.RpcClient, cursor, 1000) + if err != nil { + // TODO: wait N seconds, retry M times + return fmt.Errorf("%s: sync repos: %w", cr.RpcClient.Host, err) + } + slog.Info("got repo list", "count", len(repos.Repos)) + for _, xr := range repos.Repos { + limiter.Wait(cr.Ctx) + desc, err := atproto.RepoDescribeRepo(cr.Ctx, cr.RpcClient, xr.Did) + if err != nil { + slog.Error("repo desc", "host", cr.RpcClient.Host, "did", xr.Did, "err", err) + continue + } + for _, collection := range desc.Collections { + cr.Results <- DidCollection{Did: xr.Did, Collection: collection} + } + } + if repos.Cursor != nil { + cursor = *repos.Cursor + } else { + break + } + } + return nil +} diff --git a/cmd/collectiondir/firehose.go b/cmd/collectiondir/firehose.go new file mode 100644 index 00000000..9cceac6b --- /dev/null +++ b/cmd/collectiondir/firehose.go @@ -0,0 +1,103 @@ +package main + +import ( + "context" + "fmt" + "github.com/bluesky-social/indigo/events" + "github.com/gorilla/websocket" + "log/slog" + "net/http" + "net/url" + "time" +) + +type Firehose struct { + Log *slog.Logger + + Host string + Seq int64 + + events chan<- *events.XRPCStreamEvent +} + +func (fh *Firehose) subscribeWithRedialer(ctx context.Context, fhevents chan<- *events.XRPCStreamEvent) error { + defer close(fhevents) + d := websocket.Dialer{} + + rurl, err := url.Parse(fh.Host) + if err != nil { + rurl = new(url.URL) + rurl.Host = fh.Host + rurl.Scheme = "wss" + } else { + if rurl.Scheme == fh.Host { + rurl.Scheme = "wss" + } + if rurl.Scheme == "https" || rurl.Scheme == "wss" { + rurl.Scheme = "wss" + } else if rurl.Scheme == "http" || rurl.Scheme == "ws" { + rurl.Scheme = "ws" + } else if rurl.Scheme == "" { + rurl.Scheme = "wss" + } else { + return fmt.Errorf("host unknown scheme %#v", rurl.Scheme) + } + } + //protocol := "wss" + subscribeReposUrl := rurl.JoinPath("/xrpc/com.atproto.sync.subscribeRepos") + fh.events = fhevents + + var backoff int + for { + select { + case <-ctx.Done(): + return nil + default: + } + + header := http.Header{ + "User-Agent": []string{"bgs-rainbow-v0"}, + } + + if fh.Seq >= 0 { + subscribeReposUrl.RawQuery = fmt.Sprintf("cursor=%d", fh.Seq) + } + url := subscribeReposUrl.String() + con, res, err := d.DialContext(ctx, url, header) + if err != nil { + fh.Log.Warn("dialing failed", "url", url, "err", err, "backoff", backoff) + time.Sleep(5 * time.Second) + backoff++ + + continue + } + + fh.Log.Info("event subscription response", "code", res.StatusCode) + + if err := fh.handleConnection(ctx, con); err != nil { + fh.Log.Warn("connection failed", "host", fh.Host, "err", err) + } + } +} + +func (fh *Firehose) handleConnection(ctx context.Context, con *websocket.Conn) error { + ctx, cancel := context.WithCancel(ctx) + defer cancel() + + return events.HandleRepoStream(ctx, con, fh, fh.Log) +} + +// AddWork is part of events.Scheduler +func (fh *Firehose) AddWork(ctx context.Context, repo string, val *events.XRPCStreamEvent) error { + tsv, ok := val.GetSequence() + if ok { + fh.Seq = tsv + } + fh.events <- val + return nil +} + +// Shutdown is part of events.Scheduler +func (fh *Firehose) Shutdown() { + // unneeded in this usage +} diff --git a/cmd/collectiondir/metrics.go b/cmd/collectiondir/metrics.go new file mode 100644 index 00000000..597ce5d5 --- /dev/null +++ b/cmd/collectiondir/metrics.go @@ -0,0 +1,92 @@ +package main + +import ( + "errors" + "github.com/labstack/echo/v4" + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promauto" + "net/http" + "strconv" + "time" +) + +var firehoseReceivedCounter = promauto.NewCounter(prometheus.CounterOpts{ + Name: "collectiondir_firehose_received_total", + Help: "number of events received from upstream firehose", +}) +var firehoseCommits = promauto.NewCounter(prometheus.CounterOpts{ + Name: "collectiondir_firehose_commits", + Help: "number of #commit events received from upstream firehose", +}) +var firehoseCommitOps = promauto.NewCounterVec(prometheus.CounterOpts{ + Name: "collectiondir_firehose_commit_ops", + Help: "number of #commit events received from upstream firehose", +}, []string{"op"}) + +var firehoseDidcSet = promauto.NewCounter(prometheus.CounterOpts{ + Name: "collectiondir_firehose_didc_total", +}) + +var pebbleDup = promauto.NewCounter(prometheus.CounterOpts{ + Name: "collectiondir_pebble_dup_total", +}) + +var pebbleNew = promauto.NewCounter(prometheus.CounterOpts{ + Name: "collectiondir_pebble_new_total", +}) + +var pdsCrawledCounter = promauto.NewCounter(prometheus.CounterOpts{ + Name: "collectiondir_pds_crawled_total", +}) + +var reqDur = promauto.NewHistogramVec(prometheus.HistogramOpts{ + Name: "http_request_duration_seconds", + Help: "A histogram of latencies for requests.", + Buckets: prometheus.ExponentialBuckets(0.001, 2, 15), +}, []string{"code", "method", "path"}) + +var reqCnt = promauto.NewCounterVec(prometheus.CounterOpts{ + Name: "http_requests_total", + Help: "A counter for requests to the wrapped handler.", +}, []string{"code", "method", "path"}) + +// MetricsMiddleware defines handler function for metrics middleware +// TODO: reunify with bgs/metrics.go ? +func MetricsMiddleware(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + path := c.Path() + if path == "/metrics" || path == "/_health" { + return next(c) + } + + start := time.Now() + //requestSize := computeApproximateRequestSize(c.Request()) + + err := next(c) + + status := c.Response().Status + if err != nil { + var httpError *echo.HTTPError + if errors.As(err, &httpError) { + status = httpError.Code + } + if status == 0 || status == http.StatusOK { + status = http.StatusInternalServerError + } + } + + elapsed := float64(time.Since(start)) / float64(time.Second) + + statusStr := strconv.Itoa(status) + method := c.Request().Method + + //responseSize := float64(c.Response().Size) + + reqDur.WithLabelValues(statusStr, method, path).Observe(elapsed) + reqCnt.WithLabelValues(statusStr, method, path).Inc() + //reqSz.WithLabelValues(statusStr, method, path).Observe(float64(requestSize)) + //resSz.WithLabelValues(statusStr, method, path).Observe(responseSize) + + return err + } +} diff --git a/cmd/collectiondir/pebble.go b/cmd/collectiondir/pebble.go new file mode 100644 index 00000000..19a7d7c4 --- /dev/null +++ b/cmd/collectiondir/pebble.go @@ -0,0 +1,377 @@ +package main + +import ( + "context" + "encoding/base64" + "encoding/binary" + "errors" + "fmt" + "github.com/cockroachdb/pebble" + "log/slog" + "time" +) + +func makeCollectionInternKey(collection string) []byte { + out := make([]byte, len(collection)+1) + out[0] = 'C' + copy(out[1:], collection) + return out +} + +func parseCollectionInternKey(key []byte) string { + if key[0] != 'C' { + panic(fmt.Sprintf("collection key must start with C, got %v", key[0])) + } + return string(key[1:]) +} + +func makePrimaryPebbleRow(collectionId uint32, did string, seenMs int64) []byte { + out := make([]byte, 1+4+8+len(did)) + out[0] = 'A' + binary.BigEndian.PutUint32(out[1:], collectionId) + pos := 1 + 4 + binary.BigEndian.PutUint64(out[pos:], uint64(seenMs)) + pos += 8 + copy(out[pos:], did) + return out +} + +func parsePrimaryPebbleRow(row []byte) (collectionId uint32, did string, seenMs int64) { + if row[0] != 'A' { + panic(fmt.Sprintf("primary row key wanted A got %v", row[0])) + } + collectionId = binary.BigEndian.Uint32(row[1:5]) + seenMs = int64(binary.BigEndian.Uint64(row[5:13])) + did = string(row[13:]) + return collectionId, did, seenMs +} + +func makeByDidKey(did string, collectionId uint32) []byte { + out := make([]byte, 1+len(did)+4) + out[0] = 'D' + copy(out[1:1+len(did)], did) + pos := 1 + len(did) + binary.BigEndian.PutUint32(out[pos:], collectionId) + return out +} + +func parseByDidKey(key []byte) (did string, collectionId uint32) { + if key[0] != 'D' { + panic(fmt.Sprintf("by did key wanted D got %v", key[0])) + } + last4 := len(key) - 5 + collectionId = binary.BigEndian.Uint32(key[last4:]) + did = string(key[1 : last4+1]) + return did, collectionId +} + +// PebbleCollectionDirectory holds a DID<=>{collections} directory in pebble db. +// The primary database is (collection, seen time int64 milliseconds, did) +// Inner schema: +// C{collection} : {uint32 collectionId} +// D{did}{uint32 collectionId} : {uint64 seen ms} +// A{uint32 collectionId}{uint64 seen ms}{did} : 't' +type PebbleCollectionDirectory struct { + db *pebble.DB + + // collections can be LRU cache if it ever becomes too big + collections map[string]uint32 + collectionNames map[uint32]string // TODO: B-tree would be nice + maxCollectionId uint32 + + log *slog.Logger +} + +func (pcd *PebbleCollectionDirectory) Open(pebblePath string) error { + db, err := pebble.Open(pebblePath, &pebble.Options{}) + if err != nil { + return fmt.Errorf("%s: could not open db, %w", pebblePath, err) + } + pcd.db = db + pcd.collections = make(map[string]uint32) + pcd.collectionNames = make(map[uint32]string) + if pcd.log == nil { + pcd.log = slog.Default() + } + return pcd.ReadAllCollectionInterns(context.Background()) +} + +func (pcd *PebbleCollectionDirectory) Close() error { + err := pcd.db.Flush() + if err != nil { + pcd.log.Error("pebble flush", "err", err) + } + err = pcd.db.Close() + if err != nil { + pcd.log.Error("pebble close", "err", err) + } + return err +} + +func (pcd *PebbleCollectionDirectory) ReadAllCollectionInterns(ctx context.Context) error { + lower := []byte{'C'} + upper := []byte{'D'} + iter, err := pcd.db.NewIterWithContext(ctx, &pebble.IterOptions{ + LowerBound: lower, + UpperBound: upper, + }) + if err != nil { + return fmt.Errorf("collection iter start, %w", err) + } + defer iter.Close() + count := 0 + for iter.First(); iter.Valid(); iter.Next() { + key := iter.Key() + value, err := iter.ValueAndErr() + if err != nil { + return fmt.Errorf("collection iter, %w", err) + } + collection := parseCollectionInternKey(key) + collectionId := binary.BigEndian.Uint32(value) + count++ + pcd.collections[collection] = collectionId + pcd.collectionNames[collectionId] = collection + if collectionId > pcd.maxCollectionId { + pcd.maxCollectionId = collectionId + } + pcd.log.Debug("collection", "name", collection, "id", collectionId) + } + pcd.log.Debug("read collections", "count", count, "max", pcd.maxCollectionId) + return nil +} + +type CollectionDidTime struct { + Collection string + Did string + UnixMillis int64 +} + +func (pcd *PebbleCollectionDirectory) ReadAllPrimary(ctx context.Context, out chan<- CollectionDidTime) error { + defer close(out) + lower := []byte{'A'} + upper := []byte{'B'} + iter, err := pcd.db.NewIterWithContext(ctx, &pebble.IterOptions{ + LowerBound: lower, + UpperBound: upper, + }) + if err != nil { + return fmt.Errorf("collection iter start, %w", err) + } + defer iter.Close() + count := 0 + done := ctx.Done() + for iter.First(); iter.Valid(); iter.Next() { + key := iter.Key() + collectionId, did, seenMs := parsePrimaryPebbleRow(key) + count++ + collection := pcd.collectionNames[collectionId] + rec := CollectionDidTime{ + Collection: collection, + Did: did, + UnixMillis: seenMs, + } + select { + case <-done: + return nil + case out <- rec: + } + } + pcd.log.Debug("read primary", "count", count) + return nil +} + +func (pcd *PebbleCollectionDirectory) ReadCollection(ctx context.Context, collection, cursor string, limit int) (result []CollectionDidTime, nextCursor string, err error) { + var lower []byte + collectionId, err := pcd.CollectionToId(collection) + if err != nil { + return nil, "", fmt.Errorf("collection id err, %w", err) + } + if cursor != "" { + lower, err = base64.StdEncoding.DecodeString(cursor) + if err != nil { + return nil, "", fmt.Errorf("could not decode cursor, %w", err) + } + } else { + lower = make([]byte, 1+4) + lower[0] = 'A' + binary.BigEndian.PutUint32(lower[1:], collectionId) + } + var upper [5]byte + upper[0] = 'A' + binary.BigEndian.PutUint32(upper[1:], collectionId+1) + iter, err := pcd.db.NewIterWithContext(ctx, &pebble.IterOptions{ + LowerBound: lower, + UpperBound: upper[:], + }) + if err != nil { + return nil, "", fmt.Errorf("collection iter start, %w", err) + } + defer iter.Close() + count := 0 + done := ctx.Done() + result = make([]CollectionDidTime, 0, limit) + for iter.First(); iter.Valid(); iter.Next() { + key := iter.Key() + collectionId, did, seenMs := parsePrimaryPebbleRow(key) + count++ + collection := pcd.collectionNames[collectionId] + rec := CollectionDidTime{ + Collection: collection, + Did: did, + UnixMillis: seenMs, + } + result = append(result, rec) + breaker := false + if count >= limit { + breaker = true + } else { + select { + case <-done: + breaker = true + default: + } + } + if breaker { + prevKey := make([]byte, len(key), len(key)+1) + copy(prevKey, key) + prevKey = append(prevKey, 0) + nextCursor = base64.StdEncoding.EncodeToString(prevKey) + break + } + } + pcd.log.Debug("read primary", "count", count) + return result, nextCursor, nil +} + +func (pcd *PebbleCollectionDirectory) CollectionToId(collection string) (uint32, error) { + // easy mode: in cache + collectionId, ok := pcd.collections[collection] + if ok { + return collectionId, nil + } + + // read from db + key := makeCollectionInternKey(collection) + value, closer, err := pcd.db.Get(key) + if closer != nil { + defer closer.Close() + } + if err == nil { + collectionId = binary.BigEndian.Uint32(value) + return collectionId, nil + } + + // make new id, write to db + if errors.Is(err, pebble.ErrNotFound) { + // ok, fall through + } else if err != nil { + return 0, fmt.Errorf("pebble get err, %w", err) + } + collectionId = pcd.maxCollectionId + 1 + pcd.maxCollectionId = collectionId + var cib [4]byte + binary.BigEndian.PutUint32(cib[:], collectionId) + err = pcd.db.Set(key, cib[:], pebble.NoSync) + if err != nil { + return 0, fmt.Errorf("pebble set err, %w", err) + } + pcd.collections[collection] = collectionId + pcd.collectionNames[collectionId] = collection + return collectionId, nil +} + +var trueValue = [1]byte{'t'} + +func (pcd *PebbleCollectionDirectory) MaybeSetCollection(did, collection string) error { + collectionId, err := pcd.CollectionToId(collection) + if err != nil { + return err + } + dkey := makeByDidKey(did, collectionId) + _, closer, err := pcd.db.Get(dkey) + if closer != nil { + defer closer.Close() + } + if err == nil { + // already exists, done + pebbleDup.Inc() + return nil + } + if errors.Is(err, pebble.ErrNotFound) { + // ok, fall through + } else if err != nil { + return fmt.Errorf("pebble get err, %w", err) + } + + now := time.Now() + pkey := makePrimaryPebbleRow(collectionId, did, now.UnixMilli()) + err = pcd.db.Set(pkey, trueValue[:], pebble.NoSync) + if err != nil { + return fmt.Errorf("pebble set err, %w", err) + } + var timebytes [8]byte + binary.BigEndian.PutUint64(timebytes[:], uint64(now.UnixMilli())) + err = pcd.db.Set(dkey, timebytes[:], pebble.NoSync) + if err != nil { + return fmt.Errorf("pebble set err, %w", err) + } + pebbleNew.Inc() + return nil +} + +func (pcd *PebbleCollectionDirectory) SetFromResults(results <-chan DidCollection) { + errcount := 0 + for result := range results { + err := pcd.MaybeSetCollection(result.Did, result.Collection) + if err != nil { + errcount++ + pcd.log.Error("set collection", "err", err) + if errcount > 0 { + // TODO: signal backpressure and shutdown + return + } + } else { + errcount = 0 + } + } +} + +type CollectionStats struct { + CollectionCounts map[string]uint64 `json:"collections"` +} + +func (pcd *PebbleCollectionDirectory) GetCollectionStats() (stats CollectionStats, err error) { + ctx := context.Background() + records := make(chan CollectionDidTime, 1000) + go pcd.ReadAllPrimary(ctx, records) + + stats.CollectionCounts = make(map[string]uint64) + + for rec := range records { + stats.CollectionCounts[rec.Collection]++ + } + + return stats, nil +} + +const seqKey = "Xseq" + +func (pcd *PebbleCollectionDirectory) SetSequence(seq int64) error { + var seqb [8]byte + binary.BigEndian.PutUint64(seqb[:], uint64(seq)) + return pcd.db.Set([]byte(seqKey), seqb[:], pebble.NoSync) +} +func (pcd *PebbleCollectionDirectory) GetSequence() (int64, bool, error) { + vbytes, closer, err := pcd.db.Get([]byte(seqKey)) + if closer != nil { + defer closer.Close() + } + if errors.Is(err, pebble.ErrNotFound) { + return 0, false, nil + } + if err != nil { + return 0, false, fmt.Errorf("pebble seq err, %w", err) + } + seq := int64(binary.BigEndian.Uint64(vbytes)) + return seq, true, nil +} diff --git a/cmd/collectiondir/pebble_test.go b/cmd/collectiondir/pebble_test.go new file mode 100644 index 00000000..608a7a36 --- /dev/null +++ b/cmd/collectiondir/pebble_test.go @@ -0,0 +1,107 @@ +package main + +import ( + "context" + "encoding/csv" + "log/slog" + "strings" + "testing" + + "github.com/cockroachdb/pebble" + "github.com/cockroachdb/pebble/vfs" + "github.com/stretchr/testify/assert" +) + +type debugWriter struct { + t *testing.T +} + +func (w *debugWriter) Write(p []byte) (n int, err error) { + w.t.Helper() + w.t.Log(string(p)) + return len(p), nil +} + +// make a new pebble that writes to memory and logs to test.Log +func newMem(t *testing.T) *PebbleCollectionDirectory { + memfs := vfs.NewMem() + db, err := pebble.Open("wat", &pebble.Options{ + FS: memfs, + }) + if err != nil { + panic(err) + } + + log := slog.New(slog.NewTextHandler(&debugWriter{t: t}, &slog.HandlerOptions{Level: slog.LevelDebug})) + pcd := &PebbleCollectionDirectory{ + db: db, + collections: make(map[string]uint32), + collectionNames: make(map[uint32]string), + log: log, + } + if pcd.log == nil { + pcd.log = slog.Default() + } + return pcd +} + +// did, collection +const testDataCsv = `alice,post +alice,like +bob,post +bob,other +carol,post +eve,post +eve,like +eve,other` + +func TestPebbleCollectionDirectory(t *testing.T) { + assert := assert.New(t) + + pcd := newMem(t) + defer func() { + err := pcd.Close() + if err != nil { + t.Error(err) + } + }() + + rows, err := csv.NewReader(strings.NewReader(testDataCsv)).ReadAll() + assert.NoError(err) + for _, row := range rows { + err := pcd.MaybeSetCollection(row[0], row[1]) + assert.NoError(err) + } + stats, err := pcd.GetCollectionStats() + assert.NoError(err) + t.Log(stats) + assert.Equal(uint64(4), stats.CollectionCounts["post"]) + assert.Equal(uint64(2), stats.CollectionCounts["like"]) + assert.Equal(uint64(2), stats.CollectionCounts["other"]) + + t.Log(pcd.collections) + + wat, nextCursor, err := pcd.ReadCollection(context.Background(), "post", "", 1000) + assert.NoError(err) + assert.Equal("", nextCursor) + for _, row := range wat { + assert.Equal("post", row.Collection) + } + assert.Equal(4, len(wat)) + + wat, nextCursor, err = pcd.ReadCollection(context.Background(), "like", "", 1000) + assert.NoError(err) + assert.Equal("", nextCursor) + for _, row := range wat { + assert.Equal("like", row.Collection) + } + assert.Equal(2, len(wat)) + + wat, nextCursor, err = pcd.ReadCollection(context.Background(), "other", "", 1000) + assert.NoError(err) + assert.Equal("", nextCursor) + for _, row := range wat { + assert.Equal("other", row.Collection) + } + assert.Equal(2, len(wat)) +} diff --git a/cmd/collectiondir/serve.go b/cmd/collectiondir/serve.go new file mode 100644 index 00000000..ed868203 --- /dev/null +++ b/cmd/collectiondir/serve.go @@ -0,0 +1,867 @@ +package main + +import ( + "compress/gzip" + "context" + "encoding/csv" + "fmt" + comatproto "github.com/bluesky-social/indigo/api/atproto" + "github.com/bluesky-social/indigo/events" + "github.com/bluesky-social/indigo/xrpc" + "github.com/labstack/echo/v4" + "github.com/labstack/echo/v4/middleware" + "github.com/prometheus/client_golang/prometheus/promhttp" + "github.com/urfave/cli/v2" + "log/slog" + "net" + "net/http" + "net/url" + "os" + "os/signal" + "path/filepath" + "sort" + "strconv" + "strings" + "sync" + "syscall" + "time" +) + +var serveCmd = &cli.Command{ + Name: "serve", + Flags: []cli.Flag{ + &cli.StringFlag{ + Name: "api-listen", + Value: ":2510", + EnvVars: []string{"COLLECTIONS_API_LISTEN"}, + }, + &cli.StringFlag{ + Name: "metrics-listen", + Value: ":2511", + EnvVars: []string{"COLLECTIONS_METRICS_LISTEN"}, + }, + &cli.StringFlag{ + Name: "pebble", + Usage: "path to store pebble db", + Required: true, + }, + &cli.StringFlag{ + Name: "dau-directory", + Usage: "directory to store DAU pebble db", + Required: true, + }, + &cli.StringFlag{ + Name: "upstream", + Usage: "URL, e.g. wss://bsky.network", + EnvVars: []string{"COLLECTIONS_UPSTREAM"}, + }, + &cli.StringFlag{ + Name: "admin-token", + Usage: "admin authentication", + EnvVars: []string{"COLLECTIONS_ADMIN_TOKEN"}, + }, + &cli.Float64Flag{ + Name: "crawl-qps", + Usage: "per-PDS crawl queries-per-second limit", + Value: 100, + }, + &cli.StringFlag{ + Name: "ratelimit-header", + Usage: "secret for friend PDSes", + EnvVars: []string{"BSKY_SOCIAL_RATE_LIMIT_SKIP", "RATE_LIMIT_HEADER"}, + }, + &cli.BoolFlag{ + Name: "verbose", + }, + }, + Action: func(cctx *cli.Context) error { + var server collectionServer + return server.run(cctx) + }, +} + +type collectionServer struct { + ctx context.Context + + // the primary directory, all repos ever and their collections + pcd *PebbleCollectionDirectory + + // daily-active-user directory, new directory every 00:00:00 UTC + dauDirectory *PebbleCollectionDirectory + dauDirectoryPath string // currently open dauDirectory, {dauDirectoryDir}/{YYYY}{mm}{dd}.pebble + dauDay time.Time // YYYY-MM-DD 00:00:00 UTC + dauTomorrow time.Time + dauDirectoryDir string + + statsCache *CollectionStats + statsCacheWhen time.Time + statsCacheLock sync.Mutex + statsCacheFresh sync.Cond + statsCachePending bool + + // (did,collection) pairs from firehose + ingestFirehose chan DidCollection + // (did,collection) pairs from PDS crawl (don't apply to dauDirectory) + ingestCrawl chan DidCollection + + log *slog.Logger + + AdminToken string + ExepctedAuthHeader string + PerPDSCrawlQPS float64 + + activeCrawlHosts map[string]time.Time + activeCrawlsLock sync.Mutex + + shutdown chan struct{} + + wg sync.WaitGroup + + ratelimitHeader string + + apiServer *http.Server + //esrv *echo.Echo + metricsServer *http.Server +} + +const defaultPerPDSCrawlQPS = 100 + +func (cs *collectionServer) run(cctx *cli.Context) error { + signals := make(chan os.Signal, 1) + signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM) + cs.shutdown = make(chan struct{}) + level := slog.LevelInfo + if cctx.Bool("verbose") { + level = slog.LevelDebug + } + if cctx.IsSet("ratelimit-header") { + cs.ratelimitHeader = cctx.String("ratelimit-header") + } + log := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: level})) + cs.ingestFirehose = make(chan DidCollection, 1000) + cs.ingestCrawl = make(chan DidCollection, 1000) + cs.wg.Add(1) + go cs.ingestReceiver() + cs.log = log + cs.ctx = cctx.Context + cs.AdminToken = cctx.String("admin-token") + cs.ExepctedAuthHeader = "Bearer " + cs.AdminToken + pebblePath := cctx.String("pebble") + cs.pcd = &PebbleCollectionDirectory{ + log: cs.log, + } + err := cs.pcd.Open(pebblePath) + if err != nil { + return fmt.Errorf("%s: failed to open pebble db: %w", pebblePath, err) + } + cs.dauDirectoryDir = cctx.String("dau-directory") + if cs.dauDirectoryDir != "" { + err := cs.openDau() + if err != nil { + return err + } + } + cs.statsCacheFresh.L = &cs.statsCacheLock + errchan := make(chan error, 3) + apiAddr := cctx.String("api-listen") + cs.wg.Add(1) + go func() { + errchan <- cs.StartApiServer(cctx.Context, apiAddr) + }() + metricsAddr := cctx.String("metrics-listen") + cs.wg.Add(1) + go func() { + errchan <- cs.StartMetricsServer(cctx.Context, metricsAddr) + }() + + upstream := cctx.String("upstream") + if upstream != "" { + fh := Firehose{ + Log: log, + Host: upstream, + Seq: -1, + } + seq, seqok, err := cs.pcd.GetSequence() + if err != nil { + cs.log.Warn("db get seq", "err", err) + } else if seqok { + fh.Seq = seq + } + fhevents := make(chan *events.XRPCStreamEvent, 1000) + cs.wg.Add(1) + go cs.firehoseThread(&fh, fhevents) + cs.wg.Add(1) + go cs.handleFirehose(fhevents) + } + + select { + case <-signals: + log.Info("received shutdown signal") + go errchanlog(cs.log, "server error", errchan) + return cs.Shutdown() + case err := <-errchan: + if err != nil { + log.Error("server error", "err", err) + go errchanlog(cs.log, "server error", errchan) + return cs.Shutdown() + } + } + return nil +} + +func (cs *collectionServer) openDau() error { + now := time.Now().UTC() + ymd := now.Format("2006-01-02") + fname := fmt.Sprintf("d%s.pebble", ymd) + fpath := filepath.Join(cs.dauDirectoryDir, fname) + daud := &PebbleCollectionDirectory{ + log: cs.log, + } + err := daud.Open(fpath) + if err != nil { + return fmt.Errorf("%s: failed to open dau pebble db: %w", fpath, err) + } + cs.dauDirectory = daud + cs.dauDirectoryPath = fpath + cs.dauDay = time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC) + cs.dauTomorrow = now.AddDate(0, 0, 1) + return nil +} + +func errchanlog(log *slog.Logger, msg string, errchan <-chan error) { + for err := range errchan { + log.Error(msg, "err", err) + } +} + +func (cs *collectionServer) Shutdown() error { + close(cs.shutdown) + go func() { + cs.log.Info("metrics shutdown start") + cs.metricsServer.Shutdown(context.Background()) + cs.log.Info("metrics shutdown") + }() + cs.log.Info("api shutdown start...") + err := cs.apiServer.Shutdown(context.Background()) + //err := cs.esrv.Shutdown(context.Background()) + cs.log.Info("api shutdown, thread wait...", "err", err) + cs.wg.Wait() + cs.log.Info("threads done, db close...") + ee := cs.pcd.Close() + if ee != nil { + cs.log.Error("failed to shutdown pebble", "err", ee) + } + cs.log.Info("db done. done.") + return err +} + +// firehoseThreads is responsible for connecting to upstream firehose source +func (cs *collectionServer) firehoseThread(fh *Firehose, fhevents chan<- *events.XRPCStreamEvent) { + defer cs.wg.Done() + defer cs.log.Info("firehoseThread exit") + ctx, cancel := context.WithCancel(cs.ctx) + go func() { + <-cs.shutdown + cancel() + }() + err := fh.subscribeWithRedialer(ctx, fhevents) + if err != nil { + cs.log.Error("failed to subscribe to redialer", "err", err) + } + if fh.Seq >= 0 { + err := cs.pcd.SetSequence(fh.Seq) + if err != nil { + cs.log.Warn("db set seq", "err", err) + } + } +} + +// handleFirehose consumes XRPCStreamEvent from firehoseThread(), further parses data and applies +func (cs *collectionServer) handleFirehose(fhevents <-chan *events.XRPCStreamEvent) { + defer cs.wg.Done() + defer cs.log.Info("handleFirehose exit") + defer close(cs.ingestFirehose) + var lastSeq int64 + lastSeqSet := false + notDone := true + for notDone { + select { + case <-cs.shutdown: + cs.log.Info("firehose handler shutdown") + notDone = false + case evt, ok := <-fhevents: + if !ok { + notDone = false + cs.log.Info("firehose handler closed") + break + } + firehoseReceivedCounter.Inc() + seq, ok := evt.GetSequence() + if ok { + lastSeq = seq + lastSeqSet = true + } + if evt.RepoCommit != nil { + firehoseCommits.Inc() + cs.handleCommit(evt.RepoCommit) + } + } + } + if lastSeqSet { + cs.pcd.SetSequence(lastSeq) + } +} + +func (cs *collectionServer) handleCommit(commit *comatproto.SyncSubscribeRepos_Commit) { + for _, op := range commit.Ops { + // op.Path is collection/rkey + slash := strings.IndexRune(op.Path, '/') + if slash == -1 { + cs.log.Warn("bad op path", "repo", commit.Repo) + return + } + collection := op.Path[:slash] + firehoseCommitOps.WithLabelValues(op.Action).Inc() + if op.Action == "create" || op.Action == "update" { + firehoseDidcSet.Inc() + cs.ingestFirehose <- DidCollection{ + Did: commit.Repo, + Collection: collection, + } + } + } +} + +func (cs *collectionServer) StartMetricsServer(ctx context.Context, addr string) error { + defer cs.wg.Done() + defer cs.log.Info("metrics server exit") + cs.metricsServer = &http.Server{ + Addr: addr, + Handler: promhttp.Handler(), + } + return cs.metricsServer.ListenAndServe() +} + +func (cs *collectionServer) StartApiServer(ctx context.Context, addr string) error { + defer cs.wg.Done() + defer cs.log.Info("api server exit") + var lc net.ListenConfig + li, err := lc.Listen(ctx, "tcp", addr) + if err != nil { + return err + } + e := echo.New() + e.HideBanner = true + + e.Use(MetricsMiddleware) + e.Use(middleware.CORSWithConfig(middleware.CORSConfig{ + AllowOrigins: []string{"*"}, + AllowHeaders: []string{echo.HeaderOrigin, echo.HeaderContentType, echo.HeaderAccept, echo.HeaderAuthorization}, + })) + + e.GET("/_health", cs.healthz) + + e.GET("/v1/getDidsForCollection", cs.getDidsForCollection) + e.GET("/v1/listCollections", cs.listCollections) + + // TODO: allow public 'requestCrawl' API? + //e.GET("/xrpc/com.atproto.sync.requestCrawl", cs.crawlPds) + //e.POST("/xrpc/com.atproto.sync.requestCrawl", cs.crawlPds) + + // admin auth heador required + e.POST("/admin/pds/requestCrawl", cs.crawlPds) // same as relay + e.GET("/v1/crawlStatus", cs.crawlStatus) + + e.Listener = li + srv := &http.Server{ + Handler: e, + } + cs.apiServer = srv + return srv.Serve(li) +} + +const statsCacheDuration = time.Second * 300 + +type GetDidsForCollectionResponse struct { + Dids []string `json:"dids"` + Cursor string `json:"cursor"` +} + +func getLimit(c echo.Context, min, defaultLim, max int) int { + limstr := c.QueryParam("limit") + if limstr == "" { + return defaultLim + } + lvx, err := strconv.ParseInt(limstr, 10, 64) + if err != nil { + return defaultLim + } + lv := int(lvx) + if lv < min { + return min + } + if lv > max { + return max + } + return lv +} + +// /v1/getDidsForCollection?collection={}&cursor={} +// +// returns +// {"dids":["did:A", "..."], "cursor":"opaque text"} +func (cs *collectionServer) getDidsForCollection(c echo.Context) error { + ctx := c.Request().Context() + collection := c.QueryParam("collection") + cursor := c.QueryParam("cursor") + limit := getLimit(c, 50, 500, 1000) + they, nextCursor, err := cs.pcd.ReadCollection(ctx, collection, cursor, limit) + if err != nil { + slog.Error("ReadCollection", "collection", collection, "cursor", cursor, "limit", limit, "err", err) + return c.String(http.StatusInternalServerError, "oops") + } + cs.log.Info("getDidsForCollection", "collection", collection, "cursor", cursor, "limit", limit, "count", len(they), "nextCursor", nextCursor) + var out GetDidsForCollectionResponse + out.Dids = make([]string, len(they)) + for i, rec := range they { + out.Dids[i] = rec.Did + } + out.Cursor = nextCursor + return c.JSON(http.StatusOK, out) +} + +// return cached collection stats if they're fresh +// return new collection stats if they can be calculated quicly +// return stale cached collection stats if new stats take too long +// just wait for fresh stats if there are no cached stats +// stalenessAllowed is how old stats can be before we try to recalculate them, 0=default of 5 minutes +func (cs *collectionServer) getStatsCache(stalenessAllowed time.Duration) (*CollectionStats, error) { + if stalenessAllowed <= 0 { + stalenessAllowed = statsCacheDuration + } + var statsCache *CollectionStats + var staleCache *CollectionStats + var waiter *freshStatsWaiter + cs.statsCacheLock.Lock() + if cs.statsCache != nil { + if time.Since(cs.statsCacheWhen) < stalenessAllowed { + // has fresh! + statsCache = cs.statsCache + } else if !cs.statsCachePending { + cs.statsCachePending = true + go cs.statsBuilder() + staleCache = cs.statsCache + } else { + staleCache = cs.statsCache + } + if staleCache != nil { + waiter = &freshStatsWaiter{ + cs: cs, + freshCache: make(chan *CollectionStats), + } + go waiter.waiter() + } + } else if !cs.statsCachePending { + cs.statsCachePending = true + go cs.statsBuilder() + } + cs.statsCacheLock.Unlock() + + if statsCache != nil { + // return fresh-enough data + return statsCache, nil + } + + if staleCache == nil { + // block forever waiting for fresh data + cs.statsCacheLock.Lock() + for cs.statsCache == nil { + cs.statsCacheFresh.Wait() + } + statsCache = cs.statsCache + cs.statsCacheLock.Unlock() + return statsCache, nil + } + + // wait for up to a second for fresh data, on timeout return stale data + timeout := time.NewTimer(time.Second) + defer timeout.Stop() + select { + case <-timeout.C: + cs.statsCacheLock.Lock() + waiter.l.Lock() + waiter.obsolete = true + waiter.l.Unlock() + cs.statsCacheLock.Unlock() + return staleCache, nil + case statsCache = <-waiter.freshCache: + return statsCache, nil + } +} + +type freshStatsWaiter struct { + cs *collectionServer + l sync.Mutex + obsolete bool + freshCache chan *CollectionStats +} + +func (fsw *freshStatsWaiter) waiter() { + fsw.cs.statsCacheLock.Lock() + defer fsw.cs.statsCacheLock.Unlock() + fsw.cs.statsCacheFresh.Wait() + fsw.l.Lock() + defer fsw.l.Unlock() + if fsw.obsolete { + close(fsw.freshCache) + } else { + fsw.freshCache <- fsw.cs.statsCache + } +} + +func (cs *collectionServer) statsBuilder() { + for { + start := time.Now() + stats, err := cs.pcd.GetCollectionStats() + dt := time.Since(start) + if err == nil { + countsum := uint64(0) + for _, v := range stats.CollectionCounts { + countsum += v + } + cs.log.Info("stats built", "dt", dt, "total", countsum) + cs.statsCacheLock.Lock() + cs.statsCache = &stats + cs.statsCacheWhen = time.Now() + cs.statsCacheFresh.Broadcast() + cs.statsCachePending = false + cs.statsCacheLock.Unlock() + return + } else { + cs.log.Error("GetCollectionStats", "dt", dt, "err", err) + time.Sleep(2 * time.Second) + } + } +} + +// /v1/listCollections?c={}&cursor={}&limit={50<=limit<=1000} +// +// admin may set ?stalesec={} for a maximum number of seconds stale data is accepted +// +// returns +// {"collections":{"app.bsky.feed.post": 123456789, "some collection": 42}, "cursor":"opaque text"} +func (cs *collectionServer) listCollections(c echo.Context) error { + stalenessAllowed := statsCacheDuration + stalesecStr := c.QueryParam("stalesec") + if stalesecStr != "" && cs.isAdmin(c) { + stalesec, err := strconv.ParseInt(stalesecStr, 10, 64) + if err != nil { + return c.String(http.StatusBadRequest, "bad stalesec") + } + if stalesec == 0 { + stalenessAllowed = 1 + } else { + stalenessAllowed = time.Duration(stalesec) * time.Second + } + cs.log.Info("stalesec", "q", stalesecStr, "d", stalenessAllowed) + } + stats, err := cs.getStatsCache(stalenessAllowed) + if err != nil { + slog.Error("getStatsCache", "err", err) + return c.String(http.StatusInternalServerError, "oops") + } + cursor := c.QueryParam("cursor") + collections, hasQueryCollections := c.QueryParams()["c"] + limit := getLimit(c, 50, 500, 1000) + var out ListCollectionsResponse + if hasQueryCollections { + out.Collections = make(map[string]uint64, len(collections)) + for _, collection := range collections { + count, ok := stats.CollectionCounts[collection] + if ok { + out.Collections[collection] = count + } + } + } else { + allCollections := make([]string, 0, len(stats.CollectionCounts)) + for collection := range stats.CollectionCounts { + allCollections = append(allCollections, collection) + } + sort.Strings(allCollections) + out.Collections = make(map[string]uint64, limit) + count := 0 + for _, collection := range allCollections { + if (cursor == "") || (collection > cursor) { + out.Collections[collection] = stats.CollectionCounts[collection] + count++ + if count >= limit { + out.Cursor = collection + } + } + } + } + return c.JSON(http.StatusOK, out) +} + +type ListCollectionsResponse struct { + Collections map[string]uint64 `json:"collections"` + Cursor string `json:"cursor"` +} + +func (cs *collectionServer) ingestReceiver() { + defer cs.wg.Done() + defer cs.log.Info("ingestReceiver exit") + errcount := 0 + for { + select { + case didc, ok := <-cs.ingestFirehose: + if !ok { + cs.log.Info("ingestFirehose closed") + return + } + err := cs.pcd.MaybeSetCollection(didc.Did, didc.Collection) + if err != nil { + cs.log.Warn("pcd write", "err", err) + errcount++ + } else { + errcount = 0 + } + if cs.dauDirectory != nil { + err = cs.maybeDauWrite(didc) + if err != nil { + cs.log.Warn("dau write", "err", err) + errcount++ + } else { + errcount = 0 + } + } + case didc := <-cs.ingestCrawl: + err := cs.pcd.MaybeSetCollection(didc.Did, didc.Collection) + if err != nil { + cs.log.Warn("pcd write", "err", err) + errcount++ + } else { + errcount = 0 + } + } + if errcount > 10 { + cs.log.Error("ingestReceiver too many errors") + return // TODO: cancel parent somehow + } + } +} + +// write {dauDirectoryDir}/d{YYYY-MM-DD}.pebble stats summary to {dauDirectoryDir}/d{YYYY-MM-DD}.csv.gz +func dauStats(oldDau *PebbleCollectionDirectory, dauDay time.Time, dauDir string, log *slog.Logger) { + stats, err := oldDau.GetCollectionStats() + e2 := oldDau.Close() + if e2 != nil { + log.Error("old dau close", "err", e2) + } + if err != nil { + log.Error("old dau stats", "err", err) + } else { + fname := fmt.Sprintf("d%s.csv.gz", dauDay.Format("2006-01-02")) + outstatsPath := filepath.Join(dauDir, fname) + pcdStatsToCsvGz(stats, outstatsPath, log) + } +} + +func pcdStatsToCsvGz(stats CollectionStats, outpath string, log *slog.Logger) { + fout, err := os.Create(outpath) + if err != nil { + log.Error("dau stats open", "err", err) + return + } + defer fout.Close() + gzout := gzip.NewWriter(fout) + csvout := csv.NewWriter(gzout) + defer csvout.Flush() + defer gzout.Close() + err = csvout.Write([]string{"collection", "count"}) + if err != nil { + log.Error("dau stats header", "err", err) + return + } + var row [2]string + for collection, count := range stats.CollectionCounts { + row[0] = collection + row[1] = strconv.FormatUint(count, 10) + err = csvout.Write(row[:]) + if err != nil { + log.Error("dau stats row", "err", err) + return + } + } +} + +func (cs *collectionServer) maybeDauWrite(didc DidCollection) error { + now := time.Now() + if now.After(cs.dauTomorrow) { + go dauStats(cs.dauDirectory, cs.dauDay, cs.dauDirectoryDir, cs.log) + cs.dauDirectory = nil + err := cs.openDau() + if err != nil { + return fmt.Errorf("dau reopen, %w", err) + } + } + return cs.dauDirectory.MaybeSetCollection(didc.Did, didc.Collection) +} + +type CrawlRequest struct { + Host string `json:"hostname,omitempty"` + Hosts []string `json:"hosts,omitempty"` +} + +type CrawlRequestResponse struct { + Message string `json:"message,omitempty"` + Error string `json:"error,omitempty"` +} + +func hostOrUrlToUrl(host string) string { + xu, err := url.Parse(host) + if err != nil { + xu = new(url.URL) + xu.Host = host + xu.Scheme = "https" + return xu.String() + } else if xu.Scheme == "" { + xu.Scheme = "https" + return xu.String() + } + return host +} + +func (cs *collectionServer) isAdmin(c echo.Context) bool { + authHeader := c.Request().Header.Get("Authorization") + if authHeader == "" { + return false + } + if authHeader == cs.ExepctedAuthHeader { + return true + } + cs.log.Info("wrong auth header", "header", authHeader, "expected", cs.ExepctedAuthHeader) + return false +} + +// /v1/crawlRequest +// requires header `Authorization: Bearer {admin token}` +// +// POST {"host":"one hostname or URL", "hosts":["up to 1000 hosts", "..."]} +// OR +// POST /v1/crawlRequest?host={one host} +func (cs *collectionServer) crawlPds(c echo.Context) error { + isAdmin := cs.isAdmin(c) + if !isAdmin { + return c.JSON(http.StatusForbidden, CrawlRequestResponse{Error: "nope"}) + } + hostQ := c.QueryParam("host") + if hostQ != "" { + go cs.crawlThread(hostQ) + return c.JSON(http.StatusOK, CrawlRequestResponse{Message: "ok"}) + } + + var req CrawlRequest + err := c.Bind(&req) + if err != nil { + return c.String(http.StatusBadRequest, err.Error()) + } + if req.Host != "" { + go cs.crawlThread(req.Host) + } + for _, host := range req.Hosts { + go cs.crawlThread(host) + } + return c.JSON(http.StatusOK, CrawlRequestResponse{Message: "ok"}) +} + +func (cs *collectionServer) crawlThread(hostIn string) { + host := hostOrUrlToUrl(hostIn) + if host != hostIn { + cs.log.Info("going to crawl", "in", hostIn, "as", host) + } + httpClient := http.Client{} + rpcClient := xrpc.Client{ + Host: host, + Client: &httpClient, + } + if cs.ratelimitHeader != "" { + rpcClient.Headers = map[string]string{ + "x-ratelimit-bypass": cs.ratelimitHeader, + } + } + crawler := Crawler{ + Ctx: cs.ctx, + RpcClient: &rpcClient, + QPS: cs.PerPDSCrawlQPS, + Results: cs.ingestCrawl, + Log: cs.log, + } + start := time.Now() + ok := cs.recordCrawlStart(host, start) + if !ok { + cs.log.Info("not crawling dup", "host", host) + } + cs.log.Info("crawling", "host", host) + err := crawler.CrawlPDSRepoCollections() + cs.clearActiveCrawl(host) + pdsCrawledCounter.Inc() + if err != nil { + cs.log.Warn("crawl err", "host", host, "err", err) + } else { + dt := time.Since(start) + cs.log.Info("crawl done", "host", host, "dt", dt) + } +} + +// recordCrawlStart returns true if ok, false if duplicate +func (cs *collectionServer) recordCrawlStart(host string, start time.Time) (ok bool) { + cs.activeCrawlsLock.Lock() + defer cs.activeCrawlsLock.Unlock() + if cs.activeCrawlHosts == nil { + cs.activeCrawlHosts = make(map[string]time.Time) + cs.activeCrawlHosts[host] = start + return true + } else { + _, dup := cs.activeCrawlHosts[host] + if dup { + return false + } + cs.activeCrawlHosts[host] = start + return true + } +} + +func (cs *collectionServer) clearActiveCrawl(host string) { + cs.activeCrawlsLock.Lock() + defer cs.activeCrawlsLock.Unlock() + if cs.activeCrawlHosts == nil { + return + } + delete(cs.activeCrawlHosts, host) +} + +type CrawlStatusResponse struct { + HostStarts map[string]string `json:"host_starts"` +} + +// GET /v1/crawlStatus +func (cs *collectionServer) crawlStatus(c echo.Context) error { + authHeader := c.Request().Header.Get("Authorization") + if authHeader != cs.ExepctedAuthHeader { + return c.JSON(http.StatusForbidden, CrawlRequestResponse{Error: "nope"}) + } + var out CrawlStatusResponse + out.HostStarts = make(map[string]string) + cs.activeCrawlsLock.Lock() + defer cs.activeCrawlsLock.Unlock() + for host, start := range cs.activeCrawlHosts { + out.HostStarts[host] = start.UTC().Format(time.RFC3339) + } + return c.JSON(http.StatusOK, out) +} + +func (cs *collectionServer) healthz(c echo.Context) error { + // TODO: check database or upstream health? + return c.String(http.StatusOK, "ok") +} diff --git a/events/events.go b/events/events.go index 5619fc20..aeba2df7 100644 --- a/events/events.go +++ b/events/events.go @@ -455,6 +455,31 @@ func (evt *XRPCStreamEvent) Sequence() int64 { } } +func (evt *XRPCStreamEvent) GetSequence() (int64, bool) { + switch { + case evt == nil: + return -1, false + case evt.RepoCommit != nil: + return evt.RepoCommit.Seq, true + case evt.RepoHandle != nil: + return evt.RepoHandle.Seq, true + case evt.RepoMigrate != nil: + return evt.RepoMigrate.Seq, true + case evt.RepoTombstone != nil: + return evt.RepoTombstone.Seq, true + case evt.RepoIdentity != nil: + return evt.RepoIdentity.Seq, true + case evt.RepoAccount != nil: + return evt.RepoAccount.Seq, true + case evt.RepoInfo != nil: + return -1, false + case evt.Error != nil: + return -1, false + default: + return -1, false + } +} + func (em *EventManager) rmSubscriber(sub *Subscriber) { em.subsLk.Lock() defer em.subsLk.Unlock()