From b9ed2de2b4eb490cd3eccb238acb2d8d903efd19 Mon Sep 17 00:00:00 2001
From: Patrick O'Grady <prohb125@gmail.com>
Date: Sat, 13 Apr 2024 11:19:02 -0700
Subject: [PATCH] add mmap file err

---
 vilmo/vilmo.go      | 20 ++++++++------------
 vilmo/vilmo_test.go |  6 +++---
 2 files changed, 11 insertions(+), 15 deletions(-)

diff --git a/vilmo/vilmo.go b/vilmo/vilmo.go
index 98885dcc0f..e7626b930d 100644
--- a/vilmo/vilmo.go
+++ b/vilmo/vilmo.go
@@ -642,8 +642,9 @@ func (b *Batch) recycle() (bool, error) {
 
 	// Open old batch for writing
 	b.a.logger.Debug("continuing to build on old batch", zap.Uint64("old", oldestBatch), zap.Uint64("new", b.batch))
-	oldFile := filepath.Join(b.a.baseDir, strconv.FormatUint(oldestBatch, 10))
-	f, err := os.OpenFile(oldFile, os.O_WRONLY, 0666)
+	b.reuseFile = true
+	b.path = filepath.Join(b.a.baseDir, strconv.FormatUint(oldestBatch, 10))
+	f, err := os.OpenFile(b.path, os.O_WRONLY, 0666)
 	if err != nil {
 		return false, err
 	}
@@ -680,12 +681,12 @@ func (b *Batch) Abort() error {
 	}
 
 	// Cleanup aborted work
-	if len(b.movingPath) == 0 {
+	if !b.reuseFile {
 		if err := os.Remove(b.path); err != nil {
 			return err
 		}
 	} else {
-		if err := os.Truncate(b.movingPath, b.startingCursor); err != nil {
+		if err := os.Truncate(b.path, b.startingCursor); err != nil {
 			return err
 		}
 	}
@@ -751,7 +752,7 @@ func (b *Batch) writeChecksum() (ids.ID, error) {
 func (b *Batch) Prepare() (int64, bool) {
 	b.a.keyLock.Lock()
 
-	if len(b.movingPath) == 0 {
+	if b.reuseFile {
 		// Iterate over [alive] and update records for [keys] that were recycled
 		iter := b.t.alive.Iterator()
 		for next := iter.Next(); next != nil; next = iter.Next() {
@@ -773,7 +774,7 @@ func (b *Batch) Prepare() (int64, bool) {
 	// [keyLock] to ensure that [b.a.batches] is not referenced
 	// at the same time.
 	b.a.batches[b.batch] = b.t
-	return b.openWrites, b.pruneableBatch != nil && len(b.movingPath) == 0
+	return b.openWrites, b.pruneableBatch != nil && !b.reuseFile
 }
 
 func (b *Batch) Put(_ context.Context, key string, value []byte) error {
@@ -883,16 +884,11 @@ func (b *Batch) Write() (ids.ID, error) {
 		if err := preparedReader.reader.Close(); err != nil {
 			return ids.Empty, fmt.Errorf("%w: could not close old batch", err)
 		}
-		if len(b.movingPath) == 0 {
+		if !b.reuseFile {
 			preparedPath := filepath.Join(b.a.baseDir, strconv.FormatUint(preparedBatch, 10))
 			if err := os.Remove(preparedPath); err != nil {
 				return ids.Empty, fmt.Errorf("%w: could not remove old batch", err)
 			}
-		} else {
-			// TODO: ensure size is updated correctly
-			if err := os.Rename(b.movingPath, b.path); err != nil {
-				return ids.Empty, fmt.Errorf("%w: could not rename file", err)
-			}
 		}
 		delete(b.a.batches, preparedBatch)
 		oldestBatch := preparedBatch + 1
diff --git a/vilmo/vilmo_test.go b/vilmo/vilmo_test.go
index 666b244a34..1f8b6ea882 100644
--- a/vilmo/vilmo_test.go
+++ b/vilmo/vilmo_test.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"crypto/rand"
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
 	"strconv"
@@ -540,10 +541,9 @@ func TestMMapReuse(t *testing.T) {
 	require.NoError(err)
 	require.NoError(f.Close())
 
-	// Read from modified file
+	// Attempt to read from modified file (will fail)
 	_, err = m.ReadAt(b, 5)
-	require.NoError(err)
-	require.Equal([]byte("world"), b)
+	require.ErrorIs(io.EOF, err)
 }
 
 func BenchmarkVilmo(b *testing.B) {