at-container-registry/pkg/hold/pds/repomgr.go

package pds

// repomgr.go — RepoManager manages ATProto repository operations for the hold PDS.
//
// Originally vendored from github.com/bluesky-social/indigo/repomgr with fixes:
//   - Fixed UpdateRecord bug (was calling PutRecord internally)
//   - Added PutRecord/UpsertRecord for explicit rkey operations
//   - Added prevData support for Sync 1.1
//
// Implements the RepoOperator interface (see repo_operator.go).
// See docs/REPOMGR_MIGRATION.md for planned migration to indigo/repo directly.

import (
	"context"
	"fmt"
	"io"
	"log/slog"
	"sync"

	holddb "atcr.io/pkg/hold/db"
	atproto "github.com/bluesky-social/indigo/api/atproto"
	bsky "github.com/bluesky-social/indigo/api/bsky"
	"github.com/bluesky-social/indigo/atproto/syntax"
	"github.com/bluesky-social/indigo/models"
	"github.com/bluesky-social/indigo/repo"
	"github.com/bluesky-social/indigo/util"

	blocks "github.com/ipfs/go-block-format"
	"github.com/ipfs/go-cid"
	cbg "github.com/whyrusleeping/cbor-gen"
	"go.opentelemetry.io/otel"
)

func NewRepoManager(cs holddb.CarStore, kmgr KeyManager) *RepoManager {
	clk := syntax.NewTIDClock(0)

	return &RepoManager{
		cs:        cs,
		userLocks: make(map[models.Uid]*userLock),
		kmgr:      kmgr,
		log:       slog.Default().With("system", "repomgr"),
		clk:       clk,
	}
}

func (rm *RepoManager) SetEventHandler(cb func(context.Context, *RepoEvent), hydrateRecords bool) {
	rm.events = cb
	rm.hydrateRecords = hydrateRecords
}

type RepoManager struct {
	cs   holddb.CarStore
	kmgr KeyManager

	lklk      sync.Mutex
	userLocks map[models.Uid]*userLock

	events         func(context.Context, *RepoEvent)
	hydrateRecords bool

	log *slog.Logger
	clk *syntax.TIDClock
}

type userLock struct {
	lk    sync.Mutex
	count int
}

func (rm *RepoManager) lockUser(ctx context.Context, user models.Uid) func() {
	_, span := otel.Tracer("repoman").Start(ctx, "userLock")
	defer span.End()

	rm.lklk.Lock()

	ulk, ok := rm.userLocks[user]
	if !ok {
		ulk = &userLock{}
		rm.userLocks[user] = ulk
	}

	ulk.count++

	rm.lklk.Unlock()

	ulk.lk.Lock()

	return func() {
		rm.lklk.Lock()

		ulk.lk.Unlock()
		ulk.count--

		if ulk.count == 0 {
			delete(rm.userLocks, user)
		}
		rm.lklk.Unlock()
	}
}

func (rm *RepoManager) CreateRecord(ctx context.Context, user models.Uid, collection string, rec cbg.CBORMarshaler) (string, cid.Cid, error) {
	ctx, span := otel.Tracer("repoman").Start(ctx, "CreateRecord")
	defer span.End()

	unlock := rm.lockUser(ctx, user)
	defer unlock()

	rev, err := rm.cs.GetUserRepoRev(ctx, user)
	if err != nil {
		return "", cid.Undef, err
	}

	ds, err := rm.cs.NewDeltaSession(ctx, user, &rev)
	if err != nil {
		return "", cid.Undef, err
	}

	head := ds.BaseCid()

	r, err := repo.OpenRepo(ctx, ds, head)
	if err != nil {
		return "", cid.Undef, err
	}

	// Capture previous MST root before commit overwrites it
	var prevData *cid.Cid
	if head.Defined() {
		pd := r.DataCid()
		prevData = &pd
	}

	cc, tid, err := r.CreateRecord(ctx, collection, rec)
	if err != nil {
		return "", cid.Undef, err
	}

	nroot, nrev, err := r.Commit(ctx, rm.kmgr.SignForUser)
	if err != nil {
		return "", cid.Undef, err
	}

	rslice, err := ds.CloseWithRoot(ctx, nroot, nrev)
	if err != nil {
		return "", cid.Undef, fmt.Errorf("close with root: %w", err)
	}

	var oldroot *cid.Cid
	if head.Defined() {
		oldroot = &head
	}

	if rm.events != nil {
		rm.events(ctx, &RepoEvent{
			User:     user,
			OldRoot:  oldroot,
			NewRoot:  nroot,
			PrevData: prevData,
			Rev:      nrev,
			Since:    &rev,
			Ops: []RepoOp{{
				Kind:       EvtKindCreateRecord,
				Collection: collection,
				Rkey:       tid,
				Record:     rec,
				RecCid:     &cc,
			}},
			RepoSlice: rslice,
		})
	}

	return collection + "/" + tid, cc, nil
}

