github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/bench/microbenchmarks/http2/put_files.go (about)

     1  // Package http2 puts a given number of files with a given size into AIStore.
     2  // See /bench/http2/README.md for more info.
     3  /*
     4   * Copyright (c) 2018-2024, NVIDIA CORPORATION. All rights reserved.
     5   */
     6  package main
     7  
     8  import (
     9  	"flag"
    10  	"fmt"
    11  	"sync"
    12  
    13  	"github.com/NVIDIA/aistore/cmn"
    14  	"github.com/NVIDIA/aistore/cmn/cos"
    15  	"github.com/NVIDIA/aistore/tools"
    16  	"github.com/NVIDIA/aistore/tools/readers"
    17  )
    18  
    19  const (
    20  	kilobytes = uint64(1024)
    21  	smokeDir  = "/tmp/ais/smoke"        // smoke test dir
    22  	ProxyURL  = "http://localhost:8080" // assuming local proxy is listening on 8080
    23  )
    24  
    25  var (
    26  	files    int
    27  	workers  int
    28  	filesize uint64
    29  	bucket   cmn.Bck
    30  	url      string
    31  )
    32  
    33  func init() {
    34  	flag.StringVar(&url, "url", ProxyURL, "Proxy URL")
    35  	flag.StringVar(&bucket.Name, "bucket", "local_benchmark_bucket", "AWS or GCP bucket")
    36  	flag.StringVar(&bucket.Provider, "provider", "", "backend provider")
    37  	flag.IntVar(&files, "files", 10, "Number of files to put")
    38  	flag.IntVar(&workers, "workers", 10, "Number of workers")
    39  	flag.Uint64Var(&filesize, "filesize", 1, "Size of files to put in KB")
    40  }
    41  
    42  func worker(jobs <-chan func()) {
    43  	for j := range jobs {
    44  		j()
    45  	}
    46  }
    47  
    48  func main() {
    49  	flag.Parse()
    50  	jobs := make(chan func(), files)
    51  
    52  	for range workers {
    53  		go worker(jobs)
    54  	}
    55  
    56  	err := putSpecificFiles(filesize*kilobytes, files, bucket, jobs)
    57  	if err != nil {
    58  		fmt.Printf("%v\n", err)
    59  		return
    60  	}
    61  }
    62  
    63  func putSpecificFiles(fileSize uint64, numPuts int, bck cmn.Bck, pool chan func()) error {
    64  	var (
    65  		errCh = make(chan error, numPuts+1)
    66  		wg    = &sync.WaitGroup{}
    67  	)
    68  
    69  	cos.CreateDir(smokeDir)
    70  
    71  	for i := 1; i < numPuts+1; i++ {
    72  		r, err := readers.NewRand(int64(fileSize), cos.ChecksumXXHash)
    73  		if err != nil {
    74  			return err
    75  		}
    76  
    77  		fname := fmt.Sprintf("l%d", i)
    78  		wg.Add(1)
    79  		pool <- func() {
    80  			defer wg.Done()
    81  			tools.Put(url, bck, "__bench/"+fname, r, errCh)
    82  		}
    83  	}
    84  	close(pool)
    85  	wg.Wait()
    86  
    87  	select {
    88  	case err := <-errCh:
    89  		return err
    90  	default:
    91  		return nil
    92  	}
    93  }