milvus/internal/datacoord/index_engine_version_manager.go

package datacoord

import (
	"math"

	"github.com/blang/semver/v4"
	"github.com/samber/lo"
	"go.uber.org/zap"

	"github.com/milvus-io/milvus/internal/util/sessionutil"
	"github.com/milvus-io/milvus/pkg/v2/log"
	"github.com/milvus-io/milvus/pkg/v2/util/lock"
)

// IndexEngineVersionManager manages the index engine versions reported by all QueryNodes in the cluster.
//
// Each QueryNode registers its supported index version range [MinimalIndexVersion, CurrentIndexVersion]
// in its session. This manager aggregates versions from all QNs to determine cluster-wide compatibility:
//
//   - GetCurrent*Version(): Returns MIN of all QNs' CurrentIndexVersion.
//     This is the highest version that ALL QueryNodes can load.
//     Used when building new indexes to ensure all QNs can load them (rolling upgrade safe).
//
//   - GetMinimal*Version(): Returns MAX of all QNs' MinimalIndexVersion.
//     This is the lowest version that ANY QueryNode requires.
//     Indexes below this version may fail to load on some QNs.
//     TODO: This is not currently used in the codebase, could be used to check if the index is of too old to
//     load on any query nodes.
//
// Vector index versions come from knowhere library, while scalar index versions are defined by Milvus.
type IndexEngineVersionManager interface {
	Startup(sessions map[string]*sessionutil.Session)
	AddNode(session *sessionutil.Session)
	RemoveNode(session *sessionutil.Session)
	Update(session *sessionutil.Session)

	// Vector index version methods (from knowhere library)
	GetCurrentIndexEngineVersion() int32
	GetMinimalIndexEngineVersion() int32

	// Scalar index version methods (Milvus-defined)
	GetCurrentScalarIndexEngineVersion() int32
	GetMinimalScalarIndexEngineVersion() int32

	GetIndexNonEncoding() bool

	GetMinimalSessionVer() semver.Version
}

type versionManagerImpl struct {
	mu                  lock.Mutex
	versions            map[int64]sessionutil.IndexEngineVersion
	scalarIndexVersions map[int64]sessionutil.IndexEngineVersion
	indexNonEncoding    map[int64]bool
	sessionVersion      map[int64]semver.Version
}

func newIndexEngineVersionManager() IndexEngineVersionManager {
	return &versionManagerImpl{
		versions:            map[int64]sessionutil.IndexEngineVersion{},
		scalarIndexVersions: map[int64]sessionutil.IndexEngineVersion{},
		indexNonEncoding:    map[int64]bool{},
		sessionVersion:      map[int64]semver.Version{},
	}
}

func (m *versionManagerImpl) Startup(sessions map[string]*sessionutil.Session) {
	m.mu.Lock()
	defer m.mu.Unlock()

	sessionMap := lo.MapKeys(sessions, func(session *sessionutil.Session, _ string) int64 {
		return session.ServerID
	})

	// clean offline nodes
	for sessionID := range m.versions {
		if _, ok := sessionMap[sessionID]; !ok {
			m.removeNodeByID(sessionID)
		}
	}

	// deal with new online nodes
	for _, session := range sessions {
		m.addOrUpdate(session)
	}
}

func (m *versionManagerImpl) AddNode(session *sessionutil.Session) {
	m.mu.Lock()
	defer m.mu.Unlock()

	m.addOrUpdate(session)
}

func (m *versionManagerImpl) RemoveNode(session *sessionutil.Session) {
	m.mu.Lock()
	defer m.mu.Unlock()

	m.removeNodeByID(session.ServerID)
}

func (m *versionManagerImpl) removeNodeByID(sessionID int64) {
	delete(m.versions, sessionID)
	delete(m.scalarIndexVersions, sessionID)
	delete(m.indexNonEncoding, sessionID)
	delete(m.sessionVersion, sessionID)
}