func (rm *RepoManager) UpdateRecord(ctx context.Context, user models.Uid, collection, rkey string, rec cbg.CBORMarshaler) (cid.Cid, error) {
	ctx, span := otel.Tracer("repoman").Start(ctx, "UpdateRecord")
	defer span.End()

	unlock := rm.lockUser(ctx, user)
	defer unlock()

	rev, err := rm.cs.GetUserRepoRev(ctx, user)
	if err != nil {
		return cid.Undef, err
	}

	ds, err := rm.cs.NewDeltaSession(ctx, user, &rev)
	if err != nil {
		return cid.Undef, err
	}

	head := ds.BaseCid()
	r, err := repo.OpenRepo(ctx, ds, head)
	if err != nil {
		return cid.Undef, err
	}

	// Capture previous MST root before commit overwrites it
	var prevData *cid.Cid
	if head.Defined() {
		pd := r.DataCid()
		prevData = &pd
	}

	rpath := collection + "/" + rkey
	cc, err := r.UpdateRecord(ctx, rpath, rec)
	if err != nil {
		return cid.Undef, err
	}

	nroot, nrev, err := r.Commit(ctx, rm.kmgr.SignForUser)
	if err != nil {
		return cid.Undef, err
	}

	rslice, err := ds.CloseWithRoot(ctx, nroot, nrev)
	if err != nil {
		return cid.Undef, fmt.Errorf("close with root: %w", err)
	}

	var oldroot *cid.Cid
	if head.Defined() {
		oldroot = &head
	}

	if rm.events != nil {
		op := RepoOp{
			Kind:       EvtKindUpdateRecord,
			Collection: collection,
			Rkey:       rkey,
			RecCid:     &cc,
		}

		if rm.hydrateRecords {
			op.Record = rec
		}

		rm.events(ctx, &RepoEvent{
			User:      user,
			OldRoot:   oldroot,
			NewRoot:   nroot,
			PrevData:  prevData,
			Rev:       nrev,
			Since:     &rev,
			Ops:       []RepoOp{op},
			RepoSlice: rslice,
		})
	}

	return cc, nil
}

