colblk: ensure IndexIter releases BufferHandles

sstable/colblk/index_block.go

@@ -53,7 +53,7 @@ const (
 )
 
 // Init initializes the index block writer.
-func (w *IndexBlockWriter) Init(blockProperties ...ColumnWriter) {
+func (w *IndexBlockWriter) Init() {
  w.separators.Init()
  w.offsets.Init()
  w.lengths.InitWithDefault()
@@ -191,13 +191,15 @@ var _ block.IndexBlockIterator = (*IndexIter)(nil)
 
 // InitReader initializes an index iterator from the provided reader.
 func (i *IndexIter) InitReader(r *IndexReader) {
- *i = IndexIter{r: r, n: int(r.br.header.Rows), allocReader: i.allocReader}
+ *i = IndexIter{r: r, n: int(r.br.header.Rows), h: i.h, allocReader: i.allocReader}
 }
 
 // Init initializes an iterator from the provided block data slice.
 func (i *IndexIter) Init(
  cmp base.Compare, split base.Split, blk []byte, transforms block.IterTransforms,
 ) error {
+ i.h.Release()
+ i.h = block.BufferHandle{}
  // TODO(jackson): Handle the transforms.
  i.allocReader.Init(blk)
  i.InitReader(&i.allocReader)
@@ -206,7 +208,7 @@ func (i *IndexIter) Init(
 
 // InitHandle initializes an iterator from the provided block handle.
 func (i *IndexIter) InitHandle(
- cmp base.Compare, split base.Split, block block.BufferHandle, transforms block.IterTransforms,
+ cmp base.Compare, split base.Split, blk block.BufferHandle, transforms block.IterTransforms,
 ) error {
  // TODO(jackson): Handle the transforms.
 
@@ -216,8 +218,10 @@ func (i *IndexIter) InitHandle(
  // common to open an iterator and perform just a few seeks, so avoiding the
  // overhead can be material.)
  i.h.Release()
- i.h = block
- return i.Init(cmp, split, i.h.Get(), transforms)
+ i.h = blk
+ i.allocReader.Init(i.h.Get())
+ i.InitReader(&i.allocReader)
+ return nil
 }
 
 // RowIndex returns the index of the block entry at the iterator's current
@@ -242,7 +246,6 @@ func (i *IndexIter) Valid() bool {
 // it was initialized with.
 func (i *IndexIter) Invalidate() {
  i.r = nil
- i.h = block.BufferHandle{}
 }
 
 // IsDataInvalidated returns true when the iterator has been invalidated

sstable/colblk/index_block_test.go

@@ -9,9 +9,13 @@ import (
  "fmt"
  "strconv"
  "strings"
+ "sync"
  "testing"
 
  "github.com/cockroachdb/datadriven"
+ "github.com/cockroachdb/pebble/internal/base"
+ "github.com/cockroachdb/pebble/internal/cache"
+ "github.com/cockroachdb/pebble/internal/testkeys"
  "github.com/cockroachdb/pebble/sstable/block"
  "github.com/stretchr/testify/require"
 )
@@ -88,3 +92,58 @@ func TestIndexBlock(t *testing.T) {
  }
  })
 }
+
+// TestIndexIterInitHandle exercises initializing an IndexIter through
+// InitHandle.
+func TestIndexIterInitHandle(t *testing.T) {
+ var w IndexBlockWriter
+ w.Init()
+ bh1 := block.Handle{Offset: 0, Length: 2000}
+ bh2 := block.Handle{Offset: 2008, Length: 1000}
+ w.AddBlockHandle([]byte("a"), bh1, nil)
+ w.AddBlockHandle([]byte("b"), bh2, nil)
+ b := w.Finish(w.Rows())
+
+ c := cache.New(10 << 10)
+ defer c.Unref()
+ v := block.Alloc(len(b), nil)
+ copy(v.Get(), b)
+ v.MakeHandle(c, cache.ID(1), base.DiskFileNum(1), 0).Release()
+
+ getBlockAndIterate := func(it *IndexIter) {
+ h := c.Get(cache.ID(1), base.DiskFileNum(1), 0)
+ require.NotNil(t, h.Get())
+ require.NoError(t, it.InitHandle(
+ testkeys.Comparer.Compare,
+ testkeys.Comparer.Split,
+ block.CacheBufferHandle(h),
+ block.NoTransforms))
+ defer it.Close()
+ require.True(t, it.First())
+ bh, err := it.BlockHandleWithProperties()
+ require.NoError(t, err)
+ require.Equal(t, bh1, bh.Handle)
+ require.True(t, it.Next())
+ bh, err = it.BlockHandleWithProperties()
+ require.NoError(t, err)
+ require.Equal(t, bh2, bh.Handle)
+ require.False(t, it.Next())
+ require.False(t, it.IsDataInvalidated())
+ it.Invalidate()
+ require.True(t, it.IsDataInvalidated())
+ }
+
+ const workers = 8
+ var wg sync.WaitGroup
+ wg.Add(workers)
+ for w := 0; w < workers; w++ {
+ go func() {
+ var iter IndexIter
+ defer wg.Done()
+ for i := 0; i < 10; i++ {
+ getBlockAndIterate(&iter)
+ }
+ }()
+ }
+ wg.Wait()
+}