seaweedfs/weed/wdclient/vidmap_client.go

package wdclient

import (
	"container/list"
	"context"
	"errors"
	"fmt"
	"math/rand"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/sync/singleflight"

	"github.com/seaweedfs/seaweedfs/weed/glog"
	"github.com/seaweedfs/seaweedfs/weed/pb"
	"github.com/seaweedfs/seaweedfs/weed/util"
)

// VolumeLocationProvider is the interface for looking up volume locations
// This allows different implementations (master subscription, filer queries, etc.)
type VolumeLocationProvider interface {
	// LookupVolumeIds looks up volume locations for the given volume IDs
	// Returns a map of volume ID to locations
	LookupVolumeIds(ctx context.Context, volumeIds []string) (map[string][]Location, error)
}

// vidMapClient provides volume location caching with pluggable lookup
// It wraps the battle-tested vidMap with customizable volume lookup strategies
type vidMapClient struct {
	vidMap         *vidMap
	provider       VolumeLocationProvider
	vidLookupGroup singleflight.Group
	// recentDeletes fences (vid, url) pairs a pushed removal dropped, so a
	// provider answer still naming them is neither cached nor served until the
	// master reports the pair back or recentDeleteTTL passes. vidMapEpoch
	// keeps answers fetched before a reset or clear out of the cache.
	recentDeletesMu   sync.Mutex
	recentDeletes     map[vidLocationKey]*list.Element
	recentDeleteOrder *list.List // of recentDelete, oldest first
	vidMapEpoch       uint64
}

type recentDelete struct {
	key vidLocationKey
	at  time.Time
}

type vidLocationKey struct {
	vid uint32
	url string
}

const (
	// recentDeleteTTL outlasts a lookup in flight when the removal landed.
	// Expiring only falls back to learning by a failed read, while a mark kept
	// until the re-add would hide a live server whose re-add the master
	// dropped on a full client channel.
	recentDeleteTTL = time.Minute
	recentDeleteCap = 1 << 16
)

const (
	// DefaultVidMapCacheSize is the default number of resets a volume location
	// survives without being relearned. This provides cache history when
	// volumes move between servers.
	DefaultVidMapCacheSize = 5
)

// newVidMapClient creates a new client with the given provider and data center
func newVidMapClient(provider VolumeLocationProvider, dataCenter string, cacheSize int) *vidMapClient {
	return &vidMapClient{
		vidMap:   newVidMap(dataCenter, cacheSize),
		provider: provider,
	}
}

// GetLookupFileIdFunction returns a function that can be used to lookup file IDs
func (vc *vidMapClient) GetLookupFileIdFunction() LookupFileIdFunctionType {
	return vc.LookupFileIdWithFallback
}

// LookupFileIdWithFallback resolves a "<vid>,<cookie>" file id to a list of
// HTTP read URLs, using the cached vidMap when populated and falling back to
// the provider for a fresh lookup on miss. URLs are returned in preference
// order: same-DC local, same-DC remote-tier, then the other data centers on
// the same footing -- mirroring the cached vidMap path so both routes agree.
// Concurrent misses for the same vid are coalesced via the singleflight group
// on LookupVolumeIdsWithFallback.
func (vc *vidMapClient) LookupFileIdWithFallback(ctx context.Context, fileId string) (fullUrls []string, err error) {
	// Try cache first
	dataCenter := vc.vidMap.DataCenter
	fullUrls, err = vc.vidMap.LookupFileId(ctx, fileId)

	// Cache hit - return immediately
	if err == nil && len(fullUrls) > 0 {
		return
	}

	// Cache miss - extract volume ID from file ID (format: "volumeId,needle_id_cookie")
	if fileId == "" {
		return nil, fmt.Errorf("empty fileId")
	}
	parts := strings.Split(fileId, ",")
	if len(parts) != 2 {
		return nil, fmt.Errorf("invalid fileId %s", fileId)
	}
	volumeId := parts[0]

	// Use shared lookup logic with batching and singleflight
	vidLocations, err := vc.LookupVolumeIdsWithFallback(ctx, []string{volumeId})

	// Check for partial results first (important for multi-volume batched lookups)
	locations, found := vidLocations[volumeId]
	if !found || len(locations) == 0 {
		// Volume not found - return specific error with context from lookup if available
		if err != nil {
			return nil, fmt.Errorf("volume %s not found for fileId %s: %w", volumeId, fileId, err)
		}
		return nil, fmt.Errorf("volume %s not found for fileId %s", volumeId, fileId)
	}

	// Volume found successfully - ignore any errors about other volumes
	// (not relevant for single-volume lookup, but defensive for future batching)

	// Build HTTP URLs from locations, preferring same data center
	var sameDcUrls, otherDcUrls []string
	localUrls := make(map[string]bool)
	for _, loc := range locations {
		httpUrl := "http://" + loc.Url + "/" + fileId
		glog.V(4).Infof("lookup %s => %s, data in remote storage tier: %v", fileId, loc.Url, loc.DataInRemote)
		if !loc.DataInRemote {
			localUrls[httpUrl] = true
		}
		if dataCenter != "" && dataCenter == loc.DataCenter {
			sameDcUrls = append(sameDcUrls, httpUrl)
		} else {
			otherDcUrls = append(otherDcUrls, httpUrl)
		}
	}

	// Shuffle to distribute load across volume servers
	rand.Shuffle(len(sameDcUrls), func(i, j int) { sameDcUrls[i], sameDcUrls[j] = sameDcUrls[j], sameDcUrls[i] })
	rand.Shuffle(len(otherDcUrls), func(i, j int) { otherDcUrls[i], otherDcUrls[j] = otherDcUrls[j], otherDcUrls[i] })

	// Local replicas go first inside each data center, but never ahead of the
	// data-center preference itself. Mirrors vidMap.LookupVolumeServerUrl so
	// all client lookup paths agree.
	if len(localUrls) > 0 {
		sameDcUrls = util.ReorderToFront(localUrls, sameDcUrls)
		otherDcUrls = util.ReorderToFront(localUrls, otherDcUrls)
	}
	fullUrls = append(sameDcUrls, otherDcUrls...)
	return fullUrls, nil
}

