diff --git a/object/transformer/channel.go b/object/transformer/channel.go index e1c009b..b7a50a9 100644 --- a/object/transformer/channel.go +++ b/object/transformer/channel.go @@ -1,6 +1,8 @@ package transformer import ( + "context" + objectSDK "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object" "github.com/nspcc-dev/neo-go/pkg/util/slice" ) @@ -20,23 +22,27 @@ func NewChannelTarget(ch chan<- *objectSDK.Object) ObjectTarget { } // WriteHeader implements the ObjectTarget interface. -func (c *chanTarget) WriteHeader(object *objectSDK.Object) error { +func (c *chanTarget) WriteHeader(_ context.Context, object *objectSDK.Object) error { c.header = object return nil } // Write implements the ObjectTarget interface. -func (c *chanTarget) Write(p []byte) (n int, err error) { +func (c *chanTarget) Write(_ context.Context, p []byte) (n int, err error) { c.payload = append(c.payload, p...) return len(p), nil } // Close implements the ObjectTarget interface. -func (c *chanTarget) Close() (*AccessIdentifiers, error) { +func (c *chanTarget) Close(ctx context.Context) (*AccessIdentifiers, error) { if len(c.payload) != 0 { c.header.SetPayload(slice.Copy(c.payload)) } - c.ch <- c.header + select { + case c.ch <- c.header: + case <-ctx.Done(): + return nil, ctx.Err() + } c.header = nil c.payload = nil diff --git a/object/transformer/channel_test.go b/object/transformer/channel_test.go index 8eef6b7..2487391 100644 --- a/object/transformer/channel_test.go +++ b/object/transformer/channel_test.go @@ -1,6 +1,7 @@ package transformer import ( + "context" "crypto/rand" "testing" @@ -29,8 +30,9 @@ func TestChannelTarget(t *testing.T) { payload := make([]byte, maxSize*2+maxSize/2) _, _ = rand.Read(payload) - expectedIDs := writeObject(t, testTarget, hdr, payload) - actualIDs := writeObject(t, chTarget, hdr, payload) + ctx := context.Background() + expectedIDs := writeObject(t, ctx, testTarget, hdr, payload) + actualIDs := writeObject(t, ctx, chTarget, hdr, payload) _ = expectedIDs _ = actualIDs //require.Equal(t, expectedIDs, actualIDs) diff --git a/object/transformer/transformer.go b/object/transformer/transformer.go index efe3d98..7493387 100644 --- a/object/transformer/transformer.go +++ b/object/transformer/transformer.go @@ -1,6 +1,7 @@ package transformer import ( + "context" "crypto/ecdsa" "crypto/sha256" "fmt" @@ -52,7 +53,7 @@ func NewPayloadSizeLimiter(p Params) ObjectTarget { } } -func (s *payloadSizeLimiter) WriteHeader(hdr *object.Object) error { +func (s *payloadSizeLimiter) WriteHeader(_ context.Context, hdr *object.Object) error { s.current = fromObject(hdr) s.initialize() @@ -60,16 +61,16 @@ func (s *payloadSizeLimiter) WriteHeader(hdr *object.Object) error { return nil } -func (s *payloadSizeLimiter) Write(p []byte) (int, error) { - if err := s.writeChunk(p); err != nil { +func (s *payloadSizeLimiter) Write(ctx context.Context, p []byte) (int, error) { + if err := s.writeChunk(ctx, p); err != nil { return 0, err } return len(p), nil } -func (s *payloadSizeLimiter) Close() (*AccessIdentifiers, error) { - return s.release(true) +func (s *payloadSizeLimiter) Close(ctx context.Context) (*AccessIdentifiers, error) { + return s.release(ctx, true) } func (s *payloadSizeLimiter) initialize() { @@ -132,7 +133,7 @@ func (s *payloadSizeLimiter) initPayloadHashers() { } // nolint: funlen -func (s *payloadSizeLimiter) release(finalize bool) (*AccessIdentifiers, error) { +func (s *payloadSizeLimiter) release(ctx context.Context, finalize bool) (*AccessIdentifiers, error) { // Arg finalize is true only when called from Close method. // We finalize parent and generate linking objects only if it is more // than 1 object in split-chain. @@ -185,11 +186,11 @@ func (s *payloadSizeLimiter) release(finalize bool) (*AccessIdentifiers, error) return nil, fmt.Errorf("could not finalize object: %w", err) } - if err := s.NextTarget.WriteHeader(s.current); err != nil { + if err := s.NextTarget.WriteHeader(ctx, s.current); err != nil { return nil, fmt.Errorf("could not write header to next target: %w", err) } - if _, err := s.NextTarget.Close(); err != nil { + if _, err := s.NextTarget.Close(ctx); err != nil { return nil, fmt.Errorf("could not close next target: %w", err) } @@ -209,7 +210,7 @@ func (s *payloadSizeLimiter) release(finalize bool) (*AccessIdentifiers, error) s.initializeLinking(ids.ParentHeader) s.initializeCurrent() - if _, err := s.release(false); err != nil { + if _, err := s.release(ctx, false); err != nil { return nil, fmt.Errorf("could not release linking object: %w", err) } } @@ -224,7 +225,7 @@ func (s *payloadSizeLimiter) initializeLinking(parHdr *object.Object) { s.current.SetSplitID(s.splitID) } -func (s *payloadSizeLimiter) writeChunk(chunk []byte) error { +func (s *payloadSizeLimiter) writeChunk(ctx context.Context, chunk []byte) error { for { // statement is true if the previous write of bytes reached exactly the boundary. if s.written > 0 && s.written%s.MaxSize == 0 { @@ -233,7 +234,7 @@ func (s *payloadSizeLimiter) writeChunk(chunk []byte) error { } // we need to release current object - if _, err := s.release(false); err != nil { + if _, err := s.release(ctx, false); err != nil { return fmt.Errorf("could not release object: %w", err) } @@ -252,7 +253,7 @@ func (s *payloadSizeLimiter) writeChunk(chunk []byte) error { cut = leftToEdge } - if err := s.writeHashes(chunk[:cut]); err != nil { + if err := s.writeHashes(ctx, chunk[:cut]); err != nil { return fmt.Errorf("could not write chunk to target: %w", err) } @@ -268,8 +269,8 @@ func (s *payloadSizeLimiter) writeChunk(chunk []byte) error { } } -func (s *payloadSizeLimiter) writeHashes(chunk []byte) error { - _, err := s.NextTarget.Write(chunk) +func (s *payloadSizeLimiter) writeHashes(ctx context.Context, chunk []byte) error { + _, err := s.NextTarget.Write(ctx, chunk) if err != nil { return err } diff --git a/object/transformer/transformer_test.go b/object/transformer/transformer_test.go index 1e06bc5..270697e 100644 --- a/object/transformer/transformer_test.go +++ b/object/transformer/transformer_test.go @@ -1,6 +1,7 @@ package transformer import ( + "context" "crypto/rand" "crypto/sha256" "testing" @@ -31,7 +32,7 @@ func TestTransformer(t *testing.T) { expectedPayload := make([]byte, maxSize*2+maxSize/2) _, _ = rand.Read(expectedPayload) - ids := writeObject(t, target, hdr, expectedPayload) + ids := writeObject(t, context.Background(), target, hdr, expectedPayload) require.Equal(t, 4, len(tt.objects)) // 3 parts + linking object var actualPayload []byte @@ -85,13 +86,13 @@ func newObject(cnr cid.ID) *objectSDK.Object { return hdr } -func writeObject(t *testing.T, target ObjectTarget, header *objectSDK.Object, payload []byte) *AccessIdentifiers { - require.NoError(t, target.WriteHeader(header)) +func writeObject(t *testing.T, ctx context.Context, target ObjectTarget, header *objectSDK.Object, payload []byte) *AccessIdentifiers { + require.NoError(t, target.WriteHeader(ctx, header)) - _, err := target.Write(payload) + _, err := target.Write(ctx, payload) require.NoError(t, err) - ids, err := target.Close() + ids, err := target.Close(ctx) require.NoError(t, err) return ids @@ -112,18 +113,19 @@ func benchmarkTransformer(b *testing.B, header *objectSDK.Object, payloadSize in const maxSize = 64 * 1024 * 1024 payload := make([]byte, payloadSize) + ctx := context.Background() b.ReportAllocs() b.ResetTimer() for i := 0; i < b.N; i++ { f, _ := newPayloadSizeLimiter(maxSize, benchTarget{}) - if err := f.WriteHeader(header); err != nil { + if err := f.WriteHeader(ctx, header); err != nil { b.Fatalf("write header: %v", err) } - if _, err := f.Write(payload); err != nil { + if _, err := f.Write(ctx, payload); err != nil { b.Fatalf("write: %v", err) } - if _, err := f.Close(); err != nil { + if _, err := f.Close(ctx); err != nil { b.Fatalf("close: %v", err) } } @@ -152,15 +154,15 @@ func (s dummyEpochSource) CurrentEpoch() uint64 { type benchTarget struct{} -func (benchTarget) WriteHeader(object *objectSDK.Object) error { +func (benchTarget) WriteHeader(_ context.Context, object *objectSDK.Object) error { return nil } -func (benchTarget) Write(p []byte) (n int, err error) { +func (benchTarget) Write(_ context.Context, p []byte) (n int, err error) { return len(p), nil } -func (benchTarget) Close() (*AccessIdentifiers, error) { +func (benchTarget) Close(context.Context) (*AccessIdentifiers, error) { return nil, nil } @@ -170,17 +172,17 @@ type testTarget struct { objects []*objectSDK.Object } -func (tt *testTarget) WriteHeader(object *objectSDK.Object) error { +func (tt *testTarget) WriteHeader(_ context.Context, object *objectSDK.Object) error { tt.current = object return nil } -func (tt *testTarget) Write(p []byte) (n int, err error) { +func (tt *testTarget) Write(_ context.Context, p []byte) (n int, err error) { tt.payload = append(tt.payload, p...) return len(p), nil } -func (tt *testTarget) Close() (*AccessIdentifiers, error) { +func (tt *testTarget) Close(_ context.Context) (*AccessIdentifiers, error) { tt.current.SetPayload(tt.payload) tt.objects = append(tt.objects, tt.current) tt.current = nil diff --git a/object/transformer/types.go b/object/transformer/types.go index 6f5f159..a7e827c 100644 --- a/object/transformer/types.go +++ b/object/transformer/types.go @@ -1,7 +1,7 @@ package transformer import ( - "io" + "context" "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object" oid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object/id" @@ -31,14 +31,14 @@ type ObjectTarget interface { // that depends on the implementation. // // Must not be called after Close call. - WriteHeader(*object.Object) error + WriteHeader(context.Context, *object.Object) error // Write writes object payload chunk. // // Can be called multiple times. // // Must not be called after Close call. - io.Writer + Write(context.Context, []byte) (int, error) // Close is used to finish object writing. // @@ -48,7 +48,7 @@ type ObjectTarget interface { // Must be called no more than once. Control remains with the caller. // Re-calling can lead to undefined behavior // that depends on the implementation. - Close() (*AccessIdentifiers, error) + Close(context.Context) (*AccessIdentifiers, error) } // TargetInitializer represents ObjectTarget constructor.