// PutRecord creates a record with an explicit rkey (like CreateRecord but with specified rkey).
// This uses r.PutRecord which will FAIL if the rkey already exists (uses mst.Add, not mst.Update).
// Use UpdateRecord to modify existing records.
// Returns the collection path (e.g., "io.atcr.captain/self") and CID.
func (rm *RepoManager) PutRecord(ctx context.Context, user models.Uid, collection, rkey string, rec cbg.CBORMarshaler) (string, cid.Cid, error) {
	ctx, span := otel.Tracer("repoman").Start(ctx, "PutRecord")
	defer span.End()

	unlock := rm.lockUser(ctx, user)
	defer unlock()

	rev, err := rm.cs.GetUserRepoRev(ctx, user)
	if err != nil {
		return "", cid.Undef, err
	}

	ds, err := rm.cs.NewDeltaSession(ctx, user, &rev)
	if err != nil {
		return "", cid.Undef, err
	}

	head := ds.BaseCid()
	r, err := repo.OpenRepo(ctx, ds, head)
	if err != nil {
		return "", cid.Undef, err
	}

	// Capture previous MST root before commit overwrites it
	var prevData *cid.Cid
	if head.Defined() {
		pd := r.DataCid()
		prevData = &pd
	}

	rpath := collection + "/" + rkey
	cc, err := r.PutRecord(ctx, rpath, rec)
	if err != nil {
		return "", cid.Undef, err
	}

	nroot, nrev, err := r.Commit(ctx, rm.kmgr.SignForUser)
	if err != nil {
		return "", cid.Undef, err
	}

	rslice, err := ds.CloseWithRoot(ctx, nroot, nrev)
	if err != nil {
		return "", cid.Undef, fmt.Errorf("close with root: %w", err)
	}

	var oldroot *cid.Cid
	if head.Defined() {
		oldroot = &head
	}

	if rm.events != nil {
		op := RepoOp{
			Kind:       EvtKindCreateRecord,
			Collection: collection,
			Rkey:       rkey,
			RecCid:     &cc,
		}

		if rm.hydrateRecords {
			op.Record = rec
		}

		rm.events(ctx, &RepoEvent{
			User:      user,
			OldRoot:   oldroot,
			NewRoot:   nroot,
			PrevData:  prevData,
			Rev:       nrev,
			Since:     &rev,
			Ops:       []RepoOp{op},
			RepoSlice: rslice,
		})
	}

	return rpath, cc, nil
}

// UpsertRecord creates or updates a record with an explicit rkey.
// If the record doesn't exist, it creates it. If it exists, it updates it.
// Returns the collection path (e.g., "io.atcr.captain/self"), CID, and whether it was created (true) or updated (false).
func (rm *RepoManager) UpsertRecord(ctx context.Context, user models.Uid, collection, rkey string, rec cbg.CBORMarshaler) (string, cid.Cid, bool, error) {
	ctx, span := otel.Tracer("repoman").Start(ctx, "UpsertRecord")
	defer span.End()

	unlock := rm.lockUser(ctx, user)
	defer unlock()

	rev, err := rm.cs.GetUserRepoRev(ctx, user)
	if err != nil {
		return "", cid.Undef, false, err
	}

	ds, err := rm.cs.NewDeltaSession(ctx, user, &rev)
	if err != nil {
		return "", cid.Undef, false, err
	}

	head := ds.BaseCid()
	r, err := repo.OpenRepo(ctx, ds, head)
	if err != nil {
		return "", cid.Undef, false, err
	}

	// Capture previous MST root before commit overwrites it
	var prevData *cid.Cid
	if head.Defined() {
		pd := r.DataCid()
		prevData = &pd
	}

	rpath := collection + "/" + rkey

	// Check if record exists
	_, _, err = r.GetRecordBytes(ctx, rpath)
	recordExists := err == nil

	var cc cid.Cid
	var evtKind EventKind
	if recordExists {
		// Update existing record
		cc, err = r.UpdateRecord(ctx, rpath, rec)
		evtKind = EvtKindUpdateRecord
	} else {
		// Create new record
		cc, err = r.PutRecord(ctx, rpath, rec)
		evtKind = EvtKindCreateRecord
	}
	if err != nil {
		return "", cid.Undef, false, err
	}

	nroot, nrev, err := r.Commit(ctx, rm.kmgr.SignForUser)
	if err != nil {
		return "", cid.Undef, false, err
	}

	rslice, err := ds.CloseWithRoot(ctx, nroot, nrev)
	if err != nil {
		return "", cid.Undef, false, fmt.Errorf("close with root: %w", err)
	}

	var oldroot *cid.Cid
	if head.Defined() {
		oldroot = &head
	}

	if rm.events != nil {
		op := RepoOp{
			Kind:       evtKind,
			Collection: collection,
			Rkey:       rkey,
			RecCid:     &cc,
		}

		if rm.hydrateRecords {
			op.Record = rec
		}

		rm.events(ctx, &RepoEvent{
			User:      user,
			OldRoot:   oldroot,
			NewRoot:   nroot,
			PrevData:  prevData,
			Rev:       nrev,
			Since:     &rev,
			Ops:       []RepoOp{op},
			RepoSlice: rslice,
		})
	}

	return rpath, cc, !recordExists, nil
}

