diff --git a/Dockerfile b/Dockerfile index 073e1b0..ad33183 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,4 +1,4 @@ -FROM golang:1.22.5-alpine AS builder +FROM golang:1.22.6-alpine AS builder WORKDIR /dodo diff --git a/README.md b/README.md index 58b8992..a2ce2ca 100644 --- a/README.md +++ b/README.md @@ -99,6 +99,8 @@ docker run --rm -v ./path/config.json:/dodo/config.json -i aykhans/dodo -u https ``` ## CLI and JSON Config Parameters +If the Headers, Params, Cookies and Body fields have multiple values, each request will choose a random value from the list. + | Parameter | JSON config file | CLI Flag | CLI Short Flag | Type | Description | Default | | ----------- | ----------- | ----------- | ----------- | ----------- | ----------- | ----------- | | Config file | - | --config-file | -c | String | Path to the JSON config file | - | @@ -108,11 +110,8 @@ docker run --rm -v ./path/config.json:/dodo/config.json -i aykhans/dodo -u https | Request count | request_count | --request-count | -r | Integer | Total number of requests to send | 1000 | | Dodos count (Threads) | dodos_count | --dodos-count | -d | Integer | Number of dodos (threads) to send requests in parallel | 1 | | Timeout | timeout | --timeout | -t | Integer | Timeout for canceling each request (milliseconds) | 10000 | -| Params | params | - | - | Key-Value {string: string} | Request parameters | - | -| Headers | headers | - | - | Key-Value {string: string} | Request headers | - | -| Cookies | cookies | - | - | Key-Value {string: string} | Request cookies | - | -| Body | body | - | - | String | Request body | - | +| Params | params | - | - | Key-Value {String: [String]} | Request parameters | - | +| Headers | headers | - | - | Key-Value {String: [String]} | Request headers | - | +| Cookies | cookies | - | - | Key-Value {String: [String]} | Request cookies | - | +| Body | body | - | - | [String] | Request body | - | | Proxy | proxies | - | - | List[Key-Value {string: string}] | List of proxies (will check active proxies before sending requests) | - | - -## Examples -![dodo_example](https://raw.githubusercontent.com/aykhans/dodo/main/assets/dodo-example.gif) \ No newline at end of file diff --git a/assets/dodo-example.gif b/assets/dodo-example.gif deleted file mode 100644 index 0ecb981..0000000 Binary files a/assets/dodo-example.gif and /dev/null differ diff --git a/config.json b/config.json index 03cf494..cfa8782 100644 --- a/config.json +++ b/config.json @@ -7,7 +7,7 @@ "params": {}, "headers": {}, "cookies": {}, - "body": "", + "body": [""], "proxies": [ { "url": "http://example.com:8080", diff --git a/config/config.go b/config/config.go index 5997086..58be045 100644 --- a/config/config.go +++ b/config/config.go @@ -6,11 +6,12 @@ import ( "os" "time" + "github.com/aykhans/dodo/utils" "github.com/jedib0t/go-pretty/v6/table" ) const ( - VERSION string = "0.4.2" + VERSION string = "0.5.0" DefaultUserAgent string = "Dodo/" + VERSION ProxyCheckURL string = "https://www.google.com" DefaultMethod string = "GET" @@ -30,11 +31,11 @@ type RequestConfig struct { Timeout time.Duration DodosCount uint RequestCount uint - Params map[string]string - Headers map[string]string - Cookies map[string]string + Params map[string][]string + Headers map[string][]string + Cookies map[string][]string Proxies []Proxy - Body string + Body []string Yes bool } @@ -55,17 +56,17 @@ func (config *RequestConfig) Print() { t.AppendSeparator() t.AppendRow(table.Row{"Dodos", config.DodosCount}) t.AppendSeparator() - t.AppendRow(table.Row{"Request Count", config.RequestCount}) + t.AppendRow(table.Row{"Request", config.RequestCount}) t.AppendSeparator() - t.AppendRow(table.Row{"Params Count", len(config.Params)}) + t.AppendRow(table.Row{"Params", utils.MarshalJSON(config.Params, 3)}) t.AppendSeparator() - t.AppendRow(table.Row{"Headers Count", len(config.Headers)}) + t.AppendRow(table.Row{"Headers", utils.MarshalJSON(config.Headers, 3)}) t.AppendSeparator() - t.AppendRow(table.Row{"Cookies Count", len(config.Cookies)}) + t.AppendRow(table.Row{"Cookies", utils.MarshalJSON(config.Cookies, 3)}) t.AppendSeparator() - t.AppendRow(table.Row{"Proxies Count", len(config.Proxies)}) + t.AppendRow(table.Row{"Proxies", utils.MarshalJSON(config.Proxies, 3)}) t.AppendSeparator() - t.AppendRow(table.Row{"Body", config.Body}) + t.AppendRow(table.Row{"Body", utils.MarshalJSON(config.Body, 3)}) t.Render() } @@ -134,11 +135,11 @@ type Proxy struct { type JSONConfig struct { Config - Params map[string]string `json:"params"` - Headers map[string]string `json:"headers"` - Cookies map[string]string `json:"cookies"` - Proxies []Proxy `json:"proxies" validate:"dive"` - Body string `json:"body"` + Params map[string][]string `json:"params"` + Headers map[string][]string `json:"headers"` + Cookies map[string][]string `json:"cookies"` + Proxies []Proxy `json:"proxies" validate:"dive"` + Body []string `json:"body"` } func (config *JSONConfig) MergeConfigs(newConfig *JSONConfig) { @@ -152,7 +153,7 @@ func (config *JSONConfig) MergeConfigs(newConfig *JSONConfig) { if len(newConfig.Cookies) != 0 { config.Cookies = newConfig.Cookies } - if newConfig.Body != "" { + if len(newConfig.Body) != 0 { config.Body = newConfig.Body } if len(newConfig.Proxies) != 0 { diff --git a/go.mod b/go.mod index 6e463b0..e7edf3f 100644 --- a/go.mod +++ b/go.mod @@ -1,6 +1,6 @@ module github.com/aykhans/dodo -go 1.22.3 +go 1.22.6 require ( github.com/go-playground/validator/v10 v10.22.1 diff --git a/main.go b/main.go index da8aae0..1b5baec 100644 --- a/main.go +++ b/main.go @@ -106,6 +106,7 @@ func main() { return } else if customerrors.Is(err, customerrors.ErrNoInternet) { utils.PrintAndExit("No internet connection") + return } panic(err) } diff --git a/requests/client.go b/requests/client.go index afc7a6c..52fef8d 100644 --- a/requests/client.go +++ b/requests/client.go @@ -9,26 +9,17 @@ import ( "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" ) -// 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( +type ClientGeneratorFunc func() *fasthttp.HostClient + +// 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, @@ -36,7 +27,7 @@ func getClientDoFunc( maxConns uint, yes bool, URL *url.URL, -) ClientDoFunc { +) []*fasthttp.HostClient { isTLS := URL.Scheme == "https" if len(proxies) > 0 { @@ -71,26 +62,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{ @@ -102,7 +76,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 @@ -303,75 +277,17 @@ 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 list of fasthttp.HostClient instances. +// The function uses a local random number generator to determine the starting index and stop index for cycling through the clients. +// The returned function isn't thread-safe and should be used in a single-threaded context. +func getSharedClientFuncMultiple(clients []*fasthttp.HostClient, localRand *rand.Rand) ClientGeneratorFunc { + return utils.RandomValueCycle(clients, localRand) } -// 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 } } diff --git a/requests/request.go b/requests/request.go index 32b3730..29ca89c 100644 --- a/requests/request.go +++ b/requests/request.go @@ -1,15 +1,135 @@ package requests import ( + "context" + "math/rand" "net/url" + "time" "github.com/aykhans/dodo/config" + customerrors "github.com/aykhans/dodo/custom_errors" + "github.com/aykhans/dodo/utils" "github.com/valyala/fasthttp" ) -// newRequest creates a new fasthttp.Request object with the provided parameters. -// It sets the request URI, host header, headers, cookies, params, method, and body. +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, localRand) + } + + 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, + localRand *rand.Rand, +) RequestGeneratorFunc { + bodiesLen := len(Bodies) + getBody := func() string { return "" } + if bodiesLen == 1 { + getBody = func() string { return Bodies[0] } + } else if bodiesLen > 1 { + getBody = utils.RandomValueCycle(Bodies, localRand) + } + getHeaders := getKeyValueSetFunc(Headers, localRand) + getCookies := getKeyValueSetFunc(Cookies, localRand) + getParams := getKeyValueSetFunc(Params, localRand) + + return func() *fasthttp.Request { + return newFasthttpRequest( + URL, + getHeaders(), + getCookies(), + getParams(), + Method, + getBody(), + ) + } +} + +// 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 newFasthttpRequest( URL *url.URL, Headers map[string]string, Cookies map[string]string, @@ -70,3 +190,60 @@ func setRequestMethod(req *fasthttp.Request, method string) { func setRequestBody(req *fasthttp.Request, body string) { req.SetBody([]byte(body)) } + +// getKeyValueSetFunc generates a function that returns a map of key-value pairs based on the provided key-value set. +// The generated function will either return fixed values or random values depending on the input. +// +// Returns: +// - A function that returns a map of key-value pairs. If the input map contains multiple values for a key, +// the returned function will generate random values for that key. If the input map contains a single value +// for a key, the returned function will always return that value. If the input map is empty for a key, +// the returned function will generate an empty string for that key. +func getKeyValueSetFunc[ + KeyValueSet map[string][]string, + KeyValue map[string]string, +](keyValueSet KeyValueSet, localRand *rand.Rand) func() KeyValue { + getKeyValueSlice := []map[string]func() string{} + isRandom := false + for key, values := range keyValueSet { + valuesLen := len(values) + + // if values is empty, return a function that generates empty string + // if values has only one element, return a function that generates that element + // if values has more than one element, return a function that generates a random element + getKeyValue := func() string { return "" } + if valuesLen == 1 { + getKeyValue = func() string { return values[0] } + } else if valuesLen > 1 { + getKeyValue = utils.RandomValueCycle(values, localRand) + isRandom = true + } + + getKeyValueSlice = append( + getKeyValueSlice, + map[string]func() string{key: getKeyValue}, + ) + } + + // if isRandom is true, return a function that generates random values, + // otherwise return a function that generates fixed values to avoid unnecessary random number generation + if isRandom { + return func() KeyValue { + keyValues := make(KeyValue, len(getKeyValueSlice)) + for _, keyValue := range getKeyValueSlice { + for key, value := range keyValue { + keyValues[key] = value() + } + } + return keyValues + } + } else { + keyValues := make(KeyValue, len(getKeyValueSlice)) + for _, keyValue := range getKeyValueSlice { + for key, value := range keyValue { + keyValues[key] = value() + } + } + return func() KeyValue { return keyValues } + } +} diff --git a/requests/response.go b/requests/response.go index 82cdbc5..7955bfe 100644 --- a/requests/response.go +++ b/requests/response.go @@ -1,14 +1,12 @@ package requests import ( - "context" "fmt" "os" "time" "github.com/aykhans/dodo/utils" "github.com/jedib0t/go-pretty/v6/table" - "github.com/valyala/fasthttp" ) type Response struct { @@ -17,9 +15,7 @@ type Response struct { Time time.Duration } -type Responses []Response - -type ClientDoFunc func(ctx context.Context, request *fasthttp.Request) (*fasthttp.Response, error) +type Responses []*Response // Print prints the responses in a tabular format, including information such as // response count, minimum time, maximum time, and average time. diff --git a/requests/run.go b/requests/run.go index bdbdbb8..c1ae2bc 100644 --- a/requests/run.go +++ b/requests/run.go @@ -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,26 +40,11 @@ func Run(ctx context.Context, requestConfig *config.RequestConfig) (Responses, e requestConfig.Yes, requestConfig.URL, ) - if clientDoFunc == nil { + if clients == nil { return nil, customerrors.ErrInterrupt } - request := newRequest( - requestConfig.URL, - requestConfig.Headers, - requestConfig.Cookies, - requestConfig.Params, - requestConfig.Method, - requestConfig.Body, - ) - defer fasthttp.ReleaseRequest(request) - responses := releaseDodos( - ctx, - request, - clientDoFunc, - requestConfig.GetValidDodosCountForRequests(), - requestConfig.RequestCount, - ) + responses := releaseDodos(ctx, requestConfig, clients) if ctx.Err() != nil && len(responses) == 0 { return nil, customerrors.ErrInterrupt } @@ -59,30 +52,36 @@ func Run(ctx context.Context, requestConfig *config.RequestConfig) (Responses, e return responses, nil } -// releaseDodos sends multiple HTTP requests concurrently using multiple "dodos" (goroutines). -// It takes a mainRequest as the base request, timeout duration for each request, clientDoFunc for customizing the client behavior, -// dodosCount as the number of goroutines to be used, and requestCount as the total number of requests to be sent. -// It returns the responses received from all the requests. +// releaseDodos sends requests concurrently using multiple dodos (goroutines) and returns the aggregated responses. +// +// 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, - mainRequest *fasthttp.Request, - clientDoFunc ClientDoFunc, - dodosCount uint, - requestCount uint, + requestConfig *config.RequestConfig, + clients []*fasthttp.HostClient, ) Responses { var ( wg sync.WaitGroup streamWG sync.WaitGroup requestCountPerDodo uint - dodosCountInt = int(dodosCount) + dodosCount uint = requestConfig.GetValidDodosCountForRequests() + dodosCountInt int = int(dodosCount) + requestCount uint = uint(requestConfig.RequestCount) + responses = make([][]*Response, dodosCount) + increase = make(chan int64, requestCount) ) wg.Add(dodosCountInt) streamWG.Add(1) - responses := make([][]Response, dodosCount) - increase := make(chan int64, requestCount) - streamCtx, streamCtxCancel := context.WithCancel(context.Background()) + go streamProgress(streamCtx, &streamWG, int64(requestCount), "Dodos Working🔥", increase) for i := range dodosCount { @@ -92,16 +91,14 @@ func releaseDodos( requestCountPerDodo = ((i + 1) * requestCount / dodosCount) - (i * requestCount / dodosCount) } - dodoSpecificRequest := &fasthttp.Request{} - mainRequest.CopyTo(dodoSpecificRequest) go sendRequest( ctx, - dodoSpecificRequest, + newRequest(*requestConfig, clients, int64(i)), + requestConfig.Timeout, + requestCountPerDodo, &responses[i], increase, - requestCountPerDodo, - clientDoFunc, &wg, ) } @@ -111,29 +108,19 @@ func releaseDodos( return utils.Flatten(responses) } -// sendRequest sends an HTTP request using the provided clientDo function and handles the response. -// -// Parameters: -// - ctx: The context to control cancellation and timeout. -// - request: The HTTP request to be sent. -// - responseData: A slice to store the response data. -// - increase: A channel to signal the completion of a request. -// - requestCount: The number of requests to be sent. -// - clientDo: A function to execute the HTTP request. -// - wg: A wait group to signal the completion of the function. -// -// The function sends the specified number of requests, handles errors, and appends the response data -// to the responseData slice. +// 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, - request *fasthttp.Request, - responseData *[]Response, - increase chan<- int64, + request *Request, + timeout time.Duration, requestCount uint, - clientDo ClientDoFunc, + responseData *[]*Response, + increase chan<- int64, wg *sync.WaitGroup, ) { - defer fasthttp.ReleaseRequest(request) defer wg.Done() for range requestCount { @@ -143,14 +130,17 @@ func sendRequest( func() { 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 { return } - *responseData = append(*responseData, Response{ + *responseData = append(*responseData, &Response{ StatusCode: 0, Error: err, Time: completedTime, @@ -158,9 +148,8 @@ func sendRequest( increase <- 1 return } - defer fasthttp.ReleaseResponse(response) - *responseData = append(*responseData, Response{ + *responseData = append(*responseData, &Response{ StatusCode: response.StatusCode(), Error: nil, Time: completedTime, diff --git a/utils/convert.go b/utils/convert.go new file mode 100644 index 0000000..b97a9ff --- /dev/null +++ b/utils/convert.go @@ -0,0 +1,54 @@ +package utils + +import ( + "encoding/json" + "reflect" + "strings" +) + +func MarshalJSON(v any, maxSliceSize uint) string { + rv := reflect.ValueOf(v) + if rv.Kind() == reflect.Slice && rv.Len() == 0 { + return "[]" + } + + data, err := json.MarshalIndent(truncateLists(v, int(maxSliceSize)), "", " ") + if err != nil { + return "{}" + } + + return strings.Replace(string(data), `"..."`, "...", -1) +} + +func truncateLists(v interface{}, maxItems int) interface{} { + rv := reflect.ValueOf(v) + + switch rv.Kind() { + case reflect.Slice, reflect.Array: + if rv.Len() > maxItems { + newSlice := reflect.MakeSlice(rv.Type(), maxItems, maxItems) + reflect.Copy(newSlice, rv.Slice(0, maxItems)) + newSlice = reflect.Append(newSlice, reflect.ValueOf("...")) + return newSlice.Interface() + } + case reflect.Map: + newMap := reflect.MakeMap(rv.Type()) + for _, key := range rv.MapKeys() { + newMap.SetMapIndex(key, reflect.ValueOf(truncateLists(rv.MapIndex(key).Interface(), maxItems))) + } + return newMap.Interface() + case reflect.Struct: + newStruct := reflect.New(rv.Type()).Elem() + for i := 0; i < rv.NumField(); i++ { + newStruct.Field(i).Set(reflect.ValueOf(truncateLists(rv.Field(i).Interface(), maxItems))) + } + return newStruct.Interface() + case reflect.Ptr: + if rv.IsNil() { + return nil + } + return truncateLists(rv.Elem().Interface(), maxItems) + } + + return v +} diff --git a/utils/slice.go b/utils/slice.go index e329604..e323663 100644 --- a/utils/slice.go +++ b/utils/slice.go @@ -1,7 +1,9 @@ package utils -func Flatten[T any](nested [][]T) []T { - flattened := make([]T, 0) +import "math/rand" + +func Flatten[T any](nested [][]*T) []*T { + flattened := make([]*T, 0) for _, n := range nested { flattened = append(flattened, n...) } @@ -16,3 +18,33 @@ func Contains[T comparable](slice []T, item T) bool { } return false } + +// RandomValueCycle returns a function that cycles through the provided slice of values +// in a random order. Each call to the returned function will yield a value from the slice. +// The order of values is determined by the provided random number generator. +// +// The returned function will cycle through the values in a random order until all values +// have been returned at least once. After all values have been returned, the function will +// reset and start cycling through the values in a random order again. +// The returned function isn't thread-safe and should be used in a single-threaded context. +func RandomValueCycle[Value any](values []Value, localRand *rand.Rand) func() Value { + var ( + clientsCount int = len(values) + currentIndex int = localRand.Intn(clientsCount) + stopIndex int = currentIndex + ) + + return func() Value { + client := values[currentIndex] + currentIndex++ + if currentIndex == clientsCount { + currentIndex = 0 + } + if currentIndex == stopIndex { + currentIndex = localRand.Intn(clientsCount) + stopIndex = currentIndex + } + + return client + } +}