kratos-etcd/config/etcd_client.go

package config

import (
	"context"
	"fmt"
	"log"
	"time"

	"git.ikuban.com/server/kratos-utils/v2/initutils"
	clientv3 "go.etcd.io/etcd/client/v3"
)

func NewEtcdClient(startup *initutils.Startup) *clientv3.Client {
	if startup.RegistryType != "etcd" {
		log.Println("[WARNING] startup registry type not etcd")
		return nil
	}

	// etcd client
	etcdClient, err := clientv3.New(clientv3.Config{
		Endpoints:   []string{startup.ConfigSource},
		DialTimeout: time.Second * 5,
	})
	if err != nil {
		panic(fmt.Errorf("failed to create etcd client: %w", err))
	}

	// 检查etcd连接状态
	newCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_, err = etcdClient.Maintenance.Status(newCtx, startup.ConfigSource)
	if err != nil {
		panic(fmt.Errorf("failed to get etcd status: %w", err))
	}

	return etcdClient
}

func NewEtcdClientWithCleanup(startup *initutils.Startup) (*clientv3.Client, func()) {
	etcdClient := NewEtcdClient(startup)
	if etcdClient == nil {
		return nil, func() {}
	}
	cleanup := func() { etcdClient.Close() }
	return etcdClient, cleanup
}