diff --git a/repo/carutil/reader.go b/repo/carutil/reader.go new file mode 100644 index 00000000..cae6eb71 --- /dev/null +++ b/repo/carutil/reader.go @@ -0,0 +1,102 @@ +package carutil + +import ( + "bufio" + "encoding/binary" + "errors" + "fmt" + "io" + "sync" + + blocks "github.com/ipfs/go-block-format" + "github.com/ipfs/go-cid" + car "github.com/ipld/go-car" +) + +type Reader struct { + r *bufio.Reader + + bufs [][]byte +} + +func NewReader(r *bufio.Reader) (*Reader, cid.Cid, error) { + h, err := car.ReadHeader(r) + if err != nil { + return nil, cid.Undef, err + } + + if h.Version != 1 { + return nil, cid.Undef, fmt.Errorf("invalid version: %d", h.Version) + } + + if len(h.Roots) != 1 { + return nil, cid.Undef, fmt.Errorf("expected only 1 root in car file") + } + + return &Reader{ + r: r, + bufs: make([][]byte, 0, 10), + }, h.Roots[0], nil +} + +func (r *Reader) Free(alloc *sync.Pool) { + for _, b := range r.bufs { + alloc.Put(b) + } + r.bufs = nil +} + +const MaxAllowedSectionSize = 32 << 20 + +func (r *Reader) NextBlock(allocator *sync.Pool, allocMax uint64) (blocks.Block, error) { + data, err := ldRead(r.r, allocator, allocMax) + if err != nil { + return nil, err + } + + r.bufs = append(r.bufs, data) + + n, c, err := cid.CidFromBytes(data) + if err != nil { + return nil, err + } + + return blocks.NewBlockWithCid(data[n:], c) +} + +func ldRead(r *bufio.Reader, alloc *sync.Pool, allocMax uint64) ([]byte, error) { + if _, err := r.Peek(1); err != nil { // no more blocks, likely clean io.EOF + return nil, err + } + + l, err := binary.ReadUvarint(r) + if err != nil { + if err == io.EOF { + return nil, io.ErrUnexpectedEOF // don't silently pretend this is a clean EOF + } + return nil, err + } + + if l > uint64(MaxAllowedSectionSize) { // Don't OOM + return nil, errors.New("malformed car; header is bigger than util.MaxAllowedSectionSize") + } + + if l > allocMax { + // direct allocation, not great + buf := make([]byte, l) + if _, err := io.ReadFull(r, buf); err != nil { + return nil, err + } + + return buf, nil + } + + buf := alloc.Get().([]byte) + buf = buf[:l] + + if _, err := io.ReadFull(r, buf); err != nil { + return nil, err + } + + return buf, nil +} diff --git a/repo/repo.go b/repo/repo.go index c7ced80f..582177e9 100644 --- a/repo/repo.go +++ b/repo/repo.go @@ -1,19 +1,21 @@ package repo import ( + "bufio" "bytes" "context" "fmt" "io" + "sync" "github.com/bluesky-social/indigo/atproto/repo" "github.com/bluesky-social/indigo/atproto/syntax" lexutil "github.com/bluesky-social/indigo/lex/util" "github.com/bluesky-social/indigo/mst" + "github.com/bluesky-social/indigo/repo/carutil" "github.com/bluesky-social/indigo/util" "github.com/ipfs/go-cid" cbor "github.com/ipfs/go-ipld-cbor" - "github.com/ipld/go-car" cbg "github.com/whyrusleeping/cbor-gen" "go.opentelemetry.io/otel" ) @@ -76,17 +78,25 @@ func (uc *UnsignedCommit) BytesForSigning() ([]byte, error) { return buf.Bytes(), nil } +const repoBlockBufferSize = 128 << 10 + +var repoBlockBufferPool = &sync.Pool{ + New: func() any { + return make([]byte, repoBlockBufferSize) + }, +} + func IngestRepo(ctx context.Context, bs cbor.IpldBlockstore, r io.Reader) (cid.Cid, error) { ctx, span := otel.Tracer("repo").Start(ctx, "Ingest") defer span.End() - br, err := car.NewCarReader(r) + br, root, err := carutil.NewReader(bufio.NewReader(r)) if err != nil { return cid.Undef, fmt.Errorf("opening CAR block reader: %w", err) } for { - blk, err := br.Next() + blk, err := br.NextBlock(repoBlockBufferPool, repoBlockBufferSize) if err != nil { if err == io.EOF { break @@ -99,7 +109,7 @@ func IngestRepo(ctx context.Context, bs cbor.IpldBlockstore, r io.Reader) (cid.C } } - return br.Header.Roots[0], nil + return root, nil } func ReadRepoFromCar(ctx context.Context, r io.Reader) (*Repo, error) { -- 2.51.2 From f2631d9d3ec5c01ce74d8fe9dac3012a40851452 Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Sat, 22 Mar 2025 12:15:42 -0700 Subject: [PATCH 02/26] implement streaming repo parser --- cmd/gosky/main.go | 23 +++---- cmd/gosky/sync.go | 8 +-- repo/stream.go | 162 ++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 175 insertions(+), 18 deletions(-) create mode 100644 repo/stream.go diff --git a/cmd/gosky/main.go b/cmd/gosky/main.go index 948c4612..5b69fe10 100644 --- a/cmd/gosky/main.go +++ b/cmd/gosky/main.go @@ -634,11 +634,15 @@ var listAllRecordsCmd = &cli.Command{ var repob []byte if strings.HasPrefix(arg, "did:") { - xrpcc, err := cliutil.GetXrpcClient(cctx, true) + resp, err := identity.DefaultDirectory().LookupDID(ctx, syntax.DID(arg)) if err != nil { return err } + xrpcc := &xrpc.Client{ + Host: resp.PDSEndpoint(), + } + if arg == "" { arg = xrpcc.Auth.Did } @@ -648,6 +652,7 @@ var listAllRecordsCmd = &cli.Command{ return err } repob = rrb + fmt.Println("GOT REPO BYTES") } else { if len(arg) == 0 { return cli.Exit("must specify DID string or repo path", 127) @@ -660,11 +665,6 @@ var listAllRecordsCmd = &cli.Command{ repob = fb } - rr, err := repo.ReadRepoFromCar(ctx, bytes.NewReader(repob)) - if err != nil { - return err - } - collection := "app.bsky.feed.post" if cctx.Bool("all") { collection = "" @@ -672,24 +672,19 @@ var listAllRecordsCmd = &cli.Command{ vals := cctx.Bool("values") cids := cctx.Bool("cids") - if err := rr.ForEach(ctx, collection, func(k string, v cid.Cid) error { + if err := repo.StreamRepoRecords(ctx, bytes.NewReader(repob), collection, func(k string, cc cid.Cid, v []byte) error { if !strings.HasPrefix(k, collection) { return repo.ErrDoneIterating } fmt.Print(k) if cids { - fmt.Println(" - ", v) + fmt.Println(" - ", cc) } else { fmt.Println() } if vals { - b, err := rr.Blockstore().Get(ctx, v) - if err != nil { - return err - } - - convb, err := cborToJson(b.RawData()) + convb, err := cborToJson(v) if err != nil { return err } diff --git a/cmd/gosky/sync.go b/cmd/gosky/sync.go index 13c60f42..3b1ac1b1 100644 --- a/cmd/gosky/sync.go +++ b/cmd/gosky/sync.go @@ -9,6 +9,7 @@ import ( "github.com/bluesky-social/indigo/atproto/identity" "github.com/bluesky-social/indigo/atproto/syntax" "github.com/bluesky-social/indigo/util/cliutil" + "github.com/bluesky-social/indigo/xrpc" cli "github.com/urfave/cli/v2" ) @@ -53,11 +54,10 @@ var syncGetRepoCmd = &cli.Command{ carPath = ident.DID.String() + ".car" } - xrpcc, err := cliutil.GetXrpcClient(cctx, false) - if err != nil { - return err + xrpcc := &xrpc.Client{ + Host: ident.PDSEndpoint(), } - xrpcc.Host = ident.PDSEndpoint() + if xrpcc.Host == "" { return fmt.Errorf("no PDS endpoint for identity") } diff --git a/repo/stream.go b/repo/stream.go new file mode 100644 index 00000000..da2a0768 --- /dev/null +++ b/repo/stream.go @@ -0,0 +1,162 @@ +package repo + +import ( + "bufio" + "context" + "fmt" + "io" + "log/slog" + "sync" + + "github.com/bluesky-social/indigo/mst" + "github.com/bluesky-social/indigo/repo/carutil" + "github.com/bluesky-social/indigo/util" + block "github.com/ipfs/go-block-format" + cid "github.com/ipfs/go-cid" + "go.opentelemetry.io/otel" +) + +type waitingBlockstore struct { + lk sync.Mutex + blockWaits map[cid.Cid]chan block.Block + otherBlocks map[cid.Cid]block.Block + streamComplete bool +} + +func newWaitingBlockstore() *waitingBlockstore { + return &waitingBlockstore{ + blockWaits: make(map[cid.Cid]chan block.Block), + otherBlocks: make(map[cid.Cid]block.Block), + } +} + +func (bs *waitingBlockstore) Get(ctx context.Context, cc cid.Cid) (block.Block, error) { + bs.lk.Lock() + + if blk, ok := bs.otherBlocks[cc]; ok { + delete(bs.otherBlocks, cc) + bs.lk.Unlock() + return blk, nil + } + + if bs.streamComplete { + bs.lk.Unlock() + return nil, ErrMissingBlock + } + + bw, ok := bs.blockWaits[cc] + if ok { + bs.lk.Unlock() + return nil, fmt.Errorf("somehow already have active wait for block in question: %s", cc) + } + + bw = make(chan block.Block, 1) + + bs.blockWaits[cc] = bw + + bs.lk.Unlock() + + select { + case blk, ok := <-bw: + if !ok { + return nil, ErrMissingBlock + } + + return blk, nil + case <-ctx.Done(): + return nil, ctx.Err() + + } +} + +var ErrMissingBlock = fmt.Errorf("block was missing from archive") + +func (bs *waitingBlockstore) Put(ctx context.Context, blk block.Block) error { + bs.lk.Lock() + defer bs.lk.Unlock() + + bw, ok := bs.blockWaits[blk.Cid()] + if ok { + bw <- blk + delete(bs.blockWaits, blk.Cid()) + return nil + } + + bs.otherBlocks[blk.Cid()] = blk.(*block.BasicBlock) + return nil +} + +func (bs *waitingBlockstore) Complete() { + bs.lk.Lock() + defer bs.lk.Unlock() + bs.streamComplete = true + for _, ch := range bs.blockWaits { + close(ch) + } +} + +func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, cb func(k string, c cid.Cid, v []byte) error) error { + ctx, span := otel.Tracer("repo").Start(ctx, "RepoStream") + defer span.End() + + br, root, err := carutil.NewReader(bufio.NewReader(r)) + if err != nil { + return fmt.Errorf("opening CAR block reader: %w", err) + } + + bs := newWaitingBlockstore() + cst := util.CborStore(bs) + + var wg sync.WaitGroup + wg.Add(1) + var walkErr error + go func() { + defer wg.Done() + + var sc SignedCommit + if err := cst.Get(ctx, root, &sc); err != nil { + walkErr = fmt.Errorf("loading root from blockstore: %w", err) + return + } + + if sc.Version != ATP_REPO_VERSION && sc.Version != ATP_REPO_VERSION_2 { + walkErr = fmt.Errorf("unsupported repo version: %d", sc.Version) + return + } + // TODO: verify that signature + + t := mst.LoadMST(cst, sc.Data) + + if err := t.WalkLeavesFrom(ctx, prefix, func(k string, val cid.Cid) error { + blk, err := bs.Get(ctx, val) + if err != nil { + slog.Error("failed to get record from tree", "key", k, "cid", val, "error", err) + return nil + } + + return cb(k, val, blk.RawData()) + }); err != nil { + walkErr = fmt.Errorf("failed to walk mst: %w", err) + } + }() + + for { + blk, err := br.NextBlock(repoBlockBufferPool, repoBlockBufferSize) + if err != nil { + if err == io.EOF { + break + } + return fmt.Errorf("reading block from CAR: %w", err) + } + + if err := bs.Put(ctx, blk); err != nil { + return fmt.Errorf("copying block to store: %w", err) + } + } + + bs.Complete() + + wg.Wait() + + return walkErr +} -- 2.51.2 From 6f7f2c305190e1d603d719deaa40ad7554b222c7 Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Sat, 22 Mar 2025 12:27:26 -0700 Subject: [PATCH 03/26] okay, its a generator now --- repo/stream.go | 145 +++++++++++++++++++------------------------------ 1 file changed, 57 insertions(+), 88 deletions(-) diff --git a/repo/stream.go b/repo/stream.go index da2a0768..0b43b514 100644 --- a/repo/stream.go +++ b/repo/stream.go @@ -6,7 +6,6 @@ import ( "fmt" "io" "log/slog" - "sync" "github.com/bluesky-social/indigo/mst" "github.com/bluesky-social/indigo/repo/carutil" @@ -17,82 +16,76 @@ import ( ) type waitingBlockstore struct { - lk sync.Mutex - blockWaits map[cid.Cid]chan block.Block otherBlocks map[cid.Cid]block.Block streamComplete bool + + r *carutil.Reader } func newWaitingBlockstore() *waitingBlockstore { return &waitingBlockstore{ - blockWaits: make(map[cid.Cid]chan block.Block), otherBlocks: make(map[cid.Cid]block.Block), } } -func (bs *waitingBlockstore) Get(ctx context.Context, cc cid.Cid) (block.Block, error) { - bs.lk.Lock() +func (bs *waitingBlockstore) readUntilBlock(ctx context.Context, cc cid.Cid) (block.Block, error) { + for { + blk, err := bs.r.NextBlock(repoBlockBufferPool, repoBlockBufferSize) + if err != nil { + if err == io.EOF { + break + } + return nil, fmt.Errorf("reading block from CAR: %w", err) + } + + if blk.Cid() == cc { + return blk, nil + } + bs.otherBlocks[blk.Cid()] = blk + } + + bs.streamComplete = true + + return nil, io.EOF +} + +func (bs *waitingBlockstore) Get(ctx context.Context, cc cid.Cid) (block.Block, error) { if blk, ok := bs.otherBlocks[cc]; ok { delete(bs.otherBlocks, cc) - bs.lk.Unlock() return blk, nil } if bs.streamComplete { - bs.lk.Unlock() return nil, ErrMissingBlock } - bw, ok := bs.blockWaits[cc] - if ok { - bs.lk.Unlock() - return nil, fmt.Errorf("somehow already have active wait for block in question: %s", cc) + blk, err := bs.readUntilBlock(ctx, cc) + if err != nil { + return nil, err } - bw = make(chan block.Block, 1) - - bs.blockWaits[cc] = bw - - bs.lk.Unlock() - - select { - case blk, ok := <-bw: - if !ok { - return nil, ErrMissingBlock - } - - return blk, nil - case <-ctx.Done(): - return nil, ctx.Err() - - } + return blk, nil } var ErrMissingBlock = fmt.Errorf("block was missing from archive") func (bs *waitingBlockstore) Put(ctx context.Context, blk block.Block) error { - bs.lk.Lock() - defer bs.lk.Unlock() + return fmt.Errorf("put is not needed") + /* + bw, ok := bs.blockWaits[blk.Cid()] + if ok { + bw <- blk + delete(bs.blockWaits, blk.Cid()) + return nil + } - bw, ok := bs.blockWaits[blk.Cid()] - if ok { - bw <- blk - delete(bs.blockWaits, blk.Cid()) + bs.otherBlocks[blk.Cid()] = blk.(*block.BasicBlock) return nil - } - - bs.otherBlocks[blk.Cid()] = blk.(*block.BasicBlock) - return nil + */ } func (bs *waitingBlockstore) Complete() { - bs.lk.Lock() - defer bs.lk.Unlock() - bs.streamComplete = true - for _, ch := range bs.blockWaits { - close(ch) - } } func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, cb func(k string, c cid.Cid, v []byte) error) error { @@ -105,58 +98,34 @@ func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, cb func( } bs := newWaitingBlockstore() - cst := util.CborStore(bs) - - var wg sync.WaitGroup - wg.Add(1) - var walkErr error - go func() { - defer wg.Done() - var sc SignedCommit - if err := cst.Get(ctx, root, &sc); err != nil { - walkErr = fmt.Errorf("loading root from blockstore: %w", err) - return - } + bs.r = br - if sc.Version != ATP_REPO_VERSION && sc.Version != ATP_REPO_VERSION_2 { - walkErr = fmt.Errorf("unsupported repo version: %d", sc.Version) - return - } - // TODO: verify that signature + cst := util.CborStore(bs) - t := mst.LoadMST(cst, sc.Data) + var sc SignedCommit + if err := cst.Get(ctx, root, &sc); err != nil { + return fmt.Errorf("loading root from blockstore: %w", err) + } - if err := t.WalkLeavesFrom(ctx, prefix, func(k string, val cid.Cid) error { - blk, err := bs.Get(ctx, val) - if err != nil { - slog.Error("failed to get record from tree", "key", k, "cid", val, "error", err) - return nil - } + if sc.Version != ATP_REPO_VERSION && sc.Version != ATP_REPO_VERSION_2 { + return fmt.Errorf("unsupported repo version: %d", sc.Version) + } + // TODO: verify that signature - return cb(k, val, blk.RawData()) - }); err != nil { - walkErr = fmt.Errorf("failed to walk mst: %w", err) - } - }() + t := mst.LoadMST(cst, sc.Data) - for { - blk, err := br.NextBlock(repoBlockBufferPool, repoBlockBufferSize) + if err := t.WalkLeavesFrom(ctx, prefix, func(k string, val cid.Cid) error { + blk, err := bs.Get(ctx, val) if err != nil { - if err == io.EOF { - break - } - return fmt.Errorf("reading block from CAR: %w", err) + slog.Error("failed to get record from tree", "key", k, "cid", val, "error", err) + return nil } - if err := bs.Put(ctx, blk); err != nil { - return fmt.Errorf("copying block to store: %w", err) - } + return cb(k, val, blk.RawData()) + }); err != nil { + return fmt.Errorf("failed to walk mst: %w", err) } - bs.Complete() - - wg.Wait() - - return walkErr + return nil } -- 2.51.2 From 9c558ddd694df844978ea62d2886c6bd219c241a Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Sat, 22 Mar 2025 14:56:47 -0700 Subject: [PATCH 04/26] manage That Memory --- repo/carutil/reader.go | 14 +------------- repo/repo.go | 12 ++++++++++++ repo/stream.go | 40 +++++++++++++++++----------------------- 3 files changed, 30 insertions(+), 36 deletions(-) diff --git a/repo/carutil/reader.go b/repo/carutil/reader.go index cae6eb71..69352b5a 100644 --- a/repo/carutil/reader.go +++ b/repo/carutil/reader.go @@ -15,8 +15,6 @@ import ( type Reader struct { r *bufio.Reader - - bufs [][]byte } func NewReader(r *bufio.Reader) (*Reader, cid.Cid, error) { @@ -34,18 +32,10 @@ func NewReader(r *bufio.Reader) (*Reader, cid.Cid, error) { } return &Reader{ - r: r, - bufs: make([][]byte, 0, 10), + r: r, }, h.Roots[0], nil } -func (r *Reader) Free(alloc *sync.Pool) { - for _, b := range r.bufs { - alloc.Put(b) - } - r.bufs = nil -} - const MaxAllowedSectionSize = 32 << 20 func (r *Reader) NextBlock(allocator *sync.Pool, allocMax uint64) (blocks.Block, error) { @@ -54,8 +44,6 @@ func (r *Reader) NextBlock(allocator *sync.Pool, allocMax uint64) (blocks.Block, return nil, err } - r.bufs = append(r.bufs, data) - n, c, err := cid.CidFromBytes(data) if err != nil { return nil, err diff --git a/repo/repo.go b/repo/repo.go index 582177e9..63cfdb73 100644 --- a/repo/repo.go +++ b/repo/repo.go @@ -14,6 +14,7 @@ import ( "github.com/bluesky-social/indigo/mst" "github.com/bluesky-social/indigo/repo/carutil" "github.com/bluesky-social/indigo/util" + blocks "github.com/ipfs/go-block-format" "github.com/ipfs/go-cid" cbor "github.com/ipfs/go-ipld-cbor" cbg "github.com/whyrusleeping/cbor-gen" @@ -86,6 +87,17 @@ var repoBlockBufferPool = &sync.Pool{ }, } +func freeRepoBlock(b blocks.Block) { + bb, ok := b.(*blocks.BasicBlock) + if !ok { + return + } + + if cap(bb.RawData()) == repoBlockBufferSize { + repoBlockBufferPool.Put(bb.RawData()[:repoBlockBufferSize]) + } +} + func IngestRepo(ctx context.Context, bs cbor.IpldBlockstore, r io.Reader) (cid.Cid, error) { ctx, span := otel.Tracer("repo").Start(ctx, "Ingest") defer span.End() diff --git a/repo/stream.go b/repo/stream.go index 0b43b514..ee547067 100644 --- a/repo/stream.go +++ b/repo/stream.go @@ -15,20 +15,23 @@ import ( "go.opentelemetry.io/otel" ) -type waitingBlockstore struct { +type readStreamBlockstore struct { otherBlocks map[cid.Cid]block.Block streamComplete bool r *carutil.Reader + + lastBlockRead block.Block } -func newWaitingBlockstore() *waitingBlockstore { - return &waitingBlockstore{ +func newStreamingBlockstore(r *carutil.Reader) *readStreamBlockstore { + return &readStreamBlockstore{ otherBlocks: make(map[cid.Cid]block.Block), + r: r, } } -func (bs *waitingBlockstore) readUntilBlock(ctx context.Context, cc cid.Cid) (block.Block, error) { +func (bs *readStreamBlockstore) readUntilBlock(ctx context.Context, cc cid.Cid) (block.Block, error) { for { blk, err := bs.r.NextBlock(repoBlockBufferPool, repoBlockBufferSize) if err != nil { @@ -50,9 +53,15 @@ func (bs *waitingBlockstore) readUntilBlock(ctx context.Context, cc cid.Cid) (bl return nil, io.EOF } -func (bs *waitingBlockstore) Get(ctx context.Context, cc cid.Cid) (block.Block, error) { +func (bs *readStreamBlockstore) Get(ctx context.Context, cc cid.Cid) (block.Block, error) { + if bs.lastBlockRead != nil { + freeRepoBlock(bs.lastBlockRead) + bs.lastBlockRead = nil + } + if blk, ok := bs.otherBlocks[cc]; ok { delete(bs.otherBlocks, cc) + bs.lastBlockRead = blk return blk, nil } @@ -65,27 +74,14 @@ func (bs *waitingBlockstore) Get(ctx context.Context, cc cid.Cid) (block.Block, return nil, err } + bs.lastBlockRead = blk return blk, nil } var ErrMissingBlock = fmt.Errorf("block was missing from archive") -func (bs *waitingBlockstore) Put(ctx context.Context, blk block.Block) error { +func (bs *readStreamBlockstore) Put(ctx context.Context, blk block.Block) error { return fmt.Errorf("put is not needed") - /* - bw, ok := bs.blockWaits[blk.Cid()] - if ok { - bw <- blk - delete(bs.blockWaits, blk.Cid()) - return nil - } - - bs.otherBlocks[blk.Cid()] = blk.(*block.BasicBlock) - return nil - */ -} - -func (bs *waitingBlockstore) Complete() { } func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, cb func(k string, c cid.Cid, v []byte) error) error { @@ -97,9 +93,7 @@ func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, cb func( return fmt.Errorf("opening CAR block reader: %w", err) } - bs := newWaitingBlockstore() - - bs.r = br + bs := newStreamingBlockstore(br) cst := util.CborStore(bs) -- 2.51.2 From c38a3f597df454136cc3d87315dc17e1bdba8cb3 Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Sat, 22 Mar 2025 15:52:12 -0700 Subject: [PATCH 05/26] implement view blockstore to have better memory management --- backfill/backfill.go | 39 ++++++++++++++++----------------------- cmd/gosky/main.go | 2 +- repo/repo.go | 12 +++--------- repo/stream.go | 35 ++++++++++++++++++++--------------- 4 files changed, 40 insertions(+), 48 deletions(-) diff --git a/backfill/backfill.go b/backfill/backfill.go index dc59ae45..1821c478 100644 --- a/backfill/backfill.go +++ b/backfill/backfill.go @@ -5,6 +5,7 @@ import ( "context" "errors" "fmt" + "io" "log/slog" "net/http" "strings" @@ -301,6 +302,7 @@ func (b *Backfiller) FlushBuffer(ctx context.Context, job Job) int { type recordQueueItem struct { recordPath string nodeCid cid.Cid + data []byte } type recordResult struct { @@ -323,8 +325,8 @@ func (e *FetchRepoError) Error() string { return fmt.Sprintf("failed to get repo: %s (%d)", reason, e.StatusCode) } -// Fetches a repo CAR file over HTTP from the indicated host. If successful, parses the CAR and returns repo.Repo -func (b *Backfiller) fetchRepo(ctx context.Context, did, since, host string) (*repo.Repo, error) { +// Fetches a repo CAR file over HTTP from the indicated host. +func (b *Backfiller) fetchRepo(ctx context.Context, did, since, host string) (io.ReadCloser, error) { url := fmt.Sprintf("%s/xrpc/com.atproto.sync.getRepo?did=%s", host, did) if since != "" { @@ -366,13 +368,7 @@ func (b *Backfiller) fetchRepo(ctx context.Context, did, since, host string) (*r counter: backfillBytesProcessed.WithLabelValues(b.Name), } - defer instrumentedReader.Close() - - repo, err := repo.ReadRepoFromCar(ctx, instrumentedReader) - if err != nil { - return nil, fmt.Errorf("failed to parse repo from CAR file: %w", err) - } - return repo, nil + return &instrumentedReader, nil } // BackfillRepo backfills a repo @@ -390,7 +386,7 @@ func (b *Backfiller) BackfillRepo(ctx context.Context, job Job) (string, error) } log.Info(fmt.Sprintf("processing backfill for %s", repoDID)) - var r *repo.Repo + var r io.ReadCloser if b.tryRelayRepoFetch { rr, err := b.fetchRepo(ctx, repoDID, job.Rev(), b.RelayHost) if err != nil { @@ -426,19 +422,21 @@ func (b *Backfiller) BackfillRepo(ctx context.Context, job Job) (string, error) recordQueue := make(chan recordQueueItem, numRoutines) recordResults := make(chan recordResult, numRoutines) + var rev string // Producer routine go func() { defer close(recordQueue) - if err := r.ForEach(ctx, b.NSIDFilter, func(recordPath string, nodeCid cid.Cid) error { + rrev, err := repo.StreamRepoRecords(ctx, r, b.NSIDFilter, func(recordPath string, nodeCid cid.Cid, data []byte) error { numRecords++ - recordQueue <- recordQueueItem{recordPath: recordPath, nodeCid: nodeCid} + recordQueue <- recordQueueItem{recordPath: recordPath, nodeCid: nodeCid, data: data} return nil - }); err != nil { + }) + if err != nil { log.Error("failed to iterate records in repo", "err", err) } - }() - rev := r.SignedCommit().Rev + rev = rrev + }() // Consumer routines wg := sync.WaitGroup{} @@ -447,15 +445,10 @@ func (b *Backfiller) BackfillRepo(ctx context.Context, job Job) (string, error) go func() { defer wg.Done() for item := range recordQueue { - blk, err := r.Blockstore().Get(ctx, item.nodeCid) - if err != nil { - recordResults <- recordResult{recordPath: item.recordPath, err: fmt.Errorf("failed to get blocks for record: %w", err)} - continue - } - raw := blk.RawData() + raw := item.data - err = b.HandleCreateRecord(ctx, repoDID, rev, item.recordPath, &raw, &item.nodeCid) + err := b.HandleCreateRecord(ctx, repoDID, rev, item.recordPath, &raw, &item.nodeCid) if err != nil { recordResults <- recordResult{recordPath: item.recordPath, err: fmt.Errorf("failed to handle create record: %w", err)} continue @@ -483,7 +476,7 @@ func (b *Backfiller) BackfillRepo(ctx context.Context, job Job) (string, error) close(recordResults) resultWG.Wait() - if err := job.SetRev(ctx, r.SignedCommit().Rev); err != nil { + if err := job.SetRev(ctx, rev); err != nil { log.Error("failed to update rev after backfilling repo", "err", err) } diff --git a/cmd/gosky/main.go b/cmd/gosky/main.go index 5b69fe10..69a02145 100644 --- a/cmd/gosky/main.go +++ b/cmd/gosky/main.go @@ -672,7 +672,7 @@ var listAllRecordsCmd = &cli.Command{ vals := cctx.Bool("values") cids := cctx.Bool("cids") - if err := repo.StreamRepoRecords(ctx, bytes.NewReader(repob), collection, func(k string, cc cid.Cid, v []byte) error { + if _, err := repo.StreamRepoRecords(ctx, bytes.NewReader(repob), collection, func(k string, cc cid.Cid, v []byte) error { if !strings.HasPrefix(k, collection) { return repo.ErrDoneIterating } diff --git a/repo/repo.go b/repo/repo.go index 63cfdb73..b1200f79 100644 --- a/repo/repo.go +++ b/repo/repo.go @@ -14,7 +14,6 @@ import ( "github.com/bluesky-social/indigo/mst" "github.com/bluesky-social/indigo/repo/carutil" "github.com/bluesky-social/indigo/util" - blocks "github.com/ipfs/go-block-format" "github.com/ipfs/go-cid" cbor "github.com/ipfs/go-ipld-cbor" cbg "github.com/whyrusleeping/cbor-gen" @@ -87,14 +86,9 @@ var repoBlockBufferPool = &sync.Pool{ }, } -func freeRepoBlock(b blocks.Block) { - bb, ok := b.(*blocks.BasicBlock) - if !ok { - return - } - - if cap(bb.RawData()) == repoBlockBufferSize { - repoBlockBufferPool.Put(bb.RawData()[:repoBlockBufferSize]) +func FreeRepoBlock(b []byte) { + if cap(b) == repoBlockBufferSize { + repoBlockBufferPool.Put(b[:repoBlockBufferSize]) } } diff --git a/repo/stream.go b/repo/stream.go index ee547067..7263a801 100644 --- a/repo/stream.go +++ b/repo/stream.go @@ -20,8 +20,6 @@ type readStreamBlockstore struct { streamComplete bool r *carutil.Reader - - lastBlockRead block.Block } func newStreamingBlockstore(r *carutil.Reader) *readStreamBlockstore { @@ -54,14 +52,8 @@ func (bs *readStreamBlockstore) readUntilBlock(ctx context.Context, cc cid.Cid) } func (bs *readStreamBlockstore) Get(ctx context.Context, cc cid.Cid) (block.Block, error) { - if bs.lastBlockRead != nil { - freeRepoBlock(bs.lastBlockRead) - bs.lastBlockRead = nil - } - if blk, ok := bs.otherBlocks[cc]; ok { delete(bs.otherBlocks, cc) - bs.lastBlockRead = blk return blk, nil } @@ -74,23 +66,36 @@ func (bs *readStreamBlockstore) Get(ctx context.Context, cc cid.Cid) (block.Bloc return nil, err } - bs.lastBlockRead = blk return blk, nil } +func (bs *readStreamBlockstore) View(cc cid.Cid, cb func([]byte) error) error { + blk, err := bs.Get(context.TODO(), cc) + if err != nil { + return err + } + + if err := cb(blk.RawData()); err != nil { + return err + } + + FreeRepoBlock(blk.RawData()) + return nil +} + var ErrMissingBlock = fmt.Errorf("block was missing from archive") func (bs *readStreamBlockstore) Put(ctx context.Context, blk block.Block) error { return fmt.Errorf("put is not needed") } -func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, cb func(k string, c cid.Cid, v []byte) error) error { +func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, cb func(k string, c cid.Cid, v []byte) error) (string, error) { ctx, span := otel.Tracer("repo").Start(ctx, "RepoStream") defer span.End() br, root, err := carutil.NewReader(bufio.NewReader(r)) if err != nil { - return fmt.Errorf("opening CAR block reader: %w", err) + return "", fmt.Errorf("opening CAR block reader: %w", err) } bs := newStreamingBlockstore(br) @@ -99,11 +104,11 @@ func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, cb func( var sc SignedCommit if err := cst.Get(ctx, root, &sc); err != nil { - return fmt.Errorf("loading root from blockstore: %w", err) + return "", fmt.Errorf("loading root from blockstore: %w", err) } if sc.Version != ATP_REPO_VERSION && sc.Version != ATP_REPO_VERSION_2 { - return fmt.Errorf("unsupported repo version: %d", sc.Version) + return "", fmt.Errorf("unsupported repo version: %d", sc.Version) } // TODO: verify that signature @@ -118,8 +123,8 @@ func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, cb func( return cb(k, val, blk.RawData()) }); err != nil { - return fmt.Errorf("failed to walk mst: %w", err) + return "", fmt.Errorf("failed to walk mst: %w", err) } - return nil + return sc.Rev, nil } -- 2.51.2 From 9c21a25312c7f379f157502615485030794944c4 Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Sat, 22 Mar 2025 15:52:41 -0700 Subject: [PATCH 06/26] more memory management --- backfill/backfill.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/backfill/backfill.go b/backfill/backfill.go index 1821c478..72a3fbcd 100644 --- a/backfill/backfill.go +++ b/backfill/backfill.go @@ -454,6 +454,8 @@ func (b *Backfiller) BackfillRepo(ctx context.Context, job Job) (string, error) continue } + repo.FreeRepoBlock(raw) + backfillRecordsProcessed.WithLabelValues(b.Name).Inc() recordResults <- recordResult{recordPath: item.recordPath, err: err} } -- 2.51.2 From 348053858d0098a59043826ccf4dd9ca43e846b9 Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Sat, 22 Mar 2025 16:00:41 -0700 Subject: [PATCH 07/26] close read streams --- backfill/backfill.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/backfill/backfill.go b/backfill/backfill.go index 72a3fbcd..d2219307 100644 --- a/backfill/backfill.go +++ b/backfill/backfill.go @@ -417,6 +417,8 @@ func (b *Backfiller) BackfillRepo(ctx context.Context, job Job) (string, error) } } + defer r.Close() + numRecords := 0 numRoutines := b.ParallelRecordCreates recordQueue := make(chan recordQueueItem, numRoutines) -- 2.51.2 From 6d1151cbcf34de7fa2b4b6fabea0a488e09c0e9e Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Sat, 22 Mar 2025 22:47:51 -0700 Subject: [PATCH 08/26] some more memory tuning --- repo/carutil/reader.go | 41 ++++++++++++++++++++++++++++++++++++++--- repo/stream.go | 14 +++++++++----- 2 files changed, 47 insertions(+), 8 deletions(-) diff --git a/repo/carutil/reader.go b/repo/carutil/reader.go index 69352b5a..62b4bf45 100644 --- a/repo/carutil/reader.go +++ b/repo/carutil/reader.go @@ -8,7 +8,6 @@ import ( "io" "sync" - blocks "github.com/ipfs/go-block-format" "github.com/ipfs/go-cid" car "github.com/ipld/go-car" ) @@ -38,7 +37,7 @@ func NewReader(r *bufio.Reader) (*Reader, cid.Cid, error) { const MaxAllowedSectionSize = 32 << 20 -func (r *Reader) NextBlock(allocator *sync.Pool, allocMax uint64) (blocks.Block, error) { +func (r *Reader) NextBlock(allocator *sync.Pool, allocMax uint64) (*BasicBlock, error) { data, err := ldRead(r.r, allocator, allocMax) if err != nil { return nil, err @@ -49,7 +48,7 @@ func (r *Reader) NextBlock(allocator *sync.Pool, allocMax uint64) (blocks.Block, return nil, err } - return blocks.NewBlockWithCid(data[n:], c) + return NewBlockWithCid(data[n:], data, c), nil } func ldRead(r *bufio.Reader, alloc *sync.Pool, allocMax uint64) ([]byte, error) { @@ -88,3 +87,39 @@ func ldRead(r *bufio.Reader, alloc *sync.Pool, allocMax uint64) ([]byte, error) return buf, nil } + +type BasicBlock struct { + cid cid.Cid + data []byte + base []byte +} + +func NewBlockWithCid(data, base []byte, c cid.Cid) *BasicBlock { + return &BasicBlock{data: data, cid: c, base: base} +} + +// RawData returns the block raw contents as a byte slice. +func (b *BasicBlock) RawData() []byte { + return b.data +} + +// Cid returns the content identifier of the block. +func (b *BasicBlock) Cid() cid.Cid { + return b.cid +} + +// String provides a human-readable representation of the block CID. +func (b *BasicBlock) String() string { + return fmt.Sprintf("[Block %s]", b.Cid()) +} + +// Loggable returns a go-log loggable item. +func (b *BasicBlock) Loggable() map[string]interface{} { + return map[string]interface{}{ + "block": b.Cid().String(), + } +} + +func (b *BasicBlock) BaseBuffer() []byte { + return b.base +} diff --git a/repo/stream.go b/repo/stream.go index 7263a801..f867c408 100644 --- a/repo/stream.go +++ b/repo/stream.go @@ -16,7 +16,7 @@ import ( ) type readStreamBlockstore struct { - otherBlocks map[cid.Cid]block.Block + otherBlocks map[cid.Cid]*carutil.BasicBlock streamComplete bool r *carutil.Reader @@ -24,12 +24,12 @@ type readStreamBlockstore struct { func newStreamingBlockstore(r *carutil.Reader) *readStreamBlockstore { return &readStreamBlockstore{ - otherBlocks: make(map[cid.Cid]block.Block), + otherBlocks: make(map[cid.Cid]*carutil.BasicBlock), r: r, } } -func (bs *readStreamBlockstore) readUntilBlock(ctx context.Context, cc cid.Cid) (block.Block, error) { +func (bs *readStreamBlockstore) readUntilBlock(ctx context.Context, cc cid.Cid) (*carutil.BasicBlock, error) { for { blk, err := bs.r.NextBlock(repoBlockBufferPool, repoBlockBufferSize) if err != nil { @@ -52,6 +52,10 @@ func (bs *readStreamBlockstore) readUntilBlock(ctx context.Context, cc cid.Cid) } func (bs *readStreamBlockstore) Get(ctx context.Context, cc cid.Cid) (block.Block, error) { + return bs.get(ctx, cc) +} + +func (bs *readStreamBlockstore) get(ctx context.Context, cc cid.Cid) (*carutil.BasicBlock, error) { if blk, ok := bs.otherBlocks[cc]; ok { delete(bs.otherBlocks, cc) return blk, nil @@ -70,7 +74,7 @@ func (bs *readStreamBlockstore) Get(ctx context.Context, cc cid.Cid) (block.Bloc } func (bs *readStreamBlockstore) View(cc cid.Cid, cb func([]byte) error) error { - blk, err := bs.Get(context.TODO(), cc) + blk, err := bs.get(context.TODO(), cc) if err != nil { return err } @@ -79,7 +83,7 @@ func (bs *readStreamBlockstore) View(cc cid.Cid, cb func([]byte) error) error { return err } - FreeRepoBlock(blk.RawData()) + FreeRepoBlock(blk.BaseBuffer()) return nil } -- 2.51.2 From 65afc100118302fccc20262e9dfc428cef0980d8 Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Sat, 22 Mar 2025 23:25:24 -0700 Subject: [PATCH 09/26] maybe the pool was a bad idea, but the rest of this is fine --- repo/carutil/reader.go | 22 +++++----------------- repo/repo.go | 17 +---------------- repo/stream.go | 3 +-- 3 files changed, 7 insertions(+), 35 deletions(-) diff --git a/repo/carutil/reader.go b/repo/carutil/reader.go index 62b4bf45..34070e6e 100644 --- a/repo/carutil/reader.go +++ b/repo/carutil/reader.go @@ -6,7 +6,6 @@ import ( "errors" "fmt" "io" - "sync" "github.com/ipfs/go-cid" car "github.com/ipld/go-car" @@ -37,8 +36,8 @@ func NewReader(r *bufio.Reader) (*Reader, cid.Cid, error) { const MaxAllowedSectionSize = 32 << 20 -func (r *Reader) NextBlock(allocator *sync.Pool, allocMax uint64) (*BasicBlock, error) { - data, err := ldRead(r.r, allocator, allocMax) +func (r *Reader) NextBlock() (*BasicBlock, error) { + data, err := ldRead(r.r) if err != nil { return nil, err } @@ -51,7 +50,7 @@ func (r *Reader) NextBlock(allocator *sync.Pool, allocMax uint64) (*BasicBlock, return NewBlockWithCid(data[n:], data, c), nil } -func ldRead(r *bufio.Reader, alloc *sync.Pool, allocMax uint64) ([]byte, error) { +func ldRead(r *bufio.Reader) ([]byte, error) { if _, err := r.Peek(1); err != nil { // no more blocks, likely clean io.EOF return nil, err } @@ -68,19 +67,8 @@ func ldRead(r *bufio.Reader, alloc *sync.Pool, allocMax uint64) ([]byte, error) return nil, errors.New("malformed car; header is bigger than util.MaxAllowedSectionSize") } - if l > allocMax { - // direct allocation, not great - buf := make([]byte, l) - if _, err := io.ReadFull(r, buf); err != nil { - return nil, err - } - - return buf, nil - } - - buf := alloc.Get().([]byte) - buf = buf[:l] - + // direct allocation, not great + buf := make([]byte, l) if _, err := io.ReadFull(r, buf); err != nil { return nil, err } diff --git a/repo/repo.go b/repo/repo.go index b1200f79..56e4f477 100644 --- a/repo/repo.go +++ b/repo/repo.go @@ -6,7 +6,6 @@ import ( "context" "fmt" "io" - "sync" "github.com/bluesky-social/indigo/atproto/repo" "github.com/bluesky-social/indigo/atproto/syntax" @@ -78,20 +77,6 @@ func (uc *UnsignedCommit) BytesForSigning() ([]byte, error) { return buf.Bytes(), nil } -const repoBlockBufferSize = 128 << 10 - -var repoBlockBufferPool = &sync.Pool{ - New: func() any { - return make([]byte, repoBlockBufferSize) - }, -} - -func FreeRepoBlock(b []byte) { - if cap(b) == repoBlockBufferSize { - repoBlockBufferPool.Put(b[:repoBlockBufferSize]) - } -} - func IngestRepo(ctx context.Context, bs cbor.IpldBlockstore, r io.Reader) (cid.Cid, error) { ctx, span := otel.Tracer("repo").Start(ctx, "Ingest") defer span.End() @@ -102,7 +87,7 @@ func IngestRepo(ctx context.Context, bs cbor.IpldBlockstore, r io.Reader) (cid.C } for { - blk, err := br.NextBlock(repoBlockBufferPool, repoBlockBufferSize) + blk, err := br.NextBlock() if err != nil { if err == io.EOF { break diff --git a/repo/stream.go b/repo/stream.go index f867c408..bf928ecb 100644 --- a/repo/stream.go +++ b/repo/stream.go @@ -31,7 +31,7 @@ func newStreamingBlockstore(r *carutil.Reader) *readStreamBlockstore { func (bs *readStreamBlockstore) readUntilBlock(ctx context.Context, cc cid.Cid) (*carutil.BasicBlock, error) { for { - blk, err := bs.r.NextBlock(repoBlockBufferPool, repoBlockBufferSize) + blk, err := bs.r.NextBlock() if err != nil { if err == io.EOF { break @@ -83,7 +83,6 @@ func (bs *readStreamBlockstore) View(cc cid.Cid, cb func([]byte) error) error { return err } - FreeRepoBlock(blk.BaseBuffer()) return nil } -- 2.51.2 From 677f989a20a15d5e2b3467feb8e8d2a800239af8 Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Sat, 22 Mar 2025 23:26:09 -0700 Subject: [PATCH 10/26] rm unused --- backfill/backfill.go | 2 -- 1 file changed, 2 deletions(-) diff --git a/backfill/backfill.go b/backfill/backfill.go index d2219307..595d2ebe 100644 --- a/backfill/backfill.go +++ b/backfill/backfill.go @@ -456,8 +456,6 @@ func (b *Backfiller) BackfillRepo(ctx context.Context, job Job) (string, error) continue } - repo.FreeRepoBlock(raw) - backfillRecordsProcessed.WithLabelValues(b.Name).Inc() recordResults <- recordResult{recordPath: item.recordPath, err: err} } -- 2.51.2 From 506a8a9a26372866030e4747a8a2f29161148357 Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Sat, 22 Mar 2025 23:30:41 -0700 Subject: [PATCH 11/26] reuse bufio readers --- repo/repo.go | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/repo/repo.go b/repo/repo.go index 56e4f477..f9bf5643 100644 --- a/repo/repo.go +++ b/repo/repo.go @@ -6,6 +6,7 @@ import ( "context" "fmt" "io" + "sync" "github.com/bluesky-social/indigo/atproto/repo" "github.com/bluesky-social/indigo/atproto/syntax" @@ -77,11 +78,20 @@ func (uc *UnsignedCommit) BytesForSigning() ([]byte, error) { return buf.Bytes(), nil } +var bufrPool = &sync.Pool{ + New: func() any { + return bufio.NewReader(nil) + }, +} + func IngestRepo(ctx context.Context, bs cbor.IpldBlockstore, r io.Reader) (cid.Cid, error) { ctx, span := otel.Tracer("repo").Start(ctx, "Ingest") defer span.End() - br, root, err := carutil.NewReader(bufio.NewReader(r)) + bufr := bufrPool.Get().(*bufio.Reader) + bufr.Reset(r) + + br, root, err := carutil.NewReader(bufr) if err != nil { return cid.Undef, fmt.Errorf("opening CAR block reader: %w", err) } @@ -100,6 +110,8 @@ func IngestRepo(ctx context.Context, bs cbor.IpldBlockstore, r io.Reader) (cid.C } } + bufrPool.Put(bufr) + return root, nil } -- 2.51.2 From 2891bcbdeab941b229133a532e14886e17d80730 Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Sun, 23 Mar 2025 22:47:35 -0700 Subject: [PATCH 12/26] more alloc reductions --- cmd/gosky/main.go | 1 - events/consumer.go | 33 ++++++++++++++++++++------------- 2 files changed, 20 insertions(+), 14 deletions(-) diff --git a/cmd/gosky/main.go b/cmd/gosky/main.go index 69a02145..e1a78924 100644 --- a/cmd/gosky/main.go +++ b/cmd/gosky/main.go @@ -652,7 +652,6 @@ var listAllRecordsCmd = &cli.Command{ return err } repob = rrb - fmt.Println("GOT REPO BYTES") } else { if len(arg) == 0 { return cli.Exit("must specify DID string or repo path", 127) diff --git a/events/consumer.go b/events/consumer.go index b9180cd4..0f654ce2 100644 --- a/events/consumer.go +++ b/events/consumer.go @@ -12,6 +12,8 @@ import ( comatproto "github.com/bluesky-social/indigo/api/atproto" "github.com/prometheus/client_golang/prometheus" + cbg "github.com/whyrusleeping/cbor-gen" + "github.com/gorilla/websocket" ) @@ -162,6 +164,13 @@ func HandleRepoStream(ctx context.Context, con *websocket.Conn, sched Scheduler, return nil }) + cr := new(cbg.CborReader) + + ir := &instrumentedReader{ + addr: remoteAddr, + bytesCounter: bytesFromStreamCounter.WithLabelValues(remoteAddr), + } + lastSeq := int64(-1) for { select { @@ -182,14 +191,12 @@ func HandleRepoStream(ctx context.Context, con *websocket.Conn, sched Scheduler, // ok } - r := &instrumentedReader{ - r: rawReader, - addr: remoteAddr, - bytesCounter: bytesFromStreamCounter.WithLabelValues(remoteAddr), - } + ir.r = rawReader + + cr.SetReader(ir) var header EventHeader - if err := header.UnmarshalCBOR(r); err != nil { + if err := header.UnmarshalCBOR(cr); err != nil { return fmt.Errorf("reading header: %w", err) } @@ -200,7 +207,7 @@ func HandleRepoStream(ctx context.Context, con *websocket.Conn, sched Scheduler, switch header.MsgType { case "#commit": var evt comatproto.SyncSubscribeRepos_Commit - if err := evt.UnmarshalCBOR(r); err != nil { + if err := evt.UnmarshalCBOR(cr); err != nil { return fmt.Errorf("reading repoCommit event: %w", err) } @@ -217,7 +224,7 @@ func HandleRepoStream(ctx context.Context, con *websocket.Conn, sched Scheduler, } case "#sync": var evt comatproto.SyncSubscribeRepos_Sync - if err := evt.UnmarshalCBOR(r); err != nil { + if err := evt.UnmarshalCBOR(cr); err != nil { return fmt.Errorf("reading repoSync event: %w", err) } @@ -234,7 +241,7 @@ func HandleRepoStream(ctx context.Context, con *websocket.Conn, sched Scheduler, } case "#identity": var evt comatproto.SyncSubscribeRepos_Identity - if err := evt.UnmarshalCBOR(r); err != nil { + if err := evt.UnmarshalCBOR(cr); err != nil { return err } @@ -250,7 +257,7 @@ func HandleRepoStream(ctx context.Context, con *websocket.Conn, sched Scheduler, } case "#account": var evt comatproto.SyncSubscribeRepos_Account - if err := evt.UnmarshalCBOR(r); err != nil { + if err := evt.UnmarshalCBOR(cr); err != nil { return err } @@ -267,7 +274,7 @@ func HandleRepoStream(ctx context.Context, con *websocket.Conn, sched Scheduler, case "#info": // TODO: this might also be a LabelInfo (as opposed to RepoInfo) var evt comatproto.SyncSubscribeRepos_Info - if err := evt.UnmarshalCBOR(r); err != nil { + if err := evt.UnmarshalCBOR(cr); err != nil { return err } @@ -278,7 +285,7 @@ func HandleRepoStream(ctx context.Context, con *websocket.Conn, sched Scheduler, } case "#labels": var evt comatproto.LabelSubscribeLabels_Labels - if err := evt.UnmarshalCBOR(r); err != nil { + if err := evt.UnmarshalCBOR(cr); err != nil { return fmt.Errorf("reading Labels event: %w", err) } @@ -297,7 +304,7 @@ func HandleRepoStream(ctx context.Context, con *websocket.Conn, sched Scheduler, case EvtKindErrorFrame: var errframe ErrorFrame - if err := errframe.UnmarshalCBOR(r); err != nil { + if err := errframe.UnmarshalCBOR(cr); err != nil { return err } -- 2.51.2 From dbe374cc64df07df1303801843d7a49107fa522c Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Mon, 24 Mar 2025 11:54:34 -0700 Subject: [PATCH 13/26] try a pool --- repo/carutil/reader.go | 44 ++++++++++++++++++++++++++++++------------ repo/stream.go | 37 ++++++++++++++++++++++++++++++----- 2 files changed, 64 insertions(+), 17 deletions(-) diff --git a/repo/carutil/reader.go b/repo/carutil/reader.go index 34070e6e..7c9f9508 100644 --- a/repo/carutil/reader.go +++ b/repo/carutil/reader.go @@ -37,43 +37,63 @@ func NewReader(r *bufio.Reader) (*Reader, cid.Cid, error) { const MaxAllowedSectionSize = 32 << 20 func (r *Reader) NextBlock() (*BasicBlock, error) { - data, err := ldRead(r.r) + blk, _, err := r.NextBlockBuf(nil) if err != nil { return nil, err } + return blk, nil +} + +func (r *Reader) NextBlockBuf(buf []byte) (*BasicBlock, bool, error) { + data, usedBuf, err := ldRead(r.r, buf) + if err != nil { + return nil, false, err + } + n, c, err := cid.CidFromBytes(data) if err != nil { - return nil, err + return nil, false, err } - return NewBlockWithCid(data[n:], data, c), nil + return NewBlockWithCid(data[n:], data, c), usedBuf, nil } -func ldRead(r *bufio.Reader) ([]byte, error) { +// reads a length delimited value off of the given reader into the the given buf if its big enough, otherwise allocates a new buffer. +// returns whether or not the passed in buffer was used +func ldRead(r *bufio.Reader, buf []byte) ([]byte, bool, error) { if _, err := r.Peek(1); err != nil { // no more blocks, likely clean io.EOF - return nil, err + return nil, false, err } l, err := binary.ReadUvarint(r) if err != nil { if err == io.EOF { - return nil, io.ErrUnexpectedEOF // don't silently pretend this is a clean EOF + return nil, false, io.ErrUnexpectedEOF // don't silently pretend this is a clean EOF } - return nil, err + return nil, false, err } if l > uint64(MaxAllowedSectionSize) { // Don't OOM - return nil, errors.New("malformed car; header is bigger than util.MaxAllowedSectionSize") + return nil, false, errors.New("malformed car; header is bigger than util.MaxAllowedSectionSize") + } + + if l > uint64(len(buf)) { + // direct allocation, not great + buf := make([]byte, l) + if _, err := io.ReadFull(r, buf); err != nil { + return nil, false, err + } + + return buf, false, nil } - // direct allocation, not great - buf := make([]byte, l) + buf = buf[:l] if _, err := io.ReadFull(r, buf); err != nil { - return nil, err + return nil, false, err } - return buf, nil + return buf, true, nil } type BasicBlock struct { diff --git a/repo/stream.go b/repo/stream.go index bf928ecb..42d01d1f 100644 --- a/repo/stream.go +++ b/repo/stream.go @@ -6,6 +6,7 @@ import ( "fmt" "io" "log/slog" + "sync" "github.com/bluesky-social/indigo/mst" "github.com/bluesky-social/indigo/repo/carutil" @@ -15,34 +16,53 @@ import ( "go.opentelemetry.io/otel" ) +const bufPoolBlockSize = 512 + +var smallBlockPool = &sync.Pool{ + New: func() any { + return make([]byte, bufPoolBlockSize) + }, +} + type readStreamBlockstore struct { otherBlocks map[cid.Cid]*carutil.BasicBlock streamComplete bool r *carutil.Reader + + outOfOrder int + totalBlocks int } func newStreamingBlockstore(r *carutil.Reader) *readStreamBlockstore { return &readStreamBlockstore{ - otherBlocks: make(map[cid.Cid]*carutil.BasicBlock), + otherBlocks: make(map[cid.Cid]*carutil.BasicBlock, 20), r: r, } } func (bs *readStreamBlockstore) readUntilBlock(ctx context.Context, cc cid.Cid) (*carutil.BasicBlock, error) { for { - blk, err := bs.r.NextBlock() + buf := smallBlockPool.Get().([]byte) + blk, used, err := bs.r.NextBlockBuf(buf) if err != nil { + smallBlockPool.Put(buf) if err == io.EOF { break } return nil, fmt.Errorf("reading block from CAR: %w", err) } + if !used { + smallBlockPool.Put(buf) + } + + bs.totalBlocks++ if blk.Cid() == cc { return blk, nil } + bs.outOfOrder++ bs.otherBlocks[blk.Cid()] = blk } @@ -83,6 +103,10 @@ func (bs *readStreamBlockstore) View(cc cid.Cid, cb func([]byte) error) error { return err } + if len(blk.BaseBuffer()) == bufPoolBlockSize { + smallBlockPool.Put(blk.BaseBuffer()) + } + return nil } @@ -113,21 +137,24 @@ func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, cb func( if sc.Version != ATP_REPO_VERSION && sc.Version != ATP_REPO_VERSION_2 { return "", fmt.Errorf("unsupported repo version: %d", sc.Version) } + // TODO: verify that signature t := mst.LoadMST(cst, sc.Data) if err := t.WalkLeavesFrom(ctx, prefix, func(k string, val cid.Cid) error { - blk, err := bs.Get(ctx, val) - if err != nil { + if err := bs.View(val, func(data []byte) error { + return cb(k, val, data) + }); err != nil { slog.Error("failed to get record from tree", "key", k, "cid", val, "error", err) return nil } - return cb(k, val, blk.RawData()) + return nil }); err != nil { return "", fmt.Errorf("failed to walk mst: %w", err) } + fmt.Println("out of order blocks: ", bs.outOfOrder, bs.totalBlocks) return sc.Rev, nil } -- 2.51.2 From 0dcfa8a713aadd5907e1d1d0c4217c1f512c4882 Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Mon, 24 Mar 2025 15:57:32 -0700 Subject: [PATCH 14/26] flag in mst code to avoid caching internal stuff, makes oncethrough iteration nicer --- mst/mst.go | 31 ++++++++++++++++++++----------- mst/mst_test.go | 4 ++-- mst/mst_util.go | 6 +++--- repo/repo.go | 4 ++-- repo/stream.go | 10 +++------- 5 files changed, 30 insertions(+), 25 deletions(-) diff --git a/mst/mst.go b/mst/mst.go index c205c7b3..d566d3f8 100644 --- a/mst/mst.go +++ b/mst/mst.go @@ -137,21 +137,24 @@ type MerkleSearchTree struct { layer int pointer cid.Cid validPtr bool + + noCache bool } // NewEmptyMST reports a new empty MST using cst as its storage. func NewEmptyMST(cst cbor.IpldStore) *MerkleSearchTree { - return createMST(cst, cid.Undef, []nodeEntry{}, 0) + return createMST(cst, cid.Undef, []nodeEntry{}, 0, false) } // Typescript: MST.create(storage, entries, layer, fanout) -> MST -func createMST(cst cbor.IpldStore, ptr cid.Cid, entries []nodeEntry, layer int) *MerkleSearchTree { +func createMST(cst cbor.IpldStore, ptr cid.Cid, entries []nodeEntry, layer int, noCache bool) *MerkleSearchTree { mst := &MerkleSearchTree{ cst: cst, pointer: ptr, layer: layer, entries: entries, validPtr: ptr.Defined(), + noCache: noCache, } return mst @@ -162,7 +165,11 @@ func createMST(cst cbor.IpldStore, ptr cid.Cid, entries []nodeEntry, layer int) // This is poorly named in both implementations, because it is lazy // Typescript: MST.load(storage, cid, layer=null, fanout) -> MST func LoadMST(cst cbor.IpldStore, root cid.Cid) *MerkleSearchTree { - return createMST(cst, root, nil, -1) + return createMST(cst, root, nil, -1, false) +} + +func (mst *MerkleSearchTree) SetNoCache(v bool) { + mst.noCache = v } // === "Immutability" === @@ -173,7 +180,7 @@ func (mst *MerkleSearchTree) newTree(entries []nodeEntry) *MerkleSearchTree { if entries == nil { panic("nil entries passed to newTree") } - return createMST(mst.cst, cid.Undef, entries, mst.layer) + return createMST(mst.cst, cid.Undef, entries, mst.layer, false) } // === "Getters (lazy load)" === @@ -195,14 +202,16 @@ func (mst *MerkleSearchTree) getEntries(ctx context.Context) ([]nodeEntry, error } // NOTE(bnewbold): Typescript version computes layer in-place here, but // the entriesFromNodeData() helper does that for us in golang - entries, err := entriesFromNodeData(ctx, &nd, mst.cst) + entries, err := entriesFromNodeData(ctx, &nd, mst.cst, mst.noCache) if err != nil { return nil, err } if entries == nil { panic("got nil entries from node data decoding") } - mst.entries = entries + if !mst.noCache { + mst.entries = entries + } return entries, nil } @@ -210,7 +219,7 @@ func (mst *MerkleSearchTree) getEntries(ctx context.Context) ([]nodeEntry, error } // golang-specific helper that calls in to deserializeNodeData -func entriesFromNodeData(ctx context.Context, nd *NodeData, cst cbor.IpldStore) ([]nodeEntry, error) { +func entriesFromNodeData(ctx context.Context, nd *NodeData, cst cbor.IpldStore, noCache bool) ([]nodeEntry, error) { layer := -1 if len(nd.Entries) > 0 { // NOTE(bnewbold): can compute the layer on the first KeySuffix, because for the first entry that field is a complete key @@ -218,7 +227,7 @@ func entriesFromNodeData(ctx context.Context, nd *NodeData, cst cbor.IpldStore) layer = leadingZerosOnHashBytes(firstLeaf.KeySuffix) } - entries, err := deserializeNodeData(ctx, cst, nd, layer) + entries, err := deserializeNodeData(ctx, cst, nd, layer, noCache) if err != nil { return nil, err } @@ -468,7 +477,7 @@ func (mst *MerkleSearchTree) Add(ctx context.Context, key string, val cid.Cid, k } checkTreeInvariant(updated) - newRoot := createMST(mst.cst, cid.Undef, updated, keyZeros) + newRoot := createMST(mst.cst, cid.Undef, updated, keyZeros, mst.noCache) // NOTE(bnewbold): We do want to invalid the CID (because this node has // changed, and we are "lazy" about recomputing). Setting this flag @@ -898,7 +907,7 @@ func (mst *MerkleSearchTree) createChild(ctx context.Context) (*MerkleSearchTree return nil, err } - return createMST(mst.cst, cid.Undef, []nodeEntry{}, layer-1), nil + return createMST(mst.cst, cid.Undef, []nodeEntry{}, layer-1, mst.noCache), nil } func (mst *MerkleSearchTree) createParent(ctx context.Context) (*MerkleSearchTree, error) { @@ -907,7 +916,7 @@ func (mst *MerkleSearchTree) createParent(ctx context.Context) (*MerkleSearchTre return nil, err } - return createMST(mst.cst, cid.Undef, []nodeEntry{mkTreeEntry(mst)}, layer+1), nil + return createMST(mst.cst, cid.Undef, []nodeEntry{mkTreeEntry(mst)}, layer+1, mst.noCache), nil } // === "Finding insertion points" === diff --git a/mst/mst_test.go b/mst/mst_test.go index d38223a0..67da593b 100644 --- a/mst/mst_test.go +++ b/mst/mst_test.go @@ -36,7 +36,7 @@ func TestBasicMst(t *testing.T) { ctx := context.Background() cst := util.CborStore(blockstore.NewBlockstore(datastore.NewMapDatastore())) - mst := createMST(cst, cid.Undef, []nodeEntry{}, -1) + mst := createMST(cst, cid.Undef, []nodeEntry{}, -1, false) // NOTE: these were previously generated randomly, but the random seed behavior changed vals := map[string]cid.Cid{ @@ -364,7 +364,7 @@ func mapToCidMap(a map[string]string) map[string]cid.Cid { func cidMapToMst(t testing.TB, bs blockstore.Blockstore, m map[string]cid.Cid) *MerkleSearchTree { cst := util.CborStore(bs) - mt := createMST(cst, cid.Undef, []nodeEntry{}, -1) + mt := createMST(cst, cid.Undef, []nodeEntry{}, -1, false) for k, v := range m { nmst, err := mt.Add(context.TODO(), k, v, -1) diff --git a/mst/mst_util.go b/mst/mst_util.go index 7126e674..bdfc020a 100644 --- a/mst/mst_util.go +++ b/mst/mst_util.go @@ -66,13 +66,13 @@ func layerForEntries(entries []nodeEntry) int { } // Typescript: deserializeNodeData(storage, data, layer) -func deserializeNodeData(ctx context.Context, cst cbor.IpldStore, nd *NodeData, layer int) ([]nodeEntry, error) { +func deserializeNodeData(ctx context.Context, cst cbor.IpldStore, nd *NodeData, layer int, noCache bool) ([]nodeEntry, error) { entries := []nodeEntry{} if nd.Left != nil { // Note: like Typescript, this is actually a lazy load entries = append(entries, nodeEntry{ Kind: entryTree, - Tree: createMST(cst, *nd.Left, nil, layer-1), + Tree: createMST(cst, *nd.Left, nil, layer-1, noCache), }) } @@ -100,7 +100,7 @@ func deserializeNodeData(ctx context.Context, cst cbor.IpldStore, nd *NodeData, if e.Tree != nil { entries = append(entries, nodeEntry{ Kind: entryTree, - Tree: createMST(cst, *e.Tree, nil, layer-1), + Tree: createMST(cst, *e.Tree, nil, layer-1, noCache), Key: keyStr, }) } diff --git a/repo/repo.go b/repo/repo.go index f9bf5643..084995e3 100644 --- a/repo/repo.go +++ b/repo/repo.go @@ -91,6 +91,8 @@ func IngestRepo(ctx context.Context, bs cbor.IpldBlockstore, r io.Reader) (cid.C bufr := bufrPool.Get().(*bufio.Reader) bufr.Reset(r) + defer bufrPool.Put(bufr) + br, root, err := carutil.NewReader(bufr) if err != nil { return cid.Undef, fmt.Errorf("opening CAR block reader: %w", err) @@ -110,8 +112,6 @@ func IngestRepo(ctx context.Context, bs cbor.IpldBlockstore, r io.Reader) (cid.C } } - bufrPool.Put(bufr) - return root, nil } diff --git a/repo/stream.go b/repo/stream.go index 42d01d1f..056ce0ad 100644 --- a/repo/stream.go +++ b/repo/stream.go @@ -16,7 +16,7 @@ import ( "go.opentelemetry.io/otel" ) -const bufPoolBlockSize = 512 +const bufPoolBlockSize = 1024 var smallBlockPool = &sync.Pool{ New: func() any { @@ -29,9 +29,6 @@ type readStreamBlockstore struct { streamComplete bool r *carutil.Reader - - outOfOrder int - totalBlocks int } func newStreamingBlockstore(r *carutil.Reader) *readStreamBlockstore { @@ -57,12 +54,10 @@ func (bs *readStreamBlockstore) readUntilBlock(ctx context.Context, cc cid.Cid) smallBlockPool.Put(buf) } - bs.totalBlocks++ if blk.Cid() == cc { return blk, nil } - bs.outOfOrder++ bs.otherBlocks[blk.Cid()] = blk } @@ -142,6 +137,8 @@ func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, cb func( t := mst.LoadMST(cst, sc.Data) + t.SetNoCache(true) + if err := t.WalkLeavesFrom(ctx, prefix, func(k string, val cid.Cid) error { if err := bs.View(val, func(data []byte) error { return cb(k, val, data) @@ -155,6 +152,5 @@ func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, cb func( return "", fmt.Errorf("failed to walk mst: %w", err) } - fmt.Println("out of order blocks: ", bs.outOfOrder, bs.totalBlocks) return sc.Rev, nil } -- 2.51.2 From ca60f9503d22dadff46e4b17e951b103b0a034f9 Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Mon, 24 Mar 2025 16:26:32 -0700 Subject: [PATCH 15/26] functional nocache walk --- mst/mst.go | 46 ++++++++++++++++++++++++---------------------- mst/mst_test.go | 4 ++-- mst/mst_util.go | 6 +++--- repo/stream.go | 4 +--- 4 files changed, 30 insertions(+), 30 deletions(-) diff --git a/mst/mst.go b/mst/mst.go index d566d3f8..b0e02160 100644 --- a/mst/mst.go +++ b/mst/mst.go @@ -137,24 +137,21 @@ type MerkleSearchTree struct { layer int pointer cid.Cid validPtr bool - - noCache bool } // NewEmptyMST reports a new empty MST using cst as its storage. func NewEmptyMST(cst cbor.IpldStore) *MerkleSearchTree { - return createMST(cst, cid.Undef, []nodeEntry{}, 0, false) + return createMST(cst, cid.Undef, []nodeEntry{}, 0) } // Typescript: MST.create(storage, entries, layer, fanout) -> MST -func createMST(cst cbor.IpldStore, ptr cid.Cid, entries []nodeEntry, layer int, noCache bool) *MerkleSearchTree { +func createMST(cst cbor.IpldStore, ptr cid.Cid, entries []nodeEntry, layer int) *MerkleSearchTree { mst := &MerkleSearchTree{ cst: cst, pointer: ptr, layer: layer, entries: entries, validPtr: ptr.Defined(), - noCache: noCache, } return mst @@ -165,11 +162,7 @@ func createMST(cst cbor.IpldStore, ptr cid.Cid, entries []nodeEntry, layer int, // This is poorly named in both implementations, because it is lazy // Typescript: MST.load(storage, cid, layer=null, fanout) -> MST func LoadMST(cst cbor.IpldStore, root cid.Cid) *MerkleSearchTree { - return createMST(cst, root, nil, -1, false) -} - -func (mst *MerkleSearchTree) SetNoCache(v bool) { - mst.noCache = v + return createMST(cst, root, nil, -1) } // === "Immutability" === @@ -180,7 +173,7 @@ func (mst *MerkleSearchTree) newTree(entries []nodeEntry) *MerkleSearchTree { if entries == nil { panic("nil entries passed to newTree") } - return createMST(mst.cst, cid.Undef, entries, mst.layer, false) + return createMST(mst.cst, cid.Undef, entries, mst.layer) } // === "Getters (lazy load)" === @@ -202,16 +195,14 @@ func (mst *MerkleSearchTree) getEntries(ctx context.Context) ([]nodeEntry, error } // NOTE(bnewbold): Typescript version computes layer in-place here, but // the entriesFromNodeData() helper does that for us in golang - entries, err := entriesFromNodeData(ctx, &nd, mst.cst, mst.noCache) + entries, err := entriesFromNodeData(ctx, &nd, mst.cst) if err != nil { return nil, err } if entries == nil { panic("got nil entries from node data decoding") } - if !mst.noCache { - mst.entries = entries - } + mst.entries = entries return entries, nil } @@ -219,7 +210,7 @@ func (mst *MerkleSearchTree) getEntries(ctx context.Context) ([]nodeEntry, error } // golang-specific helper that calls in to deserializeNodeData -func entriesFromNodeData(ctx context.Context, nd *NodeData, cst cbor.IpldStore, noCache bool) ([]nodeEntry, error) { +func entriesFromNodeData(ctx context.Context, nd *NodeData, cst cbor.IpldStore) ([]nodeEntry, error) { layer := -1 if len(nd.Entries) > 0 { // NOTE(bnewbold): can compute the layer on the first KeySuffix, because for the first entry that field is a complete key @@ -227,7 +218,7 @@ func entriesFromNodeData(ctx context.Context, nd *NodeData, cst cbor.IpldStore, layer = leadingZerosOnHashBytes(firstLeaf.KeySuffix) } - entries, err := deserializeNodeData(ctx, cst, nd, layer, noCache) + entries, err := deserializeNodeData(ctx, cst, nd, layer) if err != nil { return nil, err } @@ -477,7 +468,7 @@ func (mst *MerkleSearchTree) Add(ctx context.Context, key string, val cid.Cid, k } checkTreeInvariant(updated) - newRoot := createMST(mst.cst, cid.Undef, updated, keyZeros, mst.noCache) + newRoot := createMST(mst.cst, cid.Undef, updated, keyZeros) // NOTE(bnewbold): We do want to invalid the CID (because this node has // changed, and we are "lazy" about recomputing). Setting this flag @@ -907,7 +898,7 @@ func (mst *MerkleSearchTree) createChild(ctx context.Context) (*MerkleSearchTree return nil, err } - return createMST(mst.cst, cid.Undef, []nodeEntry{}, layer-1, mst.noCache), nil + return createMST(mst.cst, cid.Undef, []nodeEntry{}, layer-1), nil } func (mst *MerkleSearchTree) createParent(ctx context.Context) (*MerkleSearchTree, error) { @@ -916,7 +907,7 @@ func (mst *MerkleSearchTree) createParent(ctx context.Context) (*MerkleSearchTre return nil, err } - return createMST(mst.cst, cid.Undef, []nodeEntry{mkTreeEntry(mst)}, layer+1, mst.noCache), nil + return createMST(mst.cst, cid.Undef, []nodeEntry{mkTreeEntry(mst)}, layer+1), nil } // === "Finding insertion points" === @@ -946,6 +937,14 @@ func (mst *MerkleSearchTree) findGtOrEqualLeafIndex(ctx context.Context, key str // key that's greater than or equal to the provided from key. // If cb returns an error, the walk is aborted and the error is returned. func (mst *MerkleSearchTree) WalkLeavesFrom(ctx context.Context, from string, cb func(key string, val cid.Cid) error) error { + return mst.walkLeavesFrom(ctx, from, false, cb) +} + +func (mst *MerkleSearchTree) WalkLeavesFromNocache(ctx context.Context, from string, cb func(key string, val cid.Cid) error) error { + return mst.walkLeavesFrom(ctx, from, true, cb) +} + +func (mst *MerkleSearchTree) walkLeavesFrom(ctx context.Context, from string, nocache bool, cb func(key string, val cid.Cid) error) error { index, err := mst.findGtOrEqualLeafIndex(ctx, from) if err != nil { return err @@ -959,7 +958,7 @@ func (mst *MerkleSearchTree) WalkLeavesFrom(ctx context.Context, from string, cb if index > 0 { prev := entries[index-1] if !prev.isUndefined() && prev.isTree() { - if err := prev.Tree.WalkLeavesFrom(ctx, from, cb); err != nil { + if err := prev.Tree.walkLeavesFrom(ctx, from, nocache, cb); err != nil { return fmt.Errorf("walk leaves %d: %w", index, err) } } @@ -971,9 +970,12 @@ func (mst *MerkleSearchTree) WalkLeavesFrom(ctx context.Context, from string, cb return err } } else { - if err := e.Tree.WalkLeavesFrom(ctx, from, cb); err != nil { + if err := e.Tree.walkLeavesFrom(ctx, from, nocache, cb); err != nil { return fmt.Errorf("walk leaves from (%d): %w", i, err) } + if nocache { + e.Tree = nil + } } } return nil diff --git a/mst/mst_test.go b/mst/mst_test.go index 67da593b..d38223a0 100644 --- a/mst/mst_test.go +++ b/mst/mst_test.go @@ -36,7 +36,7 @@ func TestBasicMst(t *testing.T) { ctx := context.Background() cst := util.CborStore(blockstore.NewBlockstore(datastore.NewMapDatastore())) - mst := createMST(cst, cid.Undef, []nodeEntry{}, -1, false) + mst := createMST(cst, cid.Undef, []nodeEntry{}, -1) // NOTE: these were previously generated randomly, but the random seed behavior changed vals := map[string]cid.Cid{ @@ -364,7 +364,7 @@ func mapToCidMap(a map[string]string) map[string]cid.Cid { func cidMapToMst(t testing.TB, bs blockstore.Blockstore, m map[string]cid.Cid) *MerkleSearchTree { cst := util.CborStore(bs) - mt := createMST(cst, cid.Undef, []nodeEntry{}, -1, false) + mt := createMST(cst, cid.Undef, []nodeEntry{}, -1) for k, v := range m { nmst, err := mt.Add(context.TODO(), k, v, -1) diff --git a/mst/mst_util.go b/mst/mst_util.go index bdfc020a..7126e674 100644 --- a/mst/mst_util.go +++ b/mst/mst_util.go @@ -66,13 +66,13 @@ func layerForEntries(entries []nodeEntry) int { } // Typescript: deserializeNodeData(storage, data, layer) -func deserializeNodeData(ctx context.Context, cst cbor.IpldStore, nd *NodeData, layer int, noCache bool) ([]nodeEntry, error) { +func deserializeNodeData(ctx context.Context, cst cbor.IpldStore, nd *NodeData, layer int) ([]nodeEntry, error) { entries := []nodeEntry{} if nd.Left != nil { // Note: like Typescript, this is actually a lazy load entries = append(entries, nodeEntry{ Kind: entryTree, - Tree: createMST(cst, *nd.Left, nil, layer-1, noCache), + Tree: createMST(cst, *nd.Left, nil, layer-1), }) } @@ -100,7 +100,7 @@ func deserializeNodeData(ctx context.Context, cst cbor.IpldStore, nd *NodeData, if e.Tree != nil { entries = append(entries, nodeEntry{ Kind: entryTree, - Tree: createMST(cst, *e.Tree, nil, layer-1, noCache), + Tree: createMST(cst, *e.Tree, nil, layer-1), Key: keyStr, }) } diff --git a/repo/stream.go b/repo/stream.go index 056ce0ad..37673567 100644 --- a/repo/stream.go +++ b/repo/stream.go @@ -137,9 +137,7 @@ func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, cb func( t := mst.LoadMST(cst, sc.Data) - t.SetNoCache(true) - - if err := t.WalkLeavesFrom(ctx, prefix, func(k string, val cid.Cid) error { + if err := t.WalkLeavesFromNocache(ctx, prefix, func(k string, val cid.Cid) error { if err := bs.View(val, func(data []byte) error { return cb(k, val, data) }); err != nil { -- 2.51.2 From ecf7051a2a6005cf59724dfb45922650e7136919 Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Fri, 2 May 2025 16:51:28 -0700 Subject: [PATCH 16/26] handle rev better --- backfill/backfill.go | 9 ++++++--- repo/stream.go | 14 ++++++++------ 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/backfill/backfill.go b/backfill/backfill.go index 595d2ebe..a55af6f7 100644 --- a/backfill/backfill.go +++ b/backfill/backfill.go @@ -425,10 +425,15 @@ func (b *Backfiller) BackfillRepo(ctx context.Context, job Job) (string, error) recordResults := make(chan recordResult, numRoutines) var rev string + // guaranteed to be called before any items are send on the recordQueue channel + setRev := func(s string) { + rev = s + } + // Producer routine go func() { defer close(recordQueue) - rrev, err := repo.StreamRepoRecords(ctx, r, b.NSIDFilter, func(recordPath string, nodeCid cid.Cid, data []byte) error { + err := repo.StreamRepoRecords(ctx, r, b.NSIDFilter, setRev, func(recordPath string, nodeCid cid.Cid, data []byte) error { numRecords++ recordQueue <- recordQueueItem{recordPath: recordPath, nodeCid: nodeCid, data: data} return nil @@ -436,8 +441,6 @@ func (b *Backfiller) BackfillRepo(ctx context.Context, job Job) (string, error) if err != nil { log.Error("failed to iterate records in repo", "err", err) } - - rev = rrev }() // Consumer routines diff --git a/repo/stream.go b/repo/stream.go index 37673567..e5f9f6ad 100644 --- a/repo/stream.go +++ b/repo/stream.go @@ -111,13 +111,13 @@ func (bs *readStreamBlockstore) Put(ctx context.Context, blk block.Block) error return fmt.Errorf("put is not needed") } -func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, cb func(k string, c cid.Cid, v []byte) error) (string, error) { +func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, setRev func(string), cb func(k string, c cid.Cid, v []byte) error) error { ctx, span := otel.Tracer("repo").Start(ctx, "RepoStream") defer span.End() br, root, err := carutil.NewReader(bufio.NewReader(r)) if err != nil { - return "", fmt.Errorf("opening CAR block reader: %w", err) + return fmt.Errorf("opening CAR block reader: %w", err) } bs := newStreamingBlockstore(br) @@ -126,15 +126,17 @@ func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, cb func( var sc SignedCommit if err := cst.Get(ctx, root, &sc); err != nil { - return "", fmt.Errorf("loading root from blockstore: %w", err) + return fmt.Errorf("loading root from blockstore: %w", err) } if sc.Version != ATP_REPO_VERSION && sc.Version != ATP_REPO_VERSION_2 { - return "", fmt.Errorf("unsupported repo version: %d", sc.Version) + return fmt.Errorf("unsupported repo version: %d", sc.Version) } // TODO: verify that signature + setRev(sc.Rev) + t := mst.LoadMST(cst, sc.Data) if err := t.WalkLeavesFromNocache(ctx, prefix, func(k string, val cid.Cid) error { @@ -147,8 +149,8 @@ func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, cb func( return nil }); err != nil { - return "", fmt.Errorf("failed to walk mst: %w", err) + return fmt.Errorf("failed to walk mst: %w", err) } - return sc.Rev, nil + return nil } -- 2.51.2 From 88d2b471cd17a390a305e18184d1a7bfbe9ade6d Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Sun, 4 May 2025 13:39:05 -0700 Subject: [PATCH 17/26] dont buffer events for enqueued backfill jobs, just for in progress ones --- backfill/gormstore.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/backfill/gormstore.go b/backfill/gormstore.go index a3816fdc..ba11821d 100644 --- a/backfill/gormstore.go +++ b/backfill/gormstore.go @@ -178,9 +178,9 @@ func (j *Gormjob) BufferOps(ctx context.Context, since *string, rev string, ops defer j.lk.Unlock() switch j.state { - case StateComplete: + case StateComplete, StateEnqueued: return false, nil - case StateInProgress, StateEnqueued: + case StateInProgress: // keep going and buffer the op default: if strings.HasPrefix(j.state, "failed") { -- 2.51.2 From 6e9ab2d59bea8f7c912bb4ffb757c61e53f51734 Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Sun, 4 May 2025 13:42:39 -0700 Subject: [PATCH 18/26] fix gosky build --- cmd/gosky/main.go | 2 +- repo/stream.go | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/cmd/gosky/main.go b/cmd/gosky/main.go index e1a78924..76e41121 100644 --- a/cmd/gosky/main.go +++ b/cmd/gosky/main.go @@ -671,7 +671,7 @@ var listAllRecordsCmd = &cli.Command{ vals := cctx.Bool("values") cids := cctx.Bool("cids") - if _, err := repo.StreamRepoRecords(ctx, bytes.NewReader(repob), collection, func(k string, cc cid.Cid, v []byte) error { + if err := repo.StreamRepoRecords(ctx, bytes.NewReader(repob), collection, nil, func(k string, cc cid.Cid, v []byte) error { if !strings.HasPrefix(k, collection) { return repo.ErrDoneIterating } diff --git a/repo/stream.go b/repo/stream.go index e5f9f6ad..bd646254 100644 --- a/repo/stream.go +++ b/repo/stream.go @@ -115,6 +115,10 @@ func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, setRev f ctx, span := otel.Tracer("repo").Start(ctx, "RepoStream") defer span.End() + if setRev == nil { + setRev = func(string) {} + } + br, root, err := carutil.NewReader(bufio.NewReader(r)) if err != nil { return fmt.Errorf("opening CAR block reader: %w", err) -- 2.51.2 From d125f0fbe36c9a016e6c65a5d0aef57bec924703 Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Sun, 4 May 2025 20:17:22 -0700 Subject: [PATCH 19/26] StreamRepoRecords needs to respect ErrDoneIterating --- repo/stream.go | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/repo/stream.go b/repo/stream.go index bd646254..6ecd5f6b 100644 --- a/repo/stream.go +++ b/repo/stream.go @@ -147,7 +147,9 @@ func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, setRev f if err := bs.View(val, func(data []byte) error { return cb(k, val, data) }); err != nil { - slog.Error("failed to get record from tree", "key", k, "cid", val, "error", err) + if err != ErrDoneIterating { + slog.Error("failed to get record from tree", "key", k, "cid", val, "error", err) + } return nil } -- 2.51.2 From d0835eae85c1350a04388b404b2029526c8d704b Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Sun, 4 May 2025 20:21:00 -0700 Subject: [PATCH 20/26] better errors --- repo/stream.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/repo/stream.go b/repo/stream.go index 6ecd5f6b..153aa9f7 100644 --- a/repo/stream.go +++ b/repo/stream.go @@ -130,7 +130,7 @@ func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, setRev f var sc SignedCommit if err := cst.Get(ctx, root, &sc); err != nil { - return fmt.Errorf("loading root from blockstore: %w", err) + return fmt.Errorf("loading root (%s) from blockstore: %w", root, err) } if sc.Version != ATP_REPO_VERSION && sc.Version != ATP_REPO_VERSION_2 { -- 2.51.2 From 8967ce944dfc4b51a671c7652424ba5f37000ff1 Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Sun, 4 May 2025 20:22:26 -0700 Subject: [PATCH 21/26] more descriptive --- repo/stream.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/repo/stream.go b/repo/stream.go index 153aa9f7..4e20a443 100644 --- a/repo/stream.go +++ b/repo/stream.go @@ -130,7 +130,7 @@ func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, setRev f var sc SignedCommit if err := cst.Get(ctx, root, &sc); err != nil { - return fmt.Errorf("loading root (%s) from blockstore: %w", root, err) + return fmt.Errorf("loading root (%s) from blockstore (other blocks: %d): %w", root, len(bs.otherBlocks), err) } if sc.Version != ATP_REPO_VERSION && sc.Version != ATP_REPO_VERSION_2 { -- 2.51.2 From 6e6ae16e77d59f2d366a69436228aadbafbd4bb0 Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Sun, 4 May 2025 23:48:17 -0700 Subject: [PATCH 22/26] dont count record stream failures as success --- backfill/backfill.go | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/backfill/backfill.go b/backfill/backfill.go index a55af6f7..08f4860a 100644 --- a/backfill/backfill.go +++ b/backfill/backfill.go @@ -417,8 +417,6 @@ func (b *Backfiller) BackfillRepo(ctx context.Context, job Job) (string, error) } } - defer r.Close() - numRecords := 0 numRoutines := b.ParallelRecordCreates recordQueue := make(chan recordQueueItem, numRoutines) @@ -431,7 +429,9 @@ func (b *Backfiller) BackfillRepo(ctx context.Context, job Job) (string, error) } // Producer routine + var streamRecordsError error go func() { + defer r.Close() defer close(recordQueue) err := repo.StreamRepoRecords(ctx, r, b.NSIDFilter, setRev, func(recordPath string, nodeCid cid.Cid, data []byte) error { numRecords++ @@ -439,7 +439,7 @@ func (b *Backfiller) BackfillRepo(ctx context.Context, job Job) (string, error) return nil }) if err != nil { - log.Error("failed to iterate records in repo", "err", err) + streamRecordsError = fmt.Errorf("failed to iterate records in repo: %w", err) } }() @@ -481,6 +481,10 @@ func (b *Backfiller) BackfillRepo(ctx context.Context, job Job) (string, error) close(recordResults) resultWG.Wait() + if streamRecordsError != nil { + return "failed to stream records", streamRecordsError + } + if err := job.SetRev(ctx, rev); err != nil { log.Error("failed to update rev after backfilling repo", "err", err) } -- 2.51.2 From 97ee6181591d7c11b2d512443e4c72e3aa555a96 Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Mon, 5 May 2025 09:50:14 -0700 Subject: [PATCH 23/26] fully ignore enqueued backfill repo events --- backfill/gormstore.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/backfill/gormstore.go b/backfill/gormstore.go index ba11821d..e97bbfc9 100644 --- a/backfill/gormstore.go +++ b/backfill/gormstore.go @@ -178,8 +178,11 @@ func (j *Gormjob) BufferOps(ctx context.Context, since *string, rev string, ops defer j.lk.Unlock() switch j.state { - case StateComplete, StateEnqueued: + case StateComplete: return false, nil + case StateEnqueued: + // if the repo is enqueue, but not actively being backfilled, just ignore events for it for now + return true, nil case StateInProgress: // keep going and buffer the op default: -- 2.51.2 From 72bd374dc5619ae22a401dcdf87da3c09ead75d4 Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Wed, 8 Oct 2025 12:50:36 -0700 Subject: [PATCH 24/26] some comments --- mst/mst.go | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mst/mst.go b/mst/mst.go index b0e02160..847c49a5 100644 --- a/mst/mst.go +++ b/mst/mst.go @@ -936,10 +936,15 @@ func (mst *MerkleSearchTree) findGtOrEqualLeafIndex(ctx context.Context, key str // WalkLeavesFrom walks the leaves of the tree, calling the cb callback on each // key that's greater than or equal to the provided from key. // If cb returns an error, the walk is aborted and the error is returned. +// NB: this method caches the tree structure in memory to make subsequent tree +// operations significantly faster func (mst *MerkleSearchTree) WalkLeavesFrom(ctx context.Context, from string, cb func(key string, val cid.Cid) error) error { return mst.walkLeavesFrom(ctx, from, false, cb) } +// WalkLeavesFromNocache works the same as WalkLeavesFrom but does not cache +// internal tree structure, intended for "once through" passes of MSTs, +// especially in streaming contexts func (mst *MerkleSearchTree) WalkLeavesFromNocache(ctx context.Context, from string, cb func(key string, val cid.Cid) error) error { return mst.walkLeavesFrom(ctx, from, true, cb) } -- 2.51.2 From 261117553bb45fac266ce503c109cfe1a3a31b51 Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Mon, 26 Jan 2026 16:46:05 -0800 Subject: [PATCH 25/26] comment new method --- repo/stream.go | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/repo/stream.go b/repo/stream.go index 4e20a443..1053d02a 100644 --- a/repo/stream.go +++ b/repo/stream.go @@ -111,6 +111,23 @@ func (bs *readStreamBlockstore) Put(ctx context.Context, blk block.Block) error return fmt.Errorf("put is not needed") } +// StreamRepoRecords parses a CAR-encoded repository from the given reader and +// iterates over its records. +// +// The prefix parameter specifies the starting point for iteration. Records are +// visited in lexicographic order starting from the first key >= prefix. If +// prefix exactly matches a full record path, that record will be included. +// Note that this does not filter to only keys with that prefix; all records +// from that point onward are visited. +// +// The setRev callback receives the repository revision string after the commit +// is loaded. It may be nil, in which case it is ignored. +// +// The cb parameter is the visitor function called for each record with the +// record's key, CID, and raw data. If cb returns an error, the error is logged +// and iteration continues; errors from cb do not stop the walk and are not +// returned from this function. To signal intentional early termination, cb can +// return ErrDoneIterating, which is handled silently without logging. func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, setRev func(string), cb func(k string, c cid.Cid, v []byte) error) error { ctx, span := otel.Tracer("repo").Start(ctx, "RepoStream") defer span.End() -- 2.51.2 From 0223d29b178df34ce53b30e72678b9df02e1c5dc Mon Sep 17 00:00:00 2001 From: whyrusleeping Date: Sat, 31 Jan 2026 18:05:09 -0800 Subject: [PATCH 26/26] swap setRev for setCommit --- backfill/backfill.go | 9 +++++---- repo/stream.go | 19 +++++++++++-------- 2 files changed, 16 insertions(+), 12 deletions(-) diff --git a/backfill/backfill.go b/backfill/backfill.go index 08f4860a..6762acb4 100644 --- a/backfill/backfill.go +++ b/backfill/backfill.go @@ -423,9 +423,10 @@ func (b *Backfiller) BackfillRepo(ctx context.Context, job Job) (string, error) recordResults := make(chan recordResult, numRoutines) var rev string - // guaranteed to be called before any items are send on the recordQueue channel - setRev := func(s string) { - rev = s + // guaranteed to be called before any items are sent on the recordQueue channel + onCommit := func(sc *repo.SignedCommit) error { + rev = sc.Rev + return nil } // Producer routine @@ -433,7 +434,7 @@ func (b *Backfiller) BackfillRepo(ctx context.Context, job Job) (string, error) go func() { defer r.Close() defer close(recordQueue) - err := repo.StreamRepoRecords(ctx, r, b.NSIDFilter, setRev, func(recordPath string, nodeCid cid.Cid, data []byte) error { + err := repo.StreamRepoRecords(ctx, r, b.NSIDFilter, onCommit, func(recordPath string, nodeCid cid.Cid, data []byte) error { numRecords++ recordQueue <- recordQueueItem{recordPath: recordPath, nodeCid: nodeCid, data: data} return nil diff --git a/repo/stream.go b/repo/stream.go index 1053d02a..cd3e1f86 100644 --- a/repo/stream.go +++ b/repo/stream.go @@ -120,20 +120,23 @@ func (bs *readStreamBlockstore) Put(ctx context.Context, blk block.Block) error // Note that this does not filter to only keys with that prefix; all records // from that point onward are visited. // -// The setRev callback receives the repository revision string after the commit -// is loaded. It may be nil, in which case it is ignored. +// The onCommit callback receives the signed commit after it is loaded. It may +// be nil, in which case it is ignored. The callback can be used to extract +// the revision string (sc.Rev), or in the future to verify the commit +// signature. If the callback returns an error, StreamRepoRecords returns +// immediately with that error. // // The cb parameter is the visitor function called for each record with the // record's key, CID, and raw data. If cb returns an error, the error is logged // and iteration continues; errors from cb do not stop the walk and are not // returned from this function. To signal intentional early termination, cb can // return ErrDoneIterating, which is handled silently without logging. -func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, setRev func(string), cb func(k string, c cid.Cid, v []byte) error) error { +func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, onCommit func(*SignedCommit) error, cb func(k string, c cid.Cid, v []byte) error) error { ctx, span := otel.Tracer("repo").Start(ctx, "RepoStream") defer span.End() - if setRev == nil { - setRev = func(string) {} + if onCommit == nil { + onCommit = func(*SignedCommit) error { return nil } } br, root, err := carutil.NewReader(bufio.NewReader(r)) @@ -154,9 +157,9 @@ func StreamRepoRecords(ctx context.Context, r io.Reader, prefix string, setRev f return fmt.Errorf("unsupported repo version: %d", sc.Version) } - // TODO: verify that signature - - setRev(sc.Rev) + if err := onCommit(&sc); err != nil { + return fmt.Errorf("commit callback: %w", err) + } t := mst.LoadMST(cst, sc.Data)