wangspfly / ratelimit Goto Github PK

package main

import (
	"fmt"
	"log"
	"strconv"
	"sync"
	"time"

	"github.com/yudeguang/ratelimit"
)

var num int //因并发问题num比实际数量小

func main() {
	log.SetFlags(log.Lshortfile | log.Ltime)
	test1()
	test2()
}

func test1() {
	fmt.Println("\r\n测试1,加载备份数据并自动定期备份以及性能测试:")
	//初始化访问规则，支持用多个规则形成的复杂规则，规则必须至少包含一条规则
	//支持从本地磁盘加载备份好的历史记录，并支持定期自动备份
	r := ratelimit.NewRule()
	r.AddRule(time.Second*10, 5)                               //每10秒只允许访问5次
	r.AddRule(time.Minute*30, 50)                              //每30分钟只允许访问50次
	r.AddRule(time.Hour*24, 500)                               //每天只允许访问500次
	err := r.LoadingAndAutoSaveToDisc("test1", time.Second*10) //设置10秒备份一次(不填写则默认60秒备份一次)，备份到程序当前文件夹下，文件名为test1.ratelimit
	if err == nil {
		log.Println("加载历史访问记录成功")
	} else {
		log.Println(err)
	}
	//构建若干个用户，模拟用户访问
	var users = make(map[string]bool)
	for i := 1; i < 1000; i++ {
		users["user_"+strconv.Itoa(i)] = true
	}
	begin := time.Now()
	//模拟多个协程访问
	chanNum := 200
	var wg sync.WaitGroup
	wg.Add(chanNum)
	for i := 0; i < chanNum; i++ {
		go func(i int, wg *sync.WaitGroup) {
			for ii := 0; ii < 50; ii++ {
				for user := range users {
					for {
						if r.AllowVisit(user) {
							num++
						} else {
							num++
							break
						}
					}
				}
			}
			wg.Done()
		}(i, &wg)
	}
	//所有线程结束，完工
	wg.Wait()

	log.Println("性能测试完成，每秒约完成:", num/int(time.Now().Sub(begin).Seconds()), "次操作")
	err = r.SaveToDiscOnce("test2") //在自动备份的同时，还支持手动备份，一般在程序要退出时调用此函数
	if err == nil {
		log.Println("完成手动数据备份")
	} else {
		log.Println(err)
	}
}
func test2() {
	fmt.Println("\r\n测试2，模拟用户访问并打印:")
	r := ratelimit.NewRule()
	r.AddRule(time.Second*10, 5)                        //每10秒只允许访问5次
	r.AddRule(time.Minute*30, 50)                       //每30分钟只允许访问50次
	r.AddRule(time.Hour*24, 500)                        //每天只允许访问500次
	r.LoadingAndAutoSaveToDisc("test2", time.Second*10) //设置10秒备份一次(不填写则默认60秒备份一次)，备份到程序当前文件夹下，文件名为test2.ratelimit
	//构建若干个用户，模拟用户访问
	users := []string{"andy", "小余", "130x"}
	for _, user := range users {
		fmt.Println("\r\n开始模拟以下用户访问:", user)
		for {
			if r.AllowVisit(user) {
				log.Println(user, "访问1次,剩余:", r.RemainingVisits(user))
			} else {
				log.Println(user, "访问过多,稍后再试")
				break
			}
			time.Sleep(time.Second * 1)
		}
	}

	//打印所有用户访问数据情况
	fmt.Println("\r\n开始打印所有用户在相关时间段内详细的剩余访问次数情况:\r\n")
	for _, user := range users {
		fmt.Println(user)
		fmt.Println("     概述:", r.RemainingVisits(user))
		fmt.Println("     具体:")
		r.PrintRemainingVisits(user)
		fmt.Println("")
	}
}