<commit_before>package middleware

import (
	"errors"
	"time"

	"github.com/achilleasa/usrv"
	"golang.org/x/net/context"
)

// Apply a throttling middleware to the input handler that throttles
// request handling to maxConcurrent requests.
//
// If the pending requests exceed the specified timeout, they will be aborted
// with ErrTimeout.
func Throttle(maxConcurrent int, timeout time.Duration) usrv.EndpointOption {

	return func(ep *usrv.Endpoint) error {

		if maxConcurrent <= 0 {
			return errors.New("maxConcurrent should be > 0")
		}

		// Allocate a buffered channel and pre-fill it with tokens
		tokens := make(chan struct{}, maxConcurrent)
		for i := 0; i < maxConcurrent; i++ {
			tokens <- struct{}{}
		}

		// Wrap original method
		originalHandler := ep.Handler
		ep.Handler = usrv.HandlerFunc(func(ctx context.Context, responseWriter usrv.ResponseWriter, request *usrv.Message) {
			fctx, cancel := context.WithTimeout(ctx, timeout)
			defer cancel()

			select {
			case <-tokens:
				// We got a token, execute request
				originalHandler.Serve(fctx, responseWriter, request)

				// Return back token
				tokens <- struct{}{}
			case <-fctx.Done():
				responseWriter.WriteError(usrv.ErrTimeout)<commit_msg>Throttle distinguishes between context cancellation and timeout errors<commit_after>package middleware

import (
	"errors"
	"time"

	"github.com/achilleasa/usrv"
	"golang.org/x/net/context"
)

// Apply a throttling middleware to the input handler that throttles
// request handling to maxConcurrent requests.
//
// If the pending requests exceed the specified timeout, they will be aborted
// with ErrTimeout.
func Throttle(maxConcurrent int, timeout time.Duration) usrv.EndpointOption {

	return func(ep *usrv.Endpoint) error {

		if maxConcurrent <= 0 {
			return errors.New("maxConcurrent should be > 0")
		}

		// Allocate a buffered channel and pre-fill it with tokens
		tokens := make(chan struct{}, maxConcurrent)
		for i := 0; i < maxConcurrent; i++ {
			tokens <- struct{}{}
		}

		// Wrap original method
		originalHandler := ep.Handler
		ep.Handler = usrv.HandlerFunc(func(ctx context.Context, responseWriter usrv.ResponseWriter, request *usrv.Message) {
			fctx, cancel := context.WithTimeout(ctx, timeout)
			defer cancel()

			select {
			case <-tokens:
				// We got a token, execute request
				originalHandler.Serve(fctx, responseWriter, request)

				// Return back token
				tokens <- struct{}{}
			case <-fctx.Done():
				if fctx.Err() == context.Canceled {
					responseWriter.WriteError(usrv.ErrCancelled)
				} else {
					responseWriter.WriteError(usrv.ErrTimeout)
				}
{'fraction_non_alphanumeric': 0.1091186570011446, 'fraction_numerical': 0.002289202594429607, 'mean_word_length': 3.6654804270462633, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 5, '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': '733391', 'n_tokens_mistral': 905, 'n_tokens_neox': 778, 'n_words': 237}