func (rm *RepoManager) DeleteRecord(ctx context.Context, user models.Uid, collection, rkey string) error {
	ctx, span := otel.Tracer("repoman").Start(ctx, "DeleteRecord")
	defer span.End()

	unlock := rm.lockUser(ctx, user)
	defer unlock()

	rev, err := rm.cs.GetUserRepoRev(ctx, user)
	if err != nil {
		return err
	}

	ds, err := rm.cs.NewDeltaSession(ctx, user, &rev)
	if err != nil {
		return err
	}

	head := ds.BaseCid()
	r, err := repo.OpenRepo(ctx, ds, head)
	if err != nil {
		return err
	}

	// Capture previous MST root before commit overwrites it
	var prevData *cid.Cid
	if head.Defined() {
		pd := r.DataCid()
		prevData = &pd
	}

	rpath := collection + "/" + rkey
	if err := r.DeleteRecord(ctx, rpath); err != nil {
		return err
	}

	nroot, nrev, err := r.Commit(ctx, rm.kmgr.SignForUser)
	if err != nil {
		return err
	}

	rslice, err := ds.CloseWithRoot(ctx, nroot, nrev)
	if err != nil {
		return fmt.Errorf("close with root: %w", err)
	}

	var oldroot *cid.Cid
	if head.Defined() {
		oldroot = &head
	}

	if rm.events != nil {
		rm.events(ctx, &RepoEvent{
			User:     user,
			OldRoot:  oldroot,
			NewRoot:  nroot,
			PrevData: prevData,
			Rev:      nrev,
			Since:    &rev,
			Ops: []RepoOp{{
				Kind:       EvtKindDeleteRecord,
				Collection: collection,
				Rkey:       rkey,
			}},
			RepoSlice: rslice,
		})
	}

	return nil

}

func (rm *RepoManager) InitNewActor(ctx context.Context, user models.Uid, handle, did, displayname string, declcid, actortype string) error {
	unlock := rm.lockUser(ctx, user)
	defer unlock()

	if did == "" {
		return fmt.Errorf("must specify DID for new actor")
	}

	if user == 0 {
		return fmt.Errorf("must specify user for new actor")
	}

	ds, err := rm.cs.NewDeltaSession(ctx, user, nil)
	if err != nil {
		return fmt.Errorf("creating new delta session: %w", err)
	}

	r := repo.NewRepo(ctx, did, ds)

	profile := &bsky.ActorProfile{
		DisplayName: &displayname,
	}

	_, err = r.PutRecord(ctx, "app.bsky.actor.profile/self", profile)
	if err != nil {
		return fmt.Errorf("setting initial actor profile: %w", err)
	}

	root, nrev, err := r.Commit(ctx, rm.kmgr.SignForUser)
	if err != nil {
		return fmt.Errorf("committing repo for actor init: %w", err)
	}

	rslice, err := ds.CloseWithRoot(ctx, root, nrev)
	if err != nil {
		return fmt.Errorf("close with root: %w", err)
	}

	if rm.events != nil {
		op := RepoOp{
			Kind:       EvtKindCreateRecord,
			Collection: "app.bsky.actor.profile",
			Rkey:       "self",
		}

		if rm.hydrateRecords {
			op.Record = profile
		}

		rm.events(ctx, &RepoEvent{
			User:      user,
			NewRoot:   root,
			Rev:       nrev,
			Ops:       []RepoOp{op},
			RepoSlice: rslice,
		})
	}

	return nil
}

func (rm *RepoManager) GetRepoRoot(ctx context.Context, user models.Uid) (cid.Cid, error) {
	unlock := rm.lockUser(ctx, user)
	defer unlock()

	return rm.cs.GetUserRepoHead(ctx, user)
}