func (m *versionManagerImpl) Update(session *sessionutil.Session) {
	m.mu.Lock()
	defer m.mu.Unlock()

	m.addOrUpdate(session)
}

func (m *versionManagerImpl) addOrUpdate(session *sessionutil.Session) {
	log.Info("addOrUpdate version", zap.Int64("nodeId", session.ServerID),
		zap.String("sessionVersion", session.Version.String()),
		zap.Int32("minimal", session.IndexEngineVersion.MinimalIndexVersion),
		zap.Int32("current", session.IndexEngineVersion.CurrentIndexVersion),
		zap.Int32("currentScalar", session.ScalarIndexEngineVersion.CurrentIndexVersion))
	m.versions[session.ServerID] = session.IndexEngineVersion
	m.scalarIndexVersions[session.ServerID] = session.ScalarIndexEngineVersion
	m.indexNonEncoding[session.ServerID] = session.IndexNonEncoding
	m.sessionVersion[session.ServerID] = session.Version
}

func (m *versionManagerImpl) GetCurrentIndexEngineVersion() int32 {
	m.mu.Lock()
	defer m.mu.Unlock()

	if len(m.versions) == 0 {
		log.Info("index versions is empty")
		return 0
	}

	current := int32(math.MaxInt32)
	for _, version := range m.versions {
		if version.CurrentIndexVersion < current {
			current = version.CurrentIndexVersion
		}
	}
	log.Info("Merged current version", zap.Int32("current", current))
	return current
}

func (m *versionManagerImpl) GetMinimalIndexEngineVersion() int32 {
	m.mu.Lock()
	defer m.mu.Unlock()

	if len(m.versions) == 0 {
		log.Info("index versions is empty")
		return 0
	}

	minimal := int32(0)
	for _, version := range m.versions {
		if version.MinimalIndexVersion > minimal {
			minimal = version.MinimalIndexVersion
		}
	}
	log.Info("Merged minimal version", zap.Int32("minimal", minimal))
	return minimal
}

func (m *versionManagerImpl) GetCurrentScalarIndexEngineVersion() int32 {
	m.mu.Lock()
	defer m.mu.Unlock()

	if len(m.scalarIndexVersions) == 0 {
		log.Info("scalar index versions is empty")
		return 0
	}

	current := int32(math.MaxInt32)
	for _, version := range m.scalarIndexVersions {
		if version.CurrentIndexVersion < current {
			current = version.CurrentIndexVersion
		}
	}
	log.Info("Merged current scalar index version", zap.Int32("current", current))
	return current
}

func (m *versionManagerImpl) GetMinimalScalarIndexEngineVersion() int32 {
	m.mu.Lock()
	defer m.mu.Unlock()

	if len(m.scalarIndexVersions) == 0 {
		log.Info("scalar index versions is empty")
		return 0
	}

	minimal := int32(0)
	for _, version := range m.scalarIndexVersions {
		if version.MinimalIndexVersion > minimal {
			minimal = version.MinimalIndexVersion
		}
	}
	log.Info("Merged minimal scalar index version", zap.Int32("minimal", minimal))
	return minimal
}

func (m *versionManagerImpl) GetIndexNonEncoding() bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	if len(m.indexNonEncoding) == 0 {
		log.Info("indexNonEncoding map is empty")
		// by default, we fall back to old index format for safety
		return false
	}
	noneEncoding := true
	for _, encoding := range m.indexNonEncoding {
		noneEncoding = noneEncoding && encoding
	}
	return noneEncoding
}

func (m *versionManagerImpl) GetMinimalSessionVer() semver.Version {
	m.mu.Lock()
	defer m.mu.Unlock()

	minVer := semver.Version{}
	first := true
	for _, version := range m.sessionVersion {
		if first {
			minVer = version
			first = false
		} else if version.LT(minVer) {
			minVer = version
		}
	}
	return minVer
}