110 lines
2.7 KiB
Go
110 lines
2.7 KiB
Go
package wire
|
||
|
||
import (
|
||
"time"
|
||
|
||
"with_you/internal/config"
|
||
"with_you/internal/grpc/runner"
|
||
"with_you/internal/pkg/redis"
|
||
|
||
"github.com/google/wire"
|
||
"github.com/google/uuid"
|
||
"go.uber.org/zap"
|
||
)
|
||
|
||
// GRPCSet gRPC 相关的 Provider Set
|
||
var GRPCSet = wire.NewSet(
|
||
ProvideGRPCServer,
|
||
ProvideRunnerHub,
|
||
ProvideTaskManager,
|
||
ProvideRunnerRegistry,
|
||
ProvideTaskDispatcher,
|
||
)
|
||
|
||
// ProvideGRPCServer 提供 gRPC 服务器
|
||
func ProvideGRPCServer(cfg *config.Config, logger *zap.Logger, hub *runner.RunnerHub) *runner.Server {
|
||
return runner.NewServerWithDeps(&cfg.GRPC, logger, hub)
|
||
}
|
||
|
||
// ProvideRunnerHub 创建 RunnerHub
|
||
func ProvideRunnerHub(logger *zap.Logger) *runner.RunnerHub {
|
||
return runner.NewRunnerHub(nil, logger)
|
||
}
|
||
|
||
// ProvideTaskManager 提供 TaskManager
|
||
func ProvideTaskManager(hub *runner.RunnerHub, logger *zap.Logger) *runner.TaskManager {
|
||
tm := runner.NewTaskManager(nil, 60*time.Second, logger)
|
||
tm.SetDispatcher(hub)
|
||
hub.SetTaskManager(tm)
|
||
return tm
|
||
}
|
||
|
||
// ProvideRunnerRegistry 提供 Runner 注册中心
|
||
// standalone 模式或 Redis 不可用时返回 nil
|
||
func ProvideRunnerRegistry(cfg *config.Config, redisClient *redis.Client) *runner.RunnerRegistry {
|
||
if cfg.Runner.Mode != "cluster" || redisClient == nil {
|
||
return nil
|
||
}
|
||
|
||
instanceID := cfg.Runner.Cluster.InstanceID
|
||
if instanceID == "" {
|
||
instanceID = uuid.New().String()
|
||
}
|
||
|
||
rdb := redisClient.GetClient()
|
||
|
||
return runner.NewRunnerRegistry(
|
||
rdb,
|
||
instanceID,
|
||
cfg.Runner.Cluster.RegistryPrefix,
|
||
cfg.Runner.Cluster.CapPrefix,
|
||
time.Duration(cfg.Runner.Cluster.RegistryTTL)*time.Second,
|
||
time.Duration(cfg.Runner.Cluster.HeartbeatInterval)*time.Second,
|
||
zap.L(),
|
||
)
|
||
}
|
||
|
||
// ProvideTaskDispatcher 提供任务分发器
|
||
// standalone 模式返回 *RunnerHub，cluster 模式返回 *TaskBus
|
||
func ProvideTaskDispatcher(
|
||
hub *runner.RunnerHub,
|
||
registry *runner.RunnerRegistry,
|
||
cfg *config.Config,
|
||
redisClient *redis.Client,
|
||
logger *zap.Logger,
|
||
) runner.TaskDispatcher {
|
||
if cfg.Runner.Mode != "cluster" || redisClient == nil || registry == nil {
|
||
logger.Info("Runner running in standalone mode")
|
||
return hub
|
||
}
|
||
|
||
instanceID := cfg.Runner.Cluster.InstanceID
|
||
if instanceID == "" {
|
||
instanceID = uuid.New().String()
|
||
}
|
||
|
||
rdb := redisClient.GetClient()
|
||
|
||
bus := runner.NewTaskBus(
|
||
hub,
|
||
rdb,
|
||
registry,
|
||
instanceID,
|
||
cfg.Runner.Cluster.TaskChannel,
|
||
cfg.Runner.Cluster.ResultChannel,
|
||
logger,
|
||
)
|
||
if err := bus.Start(); err != nil {
|
||
logger.Error("Failed to start Runner Bus, falling back to standalone", zap.Error(err))
|
||
return hub
|
||
}
|
||
|
||
// 注入 Bus 和 Registry 到 Hub
|
||
hub.SetBus(bus)
|
||
hub.SetRegistry(registry)
|
||
registry.StartHeartbeat()
|
||
|
||
logger.Info("Runner running in cluster mode", zap.String("instance_id", bus.InstanceID()))
|
||
return bus
|
||
}
|