func (rm *RepoManager) GetRepoRev(ctx context.Context, user models.Uid) (string, error) {
	unlock := rm.lockUser(ctx, user)
	defer unlock()

	return rm.cs.GetUserRepoRev(ctx, user)
}

func (rm *RepoManager) ReadRepo(ctx context.Context, user models.Uid, since string, w io.Writer) error {
	return rm.cs.ReadUserCar(ctx, user, since, true, w)
}

func (rm *RepoManager) GetRecord(ctx context.Context, user models.Uid, collection string, rkey string, maybeCid cid.Cid) (cid.Cid, cbg.CBORMarshaler, error) {
	bs, err := rm.cs.ReadOnlySession(user)
	if err != nil {
		return cid.Undef, nil, err
	}

	head, err := rm.cs.GetUserRepoHead(ctx, user)
	if err != nil {
		return cid.Undef, nil, err
	}

	r, err := repo.OpenRepo(ctx, bs, head)
	if err != nil {
		return cid.Undef, nil, err
	}

	ocid, val, err := r.GetRecord(ctx, collection+"/"+rkey)
	if err != nil {
		return cid.Undef, nil, err
	}

	if maybeCid.Defined() && ocid != maybeCid {
		return cid.Undef, nil, fmt.Errorf("record at specified key had different CID than expected")
	}

	return ocid, val, nil
}

func (rm *RepoManager) GetRecordProof(ctx context.Context, user models.Uid, collection string, rkey string) (cid.Cid, []blocks.Block, error) {
	robs, err := rm.cs.ReadOnlySession(user)
	if err != nil {
		return cid.Undef, nil, err
	}

	bs := util.NewLoggingBstore(robs)

	head, err := rm.cs.GetUserRepoHead(ctx, user)
	if err != nil {
		return cid.Undef, nil, err
	}

	r, err := repo.OpenRepo(ctx, bs, head)
	if err != nil {
		return cid.Undef, nil, err
	}

	_, _, err = r.GetRecordBytes(ctx, collection+"/"+rkey)
	if err != nil {
		return cid.Undef, nil, err
	}

	return head, bs.GetLoggedBlocks(), nil
}

func (rm *RepoManager) BatchWrite(ctx context.Context, user models.Uid, writes []*atproto.RepoApplyWrites_Input_Writes_Elem) error {
	ctx, span := otel.Tracer("repoman").Start(ctx, "BatchWrite")
	defer span.End()

	unlock := rm.lockUser(ctx, user)
	defer unlock()

	rev, err := rm.cs.GetUserRepoRev(ctx, user)
	if err != nil {
		return err
	}

	ds, err := rm.cs.NewDeltaSession(ctx, user, &rev)
	if err != nil {
		return err
	}

	head := ds.BaseCid()
	r, err := repo.OpenRepo(ctx, ds, head)
	if err != nil {
		return err
	}

	// Capture previous MST root before commit overwrites it
	var prevData *cid.Cid
	if head.Defined() {
		pd := r.DataCid()
		prevData = &pd
	}

	ops := make([]RepoOp, 0, len(writes))
	for _, w := range writes {
		switch {
		case w.RepoApplyWrites_Create != nil:
			c := w.RepoApplyWrites_Create
			var rkey string
			if c.Rkey != nil {
				rkey = *c.Rkey
			} else {
				rkey = rm.clk.Next().String()
			}

			nsid := c.Collection + "/" + rkey
			cc, err := r.PutRecord(ctx, nsid, c.Value.Val)
			if err != nil {
				return err
			}

			op := RepoOp{
				Kind:       EvtKindCreateRecord,
				Collection: c.Collection,
				Rkey:       rkey,
				RecCid:     &cc,
			}

			if rm.hydrateRecords {
				op.Record = c.Value.Val
			}

			ops = append(ops, op)
		case w.RepoApplyWrites_Update != nil:
			u := w.RepoApplyWrites_Update

			cc, err := r.PutRecord(ctx, u.Collection+"/"+u.Rkey, u.Value.Val)
			if err != nil {
				return err
			}

			op := RepoOp{
				Kind:       EvtKindUpdateRecord,
				Collection: u.Collection,
				Rkey:       u.Rkey,
				RecCid:     &cc,
			}

			if rm.hydrateRecords {
				op.Record = u.Value.Val
			}

			ops = append(ops, op)
		case w.RepoApplyWrites_Delete != nil:
			d := w.RepoApplyWrites_Delete

			if err := r.DeleteRecord(ctx, d.Collection+"/"+d.Rkey); err != nil {
				return err
			}

			ops = append(ops, RepoOp{
				Kind:       EvtKindDeleteRecord,
				Collection: d.Collection,
				Rkey:       d.Rkey,
			})
		default:
			return fmt.Errorf("no operation set in write enum")
		}
	}

	nroot, nrev, err := r.Commit(ctx, rm.kmgr.SignForUser)
	if err != nil {
		return err
	}

	rslice, err := ds.CloseWithRoot(ctx, nroot, nrev)
	if err != nil {
		return fmt.Errorf("close with root: %w", err)
	}

	var oldroot *cid.Cid
	if head.Defined() {
		oldroot = &head
	}

	if rm.events != nil {
		rm.events(ctx, &RepoEvent{
			User:      user,
			OldRoot:   oldroot,
			NewRoot:   nroot,
			PrevData:  prevData,
			RepoSlice: rslice,
			Rev:       nrev,
			Since:     &rev,
			Ops:       ops,
		})
	}

	return nil
}

