github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/pkg/loadtest/target_generator.go (about)

     1  package loadtest
     2  
     3  import (
     4  	"fmt"
     5  	"math/rand"
     6  	"net/http"
     7  	"net/url"
     8  	"strconv"
     9  	"strings"
    10  	"time"
    11  
    12  	vegeta "github.com/tsenart/vegeta/v12/lib"
    13  )
    14  
    15  const boundary = "---------------------abcdefg123456789"
    16  
    17  type TargetGenerator struct {
    18  	ServerAddress string
    19  }
    20  
    21  func randomFilepath(basename string) string {
    22  	var sb strings.Builder
    23  	const maxDepthLevel = 10
    24  	const maxDirSuffixes = 3
    25  	depth := rand.Intn(maxDepthLevel) //nolint:gosec
    26  	for i := 0; i < depth; i++ {
    27  		dirSuffix := rand.Intn(maxDirSuffixes) //nolint:gosec
    28  		sb.WriteString(fmt.Sprintf("dir%d/", dirSuffix))
    29  	}
    30  	return sb.String() + basename
    31  }
    32  
    33  func defaultTarget(method, url, body, typ string) vegeta.Target {
    34  	tgt := vegeta.Target{
    35  		Method: method, URL: url, Body: []byte(body), Header: getDefaultHeader(),
    36  	}
    37  	AddRequestType(&tgt, typ)
    38  	return tgt
    39  }
    40  
    41  func (t *TargetGenerator) GenerateCreateFileTargets(repo, branch string, num int) []vegeta.Target {
    42  	now := time.Now().UnixNano()
    43  	result := make([]vegeta.Target, num)
    44  	for i := 0; i < num; i++ {
    45  		randomContent := rand.Int() //nolint:gosec
    46  		fileContent := "--" + boundary + "\n" +
    47  			"Content-Disposition: form-data; name=\"content\"; filename=\"file\"\n" +
    48  			"Content-Type: text/plain\n\n" +
    49  			strconv.Itoa(randomContent) + "\n" + "--" + boundary + "--\n"
    50  		filename := randomFilepath(fmt.Sprintf("file_%d_%d", now, i))
    51  		tgt := vegeta.Target{
    52  			Method: "POST",
    53  			URL:    fmt.Sprintf("%s/repositories/%s/branches/%s/objects?path=%s", t.ServerAddress, repo, branch, filename),
    54  			Body:   []byte(fileContent),
    55  			Header: http.Header{
    56  				http.CanonicalHeaderKey("Accept"):          []string{"*/*"},
    57  				http.CanonicalHeaderKey("Accept-Encoding"): []string{"gzip, deflate"},
    58  				http.CanonicalHeaderKey("Content-Type"):    []string{"multipart/form-data; boundary=" + boundary},
    59  				http.CanonicalHeaderKey("Content-Length"):  []string{strconv.Itoa(len(fileContent))},
    60  			},
    61  		}
    62  		AddRequestType(&tgt, "createFile")
    63  		result[i] = tgt
    64  	}
    65  	return result
    66  }
    67  
    68  func (t *TargetGenerator) GenerateCommitTarget(repo, branch, msg string) vegeta.Target {
    69  	return defaultTarget("POST",
    70  		fmt.Sprintf("%s/repositories/%s/branches/%s/commits", t.ServerAddress, repo, branch),
    71  		fmt.Sprintf(`{"message":"%s","metadata":{}}`, msg),
    72  		"commit")
    73  }
    74  
    75  func (t *TargetGenerator) GenerateBranchTarget(repo, name string) vegeta.Target {
    76  	return defaultTarget("POST",
    77  		fmt.Sprintf("%s/repositories/%s/branches", t.ServerAddress, repo),
    78  		fmt.Sprintf(`{"name":"%s","source":"main"}`, name),
    79  		"createBranch")
    80  }
    81  
    82  func (t *TargetGenerator) GenerateMergeToMasterTarget(repo, branch string) vegeta.Target {
    83  	return defaultTarget("POST",
    84  		fmt.Sprintf("%s/repositories/%s/refs/%s/merge/main", t.ServerAddress, repo, branch),
    85  		"{}",
    86  		"merge")
    87  }
    88  
    89  func (t *TargetGenerator) GenerateListTarget(repo, branch string, amount int) vegeta.Target {
    90  	return defaultTarget("GET",
    91  		fmt.Sprintf("%s/repositories/%s/refs/%s/objects/ls?tree=%s&amount=%d&after=&", t.ServerAddress, repo, branch, randomFilepath(""), amount),
    92  		"{}",
    93  		fmt.Sprintf("list%d", amount))
    94  }
    95  
    96  func (t *TargetGenerator) GenerateDiffTarget(repo, branch string) vegeta.Target {
    97  	return defaultTarget("GET",
    98  		fmt.Sprintf("%s/repositories/%s/branches/%s/diff", t.ServerAddress, repo, branch),
    99  		"{}",
   100  		"diff")
   101  }
   102  
   103  func getDefaultHeader() http.Header {
   104  	return http.Header{
   105  		http.CanonicalHeaderKey("Accept"):       []string{"application/json"},
   106  		http.CanonicalHeaderKey("Content-Type"): []string{"application/json"},
   107  	}
   108  }
   109  
   110  func AddRequestType(tgt *vegeta.Target, typ string) {
   111  	tgt.URL = addTypeToURL(tgt.URL, typ)
   112  }
   113  
   114  func GetRequestType(res vegeta.Result) string {
   115  	return getTypeFromURL(res.URL)
   116  }
   117  
   118  func addTypeToURL(u, typ string) string {
   119  	parsedURL, err := url.Parse(u)
   120  	if err != nil {
   121  		return u
   122  	}
   123  	parsedQuery, err := url.ParseQuery(parsedURL.RawQuery)
   124  	if err != nil {
   125  		return u
   126  	}
   127  	parsedQuery.Add("loader-request-type", typ)
   128  	parsedURL.RawQuery = parsedQuery.Encode()
   129  	return parsedURL.String()
   130  }
   131  
   132  func getTypeFromURL(u string) string {
   133  	parsedURL, err := url.Parse(u)
   134  	if err != nil {
   135  		return ""
   136  	}
   137  	parsedQuery, err := url.ParseQuery(parsedURL.RawQuery)
   138  	if err != nil {
   139  		return ""
   140  	}
   141  	return parsedQuery.Get("loader-request-type")
   142  }