backend/internal/bootstrap/initializer.go

package bootstrap

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/XingfenD/yoresee_doc/internal/config"
	"github.com/XingfenD/yoresee_doc/internal/repository"
	"github.com/XingfenD/yoresee_doc/internal/service"
	"github.com/XingfenD/yoresee_doc/internal/service/mq_service"
	"github.com/XingfenD/yoresee_doc/internal/utils"
	"github.com/XingfenD/yoresee_doc/pkg/storage"
	"github.com/sirupsen/logrus"
)

type Initializer struct {
	err           error
	shutdownHooks []utils.ShutdownHook
	hookNames     map[string]struct{}
	shutdownOnce  sync.Once
	Repositories  *repository.Repositories
}

func NewInitializer() *Initializer {
	return &Initializer{
		hookNames: make(map[string]struct{}),
	}
}

func (i *Initializer) Check(step string, fn func() error) *Initializer {
	if i.err != nil {
		return i
	}
	if err := fn(); err != nil {
		i.err = fmt.Errorf("%s failed: %w", step, err)
	}
	return i
}

func (i *Initializer) CheckAllowFail(step string, fn func() error) *Initializer {
	if i.err != nil {
		return i
	}
	if err := fn(); err != nil {
		logrus.Warnf("%s failed, continue with degraded mode: %v", step, err)
	}
	return i
}

func (i *Initializer) InitConfig() *Initializer {
	return i.Check("config.InitConfig", config.InitConfig)
}

func (i *Initializer) InitPostgres() *Initializer {
	i.Check("storage.OpenPostgres", func() error {
		dbCfg := config.GlobalConfig.Database
		logCfg := config.GlobalConfig.Backend.Log
		gormLogLevel := logCfg.GormLogLevel
		if gormLogLevel == "" {
			gormLogLevel = logCfg.Level
		}

		db, err := storage.OpenPostgres(&storage.PostgresOptions{
			Host:         dbCfg.Host,
			Port:         dbCfg.Port,
			User:         dbCfg.User,
			Password:     dbCfg.Password,
			Name:         dbCfg.Name,
			MaxIdleConns: dbCfg.MaxIdleConns,
			MaxOpenConns: dbCfg.MaxOpenConns,
			GormLogLevel: gormLogLevel,
		})
		if err != nil {
			return err
		}
		storage.DB = db
		return nil
	})

	return i.addShutdownHook("Postgres", func() error {
		if storage.DB == nil {
			return nil
		}
		sqlDB, err := storage.DB.DB()
		if err != nil {
			return err
		}
		storage.DB = nil
		return sqlDB.Close()
	})
}

func (i *Initializer) InitRedis() *Initializer {
	i.Check("storage.NewRedisClient", func() error {
		cfg := config.GlobalConfig.Redis
		client, err := storage.NewRedisClient(&storage.RedisOptions{
			Host:     cfg.Host,
			Port:     cfg.Port,
			Password: cfg.Password,
			DB:       cfg.DB,
		})
		if err != nil {
			return err
		}
		storage.KVS = client
		return nil
	})

	return i.addShutdownHook("Redis", func() error {
		if storage.KVS == nil {
			return nil
		}
		client := storage.KVS
		storage.KVS = nil
		return client.Close()
	})
}

func (i *Initializer) InitConsul() *Initializer {
	return i.Check("storage.NewConsulClient", func() error {
		cfg := config.GlobalConfig.Consul
		storage.Consul = storage.NewConsulClient(&storage.ConsulOptions{
			Enabled:    cfg.Enabled,
			Address:    cfg.Address,
			Scheme:     cfg.Scheme,
			Token:      cfg.Token,
			Datacenter: cfg.Datacenter,
			Prefix:     cfg.Prefix,
		})
		return nil
	})
}

func (i *Initializer) RequireConsulEnabled() *Initializer {
	return i.Check("storage.ConsulEnabled", func() error {
		if !storage.ConsulEnabled() {
			return fmt.Errorf("consul is required but not enabled")
		}
		return nil
	})
}

