xk6-frostfs/internal/local/rawclient/rawclient.go

// Package rawclient provides a basic interface to the local storage engine.
// It can be used as a base for more complex load clients backed by local storage.
package rawclient

import (
	"context"
	"fmt"
	"time"

	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/engine"
	cid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object"
	oid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object/id"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/user"
)

// RawClient is a client to the local storage engine instance.
type RawClient struct {
	*config
	ng      *engine.StorageEngine
	ownerID *user.ID
}

// New returns a RawClient from the provided options.
func New(ng *engine.StorageEngine, opts ...Option) *RawClient {
	cfg := defaultConfig()
	for _, opt := range opts {
		opt(cfg)
	}
	client := &RawClient{cfg, ng, &user.ID{}}
	user.IDFromKey(client.ownerID, client.key.PublicKey)
	return client
}

func (c *RawClient) Put(ctx context.Context, containerID cid.ID, ownerID *user.ID, headers map[string]string, payload []byte) (oid.ID, error) {
	sz := len(payload)

	attrs := make([]object.Attribute, len(headers))
	{
		ind := 0
		for k, v := range headers {
			attrs[ind].SetKey(k)
			attrs[ind].SetValue(v)
			ind++
		}
	}

	// Note that the key is a required option, so this is never empty.
	if ownerID == nil {
		ownerID = c.ownerID
	}

	obj := object.New()
	obj.SetContainerID(containerID)
	obj.SetOwnerID(*ownerID)
	obj.SetAttributes(attrs...)
	obj.SetPayload(payload)
	obj.SetPayloadSize(uint64(sz))
	object.CalculateAndSetPayloadChecksum(obj) // needed for metabase key

	id, err := object.CalculateID(obj)
	if err != nil {
		return oid.ID{}, fmt.Errorf("calculating object id: %v", err)
	}
	obj.SetID(id)

	if err := object.CalculateAndSetSignature(c.key, obj); err != nil {
		return oid.ID{}, fmt.Errorf("calculating signature: %v", err)
	}

	var req engine.PutPrm
	req.WithObject(obj)

	start := time.Now()
	err = c.ng.Put(ctx, req)
	c.onPut(uint64(sz), err, time.Since(start))
	if err != nil {
		return oid.ID{}, err
	}

	return id, nil
}

func (c *RawClient) Get(ctx context.Context, containerID cid.ID, objectID oid.ID) (*object.Object, error) {
	var addr oid.Address
	addr.SetContainer(containerID)
	addr.SetObject(objectID)

	var req engine.GetPrm
	req.WithAddress(addr)

	start := time.Now()
	res, err := c.ng.Get(ctx, req)

	var sz uint64
	obj := res.Object()
	if obj != nil {
		sz = uint64(len(obj.Payload()))
	}

	c.onGet(sz, err, time.Since(start))
	return obj, nil
}

func (c *RawClient) Delete(ctx context.Context, containerID cid.ID, objectID oid.ID) error {
	var addr oid.Address
	addr.SetContainer(containerID)
	addr.SetObject(objectID)

	var req engine.DeletePrm
	req.WithAddress(addr)

	start := time.Now()
	_, err := c.ng.Delete(ctx, req)
	c.onDelete(err, time.Since(start))
	return err
}

func (c *RawClient) OwnerID() *user.ID {
	return c.ownerID
}