yudeguang / ratelimit Goto Github PK

package main

import (
	"fmt"
	"github.com/yudeguang/ratelimit"
	"log"
	"strconv"
	"sync"
	"time"
)

var userVisitRule visitRule

type visitRule struct {
	paidMember      *ratelimit.Rule
	freeMember      *ratelimit.Rule
	anonymousMember *ratelimit.Rule
}

func main() {
	log.SetFlags(log.Lshortfile | log.Ltime)
	example1()
	example2()
	example3()
}

// 简单规则案例
func example1() {
	//步骤一：初始化
	r := ratelimit.NewRule()
	//步骤二：增加一条或者多条规则组成复合规则，此复合规则必须至少包含一条规则
	r.AddRule(time.Second*1, 10)
	//步骤三：调用函数判断某用户是否允许访问   allow:= r.AllowVisit(user)
	for i := 0; i <= 20; i++ {
		allow := r.AllowVisit("user")
		if !allow {
			log.Println("访问量超出,其剩余访问次数情况如下:", r.RemainingVisits("user"))
		} else {
			log.Println("允许访问,其剩余访问次数情况如下:", r.RemainingVisits("user"))
		}
	}
}

// 模拟平台运营中针对不同级别用户，设定不同的访问频率控制
func example2() {
	//paidMember
	//步骤一：初始化
	userVisitRule.paidMember = ratelimit.NewRule()
	//步骤二：增加一条或者多条规则组成复合规则，此复合规则必须至少包含一条规则
	userVisitRule.paidMember.AddRule(time.Hour*24, 10000)
	userVisitRule.paidMember.AddRule(time.Hour*1, 1000)
	userVisitRule.paidMember.AddRule(time.Minute*1, 100)
	userVisitRule.paidMember.AddRule(time.Second*1, 20)
	//步骤三(可选):从本地磁盘加载历史访问数据,并指定10秒存储一次。不加载则表示不磁盘,程序中断后，访问数据将消失
	userVisitRule.paidMember.LoadingAndAutoSaveToDisc("userVisitRule_paidMember", time.Second*10)
	//freeMember
	userVisitRule.freeMember = ratelimit.NewRule()
	userVisitRule.freeMember.AddRule(time.Hour*24, 1000)
	userVisitRule.freeMember.AddRule(time.Hour*1, 100)
	userVisitRule.freeMember.AddRule(time.Minute*1, 10)
	userVisitRule.freeMember.AddRule(time.Second*1, 2)
	userVisitRule.freeMember.LoadingAndAutoSaveToDisc("userVisitRule_freemember", time.Second*10)
	//anonymousMember
	userVisitRule.anonymousMember = ratelimit.NewRule()
	userVisitRule.anonymousMember.AddRule(time.Hour*24, 100)
	userVisitRule.anonymousMember.AddRule(time.Second*1, 2)
	userVisitRule.anonymousMember.LoadingAndAutoSaveToDisc("userVisitRule_anonymousMember", time.Second*10)

	//模拟一定数量不同类型的用户，比如可用手机号或者用户名做为KEY，匿名用户，一般用IP作为KEY
	var paidMembers = []string{"17277777770", "17277777771", "17277777772", "17277777773", "17277777774", "17277777775", "17277777776"}
	var freeMembers = []string{"16277777770", "16277777771", "16277777772", "16277777773", "16277777774", "16277777775", "16277777776"}
	var anonymousMembers = []string{"192.168.0.2", "192.168.0.3", "192.168.0.4", "192.168.0.5", "192.168.0.6", "192.168.0.7", "192.168.0.8"}

	//步骤四：调用函数判断某用户是否允许访问
	/*
	   allow:= r.AllowVisit(user)
	*/
	fmt.Println("\r\n下面模拟单个用户持续访问的情况:")
	member := paidMembers[0]
	for i := 0; i <= 30; i++ {
		allow := userVisitRule.paidMember.AllowVisit(member)
		if !allow {
			log.Println(member, "访问量超出,其剩余访问次数情况如下:", userVisitRule.paidMember.RemainingVisits(member))
		} else {
			log.Println(member, "允许访问,其剩余访问次数情况如下:", userVisitRule.paidMember.RemainingVisits(member))
		}
	}
	//我们再等一段时间，看看paidMembers[0]这个用户的允许访问是否恢复
	fmt.Println("\r\n休息10秒后,该用户将恢复部分访问次数")
	time.Sleep(time.Second * 10)
	allow := userVisitRule.paidMember.AllowVisit(member)
	if !allow {
		log.Println(member, "访问量超出,其剩余访问次数情况如下:", userVisitRule.paidMember.RemainingVisits(member))
	} else {
		log.Println(member, "允许访问,其剩余访问次数情况如下:", userVisitRule.paidMember.RemainingVisits(member))
		fmt.Println("\r\n", userVisitRule.paidMember.RemainingVisits(member), "的具体含义为:")
		userVisitRule.paidMember.PrintRemainingVisits(member)
	}
	fmt.Println("\r\n手工清空", member, "后的访问记录")
	userVisitRule.paidMember.ManualEmptyVisitorRecordsOf(member)
	log.Println(member, "其剩余访问次数情况如下:", userVisitRule.paidMember.RemainingVisits(member))

	//下面模拟三种不同的用户，访问一些页面，之后再打印出来观察
	for i := 0; i <= 10; i++ {
		for _, v := range freeMembers {
			userVisitRule.paidMember.AllowVisit(v)
		}
	}
	for i := 0; i <= 10; i++ {
		for _, v := range freeMembers {
			userVisitRule.freeMember.AllowVisit(v)
		}
	}
	for i := 0; i <= 10; i++ {
		for _, v := range anonymousMembers {
			userVisitRule.anonymousMember.AllowVisit(v)
		}
	}
	//在实际的运营中，GetCurOnlineUsersVisitsDetail()函数可以自行包装以HTTP等形式输出
	fmt.Println("\r\n下面为现在付费用户剩余访问次数情况:")
	paidMemberDeatil := userVisitRule.paidMember.GetCurOnlineUsersVisitsDetail()
	for _, v := range paidMemberDeatil {
		log.Println(v)
	}
	fmt.Println("\r\n下面为现在免费用户剩余访问次数情况:")
	freeMemberDeatil := userVisitRule.freeMember.GetCurOnlineUsersVisitsDetail()
	for _, v := range freeMemberDeatil {
		log.Println(v)
	}
	fmt.Println("\r\n下面为现为匿名S用户剩余访问次数情况:")
	anonymousMemberDeatil := userVisitRule.anonymousMember.GetCurOnlineUsersVisitsDetail()
	for _, v := range anonymousMemberDeatil {
		log.Println(v)
	}

}

