xk6-frostfs/internal/s3/client.go

165 lines
3.8 KiB
Go
Raw Normal View History

package s3
import (
"bytes"
"context"
"crypto/sha256"
"encoding/hex"
"strconv"
"time"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/service/s3"
"github.com/aws/aws-sdk-go-v2/service/s3/types"
"github.com/dop251/goja"
"github.com/nspcc-dev/xk6-neofs/internal/stats"
"go.k6.io/k6/js/modules"
"go.k6.io/k6/metrics"
)
type (
Client struct {
vu modules.VU
cli *s3.Client
}
PutResponse struct {
Success bool
Error string
}
GetResponse struct {
Success bool
Error string
}
CreateBucketResponse struct {
Success bool
Error string
}
VerifyHashResponse struct {
Success bool
Error string
}
)
func (c *Client) Put(bucket, key string, payload goja.ArrayBuffer) PutResponse {
rdr := bytes.NewReader(payload.Bytes())
sz := rdr.Size()
stats.Report(c.vu, objPutTotal, 1)
start := time.Now()
_, err := c.cli.PutObject(c.vu.Context(), &s3.PutObjectInput{
Bucket: aws.String(bucket),
Key: aws.String(key),
Body: rdr,
})
if err != nil {
stats.Report(c.vu, objPutFails, 1)
return PutResponse{Success: false, Error: err.Error()}
}
stats.ReportDataSent(c.vu, float64(sz))
stats.Report(c.vu, objPutDuration, metrics.D(time.Since(start)))
return PutResponse{Success: true}
}
func (c *Client) Get(bucket, key string) GetResponse {
stats.Report(c.vu, objGetTotal, 1)
start := time.Now()
var objSize = 0
err := get(c.cli, bucket, key, func(chunk []byte) {
objSize += len(chunk)
})
if err != nil {
stats.Report(c.vu, objGetFails, 1)
return GetResponse{Success: false, Error: err.Error()}
}
stats.Report(c.vu, objGetDuration, metrics.D(time.Since(start)))
stats.ReportDataReceived(c.vu, float64(objSize))
return GetResponse{Success: true}
}
func get(
c *s3.Client,
bucket string,
key string,
onDataChunk func(chunk []byte),
) error {
var buf = make([]byte, 4*1024)
obj, err := c.GetObject(context.Background(), &s3.GetObjectInput{
Bucket: aws.String(bucket),
Key: aws.String(key),
})
if err != nil {
return err
}
for {
n, err := obj.Body.Read(buf)
if n > 0 {
onDataChunk(buf[:n])
}
if err != nil {
break
}
}
return nil
}
func (c *Client) VerifyHash(bucket, key, expectedHash string) VerifyHashResponse {
hasher := sha256.New()
err := get(c.cli, bucket, key, func(data []byte) {
hasher.Write(data)
})
if err != nil {
return VerifyHashResponse{Success: false, Error: err.Error()}
}
actualHash := hex.EncodeToString(hasher.Sum(nil))
if actualHash != expectedHash {
return VerifyHashResponse{Success: true, Error: "hash mismatch"}
}
return VerifyHashResponse{Success: true}
}
func (c *Client) CreateBucket(bucket string, params map[string]string) CreateBucketResponse {
stats.Report(c.vu, createBucketTotal, 1)
var err error
var lockEnabled bool
if lockEnabledStr, ok := params["lock_enabled"]; ok {
if lockEnabled, err = strconv.ParseBool(lockEnabledStr); err != nil {
stats.Report(c.vu, createBucketFails, 1)
return CreateBucketResponse{Success: false, Error: "invalid lock_enabled params"}
}
}
var bucketConfiguration *types.CreateBucketConfiguration
if locationConstraint, ok := params["location_constraint"]; ok {
bucketConfiguration = &types.CreateBucketConfiguration{
LocationConstraint: types.BucketLocationConstraint(locationConstraint),
}
}
start := time.Now()
_, err = c.cli.CreateBucket(c.vu.Context(), &s3.CreateBucketInput{
Bucket: aws.String(bucket),
ACL: types.BucketCannedACL(params["acl"]),
CreateBucketConfiguration: bucketConfiguration,
ObjectLockEnabledForBucket: lockEnabled,
})
if err != nil {
stats.Report(c.vu, createBucketFails, 1)
return CreateBucketResponse{Success: false, Error: err.Error()}
}
stats.Report(c.vu, createBucketDuration, metrics.D(time.Since(start)))
return CreateBucketResponse{Success: true}
}