token-bucket-test.go

token-bucket-test.go

package main

import (
	"context"
	"log"
	"time"

	"github.com/bsm/ratelimit"
	"github.com/tsenart/tb"
	// jujuRatelimit "github.com/juju/ratelimit"
	wrflyRatelimit "github.com/wrfly/testing-kit/utils/tokenbucket"
)

func main() {
	bkt := ratelimit.New(1000, time.Second)
	wrflyBkt := wrflyRatelimit.New(1000, time.Second)
	tbBkt := tb.NewBucket(1000, time.Second)
	// jujuBkt := jujuRatelimit.NewBucketWithRate(1000, 1000)

	var take, drop = 0, 0
	var start time.Time

	log.Println("5s test")

	drop = 0
	take = 0
	ctx, cancel := context.WithCancel(context.Background())
	go func() {
		start = time.Now()
		for ctx.Err() == nil {
			if bkt.Limit() {
				drop++
				continue
			}
			take++
		}
		log.Println("bsm: ", time.Since(start).Seconds())
		log.Println("take:", take)
		log.Println("drop:", drop)
		log.Println()
	}()
	time.Sleep(5 * time.Second)
	cancel()
	time.Sleep(time.Second)

	drop = 0
	take = 0
	ctx, cancel = context.WithCancel(context.Background())
	go func() {
		start = time.Now()
		for ctx.Err() == nil {
			if wrflyBkt.TakeOne() {
				take++
				continue
			}
			drop++
		}
		log.Println("wrfly: ", time.Since(start).Seconds())
		log.Println("take:", take)
		log.Println("drop:", drop)
		log.Println()
	}()
	time.Sleep(5 * time.Second)
	cancel()
	time.Sleep(time.Second)

	drop = 0
	take = 0
	ctx, cancel = context.WithCancel(context.Background())
	go func() {
		start = time.Now()
		for ctx.Err() == nil {
			if tbBkt.Take(1) == 1 {
				take++
				continue
			}
			drop++
		}
		log.Println("tb: ", time.Since(start).Seconds())
		log.Println("take:", take)
		log.Println("drop:", drop)
		log.Println()
	}()
	time.Sleep(5 * time.Second)
	cancel()
	time.Sleep(time.Second)

	end := int(1e8)
	log.Println("range test: ", end)
	start = time.Now()
	for i := 0; i < end; i++ {
		take++
		drop++
	}
	log.Println("dry run: ", time.Since(start).Seconds())
	log.Println()

	drop = 0
	take = 0
	start = time.Now()
	for i := 0; i < end; i++ {
		if bkt.Limit() {
			drop++
			continue
		}
		take++
	}
	log.Println("bsm: ", time.Since(start).Seconds())
	log.Println("take:", take)
	log.Println("drop:", drop)
	log.Println()

	drop = 0
	take = 0
	start = time.Now()
	for i := 0; i < end; i++ {
		if wrflyBkt.TakeOne() {
			take++
			continue
		}
		drop++
	}
	log.Println("wrfly: ", time.Since(start).Seconds())
	log.Println("take:", take)
	log.Println("drop:", drop)
	log.Println()

	drop = 0
	take = 0
	start = time.Now()
	for i := 0; i < end; i++ {
		if tbBkt.Take(1) == 1 {
			take++
			continue
		}
		drop++
	}
	log.Println("tb: ", time.Since(start).Seconds())
	log.Println("take:", take)
	log.Println("drop:", drop)
	log.Println()
}