airobot/lib/ai.go

package lib

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/sirupsen/logrus"
	"legu.airobot/storage"
)

type myAI struct {
	robots        []*Robot
	scenes        []*scene
	iscenes       []IScene
	tickets       Tickets //票池
	Obs           Observer
	useCount      uint32                         //计数(压入的用户数)
	useCountTotal uint32                         //总数
	lock          sync.Mutex                     //合并数据锁
	config        *storage.Config                //配置
	ReportMap     map[int]map[string]*Statistics //测试报告 key1:场景 key2:协议
}

func NewAI(aip AIParam) (*myAI, error) {
	if err := aip.Check(); err != nil {
		return nil, err
	}

	ai := &myAI{
		Obs:       NewObserver(),
		scenes:    make([]*scene, 0),
		config:    aip.Config,
		iscenes:   aip.Scenes,
		ReportMap: make(map[int]map[string]*Statistics),
	}

	if err := ai.init(); err != nil {
		return nil, err
	}

	return ai, nil
}

func (m *myAI) init() error {
	var buf bytes.Buffer
	buf.WriteString("初始化AI")

	uct := m.config.Global.UserCountTotal
	tickets, err := NewTickets(uct)
	if err != nil {
		return err
	}

	m.tickets = tickets

	buf.WriteString(fmt.Sprintf("完成 用户数量:%d", uct))
	logrus.Debug(buf.String())
	return nil
}

func (m *myAI) AppendRobot(robot *Robot) {
	defer m.lock.Unlock()
	m.lock.Lock()
	m.robots = append(m.robots, robot)
}

func (m *myAI) Start() bool {
	if len(m.config.Scenes) == 0 {
		logrus.Warn("还未设置场景")
		return false
	}
	logrus.Info("测试中...")
	go func() {
		for {
			m.tickets.Take()
			uCount := atomic.LoadUint32(&m.useCount)
			if uCount >= uint32(m.config.Global.UserCount) {
				atomic.StoreUint32(&m.useCount, 0)
				time.Sleep(time.Duration(m.config.Global.IntervalS) * time.Second)
			}
			//加入机器人
			atomic.AddUint32(&m.useCount, 1)
			robot := NewRobot(m.config)
			robot.SetScenes(m.iscenes)
			m.AppendRobot(robot)
			robot.Start()
			atomic.AddUint32(&m.useCountTotal, 1)
			m.Obs.Notify(EVENT_PROGRESS, int32(1))
		}
	}()

	go func() {
		start := time.Now()
		for {
			total := atomic.LoadUint32(&m.useCountTotal)
			if total == m.config.Global.UserCountTotal {
				elipse := time.Since(start)
				logrus.Info("开始生成测试报告")
				m.MergeResult()
				m.genReport(elipse)
				m.Obs.Notify(EVENT_PROGRESS, int32(0))
				break
			}
		}
	}()

	return true
}

func (m *myAI) Stop() {

}

func (m *myAI) MergeResult() {
	m.lock.Lock()
	defer m.lock.Unlock()
	n := make(map[int]map[string]*Statistics)
	for _, r := range m.robots {
		if len(n) == 0 {
			n = r.ReportMap
		} else {
			x := n           //已存在的
			y := r.ReportMap //将要合并的

			for i, v := range x {
				for j, w := range y {
					//场景相同
					if i == j {
						for k, a := range v {
							for l, b := range w {
								//判断协议是否相同
								if k == l {
									statis := &Statistics{}
									statis.ElapseTotal = time.Duration(int64(a.ElapseTotal + b.ElapseTotal))
									statis.AvgElapse = time.Duration(int64(a.AvgElapse+b.AvgElapse) / int64(2))
									statis.CallCount = (a.CallCount + b.CallCount)
									max := a.MaxElapse
									if max < b.MaxElapse {
										max = b.MaxElapse
									}
									statis.MaxElapse = max
									min := a.MinElapse
									if min > b.MinElapse {
										min = b.MinElapse
									}
									statis.MinElapse = min
									statis.Route = a.Route
									statis.SceneName = a.SceneName
									if n[i][k] == nil {
										n[i] = make(map[string]*Statistics)
									}
									n[i][k] = statis
								} else {
									if _, ok := n[i][k]; !ok {
										n[i][k] = a
									}
									if _, ok := n[i][l]; !ok {
										n[i][l] = b
									}
								}
							}
						}

					} else {
						if _, ok := n[i]; !ok {
							n[i] = v
						}
						if _, ok := n[j]; !ok {
							n[j] = w
						}
					}
				}
			}
		}
	}

	m.ReportMap = n
}

func (m *myAI) genReport(elipse time.Duration) {
	var msgs []string
	var i []int
	for key, _ := range m.ReportMap {
		i = append(i, key)
	}
	sort.Ints(i)
	for _, routes := range i {
		for r, d := range m.ReportMap[routes] {
			msgs = append(msgs, fmt.Sprintf("【%s】协议:%s 调用次数:%d 总耗时:%v 平均耗时:%v 最大耗时:%v 最小耗时:%v",
				d.SceneName, r, d.CallCount, d.ElapseTotal.String(), d.AvgElapse.String(), d.MaxElapse.String(), d.MinElapse.String()))
		}
	}
	record := strings.Join(msgs, "\n")
	var buf bytes.Buffer
	buf.WriteString("测试报告\n")
	buf.WriteString(fmt.Sprintf("用户总数:%d 单次投放人数:%d 投放间隔时间:%ds 共计时间:%s\n",
		m.config.Global.UserCountTotal, m.config.Global.UserCount, m.config.Global.IntervalS, elipse.String()))
	buf.WriteString(record)
	buf.WriteString("\n------------------------------------------------------------------------------------------------------\n")
	// logrus.WithField("res", buf.String()).Debug("报告内容")
	file, err := os.OpenFile(filepath.Join("./", "report.log"), os.O_TRUNC|os.O_CREATE, os.ModePerm)
	if err != nil {
		logrus.Error(err)
	}
	defer file.Close()

	if _, err := file.WriteString(buf.String()); err != nil {
		logrus.Error(err)
	}
	logrus.Info("已生成测试报告")
	m.Obs.Notify(EVENT_FLAG, true)
}