// LookupVolumeIdsWithFallback looks up volume locations, querying provider if not in cache.
// Uses singleflight to coalesce concurrent requests for the same batch of volumes.
//
// IMPORTANT: This function may return PARTIAL results with a non-nil error.
// The result map contains successfully looked up volumes, while the error aggregates
// failures for volumes that couldn't be found or had lookup errors.
//
// Callers MUST check both the result map AND the error:
//   - result != nil && err == nil: All volumes found successfully
//   - result != nil && err != nil: Some volumes found, some failed (check both)
//   - result == nil && err != nil: Complete failure (connection error, etc.)
//
// Example usage:
//
//	locs, err := mc.LookupVolumeIdsWithFallback(ctx, []string{"1", "2", "999"})
//	if len(locs) > 0 {
//	    // Process successfully found volumes
//	}
//	if err != nil {
//	    // Log/handle failed volumes
//	}
func (vc *vidMapClient) LookupVolumeIdsWithFallback(ctx context.Context, volumeIds []string) (map[string][]Location, error) {
	result := make(map[string][]Location)
	var needsLookup []string
	var lookupErrors []error

	// Check cache first and parse volume IDs once
	vidStringToUint := make(map[string]uint32, len(volumeIds))

	for _, vidString := range volumeIds {
		vid, err := strconv.ParseUint(vidString, 10, 32)
		if err != nil {
			return nil, fmt.Errorf("invalid volume id %s: %v", vidString, err)
		}
		vidStringToUint[vidString] = uint32(vid)

		locations, found := vc.vidMap.GetLocations(uint32(vid))
		if found && len(locations) > 0 {
			result[vidString] = locations
		} else {
			needsLookup = append(needsLookup, vidString)
		}
	}

	if len(needsLookup) == 0 {
		return result, nil
	}

	// Batch query all missing volumes using singleflight on the batch key
	// Sort for stable key to coalesce identical batches
	sort.Strings(needsLookup)
	batchKey := strings.Join(needsLookup, ",")

	sfResult, err := vc.doSharedLookup(ctx, batchKey, func(ctx context.Context) (interface{}, error) {
		// Double-check cache for volumes that might have been populated while waiting
		stillNeedLookup := make([]string, 0, len(needsLookup))
		batchResult := make(map[string][]Location)

		for _, vidString := range needsLookup {
			vid := vidStringToUint[vidString] // Use pre-parsed value
			if locations, found := vc.vidMap.GetLocations(vid); found && len(locations) > 0 {
				batchResult[vidString] = locations
			} else {
				stillNeedLookup = append(stillNeedLookup, vidString)
			}
		}

		if len(stillNeedLookup) == 0 {
			return batchResult, nil
		}

		// Query provider with batched volume IDs
		glog.V(2).Infof("Looking up %d volumes from provider: %v", len(stillNeedLookup), stillNeedLookup)

		fillEpoch := vc.mapEpoch()
		providerResults, err := vc.provider.LookupVolumeIds(ctx, stillNeedLookup)
		if err != nil {
			return batchResult, fmt.Errorf("provider lookup failed: %w", err)
		}

		for vidString, locations := range providerResults {
			vid, err := strconv.ParseUint(vidString, 10, 32)
			if err != nil {
				glog.Warningf("Failed to parse volume id '%s': %v", vidString, err)
				continue
			}

			var live []Location
			for _, loc := range locations {
				if vc.fillLocation(uint32(vid), loc, fillEpoch) {
					live = append(live, loc)
				}
			}

			if len(live) > 0 {
				batchResult[vidString] = live
			}
		}

		return batchResult, nil
	})

	if err != nil {
		lookupErrors = append(lookupErrors, err)
	}

	// Merge singleflight batch results
	if batchLocations, ok := sfResult.(map[string][]Location); ok {
		for vid, locs := range batchLocations {
			result[vid] = locs
		}
	}

	// Check for volumes that still weren't found
	for _, vidString := range needsLookup {
		if _, found := result[vidString]; !found {
			lookupErrors = append(lookupErrors, fmt.Errorf("volume %s not found", vidString))
		}
	}

	// Return aggregated errors
	return result, errors.Join(lookupErrors...)
}

