milvus-io--milvus/internal/util/dependency/factory.go

package dependency

import (
	"context"
	"os"
	"path/filepath"

	"github.com/cockroachdb/errors"

	"github.com/milvus-io/milvus/internal/storage"
	"github.com/milvus-io/milvus/pkg/v3/common"
	"github.com/milvus-io/milvus/pkg/v3/metrics"
	"github.com/milvus-io/milvus/pkg/v3/mlog"
	"github.com/milvus-io/milvus/pkg/v3/mq/msgstream"
	"github.com/milvus-io/milvus/pkg/v3/objectstorage"
	"github.com/milvus-io/milvus/pkg/v3/util/merr"
	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)

const (
	mqTypeDefault    = "default"
	mqTypeRocksmq    = "rocksmq"
	mqTypeKafka      = "kafka"
	mqTypePulsar     = "pulsar"
	mqTypeWoodpecker = "woodpecker"
)

type mqEnable struct {
	Rocksmq    bool
	Pulsar     bool
	Kafka      bool
	Woodpecker bool
}

// DefaultFactory is a factory that produces instances of storage.ChunkManager and message queue.
type DefaultFactory struct {
	standAlone          bool
	chunkManagerFactory storage.Factory
	msgStreamFactory    msgstream.Factory
}

// testRocksmqPath returns a unique temporary path for RocksMQ in tests.
// This avoids RocksDB LOCK file conflicts when multiple test packages run in parallel.
func testRocksmqPath() string {
	dir, _ := os.MkdirTemp("", "milvus_ut_rocksmq_*")
	return filepath.Join(dir, "rdb_data")
}

// Only for test
func NewDefaultFactory(standAlone bool) *DefaultFactory {
	params := paramtable.Get()
	return &DefaultFactory{
		standAlone:       standAlone,
		msgStreamFactory: msgstream.NewRocksmqFactory(testRocksmqPath(), &params.ServiceParam),
		chunkManagerFactory: storage.NewChunkManagerFactory("local",
			objectstorage.RootPath("/tmp/milvus")),
	}
}

// Only for test
func MockDefaultFactory(standAlone bool, params *paramtable.ComponentParam) *DefaultFactory {
	return &DefaultFactory{
		standAlone:          standAlone,
		msgStreamFactory:    msgstream.NewRocksmqFactory(testRocksmqPath(), &params.ServiceParam),
		chunkManagerFactory: storage.NewChunkManagerFactoryWithParam(params),
	}
}

// NewFactory creates a new instance of the DefaultFactory type.
// If standAlone is true, the factory will operate in standalone mode.
func NewFactory(standAlone bool) *DefaultFactory {
	return &DefaultFactory{standAlone: standAlone}
}

// Init create a msg factory(TODO only support one mq at the same time.)
// In order to guarantee backward compatibility of config file, we still support multiple mq configs.
// The initialization of MQ follows the following rules, if the mq.type is default.
// 1. standalone(local) mode: rocksmq(default) > Pulsar > Kafka
// 2. cluster mode:  Pulsar(default) > Kafka (rocksmq is unsupported in cluster mode)
func (f *DefaultFactory) Init(params *paramtable.ComponentParam) {
	// skip if using default factory
	if f.msgStreamFactory != nil {
		return
	}

	f.chunkManagerFactory = storage.NewChunkManagerFactoryWithParam(params)

	// initialize mq client or embedded mq.
	if err := f.initMQ(f.standAlone, params); err != nil {
		panic(err)
	}
}

func (f *DefaultFactory) initMQ(standalone bool, params *paramtable.ComponentParam) error {
	mqType := mustSelectMQType(standalone, params.MQCfg.Type.GetValue(), mqEnable{params.RocksmqEnable(), params.PulsarEnable(), params.KafkaEnable(), params.WoodpeckerEnable()})
	metrics.RegisterMQType(mqType)
	mlog.Info(context.TODO(), "try to init mq", mlog.Bool("standalone", standalone), mlog.String("mqType", mqType))

	switch mqType {
	case mqTypeRocksmq:
		f.msgStreamFactory = msgstream.NewRocksmqFactory(params.RocksmqCfg.Path.GetValue(), &params.ServiceParam)
	case mqTypePulsar:
		f.msgStreamFactory = msgstream.NewPmsFactory(&params.ServiceParam)
	case mqTypeKafka:
		f.msgStreamFactory = msgstream.NewKmsFactory(&params.ServiceParam)
	case mqTypeWoodpecker:
		f.msgStreamFactory = msgstream.NewWpmsFactory(&params.ServiceParam)
	}
	if f.msgStreamFactory == nil {
		return merr.WrapErrServiceInternalMsg("failed to create MQ: check the milvus log for initialization failures")
	}
	return nil
}

// Select valid mq if mq type is default.
func mustSelectMQType(standalone bool, mqType string, enable mqEnable) string {
	if mqType != mqTypeDefault {
		if err := validateMQType(standalone, mqType); err != nil {
			panic(err)
		}
		return mqType
	}

	if standalone {
		if enable.Rocksmq {
			return mqTypeRocksmq
		}
	}
	if enable.Pulsar {
		return mqTypePulsar
	}
	if enable.Kafka {
		return mqTypeKafka
	}
	if enable.Woodpecker {
		return mqTypeWoodpecker
	}

	panic(errors.Errorf("no available mq config found, %s, enable: %+v", mqType, enable))
}

// Validate mq type.
func validateMQType(standalone bool, mqType string) error {
	if mqType != mqTypeRocksmq && mqType != mqTypeKafka && mqType != mqTypePulsar && mqType != mqTypeWoodpecker {
		return merr.WrapErrParameterInvalidMsg("mq type %s is invalid", mqType)
	}
	if !standalone && mqType == mqTypeRocksmq {
		return merr.WrapErrParameterInvalidMsg("mq %s is only valid in standalone mode", mqType)
	}
	return nil
}

func (f *DefaultFactory) NewMsgStream(ctx context.Context) (msgstream.MsgStream, error) {
	return f.msgStreamFactory.NewMsgStream(ctx)
}

func (f *DefaultFactory) NewTtMsgStream(ctx context.Context) (msgstream.MsgStream, error) {
	return f.msgStreamFactory.NewTtMsgStream(ctx)
}

func (f *DefaultFactory) NewMsgStreamDisposer(ctx context.Context) func([]string, string) error {
	return f.msgStreamFactory.NewMsgStreamDisposer(ctx)
}

func (f *DefaultFactory) NewPersistentStorageChunkManager(ctx context.Context) (storage.ChunkManager, error) {
	return f.chunkManagerFactory.NewPersistentStorageChunkManager(ctx)
}

type Factory interface {
	msgstream.Factory
	Init(p *paramtable.ComponentParam)
	NewPersistentStorageChunkManager(ctx context.Context) (storage.ChunkManager, error)
}

func HealthCheck(mqType string) *common.MQClusterStatus {
	clusterStatus := &common.MQClusterStatus{MqType: mqType}
	switch mqType {
	case mqTypeRocksmq:
		// TODO: implement health checker for rocks mq
		clusterStatus.Health = true
	case mqTypePulsar:
		msgstream.PulsarHealthCheck(clusterStatus)
	case mqTypeKafka:
		msgstream.KafkaHealthCheck(clusterStatus)
	case mqTypeWoodpecker:
		// TODO: implement health checker for woodpecker
		clusterStatus.Health = true
	}
	return clusterStatus
}