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

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/go-redis/redis"
	"github.com/robfig/cron/v3"
	"github.com/zeromicro/go-queue/kq"
	"github.com/zeromicro/go-zero/core/logx"
	"github.com/zeromicro/go-zero/zrpc"
	"gitlink.org.cn/jcce-pcm/pcm-coordinator/api/internal/config"
	"gitlink.org.cn/jcce-pcm/pcm-coordinator/pkg/utils"
	"gitlink.org.cn/jcce-pcm/pcm-participant-ac/hpcacclient"
	"gitlink.org.cn/jcce-pcm/pcm-participant-ceph/cephclient"
	"gitlink.org.cn/jcce-pcm/pcm-participant-modelarts/modelartsclient"
	"gitlink.org.cn/jcce-pcm/pcm-participant-octopus/octopusclient"
	"gitlink.org.cn/jcce-pcm/pcm-participant-openstack/openstackclient"
	"gitlink.org.cn/jcce-pcm/pcm-participant-slurm/hpcthclient"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/schema"
)

type ServiceContext struct {
	Config              config.Config
	RedisClient         *redis.Client
	ScheduleHpcClient   *kq.Pusher
	ScheduleCloudClient *kq.Pusher
	ScheduleAiClient    *kq.Pusher
	Cron                *cron.Cron
	ModelArtsRpc        modelartsclient.ModelArts
	DbEngin             *gorm.DB
	ACRpc               hpcacclient.HpcAC
	THRpc               hpcthclient.HpcTH
	OctopusRpc          octopusclient.Octopus
	CephRpc             cephclient.Ceph
	OpenstackRpc        openstackclient.Openstack
	DockerClient        *client.Client
	Downloader          *s3manager.Downloader
	Uploader            *s3manager.Uploader
}

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: ", err)
		panic("InitSnowflake err")
	}
	downloader := s3manager.NewDownloader(session)
	uploader := s3manager.NewUploader(session)
	//启动Gorm支持
	dbEngin, _ := gorm.Open(mysql.Open(c.DB.DataSource), &gorm.Config{
		NamingStrategy: schema.NamingStrategy{
			SingularTable: true, // 使用单数表名，启用该选项，此时，`User` 的表名应该是 `t_user`
		},
	})
	dockerClient, err := client.NewClientWithOpts()
	if err != nil {
		logx.Error(err.Error())
		return nil
	}
	return &ServiceContext{
		Cron:    cron.New(cron.WithSeconds()),
		DbEngin: dbEngin,
		Config:  c,
		RedisClient: redis.NewClient(&redis.Options{
			Addr:     c.Redis.Host,
			Password: c.Redis.Pass,
		}),
		ScheduleHpcClient:   kq.NewPusher(c.KqProducerConf.Brokers, c.KqProducerConf.HpcTopic),
		ScheduleCloudClient: kq.NewPusher(c.KqProducerConf.Brokers, c.KqProducerConf.CloudTopic),
		ScheduleAiClient:    kq.NewPusher(c.KqProducerConf.Brokers, c.KqProducerConf.AiTopic),
		ModelArtsRpc:        modelartsclient.NewModelArts(zrpc.MustNewClient(c.ModelArtsRpcConf)),
		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)),
		DockerClient:        dockerClient,
		Downloader:          downloader,
		Uploader:            uploader,
	}
}