handling buffer overflow while creating new segmentBase

Author: Thejas-bhat

build.go

@@ -178,6 +178,8 @@ func InitSegmentBase(mem []byte, memCRC uint32, chunkMode uint32,
 	}
 	sb.updateSize()
 
+	fmt.Println("length of segment base mem", len(sb.mem))
+
 	// load the data/section starting offsets for each field
 	// by via the sectionsIndexOffset as starting point.
 	err := sb.loadFieldsNew()

new.go

@@ -17,6 +17,7 @@ package zap
 import (
 	"bytes"
 	"encoding/binary"
+	"fmt"
 	"math"
 	"sort"
 	"sync"
@@ -50,27 +51,29 @@ func (*ZapPlugin) newWithChunkMode(results []index.Document,
 	s := interimPool.Get().(*interim)
 
 	var br bytes.Buffer
-	if s.lastNumDocs > 0 {
-		// use previous results to initialize the buf with an estimate
-		// size, but note that the interim instance comes from a
-		// global interimPool, so multiple scorch instances indexing
-		// different docs can lead to low quality estimates
-		estimateAvgBytesPerDoc := int(float64(s.lastOutSize/s.lastNumDocs) *
-			NewSegmentBufferNumResultsFactor)
-		estimateNumResults := int(float64(len(results)+NewSegmentBufferNumResultsBump) *
-			NewSegmentBufferAvgBytesPerDocFactor)
-		br.Grow(estimateAvgBytesPerDoc * estimateNumResults)
-	}
-
+	// if s.lastNumDocs > 0 {
+	// 	// use previous results to initialize the buf with an estimate
+	// 	// size, but note that the interim instance comes from a
+	// 	// global interimPool, so multiple scorch instances indexing
+	// 	// different docs can lead to low quality estimates
+	// 	estimateAvgBytesPerDoc := int(float64(s.lastOutSize/s.lastNumDocs) *
+	// 		NewSegmentBufferNumResultsFactor)
+	// 	estimateNumResults := int(float64(len(results)+NewSegmentBufferNumResultsBump) *
+	// 		NewSegmentBufferAvgBytesPerDocFactor)
+	// 	br.Grow(estimateAvgBytesPerDoc * estimateNumResults)
+	// 	fmt.Println("=============")
+	// }
+	// fmt.Println("br size", br.Cap())
+	br.Grow(255)
 	s.results = results
 	s.chunkMode = chunkMode
 	s.w = NewCountHashWriter(&br)
-
+	fmt.Println("buffer initial capacity:", br.Cap())
 	storedIndexOffset, dictOffsets, sectionsIndexOffset, err := s.convert()
 	if err != nil {
 		return nil, uint64(0), err
 	}
-
+	fmt.Println("buffer capacity after write:", br.Cap(), "length of buffer:", len(br.Bytes()))
 	sb, err := InitSegmentBase(br.Bytes(), s.w.Sum32(), chunkMode,
 		s.FieldsMap, s.FieldsInv, uint64(len(results)),
 		storedIndexOffset, dictOffsets, sectionsIndexOffset)
@@ -225,9 +228,11 @@ func (s *interim) convert() (uint64, []uint64, uint64, error) {
 
 	// after persisting the sections to the writer, account corresponding
 	for _, opaque := range s.opaque {
+		fmt.Println("type of opaque", opaque.Type())
 		opaqueIO, ok := opaque.(segment.DiskStatsReporter)
 		if ok {
 			s.incrementBytesWritten(opaqueIO.BytesWritten())
+			fmt.Println("bytes written", opaqueIO.BytesWritten())
 		}
 	}
 
@@ -242,6 +247,7 @@ func (s *interim) convert() (uint64, []uint64, uint64, error) {
 		return 0, nil, 0, err
 	}
 
+	fmt.Println("offset values", storedIndexOffset, sectionsIndexOffset)
 	return storedIndexOffset, dictOffsets, sectionsIndexOffset, nil
 }
 

section.go

@@ -51,6 +51,7 @@ type section interface {
 type resetable interface {
 	Reset() error
 	Set(key string, value interface{})
+	Type() string
 }
 
 // -----------------------------------------------------------------------------

section_faiss_vector_index.go

@@ -59,7 +59,9 @@ func (v *faissVectorIndexSection) Process(opaque map[int]resetable, docNum uint3
 
 func (v *faissVectorIndexSection) Persist(opaque map[int]resetable, w *CountHashWriter) (n int64, err error) {
 	vo := v.getvectorIndexOpaque(opaque)
+	fmt.Println("the writer count", w.Count(), w.Sum32())
 	vo.writeVectorIndexes(w)
+	fmt.Println("the writer count after", w.Count(), w.Sum32())
 	return 0, nil
 }
 
@@ -755,6 +757,10 @@ func (v *vectorIndexOpaque) BytesRead() uint64 {
 func (v *vectorIndexOpaque) ResetBytesRead(uint64) {
 }
 
+func (v *vectorIndexOpaque) Type() string {
+	return "vector"
+}
+
 // cleanup stuff over here for reusability
 func (v *vectorIndexOpaque) Reset() (err error) {
 	// tracking the number of vecs and fields processed and tracked in this

section_inverted_text_index.go

@@ -1014,3 +1014,7 @@ func (i *invertedIndexOpaque) Set(key string, val interface{}) {
 		i.numDocs = val.(uint64)
 	}
 }
+
+func (v *invertedIndexOpaque) Type() string {
+	return "text"
+}

segment.go

@@ -322,8 +322,20 @@ func (s *SegmentBase) loadFieldsNew() error {
 		return s.loadFields()
 	}
 
+	fmt.Println("pos:", pos, "pos+binary.MaxVarintLen64:", pos+binary.MaxVarintLen64, "capacity of buffer:", cap(s.mem))
+
+	seek := pos + binary.MaxVarintLen64
+	if seek > uint64(cap(s.mem)) {
+		// handling a buffer overflow case.
+		// a rare case where the backing buffer is not large enough to be read directly via
+		// a pos+binary.MaxVarinLen64 seek. For eg, this can happen when there is only
+		// one field to be indexed in the entire batch of data and while writing out
+		// these fields metadata, you write 1 + 8 bytes whereas the MaxVarintLen64 = 10.
+		seek = uint64(cap(s.mem))
+	}
+
 	// read the number of fields
-	numFields, sz := binary.Uvarint(s.mem[pos : pos+binary.MaxVarintLen64])
+	numFields, sz := binary.Uvarint(s.mem[pos:seek])
 	pos += uint64(sz)
 	s.incrementBytesRead(uint64(sz))
 

write.go

@@ -16,6 +16,7 @@ package zap
 
 import (
 	"encoding/binary"
+	"fmt"
 	"io"
 
 	"github.com/RoaringBitmap/roaring"
@@ -53,7 +54,7 @@ func writeRoaringWithLen(r *roaring.Bitmap, w io.Writer,
 func persistFieldsSection(fieldsInv []string, w *CountHashWriter, dictLocs []uint64, opaque map[int]resetable) (uint64, error) {
 	var rv uint64
 	fieldsOffsets := make([]uint64, 0, len(fieldsInv))
-
+	fmt.Println("total number of fields:", len(fieldsInv))
 	for fieldID, fieldName := range fieldsInv {
 		// record start of this field
 		fieldsOffsets = append(fieldsOffsets, uint64(w.Count()))
@@ -70,6 +71,8 @@ func persistFieldsSection(fieldsInv []string, w *CountHashWriter, dictLocs []uin
 			return 0, err
 		}
 
+		fmt.Println(" ->field being written out:", fieldName, "number of sections:", len(segmentSections))
+
 		// write out the number of field-specific indexes
 		// FIXME hard-coding to 2, and not attempting to support sparseness well
 		_, err = writeUvarints(w, uint64(len(segmentSections)))