From c0f400f47dddde3eeb4ef85e67352d1b632f4a47 Mon Sep 17 00:00:00 2001
From: Anton Kolesnikov <anton.e.kolesnikov@gmail.com>
Date: Thu, 17 Oct 2024 17:06:52 +0800
Subject: [PATCH] fix(v2): allocate dedicated buffer for tsdb

---
 pkg/experiment/query_backend/block/dataset.go      | 9 ++++++---
 pkg/experiment/query_backend/block/section_tsdb.go | 8 +++++---
 pkg/objstore/read_only_file.go                     | 3 +++
 3 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/pkg/experiment/query_backend/block/dataset.go b/pkg/experiment/query_backend/block/dataset.go
index 47bfc470a1..c0eac51330 100644
--- a/pkg/experiment/query_backend/block/dataset.go
+++ b/pkg/experiment/query_backend/block/dataset.go
@@ -55,15 +55,18 @@ func WithDatasetMaxSizeLoadInMemory(size int) DatasetOption {
 // is only initialized once. While it is possible to open the dataset
 // repeatedly after close, the caller must pass the failure reason to
 // the CloseWithError call, preventing further use, if applicable.
-func (s *Dataset) Open(ctx context.Context, sections ...Section) (err error) {
+func (s *Dataset) Open(ctx context.Context, sections ...Section) error {
 	return s.refs.IncErr(func() error {
-		return s.open(ctx, sections...)
+		if err := s.open(ctx, sections...); err != nil {
+			return fmt.Errorf("%w (%s)", err, s.obj.meta.Id)
+		}
+		return nil
 	})
 }
 
 func (s *Dataset) open(ctx context.Context, sections ...Section) (err error) {
 	if s.err != nil {
-		// The tenant dataset has been already closed with an error.
+		// The tenant dataset has already been closed with an error.
 		return s.err
 	}
 	if err = s.obj.Open(ctx); err != nil {
diff --git a/pkg/experiment/query_backend/block/section_tsdb.go b/pkg/experiment/query_backend/block/section_tsdb.go
index ac93674014..7aa3a70b59 100644
--- a/pkg/experiment/query_backend/block/section_tsdb.go
+++ b/pkg/experiment/query_backend/block/section_tsdb.go
@@ -39,11 +39,13 @@ type tsdbBuffer struct {
 }
 
 func (b *tsdbBuffer) Close() (err error) {
-	if b.buf != nil {
-		bufferpool.Put(b.buf)
-	}
 	if b.index != nil {
 		err = b.index.Close()
+		b.index = nil
+	}
+	if b.buf != nil {
+		bufferpool.Put(b.buf)
+		b.buf = nil
 	}
 	return err
 }
diff --git a/pkg/objstore/read_only_file.go b/pkg/objstore/read_only_file.go
index 350010262a..d9151bac32 100644
--- a/pkg/objstore/read_only_file.go
+++ b/pkg/objstore/read_only_file.go
@@ -54,6 +54,9 @@ func download(ctx context.Context, name string, src BucketReader, dir string) (f
 	if err != nil {
 		return nil, err
 	}
+	defer func() {
+		_ = dst.Close()
+	}()
 	buf := bufferpool.GetBuffer(32 << 10)
 	defer bufferpool.Put(buf)
 	buf.B = buf.B[:cap(buf.B)]