<commit_before>package memory

import (
	"context"
	"sync"
	"time"
)

// TimeLimiter provides a means limit the amount of time spent working.TimeLimiter
type TimeLimiter struct {
	sync.Mutex
	ctx          context.Context
	lockDuration time.Duration
	window       time.Duration
	limit        time.Duration
	accountCh    chan time.Duration
	queryCh      chan chan struct{}
}

// NewTimeLimiter creates a new TimeLimiter.  A background thread will run until the
// provided context is done.  When the amount of time spent on task (the time is determined
// by calls to "Add()") every "window" duration is more then "limit",  then calls to
// Wait() will block until the start if the next window period.
func NewTimeLimiter(ctx context.Context, window, limit time.Duration) *TimeLimiter {
	l := &TimeLimiter{
		ctx:       ctx,
		window:    window,
		limit:     limit,
		accountCh: make(chan time.Duration),
		queryCh:   make(chan chan struct{}),
	}
	go l.run()
	return l
}

func (l *TimeLimiter) run() {
	ticker := time.NewTicker(l.window)
	done := l.ctx.Done()
	var blockedQueries []chan struct{}
	for {
		select {
		case <-done:
			//context done. shutting down
			for _, ch := range blockedQueries {
				close(ch)
			}
			blockedQueries = nil
			return
		case <-ticker.C:
			// reset lockDuration to 0
			l.lockDuration = time.Duration(0)
			for _, ch := range blockedQueries {
				close(ch)
			}
			blockedQueries = nil
		case d := <-l.accountCh:
			l.lockDuration += d
		case respCh := <-l.queryCh:
			if l.lockDuration < l.limit {
				close(respCh)
			} else {
				// rate limit exceeded.  On the next tick respCh will be closed
				// notifying the caller that they can continue.
				blockedQueries = append(blockedQueries, respCh)
			}
		}
	}
}

// Add increments the counter of time spent doing something by "d"
func (l *TimeLimiter) Add(d time.Duration) {
	l.accountCh <- d
}

// Wait will return immediately if we are not rate limited, otherwise it will
// block until we are no longer limited.  The longest we will block for is
// the size of the defined time window.
func (l *TimeLimiter) Wait() {
	respCh := make(chan struct{})
	l.queryCh <- respCh

	// if we have not exceeded our locking quota then respCh will be
	// immediately closed. Otherwise it wont be closed until the next tick (duration of "l.window")
	// and we will block until then.
	<-respCh
}<commit_msg>typos, simplifications, clarifications and an allocation tweak<commit_after>package memory

import (
	"context"
	"sync"
	"time"
)

// TimeLimiter limits the rate of a set of operations.
// It does this by slowing down further operations as soon
// as one Add() is called informing it the per-window allowed budget has been exceeded.
// Limitations:
// * concurrently running operations can all exceed the budget,
//   so it works best for serial operations.
// * for serial operations, the last operation is allowed to exceed the budget
// * when an operation takes very long (e.g. 10 seconds, with a 100ms limit per second), it
//   is counted as exceeding the 100ms budget, but no other provisions are being made.
//
// Thus, TimeLimiter is designed for, and works best with, serially running operations,
// each of which takes a fraction of the limit.
type TimeLimiter struct {
	sync.Mutex
	ctx       context.Context
	timeSpent time.Duration
	window    time.Duration
	limit     time.Duration
	addCh     chan time.Duration
	queryCh   chan chan struct{}
}

// NewTimeLimiter creates a new TimeLimiter.  A background goroutine will run until the
// provided context is done.  When the amount of time spent on task (the time is determined
// by calls to "Add()") every "window" duration is more then "limit",  then calls to
// Wait() will block until the start if the next window period.
func NewTimeLimiter(ctx context.Context, window, limit time.Duration) *TimeLimiter {
	l := &TimeLimiter{
		ctx:     ctx,
		window:  window,
		limit:   limit,
		addCh:   make(chan time.Duration),
		queryCh: make(chan chan struct{}),
	}
	go l.run()
	return l
}

func (l *TimeLimiter) run() {
	ticker := time.NewTicker(l.window)
	done := l.ctx.Done()
	var blockedQueries []chan struct{}
	for {
		select {
		case <-done:
			//context done. shutting down
			for _, ch := range blockedQueries {
				close(ch)
			}
			return
		case <-ticker.C:
			l.timeSpent = 0
			for _, ch := range blockedQueries {
				close(ch)
			}
			blockedQueries = blockedQueries[:0]
		case d := <-l.addCh:
			l.timeSpent += d
		case respCh := <-l.queryCh:
			if l.timeSpent < l.limit {
				close(respCh)
			} else {
				// rate limit exceeded.  On the next tick respCh will be closed
				// notifying the caller that they can continue.
				blockedQueries = append(blockedQueries, respCh)
			}
		}
	}
}

// Add increments the "time spent" counter by "d"
func (l *TimeLimiter) Add(d time.Duration) {
	l.addCh <- d
}

// Wait returns when we are not rate limited, which may be
// anywhere between immediately or after the window.
func (l *TimeLimiter) Wait() {
	respCh := make(chan struct{})
	l.queryCh <- respCh

	// if we have not exceeded our locking quota then respCh will be
	// immediately closed. Otherwise it wont be closed until the next tick (duration of "l.window")
	// and we will block until then.
	<-respCh
}
{'fraction_non_alphanumeric': 0.09627092846270928, 'fraction_numerical': 0.00228310502283105, 'mean_word_length': 3.1166797180892716, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 3, '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': '2849365', 'n_tokens_mistral': 1767, 'n_tokens_neox': 1601, 'n_words': 654}