// doSharedLookup runs fn once for all concurrent callers of key. fn gets a
// context no single caller can cancel, since one caller leaving must not fail
// the others; each caller stops waiting when its own context ends.
func (vc *vidMapClient) doSharedLookup(ctx context.Context, key string, fn func(context.Context) (interface{}, error)) (interface{}, error) {
	resultCh := vc.vidLookupGroup.DoChan(key, func() (interface{}, error) {
		return fn(context.WithoutCancel(ctx))
	})
	select {
	case res := <-resultCh:
		return res.Val, res.Err
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

// LookupVolumeIdsAuthoritative bypasses the asynchronously maintained vid map.
// Use it when a stale positive result is unsafe, such as before serving cache
// payload from a Volume that the master may have retired. The provider still
// batches all requested IDs into one lookup, returns a nil map when it got no
// answer, and reports the volumes the master does not serve as errors.
func (vc *vidMapClient) LookupVolumeIdsAuthoritative(ctx context.Context, volumeIds []string) (map[string][]Location, error) {
	return vc.provider.LookupVolumeIds(ctx, volumeIds)
}

// Public methods for external access
//
// The vidMap itself is never replaced, so these all read the one map under its
// own lock. Resets bump its generation instead of swapping in a fresh instance.

// GetLocations safely retrieves volume locations
func (vc *vidMapClient) GetLocations(vid uint32) (locations []Location, found bool) {
	return vc.vidMap.GetLocations(vid)
}

// GetLocationsClone safely retrieves a clone of volume locations
func (vc *vidMapClient) GetLocationsClone(vid uint32) (locations []Location, found bool) {
	return vc.vidMap.GetLocationsClone(vid)
}

// GetVidLocations safely retrieves volume locations by string ID
func (vc *vidMapClient) GetVidLocations(vid string) (locations []Location, err error) {
	return vc.vidMap.GetVidLocations(vid)
}

// LookupFileId safely looks up URLs for a file ID
func (vc *vidMapClient) LookupFileId(ctx context.Context, fileId string) (fullUrls []string, err error) {
	return vc.vidMap.LookupFileId(ctx, fileId)
}

// LookupVolumeServerUrl safely looks up volume server URLs
func (vc *vidMapClient) LookupVolumeServerUrl(vid string) (serverUrls []string, err error) {
	return vc.vidMap.LookupVolumeServerUrl(vid)
}

// HasVolumeServer reports whether addr is currently a known volume server
// (hosts at least one volume or EC shard) in the cached vid map. Used by
// admission paths that must only contact peers learned from the master.
func (vc *vidMapClient) HasVolumeServer(addr pb.ServerAddress) bool {
	return vc.vidMap.hasVolumeServer(addr)
}

// GetDataCenter safely retrieves the data center
func (vc *vidMapClient) GetDataCenter() string {
	return vc.vidMap.DataCenter
}

// Thread-safe helpers for vidMap operations

// addLocation adds a volume location
func (vc *vidMapClient) addLocation(vid uint32, location Location) {
	vc.vidMap.addLocation(vid, location)
}

// deleteLocation removes a volume location
func (vc *vidMapClient) deleteLocation(vid uint32, location Location) {
	vc.vidMap.deleteLocation(vid, location)
}

// addEcLocation adds an EC volume location
func (vc *vidMapClient) addEcLocation(vid uint32, location Location) {
	vc.vidMap.addEcLocation(vid, location)
}

// deleteEcLocation removes an EC volume location
func (vc *vidMapClient) deleteEcLocation(vid uint32, location Location) {
	vc.vidMap.deleteEcLocation(vid, location)
}

// resetVidMap starts a new generation, as when the master changes: what the
// previous one told us stays readable until it is relearned or expires.
func (vc *vidMapClient) resetVidMap() {
	vc.recentDeletesMu.Lock()
	vc.vidMapEpoch++
	vc.recentDeletesMu.Unlock()
	vc.vidMap.reset()
}

// clearVidMap drops every location, for a cache whose feed restarted without
// replaying the removals it missed.
func (vc *vidMapClient) clearVidMap() {
	vc.recentDeletesMu.Lock()
	vc.vidMapEpoch++
	vc.recentDeletesMu.Unlock()
	vc.vidMap.clear()
}

func (vc *vidMapClient) mapEpoch() uint64 {
	vc.recentDeletesMu.Lock()
	defer vc.recentDeletesMu.Unlock()
	return vc.vidMapEpoch
}

func (vc *vidMapClient) tombstone(vid uint32, loc Location) {
	key := vidLocationKey{vid: vid, url: loc.Url}
	now := time.Now()
	vc.recentDeletesMu.Lock()
	defer vc.recentDeletesMu.Unlock()
	if vc.recentDeletes == nil {
		vc.recentDeletes = make(map[vidLocationKey]*list.Element)
		vc.recentDeleteOrder = list.New()
	}
	vc.expireRecentDeletes(now)
	if e, found := vc.recentDeletes[key]; found {
		e.Value = recentDelete{key: key, at: now}
		vc.recentDeleteOrder.MoveToBack(e)
		return
	}
	vc.recentDeletes[key] = vc.recentDeleteOrder.PushBack(recentDelete{key: key, at: now})
	if len(vc.recentDeletes) > recentDeleteCap {
		vc.removeRecentDelete(vc.recentDeleteOrder.Front())
	}
}

// expireRecentDeletes drops marks older than recentDeleteTTL. Callers must
// hold recentDeletesMu.
func (vc *vidMapClient) expireRecentDeletes(now time.Time) {
	if vc.recentDeleteOrder == nil {
		return
	}
	for e := vc.recentDeleteOrder.Front(); e != nil && now.Sub(e.Value.(recentDelete).at) >= recentDeleteTTL; e = vc.recentDeleteOrder.Front() {
		vc.removeRecentDelete(e)
	}
}

func (vc *vidMapClient) removeRecentDelete(e *list.Element) {
	delete(vc.recentDeletes, e.Value.(recentDelete).key)
	vc.recentDeleteOrder.Remove(e)
}

func (vc *vidMapClient) untombstone(vid uint32, loc Location) {
	key := vidLocationKey{vid: vid, url: loc.Url}
	vc.recentDeletesMu.Lock()
	defer vc.recentDeletesMu.Unlock()
	if e, found := vc.recentDeletes[key]; found {
		vc.removeRecentDelete(e)
	}
}

// fillLocation reports whether a provider answer may be served, and caches it
// unless a reset or clear happened since epoch. Holding recentDeletesMu across
// the check and the add keeps a concurrent removal from being undone.
func (vc *vidMapClient) fillLocation(vid uint32, loc Location, epoch uint64) bool {
	vc.recentDeletesMu.Lock()
	defer vc.recentDeletesMu.Unlock()
	vc.expireRecentDeletes(time.Now())
	if _, found := vc.recentDeletes[vidLocationKey{vid: vid, url: loc.Url}]; found {
		return false
	}
	if vc.vidMapEpoch == epoch {
		vc.vidMap.addLocation(vid, loc)
	}
	return true
}

// InvalidateCache removes all cached locations for a volume ID
func (vc *vidMapClient) InvalidateCache(fileId string) {
	parts := strings.Split(fileId, ",")
	vidString := parts[0]
	vid, err := strconv.ParseUint(vidString, 10, 32)
	if err != nil {
		return
	}
	vc.vidMap.deleteVid(uint32(vid))
}