Refactor request generation to use local random number generator (fix #26, #27)

This commit is contained in:
Aykhan Shahsuvarov 2024-09-14 19:42:40 +04:00
parent 8ad0bb5697
commit 891f1f1333
3 changed files with 179 additions and 215 deletions

View File

@ -3,34 +3,22 @@ package requests
import (
"context"
"fmt"
"math/rand"
"net/url"
"sync"
"time"
"github.com/aykhans/dodo/config"
customerrors "github.com/aykhans/dodo/custom_errors"
"github.com/aykhans/dodo/readers"
"github.com/aykhans/dodo/utils"
"github.com/valyala/fasthttp"
"github.com/valyala/fasthttp/fasthttpproxy"
)
type ClientDoFunc func(ctx context.Context, request *fasthttp.Request) (*fasthttp.Response, error)
type ClientGeneratorFunc func() *fasthttp.HostClient
// getClientDoFunc returns a ClientDoFunc function that can be used to make HTTP requests.
//
// The function first checks if there are any proxies available. If there are, it retrieves the active proxy clients
// using the getActiveProxyClients function. If the context is canceled during this process, it returns nil.
// It then checks the number of active proxy clients and prompts the user to continue if there are none.
// If the user chooses to continue, it creates a fasthttp.HostClient with the appropriate settings and returns
// a ClientDoFunc function using the getSharedClientDoFunc function.
// If there is only one active proxy client, it uses that client to create the ClientDoFunc function.
// If there are multiple active proxy clients, it uses the getSharedRandomClientDoFunc function to create the ClientDoFunc function.
//
// If there are no proxies available, it creates a fasthttp.HostClient with the appropriate settings and returns
// a ClientDoFunc function using the getSharedClientDoFunc function.
func getClientDoFunc(
// getClients initializes and returns a slice of fasthttp.HostClient based on the provided parameters.
// It can either return clients with proxies or a single client without proxies.
func getClients(
ctx context.Context,
timeout time.Duration,
proxies []config.Proxy,
@ -38,7 +26,7 @@ func getClientDoFunc(
maxConns uint,
yes bool,
URL *url.URL,
) ClientDoFunc {
) []*fasthttp.HostClient {
isTLS := URL.Scheme == "https"
if len(proxies) > 0 {
@ -73,26 +61,9 @@ func getClientDoFunc(
}
}
fmt.Println()
if activeProxyClientsCount == 0 {
client := &fasthttp.HostClient{
MaxConns: int(maxConns),
IsTLS: isTLS,
Addr: URL.Host,
MaxIdleConnDuration: timeout,
MaxConnDuration: timeout,
WriteTimeout: timeout,
ReadTimeout: timeout,
}
return getSharedClientDoFunc(client, timeout)
} else if activeProxyClientsCount == 1 {
client := activeProxyClients[0]
return getSharedClientDoFunc(client, timeout)
if activeProxyClientsCount > 0 {
return activeProxyClients
}
return getSharedRandomClientDoFunc(
activeProxyClients,
activeProxyClientsCount,
timeout,
)
}
client := &fasthttp.HostClient{
@ -104,7 +75,7 @@ func getClientDoFunc(
WriteTimeout: timeout,
ReadTimeout: timeout,
}
return getSharedClientDoFunc(client, timeout)
return []*fasthttp.HostClient{client}
}
// getActiveProxyClients divides the proxies into slices based on the number of dodos and
@ -305,75 +276,34 @@ func getDialFunc(proxy *config.Proxy, timeout time.Duration) (fasthttp.DialFunc,
return dialer, nil
}
// getSharedRandomClientDoFunc is equivalent to getSharedClientDoFunc but uses a random client from the provided slice.
func getSharedRandomClientDoFunc(
clients []*fasthttp.HostClient,
clientsCount uint,
timeout time.Duration,
) ClientDoFunc {
clientsCountInt := int(clientsCount)
return func(ctx context.Context, request *fasthttp.Request) (*fasthttp.Response, error) {
client := clients[rand.Intn(clientsCountInt)]
defer client.CloseIdleConnections()
response := fasthttp.AcquireResponse()
ch := make(chan error)
go func() {
err := client.DoTimeout(request, response, timeout)
ch <- err
}()
select {
case err := <-ch:
if err != nil {
fasthttp.ReleaseResponse(response)
return nil, err
}
return response, nil
case <-time.After(timeout):
fasthttp.ReleaseResponse(response)
return nil, customerrors.ErrTimeout
case <-ctx.Done():
return nil, customerrors.ErrInterrupt
// getSharedClientFuncMultiple returns a ClientGeneratorFunc that cycles through
// a provided list of fasthttp.HostClient instances. Each call to the returned
// function will return the next client in the list, cycling back to the first
// client after reaching the end of the slice.
//
// The returned function isn't thread-safe and should be used in a single-threaded context.
func getSharedClientFuncMultiple(clients []*fasthttp.HostClient) ClientGeneratorFunc {
var (
currentIndex int = 0
clientsCount int = len(clients)
)
return func() *fasthttp.HostClient {
client := clients[currentIndex]
if currentIndex == clientsCount-1 {
currentIndex = 0
} else {
currentIndex++
}
return client
}
}
// getSharedClientDoFunc is a function that returns a ClientDoFunc, which is a function type used for making HTTP requests using a shared client.
// It takes a client of type *fasthttp.HostClient and a timeout of type time.Duration as input parameters.
// The returned ClientDoFunc function can be used to make an HTTP request with the given client and timeout.
// It takes a context.Context and a *fasthttp.Request as input parameters and returns a *fasthttp.Response and an error.
// The function internally creates a new response using fasthttp.AcquireResponse() and a channel to handle errors.
// It then spawns a goroutine to execute the client.DoTimeout() method with the given request, response, and timeout.
// The function uses a select statement to handle three cases:
// - If an error is received from the channel, it checks if the error is not nil. If it's not nil, it releases the response and returns nil and the error.
// Otherwise, it returns the response and nil.
// - If the timeout duration is reached, it releases the response and returns nil and a custom timeout error.
// - If the context is canceled, it returns nil and a custom interrupt error.
//
// The function ensures that idle connections are closed by calling client.CloseIdleConnections() using a defer statement.
func getSharedClientDoFunc(
client *fasthttp.HostClient,
timeout time.Duration,
) ClientDoFunc {
return func(ctx context.Context, request *fasthttp.Request) (*fasthttp.Response, error) {
defer client.CloseIdleConnections()
response := fasthttp.AcquireResponse()
ch := make(chan error)
go func() {
err := client.DoTimeout(request, response, timeout)
ch <- err
}()
select {
case err := <-ch:
if err != nil {
fasthttp.ReleaseResponse(response)
return nil, err
}
return response, nil
case <-time.After(timeout):
fasthttp.ReleaseResponse(response)
return nil, customerrors.ErrTimeout
case <-ctx.Done():
return nil, customerrors.ErrInterrupt
}
// getSharedClientFuncSingle returns a ClientGeneratorFunc that always returns the provided fasthttp.HostClient instance.
// This can be useful for sharing a single client instance across multiple requests.
func getSharedClientFuncSingle(client *fasthttp.HostClient) ClientGeneratorFunc {
return func() *fasthttp.HostClient {
return client
}
}

View File

@ -4,51 +4,117 @@ import (
"context"
"math/rand"
"net/url"
"time"
"github.com/aykhans/dodo/config"
customerrors "github.com/aykhans/dodo/custom_errors"
"github.com/valyala/fasthttp"
)
// getRequests generates a list of HTTP requests based on the provided parameters.
//
// Parameters:
// - ctx: The context to control cancellation and deadlines.
// - URL: The base URL for the requests.
// - Headers: A map of headers to include in each request.
// - Cookies: A map of cookies to include in each request.
// - Params: A map of query parameters to include in each request.
// - Method: The HTTP method to use for the requests (e.g., GET, POST).
// - Bodies: A list of request bodies to cycle through for each request.
// - RequestCount: The number of requests to generate.
//
// Returns:
// - A list of fasthttp.Request objects based on the provided parameters.
// - An error if the context is canceled.
func getRequests(
ctx context.Context,
type RequestGeneratorFunc func() *fasthttp.Request
// Request represents an HTTP request to be sent using the fasthttp client.
// It isn't thread-safe and should be used by a single goroutine.
type Request struct {
getClient ClientGeneratorFunc
getRequest RequestGeneratorFunc
}
// Send sends the HTTP request using the fasthttp client with a specified timeout.
// It returns the HTTP response or an error if the request fails or times out.
func (r *Request) Send(ctx context.Context, timeout time.Duration) (*fasthttp.Response, error) {
client := r.getClient()
request := r.getRequest()
defer client.CloseIdleConnections()
defer fasthttp.ReleaseRequest(request)
response := fasthttp.AcquireResponse()
ch := make(chan error)
go func() {
err := client.DoTimeout(request, response, timeout)
ch <- err
}()
select {
case err := <-ch:
if err != nil {
fasthttp.ReleaseResponse(response)
return nil, err
}
return response, nil
case <-time.After(timeout):
fasthttp.ReleaseResponse(response)
return nil, customerrors.ErrTimeout
case <-ctx.Done():
return nil, customerrors.ErrInterrupt
}
}
// newRequest creates a new Request instance based on the provided configuration and clients.
// It initializes a random number generator using the current time and a unique identifier (uid).
// Depending on the number of clients provided, it sets up a function to select the appropriate client.
// It also sets up a function to generate the request based on the provided configuration.
func newRequest(
requestConfig config.RequestConfig,
clients []*fasthttp.HostClient,
uid int64,
) *Request {
localRand := rand.New(rand.NewSource(time.Now().UnixNano() + uid))
clientsCount := len(clients)
if clientsCount < 1 {
panic("no clients")
}
getClient := ClientGeneratorFunc(nil)
if clientsCount == 1 {
getClient = getSharedClientFuncSingle(clients[0])
} else {
getClient = getSharedClientFuncMultiple(clients)
}
getRequest := getRequestGeneratorFunc(
requestConfig.URL,
requestConfig.Headers,
requestConfig.Cookies,
requestConfig.Params,
requestConfig.Method,
requestConfig.Body,
localRand,
)
requests := &Request{
getClient: getClient,
getRequest: getRequest,
}
return requests
}
// getRequestGeneratorFunc returns a RequestGeneratorFunc which generates HTTP requests
// with the specified parameters.
// The function uses a local random number generator to select bodies, headers, cookies, and parameters
// if multiple options are provided.
func getRequestGeneratorFunc(
URL *url.URL,
Headers map[string][]string,
Cookies map[string][]string,
Params map[string][]string,
Method string,
Bodies []string,
RequestCount uint,
) ([]*fasthttp.Request, error) {
requests := make([]*fasthttp.Request, 0, RequestCount)
localRand *rand.Rand,
) RequestGeneratorFunc {
bodiesLen := len(Bodies)
getBody := func() string { return "" }
if bodiesLen == 1 {
getBody = func() string { return Bodies[0] }
} else if bodiesLen > 1 {
currentIndex := 0
currentIndex := localRand.Intn(bodiesLen)
stopIndex := bodiesLen - 1
getBody = func() string {
body := Bodies[currentIndex%bodiesLen]
if currentIndex == stopIndex {
currentIndex = rand.Intn(bodiesLen)
currentIndex = localRand.Intn(bodiesLen)
stopIndex = currentIndex - 1
} else {
currentIndex = (currentIndex + 1) % bodiesLen
@ -56,15 +122,12 @@ func getRequests(
return body
}
}
getHeaders := getKeyValueSetFunc(Headers)
getCookies := getKeyValueSetFunc(Cookies)
getParams := getKeyValueSetFunc(Params)
getHeaders := getKeyValueSetFunc(Headers, localRand)
getCookies := getKeyValueSetFunc(Cookies, localRand)
getParams := getKeyValueSetFunc(Params, localRand)
for range RequestCount {
if ctx.Err() != nil {
return nil, customerrors.ErrInterrupt
}
request := newRequest(
return func() *fasthttp.Request {
return newFasthttpRequest(
URL,
getHeaders(),
getCookies(),
@ -72,15 +135,12 @@ func getRequests(
Method,
getBody(),
)
requests = append(requests, request)
}
return requests, nil
}
// newRequest creates a new fasthttp.Request object with the provided parameters.
// newFasthttpRequest creates a new fasthttp.Request object with the provided parameters.
// It sets the request URI, host header, headers, cookies, params, method, and body.
func newRequest(
func newFasthttpRequest(
URL *url.URL,
Headers map[string]string,
Cookies map[string]string,
@ -153,7 +213,7 @@ func setRequestBody(req *fasthttp.Request, body string) {
func getKeyValueSetFunc[
KeyValueSet map[string][]string,
KeyValue map[string]string,
](keyValueSet KeyValueSet) func() KeyValue {
](keyValueSet KeyValueSet, localRand *rand.Rand) func() KeyValue {
getKeyValueSlice := []map[string]func() string{}
isRandom := false
for key, values := range keyValueSet {
@ -166,13 +226,13 @@ func getKeyValueSetFunc[
if valuesLen == 1 {
getKeyValue = func() string { return values[0] }
} else if valuesLen > 1 {
currentIndex := 0
currentIndex := localRand.Intn(valuesLen)
stopIndex := valuesLen - 1
getKeyValue = func() string {
value := values[currentIndex%valuesLen]
if currentIndex == stopIndex {
currentIndex = rand.Intn(valuesLen)
currentIndex = localRand.Intn(valuesLen)
stopIndex = currentIndex - 1
} else {
currentIndex = (currentIndex + 1) % valuesLen

View File

@ -11,10 +11,18 @@ import (
"github.com/valyala/fasthttp"
)
// Run executes the HTTP requests based on the provided request configuration.
// It checks for internet connection and returns an error if there is no connection.
// If the context is canceled while checking proxies, it returns the ErrInterrupt.
// If the context is canceled while sending requests, it returns the response objects obtained so far.
// Run executes the main logic for processing requests based on the provided configuration.
// It first checks for an internet connection with a timeout context. If no connection is found,
// it returns an error. Then, it initializes clients based on the request configuration and
// releases the dodos. If the context is canceled and no responses are collected, it returns an interrupt error.
//
// Parameters:
// - ctx: The context for managing request lifecycle and cancellation.
// - requestConfig: The configuration for the request, including timeout, proxies, and other settings.
//
// Returns:
// - Responses: A collection of responses from the executed requests.
// - error: An error if the operation fails, such as no internet connection or an interrupt.
func Run(ctx context.Context, requestConfig *config.RequestConfig) (Responses, error) {
checkConnectionCtx, checkConnectionCtxCancel := context.WithTimeout(ctx, 8*time.Second)
if !checkConnection(checkConnectionCtx) {
@ -23,7 +31,7 @@ func Run(ctx context.Context, requestConfig *config.RequestConfig) (Responses, e
}
checkConnectionCtxCancel()
clientDoFunc := getClientDoFunc(
clients := getClients(
ctx,
requestConfig.Timeout,
requestConfig.Proxies,
@ -32,30 +40,8 @@ func Run(ctx context.Context, requestConfig *config.RequestConfig) (Responses, e
requestConfig.Yes,
requestConfig.URL,
)
if clientDoFunc == nil {
return nil, customerrors.ErrInterrupt
}
requests, err := getRequests(
ctx,
requestConfig.URL,
requestConfig.Headers,
requestConfig.Cookies,
requestConfig.Params,
requestConfig.Method,
requestConfig.Body,
requestConfig.RequestCount,
)
if err != nil {
return nil, err
}
responses := releaseDodos(
ctx,
requests,
clientDoFunc,
requestConfig.GetValidDodosCountForRequests(),
)
responses := releaseDodos(ctx, requestConfig, clients)
if ctx.Err() != nil && len(responses) == 0 {
return nil, customerrors.ErrInterrupt
}
@ -63,59 +49,55 @@ func Run(ctx context.Context, requestConfig *config.RequestConfig) (Responses, e
return responses, nil
}
// releaseDodos sends HTTP requests concurrently using multiple "dodos" (goroutines).
// releaseDodos sends requests concurrently using multiple dodos (goroutines) and returns the aggregated responses.
//
// Parameters:
// - ctx: The context to control the lifecycle of the requests.
// - requests: A slice of HTTP requests to be sent.
// - clientDoFunc: A function to execute the HTTP requests.
// - dodosCount: The number of dodos (goroutines) to use for sending the requests.
//
// Returns:
// - A slice of Response objects containing the results of the requests.
//
// The function divides the requests into equal parts based on the number of dodos.
// It then sends each part concurrently using a separate goroutine.
// The function performs the following steps:
// 1. Initializes wait groups and other necessary variables.
// 2. Starts a goroutine to stream progress updates.
// 3. Distributes the total request count among the dodos.
// 4. Starts a goroutine for each dodo to send requests concurrently.
// 5. Waits for all dodos to complete their requests.
// 6. Cancels the progress streaming context and waits for the progress goroutine to finish.
// 7. Flattens and returns the aggregated responses.
func releaseDodos(
ctx context.Context,
requests []*fasthttp.Request,
clientDoFunc ClientDoFunc,
dodosCount uint,
requestConfig *config.RequestConfig,
clients []*fasthttp.HostClient,
) Responses {
var (
wg sync.WaitGroup
streamWG sync.WaitGroup
requestCountPerDodo uint
dodosCount uint = requestConfig.GetValidDodosCountForRequests()
dodosCountInt int = int(dodosCount)
totalRequestCount uint = uint(len(requests))
requestCount uint = 0
requestCount uint = uint(requestConfig.RequestCount)
responses = make([][]*Response, dodosCount)
increase = make(chan int64, totalRequestCount)
increase = make(chan int64, requestCount)
)
wg.Add(dodosCountInt)
streamWG.Add(1)
streamCtx, streamCtxCancel := context.WithCancel(context.Background())
go streamProgress(streamCtx, &streamWG, int64(totalRequestCount), "Dodos Working🔥", increase)
go streamProgress(streamCtx, &streamWG, int64(requestCount), "Dodos Working🔥", increase)
for i := range dodosCount {
if i+1 == dodosCount {
requestCountPerDodo = totalRequestCount - (i * totalRequestCount / dodosCount)
requestCountPerDodo = requestCount - (i * requestCount / dodosCount)
} else {
requestCountPerDodo = ((i + 1) * totalRequestCount / dodosCount) -
(i * totalRequestCount / dodosCount)
requestCountPerDodo = ((i + 1) * requestCount / dodosCount) -
(i * requestCount / dodosCount)
}
go sendRequest(
ctx,
requests[requestCount:requestCount+requestCountPerDodo],
newRequest(*requestConfig, clients, int64(i)),
requestConfig.Timeout,
requestCountPerDodo,
&responses[i],
increase,
clientDoFunc,
&wg,
)
requestCount += requestCountPerDodo
}
wg.Wait()
streamCtxCancel()
@ -123,40 +105,33 @@ func releaseDodos(
return utils.Flatten(responses)
}
// sendRequest sends multiple HTTP requests concurrently and collects their responses.
//
// Parameters:
// - ctx: The context to control cancellation and timeout.
// - requests: A slice of pointers to fasthttp.Request objects to be sent.
// - responseData: A pointer to a slice of *Response objects to store the results.
// - increase: A channel to signal the completion of each request.
// - clientDo: A function to execute the HTTP request.
// - wg: A wait group to synchronize the completion of the requests.
//
// The function iterates over the provided requests, sending each one using the clientDo function.
// It measures the time taken for each request and appends the response data to responseData.
// If an error occurs, it appends an error response. The function signals completion through the increase channel
// and ensures proper resource cleanup by releasing requests and responses.
// sendRequest sends a specified number of HTTP requests concurrently with a given timeout.
// It appends the responses to the provided responseData slice and sends the count of completed requests
// to the increase channel. The function terminates early if the context is canceled or if a custom
// interrupt error is encountered.
func sendRequest(
ctx context.Context,
requests []*fasthttp.Request,
request *Request,
timeout time.Duration,
requestCount uint,
responseData *[]*Response,
increase chan<- int64,
clientDo ClientDoFunc,
wg *sync.WaitGroup,
) {
defer wg.Done()
for _, request := range requests {
for range requestCount {
if ctx.Err() != nil {
return
}
func() {
defer fasthttp.ReleaseRequest(request)
startTime := time.Now()
response, err := clientDo(ctx, request)
response, err := request.Send(ctx, timeout)
completedTime := time.Since(startTime)
if response != nil {
defer fasthttp.ReleaseResponse(response)
}
if err != nil {
if err == customerrors.ErrInterrupt {
@ -170,7 +145,6 @@ func sendRequest(
increase <- 1
return
}
defer fasthttp.ReleaseResponse(response)
*responseData = append(*responseData, &Response{
StatusCode: response.StatusCode(),