diff --git a/pkg/services/object/qos.go b/pkg/services/object/qos.go index 145a316e..01eb1ea8 100644 --- a/pkg/services/object/qos.go +++ b/pkg/services/object/qos.go @@ -3,6 +3,8 @@ package object import ( "context" + "git.frostfs.info/TrueCloudLab/frostfs-node/internal/assert" + "git.frostfs.info/TrueCloudLab/frostfs-qos/tagging" "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/api/object" "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/api/session" ) @@ -120,13 +122,24 @@ type qosSendRecv[TReq qosVerificationHeader, TResp any] interface { type qosWriteStream[TReq qosVerificationHeader, TResp any] struct { s qosSendRecv[TReq, TResp] adj AdjustIOTag + + ioTag string + ioTagDefined bool } func (q *qosWriteStream[TReq, TResp]) CloseAndRecv(ctx context.Context) (TResp, error) { + if q.ioTagDefined { + ctx = tagging.ContextWithIOTag(ctx, q.ioTag) + } return q.s.CloseAndRecv(ctx) } func (q *qosWriteStream[TReq, TResp]) Send(ctx context.Context, req TReq) error { - ctx = q.adj.AdjustIncomingTag(ctx, req.GetVerificationHeader().GetBodySignature().GetKey()) + if !q.ioTagDefined { + ctx = q.adj.AdjustIncomingTag(ctx, req.GetVerificationHeader().GetBodySignature().GetKey()) + q.ioTag, q.ioTagDefined = tagging.IOTagFromContext(ctx) + } + assert.True(q.ioTagDefined, "io tag undefined after incoming tag adjustment") + ctx = tagging.ContextWithIOTag(ctx, q.ioTag) return q.s.Send(ctx, req) }