at-container-registry/pkg/appview/storage/drain.go

package storage

import (
	"context"
	"encoding/json"
	"log/slog"
	"strings"
	"sync"
	"time"

	"atcr.io/pkg/atproto"
	"atcr.io/pkg/auth"
)

// drainLocks prevents concurrent drain operations per DID.
// If a drain is already running for a user (from login or push), skip.
var drainLocks sync.Map

// MigrateManifestsForSuccessor rewrites manifest records and profile
// when any of the user's holds have a successor. Best-effort, runs in background.
//
// Steps:
// 1. Get user's sailor profile
// 2. Collect candidate holds: profile.DefaultHold + distinct holds from manifests DB
// 3. For each hold with a successor: rewrite PDS manifest records and local DB
// 4. If profile.DefaultHold itself had a successor, update the profile too
func MigrateManifestsForSuccessor(
	ctx context.Context,
	client *atproto.Client,
	authorizer auth.HoldAuthorizer,
	db HoldDIDLookup,
	did string,
) {
	// Lock per DID — skip if already running
	if _, loaded := drainLocks.LoadOrStore(did, true); loaded {
		return
	}
	defer drainLocks.Delete(did)

	// 1. Get user's profile
	profile, err := GetProfile(ctx, client)
	if err != nil {
		slog.Debug("Drain: failed to get profile", "component", "storage/drain", "did", did, "error", err)
		return
	}

	// 2. Collect candidate holds to check for successors.
	// Start with profile.DefaultHold, then add any distinct holds from the DB.
	candidates := make(map[string]bool)
	if profile != nil && profile.DefaultHold != "" {
		candidates[profile.DefaultHold] = true
	}
	if db != nil {
		manifestHolds, err := db.GetDistinctManifestHoldDIDs(did)
		if err != nil {
			slog.Warn("Drain: failed to get distinct manifest holds", "component", "storage/drain", "did", did, "error", err)
		}
		for _, h := range manifestHolds {
			candidates[h] = true
		}
	}

	if len(candidates) == 0 {
		return
	}

	// 3. Check each candidate for a successor and drain if found
	for oldHold := range candidates {
		captain, err := authorizer.GetCaptainRecord(ctx, oldHold)
		if err != nil {
			slog.Debug("Drain: failed to get captain record", "component", "storage/drain", "did", did, "hold", oldHold, "error", err)
			continue
		}
		if captain == nil || captain.Successor == "" {
			continue
		}
		newHold := captain.Successor

		slog.Info("Starting hold drain", "component", "storage/drain", "did", did, "from", oldHold, "to", newHold)

		drainHold(ctx, client, db, did, oldHold, newHold)

		// 4. If profile.DefaultHold pointed to this old hold, update it
		if profile != nil && profile.DefaultHold == oldHold {
			profile.DefaultHold = newHold
			profile.UpdatedAt = time.Now()
			if err := UpdateProfile(ctx, client, profile); err != nil {
				slog.Warn("Drain: failed to update profile", "component", "storage/drain", "did", did, "error", err)
			} else {
				slog.Info("Drain: updated profile defaultHold", "component", "storage/drain", "did", did, "newHold", newHold)
			}
		}
	}
}

// drainHold rewrites all PDS manifest records and local DB rows from oldHold to newHold.
func drainHold(
	ctx context.Context,
	client *atproto.Client,
	db HoldDIDLookup,
	did, oldHold, newHold string,
) {
	// Walk PDS manifest records, rewrite holdDid
	cursor := ""
	rewritten := 0
	for {
		records, nextCursor, err := client.ListRecordsWithCursor(ctx, atproto.ManifestCollection, 100, cursor)
		if err != nil {
			slog.Warn("Drain: failed to list manifest records", "component", "storage/drain", "did", did, "error", err)
			break
		}

		for _, rec := range records {
			var manifest atproto.ManifestRecord
			if err := json.Unmarshal(rec.Value, &manifest); err != nil {
				slog.Debug("Drain: failed to unmarshal manifest", "component", "storage/drain", "uri", rec.URI, "error", err)
				continue
			}

			// Check if this manifest points to the old hold (via DID or legacy endpoint)
			needsRewrite := false
			if manifest.HoldDID == oldHold {
				needsRewrite = true
			} else if manifest.HoldEndpoint != "" {
				if resolvedDID, resolveErr := atproto.ResolveHoldDID(ctx, manifest.HoldEndpoint); resolveErr == nil && resolvedDID == oldHold {
					needsRewrite = true
				}
			}

			if !needsRewrite {
				continue
			}

			// Rewrite to new hold
			manifest.HoldDID = newHold
			manifest.HoldEndpoint = "" // Clear legacy field

			// Extract rkey from AT URI (at://did/collection/rkey)
			uriParts := strings.Split(rec.URI, "/")
			if len(uriParts) < 2 {
				continue
			}
			rkey := uriParts[len(uriParts)-1]

			if _, err := client.PutRecord(ctx, atproto.ManifestCollection, rkey, &manifest); err != nil {
				slog.Warn("Drain: failed to rewrite manifest", "component", "storage/drain", "uri", rec.URI, "error", err)
				continue
			}
			rewritten++
		}

		if nextCursor == "" {
			break
		}
		cursor = nextCursor
	}

	// Update appview's local manifests table
	if db != nil {
		dbUpdated, err := db.UpdateManifestHoldDID(did, oldHold, newHold)
		if err != nil {
			slog.Warn("Drain: failed to update local DB", "component", "storage/drain", "did", did, "error", err)
		} else if dbUpdated > 0 {
			slog.Info("Drain: updated local DB manifests", "component", "storage/drain", "did", did, "rows", dbUpdated)
		}
	}

	slog.Info("Hold drain complete", "component", "storage/drain", "did", did, "from", oldHold, "to", newHold, "rewritten", rewritten)
}