From d3b36cfd2cc7dc939cf4481880fc73670890dfad Mon Sep 17 00:00:00 2001
From: Radu Berinde <radu@cockroachlabs.com>
Date: Fri, 18 Oct 2024 07:32:57 -0700
Subject: [PATCH] sstable: clarify BufferHandle ownership by iterators

---
 sstable/block/block.go         | 16 ++++++++++++----
 sstable/reader.go              |  3 +++
 sstable/reader_iter_two_lvl.go |  7 +++++--
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/sstable/block/block.go b/sstable/block/block.go
index 8466b217ff..57a8b31791 100644
--- a/sstable/block/block.go
+++ b/sstable/block/block.go
@@ -164,13 +164,17 @@ const MetadataSize = 328
 const _ uint = -(MetadataSize % 8)
 
 // DataBlockIterator is a type constraint for implementations of block iterators
-// over data blocks.
+// over data blocks. It's implemented by *rowblk.Iter and *colblk.DataBlockIter.
 type DataBlockIterator interface {
 	base.InternalIterator
 
 	// Handle returns the handle to the block.
 	Handle() BufferHandle
 	// InitHandle initializes the block from the provided buffer handle.
+	//
+	// The iterator takes ownership of the BufferHandle and releases it when it is
+	// closed (or re-initialized with another handle). This happens even in error
+	// cases.
 	InitHandle(base.Compare, base.Split, BufferHandle, IterTransforms) error
 	// Valid returns true if the iterator is currently positioned at a valid KV.
 	Valid() bool
@@ -198,13 +202,17 @@ type DataBlockIterator interface {
 	IsDataInvalidated() bool
 }
 
-// IndexBlockIterator is an interface for implementations of block
-// iterators over index blocks. It's currently satisifed by the
-// *rowblk.IndexIter type.
+// IndexBlockIterator is an interface for implementations of block iterators
+// over index blocks. It's implemented by *rowblk.IndexIter and
+// *colblk.IndexBlockIter.
 type IndexBlockIterator interface {
 	// Init initializes the block iterator from the provided block.
 	Init(base.Compare, base.Split, []byte, IterTransforms) error
 	// InitHandle initializes an iterator from the provided block handle.
+	//
+	// The iterator takes ownership of the BufferHandle and releases it when it is
+	// closed (or re-initialized with another handle). This happens even in error
+	// cases.
 	InitHandle(base.Compare, base.Split, BufferHandle, IterTransforms) error
 	// Valid returns true if the iterator is currently positioned at a valid
 	// block handle.
diff --git a/sstable/reader.go b/sstable/reader.go
index ea8dcf3f8e..132ea82ba3 100644
--- a/sstable/reader.go
+++ b/sstable/reader.go
@@ -861,6 +861,9 @@ func estimateDiskUsage[I any, PI indexBlockIterator[I]](
 	if err != nil {
 		return 0, err
 	}
+	// We are using InitHandle below but we never Close those iterators, which
+	// allows us to release the index handle ourselves.
+	// TODO(radu): clean this up.
 	defer indexH.Release()
 
 	// Iterators over the bottom-level index blocks containing start and end.
diff --git a/sstable/reader_iter_two_lvl.go b/sstable/reader_iter_two_lvl.go
index e9e9a97582..4e0b11ea34 100644
--- a/sstable/reader_iter_two_lvl.go
+++ b/sstable/reader_iter_two_lvl.go
@@ -69,10 +69,13 @@ func (i *twoLevelIterator[I, PI, D, PD]) loadSecondLevelIndexBlock(dir int8) loa
 		// blockIntersects
 	}
 	indexBlock, err := i.secondLevel.reader.readIndexBlock(i.secondLevel.ctx, i.secondLevel.readBlockEnv, i.secondLevel.indexFilterRH, bhp.Handle)
-	if err == nil {
-		err = PI(&i.secondLevel.index).InitHandle(i.secondLevel.cmp, i.secondLevel.reader.Split, indexBlock, i.secondLevel.transforms)
+	if err != nil {
+		i.secondLevel.err = err
+		return loadBlockFailed
 	}
+	err = PI(&i.secondLevel.index).InitHandle(i.secondLevel.cmp, i.secondLevel.reader.Split, indexBlock, i.secondLevel.transforms)
 	if err != nil {
+		PI(&i.secondLevel.index).Invalidate()
 		i.secondLevel.err = err
 		return loadBlockFailed
 	}