box-cost/boxcost/api/print.go

package api

import (
	"box-cost/db/model"
	"box-cost/db/repo"
	randc "crypto/rand"
	"encoding/json"
	"fmt"
	"log"
	"math/big"
	"math/rand"
	"sort"
	"sync"
	"time"
	"unsafe"

	"github.com/gin-gonic/gin"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
)

func Printr(c *gin.Context, apictx *ApiSession) (interface{}, error) {
	loc, _ := time.LoadLocation("Local")
	date := time.Now().Format("2006-01-02")
	startDate := date + " 00:00:00"
	startTime, _ := time.ParseInLocation("2006-01-02 15:04:05", startDate, loc)
	endDate := date + " 23:59:59"
	endTime, _ := time.ParseInLocation("2006-01-02 15:04:05", endDate, loc)
	fmt.Println(startTime.Unix())
	fmt.Println(endTime.Unix())

	return "success", nil

}

// 生成百万数据
func GenData(_ *gin.Context, apictx *ApiSession) (interface{}, error) {
	var wg sync.WaitGroup
	for i := 0; i < 100; i++ {
		wg.Add(1)
		go insertData(apictx, &wg)
	}
	wg.Wait()

	return true, nil
}

// 聚合查询百万数据
func SearchData(_ *gin.Context, apictx *ApiSession) (interface{}, error) {
	dataStatResult := make(chan bson.M)
	var output ResultSlice
	var wg sync.WaitGroup
	// 82个gorutine 执行聚合
	for i := 18; i < 100; i++ {
		wg.Add(1)
		go DataAggregate(apictx, i, dataStatResult, &wg)
	}

	// 从管道中获取聚合结果 保存在数组中
	for value := range dataStatResult {
		output = append(output, OutPut{
			Age:       value["age"].(int32),
			Sex:       value["sex"].(int32),
			Total:     value["total"].(int32),
			AvgSalary: value["avgSalary"].(float64),
		})
		if len(output) == 164 { // 如果大于164，不会跳出；如果小于164，wg.done != wg.add ，wg.wait阻塞
			break
		}
	}

	wg.Wait()
	//倒序排列
	sort.Sort(output)
	for _, v := range output {
		result, err := json.Marshal(&v)
		if err != nil {
			fmt.Println("json.marshal failed, err:", err)
			return nil, err
		}
		fmt.Println(string(result))
	}
	return output, nil
}

func insertData(apictx *ApiSession, wg *sync.WaitGroup) {
	collectName := "aggregate-test"
	rowNum := 10000
	for i := 0; i < rowNum; i++ {
		repo.RepoAddDoc(apictx.CreateRepoCtx(), collectName, &model.AggregateTest{
			Name:   randName(6),
			Age:    randAge(),
			Sex:    randSex(),
			Salary: randSalary(),
		})

	}
	wg.Done()
}

func randSex() *int {
	result, _ := randc.Int(randc.Reader, big.NewInt(2))
	sex := int(result.Int64())
	return &sex
}
func randAge() int {
	result, _ := randc.Int(randc.Reader, big.NewInt(82))
	return int(result.Int64() + 18)
}

func randSalary() int {
	result, _ := randc.Int(randc.Reader, big.NewInt(10000))

	return int(result.Int64() + 2000)
}

const letters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"

var src = rand.NewSource(time.Now().UnixNano())

const (
	// 6 bits to represent a letter index
	letterIdBits = 6
	// All 1-bits as many as letterIdBits
	letterIdMask = 1<<letterIdBits - 1
	letterIdMax  = 63 / letterIdBits
)

func randName(n int) string {
	b := make([]byte, n)
	// A rand.Int63() generates 63 random bits, enough for letterIdMax letters!
	for i, cache, remain := n-1, src.Int63(), letterIdMax; i >= 0; {
		if remain == 0 {
			cache, remain = src.Int63(), letterIdMax
		}
		if idx := int(cache & letterIdMask); idx < len(letters) {
			b[i] = letters[idx]
			i--
		}
		cache >>= letterIdBits
		remain--
	}
	return *(*string)(unsafe.Pointer(&b))
}

func genPipeline(age int) (bson.D, bson.D, bson.D) {

	matchStage := bson.D{
		{"$match", bson.D{
			{"age",
				bson.D{
					{"$eq", age},
				}},
		}},
	}
	groupStage := bson.D{
		{"$group", bson.D{
			{"_id", bson.D{
				{"age", "$age"},
				{"sex", "$sex"},
			}},
			{"age", bson.D{
				{"$first", "$age"},
			}},
			{"sex", bson.D{
				{"$first", "$sex"},
			}},
			{"total", bson.D{
				{"$sum", 1},
			}},
			{"avgSalary", bson.D{
				{"$avg", "$salary"},
			}},
		}},
	}
	projectStage := bson.D{
		{"$project", bson.D{
			{"_id", 0},
			{"age", 1},
			{"sex", 1},
			{"total", 1},
			{"avgSalary", 1},
		}},
	}

	return matchStage, groupStage, projectStage

}

func DataAggregate(apictx *ApiSession, age int, resultChan chan bson.M, wg *sync.WaitGroup) {
	matchStage, groupStage, projectStage := genPipeline(age)

	// opts := options.Aggregate().SetMaxTime(15 * time.Second)
	cursor, err := apictx.Svc.Mongo.GetCollection("aggregate-test").Aggregate(apictx.CreateRepoCtx().Ctx, mongo.Pipeline{matchStage, groupStage, projectStage})
	if err != nil {
		log.Println(err)
	}
	//打印文档内容
	var results []bson.M
	if err = cursor.All(apictx.CreateRepoCtx().Ctx, &results); err != nil {
		log.Println(err)
	}
	log.Printf("%#v\n", results)
	for _, result := range results {
		resultChan <- result
	}

	wg.Done()
}

type OutPut struct {
	Age       int32   `json:"age"`
	Sex       int32   `json:"sex"`
	Total     int32   `json:"total"`
	AvgSalary float64 `json:"avg_salary"`
}

type ResultSlice []OutPut

func (a ResultSlice) Len() int { // 重写 Len() 方法
	return len(a)
}
func (a ResultSlice) Swap(i, j int) { // 重写 Swap() 方法
	a[i], a[j] = a[j], a[i]
}
func (a ResultSlice) Less(i, j int) bool { // 重写 Less() 方法， 从大到小排序
	return a[j].Age < a[i].Age
}