diff --git a/pkg/common/discoveryregister/etcd/etcd.go b/pkg/common/discoveryregister/etcd/etcd.go index 904f83dd0..51fafd0b6 100644 --- a/pkg/common/discoveryregister/etcd/etcd.go +++ b/pkg/common/discoveryregister/etcd/etcd.go @@ -3,6 +3,7 @@ package etcd import ( "context" "fmt" + "github.com/pkg/errors" clientv3 "go.etcd.io/etcd/client/v3" "go.etcd.io/etcd/client/v3/naming/endpoints" "go.etcd.io/etcd/client/v3/naming/resolver" @@ -173,3 +174,48 @@ func (r *SvcDiscoveryRegistryImpl) Close() { _ = r.client.Close() } } + +// Check verifies if etcd is running by checking the existence of the root node and optionally creates it +func Check(ctx context.Context, etcdServers []string, etcdRoot string, createIfNotExist bool, options ...ZkOption) error { + // Configure the etcd client with default settings + cfg := clientv3.Config{ + Endpoints: etcdServers, + } + + // Apply provided options to the config + for _, opt := range options { + opt(&cfg) + } + + client, err := clientv3.New(cfg) + if err != nil { + return errors.Wrap(err, "failed to connect to etcd") + } + defer client.Close() + + // Create a child context with a default timeout or use the provided context + opCtx, cancel := context.WithCancel(ctx) + defer cancel() + + // Check if the root node exists + resp, err := client.Get(opCtx, etcdRoot) + if err != nil { + return errors.Wrap(err, "failed to get the root node from etcd") + } + + // If root node does not exist and createIfNotExist is true, create the root node + if len(resp.Kvs) == 0 { + if createIfNotExist { + _, err := client.Put(opCtx, etcdRoot, "") + if err != nil { + return errors.Wrap(err, "failed to create the root node in etcd") + } + fmt.Printf("Root node %s did not exist, but has been created.\n", etcdRoot) + } else { + return fmt.Errorf("root node %s does not exist in etcd", etcdRoot) + } + } else { + fmt.Printf("Etcd is running and the root node %s exists.\n", etcdRoot) + } + return nil +} diff --git a/tools/check-component/main.go b/tools/check-component/main.go index 7fe64d3c5..5a936f775 100644 --- a/tools/check-component/main.go +++ b/tools/check-component/main.go @@ -20,6 +20,7 @@ import ( "fmt" "github.com/openimsdk/open-im-server/v3/pkg/common/cmd" "github.com/openimsdk/open-im-server/v3/pkg/common/config" + "github.com/openimsdk/open-im-server/v3/pkg/common/discoveryregister/etcd" "github.com/openimsdk/tools/db/mongoutil" "github.com/openimsdk/tools/db/redisutil" "github.com/openimsdk/tools/discovery/zookeeper" @@ -43,6 +44,25 @@ func CheckZookeeper(ctx context.Context, config *config.ZooKeeper) error { return zookeeper.Check(ctx, config.Address, config.Schema, zookeeper.WithUserNameAndPassword(config.Username, config.Password)) } +func CheckEtcd(ctx context.Context, config *config.Etcd) error { + etcd.Check(ctx, config.Address, "/check_openim_component", + true, + etcd.WithDialTimeout(10*time.Second), + etcd.WithMaxCallSendMsgSize(20*1024*1024), + etcd.WithUsernameAndPassword(config.Username, config.Password)) + return nil +} + +func CheckDiscovery(ctx context.Context, config *config.Discovery) error { + switch config.Enable { + case "etcd": + return CheckEtcd(ctx, &config.Etcd) + + } + + return nil +} + func CheckMongo(ctx context.Context, config *config.Mongo) error { return mongoutil.Check(ctx, config.Build()) } @@ -59,14 +79,14 @@ func CheckKafka(ctx context.Context, conf *config.Kafka) error { return kafka.Check(ctx, conf.Build(), []string{conf.ToMongoTopic, conf.ToRedisTopic, conf.ToPushTopic}) } -func initConfig(configDir string) (*config.Mongo, *config.Redis, *config.Kafka, *config.Minio, *config.ZooKeeper, error) { +func initConfig(configDir string) (*config.Mongo, *config.Redis, *config.Kafka, *config.Minio, *config.Discovery, error) { var ( - mongoConfig = &config.Mongo{} - redisConfig = &config.Redis{} - kafkaConfig = &config.Kafka{} - minioConfig = &config.Minio{} - zookeeperConfig = &config.ZooKeeper{} - thirdConfig = &config.Third{} + mongoConfig = &config.Mongo{} + redisConfig = &config.Redis{} + kafkaConfig = &config.Kafka{} + minioConfig = &config.Minio{} + discovery = &config.Discovery{} + thirdConfig = &config.Third{} ) err := config.LoadConfig(filepath.Join(configDir, cmd.MongodbConfigFileName), cmd.ConfigEnvPrefixMap[cmd.MongodbConfigFileName], mongoConfig) if err != nil { @@ -96,11 +116,11 @@ func initConfig(configDir string) (*config.Mongo, *config.Redis, *config.Kafka, } else { minioConfig = nil } - err = config.LoadConfig(filepath.Join(configDir, cmd.ZookeeperConfigFileName), cmd.ConfigEnvPrefixMap[cmd.ZookeeperConfigFileName], zookeeperConfig) + err = config.LoadConfig(filepath.Join(configDir, cmd.DiscoveryConfigFilename), cmd.ConfigEnvPrefixMap[cmd.DiscoveryConfigFilename], discovery) if err != nil { return nil, nil, nil, nil, nil, err } - return mongoConfig, redisConfig, kafkaConfig, minioConfig, zookeeperConfig, nil + return mongoConfig, redisConfig, kafkaConfig, minioConfig, discovery, nil } func main() { @@ -127,13 +147,10 @@ func main() { } } -func performChecks(ctx context.Context, mongoConfig *config.Mongo, redisConfig *config.Redis, kafkaConfig *config.Kafka, minioConfig *config.Minio, zookeeperConfig *config.ZooKeeper, maxRetry int) error { +func performChecks(ctx context.Context, mongoConfig *config.Mongo, redisConfig *config.Redis, kafkaConfig *config.Kafka, minioConfig *config.Minio, discovery *config.Discovery, maxRetry int) error { checksDone := make(map[string]bool) checks := map[string]func() error{ - "Zookeeper": func() error { - return CheckZookeeper(ctx, zookeeperConfig) - }, "Mongo": func() error { return CheckMongo(ctx, mongoConfig) }, @@ -144,12 +161,20 @@ func performChecks(ctx context.Context, mongoConfig *config.Mongo, redisConfig * return CheckKafka(ctx, kafkaConfig) }, } - if minioConfig != nil { checks["MinIO"] = func() error { return CheckMinIO(ctx, minioConfig) } } + if discovery.Enable == "etcd" { + checks["Etcd"] = func() error { + return CheckEtcd(ctx, &discovery.Etcd) + } + } else if discovery.Enable == "zookeeper" { + checks["Zookeeper"] = func() error { + return CheckZookeeper(ctx, &discovery.ZooKeeper) + } + } for i := 0; i < maxRetry; i++ { allSuccess := true