backend/internal/wire/grpc.go

package wire

import (
	"time"

	"with_you/internal/config"
	"with_you/internal/grpc/runner"
	"with_you/internal/pkg/redis"

	"github.com/google/uuid"
	"github.com/google/wire"
	"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
}