func (i *Initializer) InitMinio() *Initializer {
	return i.Check("storage.NewMinioClient", func() error {
		cfg := config.GlobalConfig.Minio
		client, err := storage.NewMinioClient(&storage.MinioOptions{
			Endpoint:  cfg.Endpoint,
			AccessKey: cfg.AccessKey,
			SecretKey: cfg.SecretKey,
			UseSSL:    cfg.UseSSL,
		})
		if err != nil {
			return err
		}
		if err = storage.EnsureBucket(context.Background(), client, cfg.Bucket); err != nil {
			return err
		}
		if err = storage.EnsureBucketPublicRead(context.Background(), client, cfg.Bucket); err != nil {
			return err
		}
		storage.MinioClient = client
		return nil
	})
}

func (i *Initializer) InitElasticsearch() *Initializer {
	i.Check("storage.NewElasticsearchClient", func() error {
		cfg := config.GlobalConfig.Elasticsearch
		client, err := storage.NewElasticsearchClient(&storage.ElasticsearchOptions{
			Enabled:   cfg.Enabled,
			Addresses: cfg.Addresses,
			Username:  cfg.Username,
			Password:  cfg.Password,
			Timeout:   cfg.Timeout,
		})
		if err != nil {
			return err
		}
		storage.ES = client
		return nil
	})

	return i.addShutdownHook("Elasticsearch", func() error {
		storage.ES = nil
		return nil
	})
}

func (i *Initializer) InitElasticsearchAllowFail() *Initializer {
	i.CheckAllowFail("storage.NewElasticsearchClient", func() error {
		cfg := config.GlobalConfig.Elasticsearch
		client, err := storage.NewElasticsearchClient(&storage.ElasticsearchOptions{
			Enabled:   cfg.Enabled,
			Addresses: cfg.Addresses,
			Username:  cfg.Username,
			Password:  cfg.Password,
			Timeout:   cfg.Timeout,
		})
		if err != nil {
			return err
		}
		storage.ES = client
		return nil
	})
	if storage.ES != nil {
		i.addShutdownHook("Elasticsearch", func() error {
			storage.ES = nil
			return nil
		})
	}
	return i
}

func (i *Initializer) InitMQ() *Initializer {
	i.Check("mq_service.Init", func() error {
		return mq_service.MQSvc.Init(&config.GlobalConfig.MQConfig)
	})
	return i.addShutdownHook("MQ", mq_service.MQSvc.Close)
}

func (i *Initializer) InitRepository() *Initializer {
	return i.Check("repository.NewRepositories", func() error {
		i.Repositories = repository.NewRepositories(storage.DB, storage.KVS)
		return nil
	})
}

func (i *Initializer) InitService() *Initializer {
	return i.Check("service.Init", func() error {
		return service.Init(config.GlobalConfig, i.Repositories)
	})
}

func (i *Initializer) Err() error {
	return i.err
}

func (i *Initializer) Shutdown() {
	i.shutdownOnce.Do(func() {
		if len(i.shutdownHooks) == 0 {
			return
		}
		reversed := make([]utils.ShutdownHook, 0, len(i.shutdownHooks))
		for idx := len(i.shutdownHooks) - 1; idx >= 0; idx-- {
			reversed = append(reversed, i.shutdownHooks[idx])
		}
		utils.RunShutdownHooks(reversed...)
	})
}

func (i *Initializer) ShutdownOnSignal(delay time.Duration) {
	utils.WaitForShutdownSignal()
	if delay > 0 {
		time.Sleep(delay)
	}
	i.Shutdown()
}

func (i *Initializer) addShutdownHook(name string, fn func() error) *Initializer {
	if i.err != nil || fn == nil || name == "" {
		return i
	}
	if _, exists := i.hookNames[name]; exists {
		return i
	}
	i.hookNames[name] = struct{}{}
	i.shutdownHooks = append(i.shutdownHooks, utils.ShutdownHook{
		Name: name,
		Fn:   fn,
	})
	return i
}