Files
versitygw/cmd/cuobjtest/main.go
T

520 lines
14 KiB
Go

// Copyright 2026 Versity Software
// This file is licensed under the Apache License, Version 2.0
// (the "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
package main
import (
"bytes"
"context"
crand "crypto/rand"
"errors"
"flag"
"fmt"
"hash/crc32"
"io"
"math"
"os"
"sort"
"strconv"
"strings"
"time"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/credentials"
s3lib "github.com/aws/aws-sdk-go-v2/service/s3"
s3types "github.com/aws/aws-sdk-go-v2/service/s3/types"
"github.com/versity/versitygw/cuobjclient"
)
var (
endpoint = flag.String("endpoint", "http://localhost:7070", "cuserver S3 endpoint URL")
bucket = flag.String("bucket", "cuobjtest", "S3 bucket (created if absent)")
key = flag.String("key", "cuobjtest-object", "S3 object key prefix")
sizeStr = flag.String("size", "64MiB", "Transfer size per iteration: e.g. 4MiB, 256MiB, 1GiB")
access = flag.String("access", "admin", "S3 access key")
secret = flag.String("secret", "", "S3 secret key (default: $AWS_SECRET_ACCESS_KEY); prefer the environment variable over this flag, which is visible in shell history and process listings")
region = flag.String("region", "us-east-1", "S3 region")
iterations = flag.Int("n", 3, "Number of benchmark iterations")
putOnly = flag.Bool("put-only", false, "Run PUT only (skip GET and checksum verification)")
getOnly = flag.Bool("get-only", false, "Run GET only (assumes object already exists)")
stdS3 = flag.Bool("std-s3", false, "Use standard S3 PUT/GET (no cuObject RDMA transfer)")
)
func main() {
flag.Parse()
if *secret == "" {
*secret = os.Getenv("AWS_SECRET_ACCESS_KEY")
}
if *secret == "" {
fmt.Fprintln(os.Stderr, "cuobjtest: -secret or AWS_SECRET_ACCESS_KEY is required")
flag.Usage()
os.Exit(1)
}
if *putOnly && *getOnly {
fatalf("cuobjtest: -put-only and -get-only are mutually exclusive")
}
if *iterations < 1 {
fatalf("cuobjtest: -n must be >= 1")
}
size, err := parseSize(*sizeStr)
if err != nil {
fatalf("cuobjtest: -size: %v", err)
}
if size <= 0 || size > cuobjclient.MaxTransferSize {
fatalf("cuobjtest: -size must be between 1 B and %d bytes", cuobjclient.MaxTransferSize)
}
if err := runBenchmark(size); err != nil {
fatalf("cuobjtest: %v", err)
}
}
type iterationResult struct {
index int
putDur time.Duration
getDur time.Duration
status string
err error
}
type benchmarkStats struct {
putDurs []time.Duration
getDurs []time.Duration
mismatch int
}
func runBenchmark(size int) error {
base := newS3Client(*endpoint, *access, *secret, *region)
if err := ensureBucket(base, *bucket); err != nil {
return err
}
printRunHeader(size)
stats := benchmarkStats{}
var results []iterationResult
var benchmarkDur time.Duration
if *stdS3 {
results, benchmarkDur = runConcurrentStdS3(base, size)
} else {
results, benchmarkDur = runConcurrent(base, size)
}
sort.Slice(results, func(i, j int) bool {
return results[i].index < results[j].index
})
for _, r := range results {
if r.err != nil {
return fmt.Errorf("iter %d: %w", r.index, r.err)
}
accumulateStats(&stats, r)
printIterationResult(size, r)
}
return printSummary(stats, benchmarkDur, size)
}
func printRunHeader(size int) {
mode := runModeLabel()
if *stdS3 {
mode = "standard-s3"
}
fmt.Printf("\ncuobjtest endpoint=%-28s size=%s iterations=%d\n", *endpoint, formatSize(size), *iterations)
fmt.Printf(" mode=%s\n\n", mode)
if *putOnly {
fmt.Printf("%-4s %-14s %-10s\n", "#", "PUT latency", "PUT GB/s")
} else if *getOnly {
fmt.Printf("%-4s %-14s %-10s\n", "#", "GET latency", "GET GB/s")
} else {
fmt.Printf("%-4s %-14s %-10s %-14s %-10s %s\n",
"#", "PUT latency", "PUT GB/s", "GET latency", "GET GB/s", "checksum")
}
fmt.Printf("%s\n", strings.Repeat("-", 70))
}
func accumulateStats(stats *benchmarkStats, r iterationResult) {
if !*getOnly {
stats.putDurs = append(stats.putDurs, r.putDur)
}
if !*putOnly {
stats.getDurs = append(stats.getDurs, r.getDur)
}
if r.status != "OK" {
stats.mismatch++
}
}
func printSummary(stats benchmarkStats, benchmarkDur time.Duration, size int) error {
fmt.Printf("%s\n", strings.Repeat("-", 70))
putCount := len(stats.putDurs)
getCount := len(stats.getDurs)
putBytes := int64(putCount * size)
getBytes := int64(getCount * size)
totalBytes := putBytes + getBytes
// PUT/GET throughput is measured against the time actually spent in
// PUT/GET calls, not the combined benchmarkDur (which also includes
// the other operation and buffer preparation).
putDur := sumDurs(stats.putDurs)
getDur := sumDurs(stats.getDurs)
fmt.Printf("OVERALL elapsed=%s", fmtDur(benchmarkDur))
if putCount > 0 {
fmt.Printf(" PUT agg=%s (%d ops, %.2f ops/s)",
fmt.Sprintf("%.3f GB/s", float64(putBytes)/putDur.Seconds()/1e9),
putCount,
float64(putCount)/putDur.Seconds())
}
if getCount > 0 {
fmt.Printf(" GET agg=%s (%d ops, %.2f ops/s)",
fmt.Sprintf("%.3f GB/s", float64(getBytes)/getDur.Seconds()/1e9),
getCount,
float64(getCount)/getDur.Seconds())
}
if putCount > 0 && getCount > 0 {
fmt.Printf(" COMBINED=%s",
fmt.Sprintf("%.3f GB/s", float64(totalBytes)/benchmarkDur.Seconds()/1e9))
}
fmt.Printf("\n")
if *iterations > 1 {
if !*getOnly {
pa, plo, phi := durStats(stats.putDurs)
fmt.Printf("PUT avg=%-10s min=%-10s max=%-10s avg %s\n",
fmtDur(pa), fmtDur(plo), fmtDur(phi), fmtGBps(size, pa))
}
if !*putOnly {
ga, glo, ghi := durStats(stats.getDurs)
fmt.Printf("GET avg=%-10s min=%-10s max=%-10s avg %s\n",
fmtDur(ga), fmtDur(glo), fmtDur(ghi), fmtGBps(size, ga))
}
}
if *getOnly {
fmt.Printf("\nGET-only run complete.\n")
return nil
}
if *putOnly {
fmt.Printf("\nPUT-only run complete.\n")
return nil
}
if stats.mismatch > 0 {
return fmt.Errorf("%d/%d iterations had checksum mismatches", stats.mismatch, *iterations)
}
fmt.Printf("\nAll checksums OK.\n")
return nil
}
func runConcurrent(base *s3lib.Client, size int) ([]iterationResult, time.Duration) {
session, err := cuobjclient.NewSession(size)
if err != nil {
return []iterationResult{{index: 1, err: fmt.Errorf("init cuObj session: %w", err)}}, 0
}
defer session.Close()
putHost := make([]byte, size)
getHost := make([]byte, size)
out := make([]iterationResult, 0, *iterations)
benchmarkStart := time.Now()
for i := 1; i <= *iterations; i++ {
res := iterationResult{index: i, status: "OK"}
iterKey := *key
if !*getOnly {
if _, err := crand.Read(putHost); err != nil {
res.err = fmt.Errorf("fill PUT buffer: %w", err)
out = append(out, res)
continue
}
putCRC := crc32.ChecksumIEEE(putHost)
putStart := time.Now()
err := session.Upload(base, *bucket, iterKey, putHost)
res.putDur = time.Since(putStart)
if err != nil {
res.err = fmt.Errorf("PUT: %w", err)
out = append(out, res)
continue
}
if *putOnly {
out = append(out, res)
continue
}
getStart := time.Now()
err = session.Download(base, *bucket, iterKey, getHost)
res.getDur = time.Since(getStart)
if err != nil {
res.err = fmt.Errorf("GET: %w", err)
out = append(out, res)
continue
}
getCRC := crc32.ChecksumIEEE(getHost)
if getCRC != putCRC {
res.status = fmt.Sprintf("MISMATCH put=%08x get=%08x", putCRC, getCRC)
}
out = append(out, res)
continue
}
getStart := time.Now()
err := session.Download(base, *bucket, iterKey, getHost)
res.getDur = time.Since(getStart)
if err != nil {
res.err = fmt.Errorf("GET: %w", err)
}
out = append(out, res)
}
return out, time.Since(benchmarkStart)
}
// runConcurrentStdS3 exercises plain S3 PUT/GET (request body carries the
// actual data, no RDMA headers) so it can be benchmarked against the
// cuObject RDMA path in runConcurrent.
func runConcurrentStdS3(base *s3lib.Client, size int) ([]iterationResult, time.Duration) {
putHost := make([]byte, size)
getHost := make([]byte, size)
out := make([]iterationResult, 0, *iterations)
benchmarkStart := time.Now()
for i := 1; i <= *iterations; i++ {
res := iterationResult{index: i, status: "OK"}
iterKey := *key
if !*getOnly {
if _, err := crand.Read(putHost); err != nil {
res.err = fmt.Errorf("fill PUT buffer: %w", err)
out = append(out, res)
continue
}
putCRC := crc32.ChecksumIEEE(putHost)
putStart := time.Now()
err := s3PutObject(base, *bucket, iterKey, putHost)
res.putDur = time.Since(putStart)
if err != nil {
res.err = fmt.Errorf("PUT: %w", err)
out = append(out, res)
continue
}
if *putOnly {
out = append(out, res)
continue
}
getStart := time.Now()
err = s3GetObject(base, *bucket, iterKey, getHost)
res.getDur = time.Since(getStart)
if err != nil {
res.err = fmt.Errorf("GET: %w", err)
out = append(out, res)
continue
}
getCRC := crc32.ChecksumIEEE(getHost)
if getCRC != putCRC {
res.status = fmt.Sprintf("MISMATCH put=%08x get=%08x", putCRC, getCRC)
}
out = append(out, res)
continue
}
getStart := time.Now()
err := s3GetObject(base, *bucket, iterKey, getHost)
res.getDur = time.Since(getStart)
if err != nil {
res.err = fmt.Errorf("GET: %w", err)
}
out = append(out, res)
}
return out, time.Since(benchmarkStart)
}
// s3PutObject issues a standard S3 PUT carrying data in the request body.
func s3PutObject(base *s3lib.Client, bucket, key string, data []byte) error {
_, err := base.PutObject(context.Background(), &s3lib.PutObjectInput{
Bucket: aws.String(bucket),
Key: aws.String(key),
Body: bytes.NewReader(data),
ContentLength: aws.Int64(int64(len(data))),
})
return err
}
// s3GetObject issues a standard S3 GET and reads the full body into dst.
func s3GetObject(base *s3lib.Client, bucket, key string, dst []byte) error {
out, err := base.GetObject(context.Background(), &s3lib.GetObjectInput{
Bucket: aws.String(bucket),
Key: aws.String(key),
})
if err != nil {
return err
}
defer out.Body.Close()
// io.ReadFull only guarantees len(dst) bytes were read; it doesn't
// notice a larger object whose extra bytes are simply left unread.
// Compare against the reported object size so a too-small dst is
// caught instead of silently benchmarking a truncated prefix.
if out.ContentLength != nil && *out.ContentLength != int64(len(dst)) {
return fmt.Errorf("object size %d does not match expected size %d", *out.ContentLength, len(dst))
}
_, err = io.ReadFull(out.Body, dst)
return err
}
func newS3Client(endpoint, access, secret, region string) *s3lib.Client {
return s3lib.New(s3lib.Options{
BaseEndpoint: aws.String(endpoint),
Region: region,
Credentials: aws.NewCredentialsCache(credentials.NewStaticCredentialsProvider(access, secret, "")),
UsePathStyle: true,
})
}
func ensureBucket(c *s3lib.Client, bucket string) error {
_, err := c.HeadBucket(context.Background(), &s3lib.HeadBucketInput{Bucket: aws.String(bucket)})
if err == nil {
return nil
}
var notFound *s3types.NotFound
if !errors.As(err, &notFound) {
return fmt.Errorf("head bucket %q: %w", bucket, err)
}
_, err = c.CreateBucket(context.Background(), &s3lib.CreateBucketInput{
Bucket: aws.String(bucket),
CreateBucketConfiguration: &s3types.CreateBucketConfiguration{},
})
if err != nil {
return fmt.Errorf("create bucket %q: %w", bucket, err)
}
fmt.Printf("cuobjtest: created bucket %q\n", bucket)
return nil
}
// printIterationResult renders one benchmark iteration according to mode.
func printIterationResult(size int, r iterationResult) {
if *putOnly {
fmt.Printf("%-4d %-14s %-10s\n", r.index, fmtDur(r.putDur), fmtGBps(size, r.putDur))
return
}
if *getOnly {
fmt.Printf("%-4d %-14s %-10s\n", r.index, fmtDur(r.getDur), fmtGBps(size, r.getDur))
return
}
fmt.Printf("%-4d %-14s %-10s %-14s %-10s %s\n",
r.index,
fmtDur(r.putDur), fmtGBps(size, r.putDur),
fmtDur(r.getDur), fmtGBps(size, r.getDur),
r.status)
}
// fmtDur renders a human-friendly duration.
func fmtDur(d time.Duration) string {
ms := d.Seconds() * 1000
if ms >= 1000 {
return fmt.Sprintf("%.2f s", d.Seconds())
}
return fmt.Sprintf("%.2f ms", ms)
}
// fmtGBps formats bytes-per-second throughput as decimal GB/s.
func fmtGBps(bytes int, d time.Duration) string {
if d <= 0 {
return "-"
}
return fmt.Sprintf("%.3f GB/s", float64(bytes)/d.Seconds()/1e9)
}
// formatSize renders a byte size using binary units when evenly divisible.
func formatSize(n int) string {
switch {
case n >= 1<<30 && n%(1<<30) == 0:
return fmt.Sprintf("%d GiB", n>>30)
case n >= 1<<20 && n%(1<<20) == 0:
return fmt.Sprintf("%d MiB", n>>20)
case n >= 1<<10 && n%(1<<10) == 0:
return fmt.Sprintf("%d KiB", n>>10)
default:
return fmt.Sprintf("%d B", n)
}
}
// durStats returns average, minimum, and maximum durations.
func durStats(ds []time.Duration) (avg, min, max time.Duration) {
min = time.Duration(math.MaxInt64)
var sum time.Duration
for _, d := range ds {
sum += d
if d < min {
min = d
}
if d > max {
max = d
}
}
avg = sum / time.Duration(len(ds))
return
}
// sumDurs returns the sum of the given durations.
func sumDurs(ds []time.Duration) time.Duration {
var sum time.Duration
for _, d := range ds {
sum += d
}
return sum
}
// parseSize parses byte sizes like 64MiB, 1GiB, or 1000000.
func parseSize(s string) (int, error) {
units := []struct {
suffix string
mult int64
}{
{"GiB", 1 << 30}, {"GB", 1_000_000_000},
{"MiB", 1 << 20}, {"MB", 1_000_000},
{"KiB", 1 << 10}, {"KB", 1_000},
{"G", 1 << 30}, {"M", 1 << 20}, {"K", 1 << 10},
{"B", 1},
}
upper := strings.ToUpper(strings.TrimSpace(s))
for _, u := range units {
if strings.HasSuffix(upper, strings.ToUpper(u.suffix)) {
numStr := strings.TrimSpace(s[:len(s)-len(u.suffix)])
n, err := strconv.ParseFloat(numStr, 64)
if err != nil {
return 0, fmt.Errorf("invalid size %q", s)
}
return int(n * float64(u.mult)), nil
}
}
n, err := strconv.Atoi(strings.TrimSpace(s))
if err != nil {
return 0, fmt.Errorf("invalid size %q", s)
}
return n, nil
}
// fatalf prints an error and exits with status code 1.
func fatalf(format string, args ...any) {
fmt.Fprintf(os.Stderr, format+"\n", args...)
os.Exit(1)
}