<commit_before>package ratelimit

import (
	"math"
	"time"
)

const resolution = 1 * time.Millisecond

type Bucket struct {
	ts       time.Time
	rate     int64
	capacity int64
	avail    int64
}

func NewBucket(rate int64, capacity int64) *Bucket {
	b := &Bucket{
		ts:    time.Now().Round(resolution),
		avail: capacity,
	}
	b.Set(rate, capacity)

	return b
}

func (self *Bucket) Set(rate int64, capacity int64) {
	if rate <= 0 {
		panic("rate <= 0")
	}
	if capacity <= 0 {
		panic("capacity <= 0")
	}

	// convert rate from tokens per second to tokens per resolution unit
	conv := float64(time.Second) / float64(resolution)
	self.rate = int64(math.Ceil(float64(rate) / conv))

	self.capacity = capacity
	if self.avail > self.capacity {
		self.avail = self.capacity
	}
}

func (self *Bucket) Wait(count int64) {
	if count <= self.avail {
		self.avail -= count
		return
	}

	// refill the bucket
	since := time.Since(self.ts)
	diff := int64(since / resolution)
	self.avail += diff * self.rate
	if self.avail > self.capacity {
		self.avail = self.capacity
	}
	self.ts = self.ts.Add(time.Duration(diff) * resolution)

	// re-check
	if count <= self.avail {
		self.avail -= count
		return
	}

	// we need to wait
	count -= self.avail
	self.avail = 0

	dur := time.Duration(math.Ceil(float64(count)/float64(self.rate))) * resolution
	if dur < resolution {
		dur = resolution
	}
	time.Sleep(dur)
	self.Wait(count)
}

func (self *Bucket) Fill() {
	self.avail = self.capacity
}<commit_msg>add Use and Check<commit_after>package ratelimit

import (
	"math"
	"time"
)

const resolution = 1 * time.Millisecond

type Bucket struct {
	ts       time.Time
	rate     int64
	capacity int64
	avail    int64
}

func NewBucket(rate int64, capacity int64) *Bucket {
	b := &Bucket{
		ts:    time.Now().Round(resolution),
		avail: capacity,
	}
	b.Set(rate, capacity)

	return b
}

func (self *Bucket) Set(rate int64, capacity int64) {
	if rate <= 0 {
		panic("rate <= 0")
	}
	if capacity <= 0 {
		panic("capacity <= 0")
	}

	// convert rate from tokens per second to tokens per resolution unit
	conv := float64(time.Second) / float64(resolution)
	self.rate = int64(math.Ceil(float64(rate) / conv))

	self.capacity = capacity
	if self.avail > self.capacity {
		self.avail = self.capacity
	}
}

func (self *Bucket) Wait(count int64) {
	enough := self.Use(count)
	if enough {
		return
	}

	// we need to wait
	count -= self.avail
	self.Use(self.avail)

	dur := time.Duration(math.Ceil(float64(count)/float64(self.rate))) * resolution
	if dur < resolution {
		dur = resolution
	}

	time.Sleep(dur)

	self.Wait(count)
}

func (self *Bucket) Use(count int64) bool {
	if count <= self.avail {
		self.avail -= count
		return true
	}

	// refill the bucket
	since := time.Since(self.ts)
	diff := int64(since / resolution)
	self.avail += diff * self.rate
	if self.avail > self.capacity {
		self.avail = self.capacity
	}
	self.ts = self.ts.Add(time.Duration(diff) * resolution)

	// re-check
	if count <= self.avail {
		self.avail -= count
		return true
	}

	return false
}

func (self *Bucket) Check(count int64) bool {
	if count <= self.avail {
		return true
	}
	return false
}

func (self *Bucket) Fill() {
	self.avail = self.capacity
}
{'fraction_non_alphanumeric': 0.13474734872114785, 'fraction_numerical': 0.023393636930754833, 'mean_word_length': 3.2364597093791283, 'pattern_counts': {'":': 0, '<': 18, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2167321', 'n_tokens_mistral': 1274, 'n_tokens_neox': 1177, 'n_words': 306}