pcm-coordinator/api/internal/svc/servicecontext.go

/*

 Copyright (c) [2023] [pcm]
 [pcm-coordinator] is licensed under Mulan PSL v2.
 You can use this software according to the terms and conditions of the Mulan PSL v2.
 You may obtain a copy of Mulan PSL v2 at:
          http://license.coscl.org.cn/MulanPSL2
 THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 EITHER EXPaRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 See the Mulan PSL v2 for more details.

*/

package svc

import (
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"github.com/docker/docker/client"
	"github.com/redis/go-redis/v9"
	"github.com/robfig/cron/v3"
	"github.com/zeromicro/go-zero/core/logx"
	"github.com/zeromicro/go-zero/zrpc"
	"gitlink.org.cn/jcce-pcm/pcm-ac/hpcacclient"
	"gitlink.org.cn/jcce-pcm/pcm-coordinator/api/internal/config"
	"gitlink.org.cn/jcce-pcm/pcm-coordinator/pkg/tracker"
	"gitlink.org.cn/jcce-pcm/pcm-coordinator/pkg/utils"
	"gitlink.org.cn/jcce-pcm/pcm-coordinator/rpc/client/participantservice"
	"gitlink.org.cn/jcce-pcm/pcm-participant-ceph/cephclient"
	"gitlink.org.cn/jcce-pcm/pcm-participant-kubernetes/kubernetesclient"
	"gitlink.org.cn/jcce-pcm/pcm-participant-modelarts/client/imagesservice"
	"gitlink.org.cn/jcce-pcm/pcm-participant-modelarts/client/modelartsservice"
	"gitlink.org.cn/jcce-pcm/pcm-participant-octopus/octopusclient"
	"gitlink.org.cn/jcce-pcm/pcm-participant-openstack/openstackclient"
	"gitlink.org.cn/jcce-pcm/pcm-slurm/slurmclient"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/schema"
	"time"
)

type ServiceContext struct {
	Config          config.Config
	RedisClient     *redis.Client
	Cron            *cron.Cron
	ModelArtsRpc    modelartsservice.ModelArtsService
	ModelArtsImgRpc imagesservice.ImagesService
	DbEngin         *gorm.DB
	ACRpc           hpcacclient.HpcAC
	THRpc           slurmclient.Slurm
	OctopusRpc      octopusclient.Octopus
	CephRpc         cephclient.Ceph
	OpenstackRpc    openstackclient.Openstack
	DockerClient    *client.Client
	Downloader      *s3manager.Downloader
	Uploader        *s3manager.Uploader
	K8sRpc          kubernetesclient.Kubernetes
	PromClient      map[int64]tracker.Prometheus
	ParticipantRpc  participantservice.ParticipantService
}

func NewServiceContext(c config.Config) *ServiceContext {
	// 创建s3 session
	session, _ := session.NewSession(&aws.Config{
		Credentials:      credentials.NewStaticCredentials(c.MinioConf.AccessKey, c.MinioConf.Secret, ""), //使用静态凭据,硬编码
		Endpoint:         aws.String(c.MinioConf.Endpoint),                                                //配置端点
		Region:           aws.String("default"),                                                           //配置区域
		DisableSSL:       aws.Bool(false),                                                                 //是否禁用https,这里表示不禁用,即使用HTTPS
		S3ForcePathStyle: aws.Bool(true),                                                                  //使用路径样式而非虚拟主机样式,区别请参考:https://docs.aws.amazon.com/AmazonS3/latest/dev/VirtualHosting.html
	})
	//添加snowflake支持
	err := utils.InitSnowflake(c.SnowflakeConf.MachineId)
	if err != nil {
		logx.Errorf("InitSnowflake err: %v", err)
		panic("InitSnowflake err")
	}
	downloader := s3manager.NewDownloader(session)
	uploader := s3manager.NewUploader(session)
	//启动Gorm支持
	dbEngin, err := gorm.Open(mysql.Open(c.DB.DataSource), &gorm.Config{
		NamingStrategy: schema.NamingStrategy{
			SingularTable: true, // 使用单数表名，启用该选项，此时，`User` 的表名应该是 `t_user`
		},
	})
	if err != nil {
		logx.Errorf("数据库连接失败, err%v", err)
		panic(err)
	}
	sqlDB, err := dbEngin.DB()
	// SetMaxIdleConns 设置空闲连接池中连接的最大数量
	sqlDB.SetMaxIdleConns(10)

	// SetMaxOpenConns 设置打开数据库连接的最大数量。
	sqlDB.SetMaxOpenConns(50)

	// SetConnMaxLifetime 设置了连接可复用的最大时间。
	sqlDB.SetConnMaxLifetime(time.Hour)

	dockerClient, err := client.NewClientWithOpts()
	if err != nil {
		logx.Error(err.Error())
		return nil
	}
	redisClient := redis.NewClient(&redis.Options{
		Addr:     c.Redis.Host,
		Password: c.Redis.Pass,
	})
	return &ServiceContext{
		Cron:            cron.New(cron.WithSeconds()),
		DbEngin:         dbEngin,
		Config:          c,
		RedisClient:     redisClient,
		ModelArtsRpc:    modelartsservice.NewModelArtsService(zrpc.MustNewClient(c.ModelArtsRpcConf)),
		ModelArtsImgRpc: imagesservice.NewImagesService(zrpc.MustNewClient(c.ModelArtsImgRpcConf)),
		CephRpc:         cephclient.NewCeph(zrpc.MustNewClient(c.CephRpcConf)),
		ACRpc:           hpcacclient.NewHpcAC(zrpc.MustNewClient(c.ACRpcConf)),
		OctopusRpc:      octopusclient.NewOctopus(zrpc.MustNewClient(c.OctopusRpcConf)),
		OpenstackRpc:    openstackclient.NewOpenstack(zrpc.MustNewClient(c.OpenstackRpcConf)),
		K8sRpc:          kubernetesclient.NewKubernetes(zrpc.MustNewClient(c.K8sNativeConf)),
		PromClient:      make(map[int64]tracker.Prometheus),
		ParticipantRpc:  participantservice.NewParticipantService(zrpc.MustNewClient(c.PcmCoreRpcConf)),
		DockerClient:    dockerClient,
		Downloader:      downloader,
		Uploader:        uploader,
	}
}