// BulkUpsert writes multiple records in a single delta session and commit.
// Each record is upserted: created if new, updated if it already exists.
func (rm *RepoManager) BulkUpsert(ctx context.Context, user models.Uid, records []BulkRecord) error {
	ctx, span := otel.Tracer("repoman").Start(ctx, "BulkUpsert")
	defer span.End()

	unlock := rm.lockUser(ctx, user)
	defer unlock()

	rev, err := rm.cs.GetUserRepoRev(ctx, user)
	if err != nil {
		return err
	}

	ds, err := rm.cs.NewDeltaSession(ctx, user, &rev)
	if err != nil {
		return err
	}

	head := ds.BaseCid()
	r, err := repo.OpenRepo(ctx, ds, head)
	if err != nil {
		return err
	}

	// Capture previous MST root before commit overwrites it
	var prevData *cid.Cid
	if head.Defined() {
		pd := r.DataCid()
		prevData = &pd
	}

	ops := make([]RepoOp, 0, len(records))
	for _, rec := range records {
		rpath := rec.Collection + "/" + rec.Rkey

		// Check if record exists to determine create vs update
		_, _, getErr := r.GetRecordBytes(ctx, rpath)
		recordExists := getErr == nil

		var cc cid.Cid
		var evtKind EventKind
		if recordExists {
			cc, err = r.UpdateRecord(ctx, rpath, rec.Data)
			evtKind = EvtKindUpdateRecord
		} else {
			cc, err = r.PutRecord(ctx, rpath, rec.Data)
			evtKind = EvtKindCreateRecord
		}
		if err != nil {
			return fmt.Errorf("failed to write %s: %w", rpath, err)
		}

		ops = append(ops, RepoOp{
			Kind:       evtKind,
			Collection: rec.Collection,
			Rkey:       rec.Rkey,
			RecCid:     &cc,
		})
	}

	nroot, nrev, err := r.Commit(ctx, rm.kmgr.SignForUser)
	if err != nil {
		return err
	}

	rslice, err := ds.CloseWithRoot(ctx, nroot, nrev)
	if err != nil {
		return fmt.Errorf("close with root: %w", err)
	}

	var oldroot *cid.Cid
	if head.Defined() {
		oldroot = &head
	}

	if rm.events != nil {
		rm.events(ctx, &RepoEvent{
			User:      user,
			OldRoot:   oldroot,
			NewRoot:   nroot,
			PrevData:  prevData,
			Rev:       nrev,
			Since:     &rev,
			Ops:       ops,
			RepoSlice: rslice,
		})
	}

	return nil
}