Merge pull request #6 from ahrav/optimize-seek

Optimize the seek method

Author: bill-rich

disk_buffer_reader.go

@@ -98,44 +98,46 @@ func (dbr *DiskBufferReader) Reset() error {
 
 // Seek sets the offset for the next Read or Write to offset.
 func (dbr *DiskBufferReader) Seek(offset int64, whence int) (int64, error) {
+	newIndex := dbr.index
+
 	switch whence {
 	case io.SeekStart:
-		switch {
-		case offset < 0:
-			return 0, fmt.Errorf("can not seek to before start of reader")
-		case offset > dbr.bytesRead:
-			trashBytes := make([]byte, offset-dbr.bytesRead)
-			dbr.Read(trashBytes)
-		}
-		dbr.index = offset
+		newIndex = offset
 	case io.SeekCurrent:
-		switch {
-		case dbr.index+offset < 0:
-			return 0, fmt.Errorf("can not seek to before start of reader")
-		case offset > 0:
-			trashBytes := make([]byte, offset)
-			dbr.Read(trashBytes)
-		}
-		dbr.index += offset
+		newIndex += offset
 	case io.SeekEnd:
-		trashBytes := make([]byte, 1024)
-		for {
-			_, err := dbr.Read(trashBytes)
-			if err != nil {
-				if errors.Is(err, io.EOF) {
-					break
-				}
-				return dbr.index, err
-			}
+		newIndex = dbr.bytesRead + offset
+	}
+
+	if newIndex < 0 {
+		return 0, fmt.Errorf("can not seek to before start of reader")
+	}
+
+	// If seeking past the bytes read and recording is on, fill the gap by reading the necessary bytes.
+	if newIndex > dbr.bytesRead && dbr.recording {
+		_, err := dbr.tmpFile.Seek(0, io.SeekEnd)
+		if err != nil {
+			return 0, err
 		}
-		if dbr.index+offset < 0 {
-			return 0, fmt.Errorf("can not seek to before start of reader")
+
+		bytesToRead := int(newIndex - dbr.bytesRead)
+		trashBytes := make([]byte, bytesToRead)
+
+		n, err := dbr.reader.Read(trashBytes)
+		if err != nil && !errors.Is(err, io.EOF) {
+			return 0, err
+		}
+
+		m, err := dbr.tmpFile.Write(trashBytes[:n])
+		if err != nil {
+			return 0, err
 		}
-		dbr.index += offset
-		return dbr.index, nil
+
+		dbr.bytesRead += int64(m)
 	}
 
-	return dbr.index, nil
+	dbr.index = newIndex
+	return newIndex, nil
 }
 
 // ReadAt reads len(p) bytes into p starting at offset off in the underlying input source.

disk_buffer_reader_test.go

@@ -180,3 +180,31 @@ func TestReadAllLarge(t *testing.T) {
 		t.Fatalf("Unexpected error. Got: %s, expected: %s", testErr, baseErr)
 	}
 }
+
+func BenchmarkSeek(b *testing.B) {
+	// Use a fixed data source for consistent benchmarking.
+	data := make([]byte, 100000) // Example: 100KB of data
+	for i := range data {
+		data[i] = byte(i % 256)
+	}
+
+	tmpReader := bytes.NewBuffer(data)
+
+	dbr, err := New(tmpReader)
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer dbr.Close()
+
+	whenceOptions := []int{io.SeekStart, io.SeekCurrent, io.SeekEnd}
+	offset := int64(100) // Example offset value
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		whence := whenceOptions[i%len(whenceOptions)] // Vary whence to cover different cases
+		_, err := dbr.Seek(offset, whence)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}