frostfs-sdk-go/client/object_put_transformer.go

package client

import (
	"context"

	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object/transformer"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

func (c *Client) objectPutInitTransformer(prm PrmObjectPutInit) (*objectWriterTransformer, error) {
	var w objectWriterTransformer
	w.it = internalTarget{
		client: c,
		prm:    prm,
	}
	key := &c.prm.key
	if prm.key != nil {
		key = prm.key
	}
	w.ot = transformer.NewPayloadSizeLimiter(transformer.Params{
		Key:                    key,
		NextTargetInit:         func() transformer.ObjectTarget { return &w.it },
		MaxSize:                prm.maxSize,
		WithoutHomomorphicHash: prm.withoutHomomorphicHash,
		NetworkState:           prm.epochSource,
	})
	return &w, nil
}

type objectWriterTransformer struct {
	ot  transformer.ObjectTarget
	it  internalTarget
	err error
}

func (x *objectWriterTransformer) WriteHeader(ctx context.Context, hdr object.Object) bool {
	x.err = x.ot.WriteHeader(ctx, &hdr)
	return x.err == nil
}

func (x *objectWriterTransformer) WritePayloadChunk(ctx context.Context, chunk []byte) bool {
	_, x.err = x.ot.Write(ctx, chunk)
	return x.err == nil
}

func (x *objectWriterTransformer) Close(ctx context.Context) (*ResObjectPut, error) {
	ai, err := x.ot.Close(ctx)
	if err != nil {
		return nil, err
	}

	if ai != nil && ai.ParentID != nil {
		x.it.res.obj = *ai.ParentID
	}
	return x.it.res, nil
}

type internalTarget struct {
	current   *object.Object
	client    *Client
	res       *ResObjectPut
	prm       PrmObjectPutInit
	payload   []byte
	useStream bool
}

func (it *internalTarget) WriteHeader(_ context.Context, object *object.Object) error {
	it.current = object
	return nil
}

func (it *internalTarget) Write(_ context.Context, p []byte) (n int, err error) {
	it.payload = append(it.payload, p...)
	return len(p), nil
}

func (it *internalTarget) Close(ctx context.Context) (*transformer.AccessIdentifiers, error) {
	it.current.SetPayload(it.payload)

	putSingleImplemented, err := it.tryPutSingle(ctx)
	if putSingleImplemented {
		return nil, err
	}
	it.useStream = true
	return nil, it.putAsStream(ctx)
}

func (it *internalTarget) putAsStream(ctx context.Context) error {
	wrt, err := it.client.objectPutInitRaw(ctx, it.prm)
	if err != nil {
		return err
	}
	if wrt.WriteHeader(ctx, *it.current) {
		wrt.WritePayloadChunk(ctx, it.current.Payload())
	}
	it.res, err = wrt.Close(ctx)
	it.current = nil
	it.payload = nil
	return err
}

func (it *internalTarget) tryPutSingle(ctx context.Context) (bool, error) {
	if it.useStream {
		return false, nil
	}
	var prm PrmObjectPutSingle
	prm.SetCopiesNumber(it.prm.copyNum)
	prm.SetObject(it.current.ToV2())
	prm.UseKey(prm.key)
	prm.meta = it.prm.meta

	res, err := it.client.ObjectPutSingle(ctx, prm)
	if err != nil && status.Code(err) == codes.Unimplemented {
		return false, err
	}

	if err == nil {
		id, _ := it.current.ID()
		it.res = &ResObjectPut{
			statusRes: res.statusRes,
			obj:       id,
		}
	}
	it.current = nil
	it.payload = nil
	return true, err
}