wheat-cache/pkg/middle/middleware.go

package middle

import (
	"context"

	_ "gitee.com/timedb/wheatCache/conf"
	"gitee.com/timedb/wheatCache/pkg/event"
	"gitee.com/timedb/wheatCache/pkg/logx"
	middleMsg "gitee.com/timedb/wheatCache/pkg/middle-msg"
	"gitee.com/timedb/wheatCache/plugins"
	"gitee.com/timedb/wheatCache/plugins/config"
	"github.com/spf13/viper"
)

type MiddleWare struct {
	eventDriver   event.DriverInterface
	eventConsumer event.ConsumerInterface
	eventProduce  event.ProduceInterface
	plugins       map[string][]plugins.MiddleToolsInterface
	consumerCount int
	driverCount   int
}

func NewMiddleWare() *MiddleWare {
	oneMiddle.Do(func() {
		consumerCount, driverCount := loadConfigAndDefault()

		driver := event.NewDriver(driverCount)
		middleWareDriver = &MiddleWare{
			eventDriver:   driver,
			eventConsumer: event.NewConsumer(driver),
			eventProduce:  event.NewProduce(driver),
			driverCount:   driverCount,
			consumerCount: consumerCount,
		}
		middleWareDriver.loadPlugins()

		middleWareDriver.startWork()

	})
	return middleWareDriver
}

func (m *MiddleWare) GetEventDriver() event.DriverInterface {
	return m.eventDriver
}

func (m *MiddleWare) loadPlugins() {
	plug := viper.GetStringMapStringSlice("plugins-control")

	pluginsMap := config.GetMiddlewareMap()

	pluginsContext := make(map[string][]plugins.MiddleToolsInterface)

	for middleMsg, pluNames := range plug {
		pulgSingle := make([]plugins.MiddleToolsInterface, 0)
		for _, name := range pluNames {
			pulgSingle = append(pulgSingle, pluginsMap[name])
		}

		pluginsContext[middleMsg] = pulgSingle
	}

	m.plugins = pluginsContext
}

func loadConfigAndDefault() (int, int) {
	// 加载 consumerCount
	consumerCount := viper.GetInt("middle-driver.middleConsumerCount")
	if consumerCount == 0 {
		consumerCount = defaultConsumerCount
	}
	driverCount := viper.GetInt("middle-driver.driverCount")
	if driverCount == 0 {
		driverCount = defaultDriverCount
	}
	return consumerCount, driverCount
}

func (m *MiddleWare) startWork() {

	for i := 0; i < m.consumerCount; i++ {
		go func() {
			ctx := context.Background()
			for {
				workEvent := m.eventConsumer.Receive(ctx)
				plugs := m.plugins[workEvent.GetEventName()]
				msg, ok := workEvent.GetValue(middleMsg.MiddleMsgKey)
				if !ok {
					logx.With(ctx, m.eventProduce).Error("get event value err，not key:%s", middleMsg.MiddleMsgKey)
					continue
				}
				for _, val := range plugs {
					val.Exec(msg)
				}
			}
		}()
	}
}