// 模拟1000个用户，累计进行总共约1亿次性能测试
func example3() {
	var Visits int //因并发问题num比实际数量稍小
	fmt.Println("\r\n性能测试，预计耗时1分钟，请耐心等待:")
	//步骤一：初始化
	r := ratelimit.NewRule()
	//步骤二：增加一条或者多条规则组成复合规则，规则必须至少包含一条规则
	//此处对于性能测试，为方便准确计数，只需要添加一条规则
	r.AddRule(time.Second*10, 1000) //每10秒只允许访问1000次
	/*
		r.AddRule(time.Second*10, 10)   //每10秒只允许访问10次
		r.AddRule(time.Minute*30, 1000) //每30分钟只允许访问1000次
		r.AddRule(time.Hour*24, 5000)   //每天只允许访问500次
	*/
	//步骤三(可选):从本地磁盘加载历史访问数据
	r.LoadingAndAutoSaveToDisc("example2", time.Second*10) //设置10秒备份一次(不填写则默认60秒备份一次)，备份到程序当前文件夹下，文件名为test1.ratelimit
	log.Println("性能测试正式开始")
	//步骤四：调用函数判断某用户是否允许访问
	/*
	   allow:= r.AllowVisit(user)
	*/
	//构建若干个用户，模拟用户访问
	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 < 5000; ii++ {
				for user := range users {
					for {
						Visits++
						if !r.AllowVisit(user) {
							break
						}
					}
				}
			}
			wg.Done()
		}(i, &wg)
	}
	//所有线程结束，完工
	wg.Wait()
	t := int(time.Now().Sub(begin).Seconds())
	log.Println("性能测试完成:共计访问", Visits, "次,", "耗时", t, "秒,即每秒约完成", Visits/t, "次操作")
	//步骤五:程序退出前主动手动存盘
	err := r.SaveToDiscOnce() //在自动备份的同时，还支持手动备份，一般在程序要退出时调用此函数
	if err == nil {
		log.Println("完成手动数据备份")
	} else {
		log.Println(err)
	}
}